build.envd

# syntax=v1
# You can install envd with pip: `pip install envd`.
# Make sure you already have docker installed.
# To start the development of mosec, just run `envd up`, then you will get an
# isolated and clean Python & Rust development environment.
# To access GPU, you can run `envd up :gpu`. Make sure you have nvidia-docker
# installed.

def jax():
    install.python_packages(
        name=[
            "jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html",
            "chex",
        ]
    )


def build():
    base(dev=True, image="ubuntu:22.04")
    install.apt_packages(name=["build-essential"])
    install.uv()
    shell("fish")
    install.rust()
    install.codex()
    runtime.mount("~/.codex", "~/.codex")
    runtime.init(["make install"])


def gpu():
    build()
    install.cuda(version="11.6.2")