From fc9fd5a6c3afd0d22b95a4551318474c061edaf4 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 18 Mar 2026 17:03:09 +0100 Subject: [PATCH] chore: configure devcontainer --- .devcontainer/Dockerfile | 29 ++++++++++++++++++++ .devcontainer/devcontainer.json | 47 +++++++++++++++++++++++++++++++++ .devcontainer/post-create.sh | 15 +++++++++++ pyproject.toml | 7 ++++- tests/test_jax_init.py | 22 +++++++++++++++ uv.lock | 19 +++++++++++-- 6 files changed, 136 insertions(+), 3 deletions(-) create mode 100644 .devcontainer/Dockerfile create mode 100644 .devcontainer/devcontainer.json create mode 100644 .devcontainer/post-create.sh create mode 100644 tests/test_jax_init.py diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 0000000..20ca330 --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,29 @@ +FROM mcr.microsoft.com/devcontainers/python:3.12 + +# Install uv for dependency management +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + +# Install system dependencies for JAX, OpenGL, and Mujoco +RUN apt-get update && apt-get install -y --no-install-recommends \ + libgl1 \ + libosmesa6-dev \ + libglew-dev \ + libglfw3 \ + patchelf \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Set environment variables for hardware acceleration and rendering +ENV LD_LIBRARY_PATH=/usr/lib/nvidia +ENV NVIDIA_VISIBLE_DEVICES=all +ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics + +# Set the standardized workspace path +ENV WORKSPACE_PATH=/workspaces/project +ENV PATH="$WORKSPACE_PATH/.venv/bin:$PATH" + +# Ensure the workspace directory exists and is owned by the vscode user +# This prevents permission errors when the postCreateCommand runs uv sync +RUN mkdir -p $WORKSPACE_PATH && chown vscode:vscode $WORKSPACE_PATH + +USER vscode +WORKDIR $WORKSPACE_PATH diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..b2e83cb --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,47 @@ +{ + "name": "Brittle Star JAX/CUDA", + "build": { + "dockerfile": "Dockerfile", + "context": ".." + }, + // Standardize workspace folder to /workspaces/project for IDE parity + "workspaceFolder": "/workspaces/project", + "customizations": { + "vscode": { + "extensions": [ + "ms-python.python", + "charliermarsh.ruff", + "ms-python.vscode-pylance", + "ms-toolsai.jupyter" + ], + "settings": { + "python.defaultInterpreterPath": "/workspaces/project/.venv/bin/python", + "python.analysis.localRoot": "/workspaces/project", + "python.analysis.extraPaths": [ + "/workspaces/project/.venv/lib/python3.12/site-packages" + ] + } + }, + "jetbrains": { + "plugins": [ + "Pythonid", + "fleet.python", + "com.koxudaxi.ruff" + ] + } + }, + "remoteUser": "vscode", + "runArgs": [ + "--gpus", + "all" + ], + // Ensure the .venv persists using a named volume for performance and parity + "mounts": [ + "source=brittle_star-venv,target=/workspaces/project/.venv,type=volume" + ], + // Invoke the hardware-aware sync script + "postCreateCommand": "bash .devcontainer/post-create.sh", + "features": { + "ghcr.io/devcontainers/features/common-utils:1": {} + } +} diff --git a/.devcontainer/post-create.sh b/.devcontainer/post-create.sh new file mode 100644 index 0000000..9fd5e3e --- /dev/null +++ b/.devcontainer/post-create.sh @@ -0,0 +1,15 @@ +#!/bin/bash +set -e + +# Detect if a GPU is available via nvidia-smi. +# This works for Linux hosts and Windows (WSL2) with NVIDIA Container Toolkit. +# On Mac (Apple Silicon) or systems without NVIDIA GPUs, this will skip the 'cuda' extra. +if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null; then + echo "GPU detected. Syncing with 'cuda' extra..." + uv sync --frozen --extra cuda +else + echo "No GPU detected or nvidia-smi failed. Syncing without 'cuda' extra..." + uv sync --frozen +fi + +echo "Environment synced successfully." diff --git a/pyproject.toml b/pyproject.toml index a44dd05..9343f2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ dependencies = [ "biorobot==0.4.2", "evosax==0.2.0", "ipykernel==7.2.0", - "jax[cuda13]==0.9.0.1", + "jax==0.9.0.1", "matplotlib==3.10.8", "mediapy==1.2.6", "pyopengl>=3.1.10", @@ -16,6 +16,11 @@ dependencies = [ "wandb==0.24.2", ] +[project.optional-dependencies] +cuda = [ + "jax[cuda13]==0.9.0.1", +] + [dependency-groups] dev = [ "ruff>=0.15.2", diff --git a/tests/test_jax_init.py b/tests/test_jax_init.py new file mode 100644 index 0000000..3d19422 --- /dev/null +++ b/tests/test_jax_init.py @@ -0,0 +1,22 @@ +import jax +import sys + +def verify_jax(): + print(f"Python version: {sys.version}") + print(f"JAX version: {jax.__version__}") + + devices = jax.devices() + print(f"Available devices: {devices}") + + gpu_found = any(d.device_kind == 'gpu' for d in devices) + if gpu_found: + print("SUCCESS: GPU detected!") + else: + print("INFO: Only CPU detected (expected if not in GPU-enabled environment/container).") + +if __name__ == "__main__": + try: + verify_jax() + except Exception as e: + print(f"ERROR during JAX initialization: {e}") + sys.exit(1) diff --git a/uv.lock b/uv.lock index 164ef5f..7585030 100644 --- a/uv.lock +++ b/uv.lock @@ -14,7 +14,7 @@ dependencies = [ { name = "biorobot" }, { name = "evosax" }, { name = "ipykernel" }, - { name = "jax", extra = ["cuda13"] }, + { name = "jax" }, { name = "matplotlib" }, { name = "mediapy" }, { name = "pyopengl" }, @@ -22,6 +22,11 @@ dependencies = [ { name = "wandb" }, ] +[package.optional-dependencies] +cuda = [ + { name = "jax", extra = ["cuda13"] }, +] + [package.dev-dependencies] dev = [ { name = "ruff" }, @@ -32,13 +37,15 @@ requires-dist = [ { name = "biorobot", specifier = "==0.4.2" }, { name = "evosax", specifier = "==0.2.0" }, { name = "ipykernel", specifier = "==7.2.0" }, - { name = "jax", extras = ["cuda13"], specifier = "==0.9.0.1" }, + { name = "jax", specifier = "==0.9.0.1" }, + { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" }, { name = "matplotlib", specifier = "==3.10.8" }, { name = "mediapy", specifier = "==1.2.6" }, { name = "pyopengl", specifier = ">=3.1.10" }, { name = "pyopengl-accelerate", specifier = ">=3.1.10" }, { name = "wandb", specifier = "==0.24.2" }, ] +provides-extras = ["cuda"] [package.metadata.requires-dev] dev = [{ name = "ruff", specifier = ">=0.15.2" }] @@ -511,6 +518,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/3f/efeb7c6801c46e11bd666a5180f0d615f74f72264212f74f39586c6fda9d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_x86_64.whl", hash = "sha256:ce6724bb7cb3d0543dcba17206dce909f94176e68220b8eafee72e9f92bcf542", size = 243522, upload-time = "2026-01-28T05:58:03.517Z" }, { url = "https://files.pythonhosted.org/packages/cf/b9/b04c3aa0aad2870cfe799f32f8b59789c98e1816bbce9e83f4823c5b840b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win32.whl", hash = "sha256:fca724a21a372731edb290841edd28a9fb1ee490f833392752844ac807c0086a", size = 552682, upload-time = "2026-01-28T05:58:05.649Z" }, { url = "https://files.pythonhosted.org/packages/bd/e1/6d6816b296a529ac9b897ad228b1e084eb1f92319e96371880eebdc874a6/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win_amd64.whl", hash = "sha256:823c0bd7770977d4b10e0ed0aef2f3682276b7c88b8b65cfc540afce5951392f", size = 559464, upload-time = "2026-01-28T05:58:07.261Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a8/d4dab8a58fc2e6981fc7a58c4e56ba9d777fb24931cec6a22152edbb3540/glfw-2.10.0-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:a0d1f29f206219cc291edfb6cace663a86da2470632551c998e3db82d48ea177", size = 105288, upload-time = "2026-03-10T17:21:19.929Z" }, + { url = "https://files.pythonhosted.org/packages/14/61/68d35e001872a7705112418da236fa2418d4f2e5419f8b2837f9b81bb3da/glfw-2.10.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:d28d6f3ef217e64e35dc6fd0a7acb4cec9bfe7cd14dd9b35a7228a87002de154", size = 102139, upload-time = "2026-03-10T17:21:21.645Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e1/ca5984081aaae07c9d371cb11dc4e4ff603510678ed9b73e58b6c351fe63/glfw-2.10.0-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:f968b522bb6a0e04aaf4dcac30a476d7229308bb2bac406a60587debb5a61e29", size = 229998, upload-time = "2026-03-10T17:21:23.549Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c4/82ac75fdcfba2896da7a573c0fc7f8ceb8f77ead6866d500d06c32f1c464/glfw-2.10.0-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:68cf3752bdadb6f4bc0a876247c28c88c7251ac39f8af076ed938fdfd71e72dd", size = 241944, upload-time = "2026-03-10T17:21:26.102Z" }, + { url = "https://files.pythonhosted.org/packages/e3/96/9f691823cca5eb6a08f346bd0ff03b78032db9370b509a1e9c8976fb20a5/glfw-2.10.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:44d98de5dbf8f727e0cb29f9b29d29528ea7570f2e6f42f8430a69df05f12b48", size = 231009, upload-time = "2026-03-10T17:21:28.481Z" }, + { url = "https://files.pythonhosted.org/packages/3f/93/977b9e679e356871d428ae7a1139ec767dd5177bed58a6344b4d2199e00f/glfw-2.10.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:cca5158d62189e08792b1ae54f92307a282921a0e7783315b467e21b0a381c88", size = 243480, upload-time = "2026-03-10T17:21:30.538Z" }, + { url = "https://files.pythonhosted.org/packages/f9/bd/cea9569c8f2188b0a104472951420434a3e1f5cf26f5836ef9d7227a1a30/glfw-2.10.0-py2.py3-none-win32.whl", hash = "sha256:5e024509989740e8e7b86cc4aab508195495f79879072b0e1f68bd036a2916ad", size = 552641, upload-time = "2026-03-10T17:21:32.653Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9b/4366ad3e1c0688146c70aa6143584d6a8d88583b9390f106250e25a3d5cd/glfw-2.10.0-py2.py3-none-win_amd64.whl", hash = "sha256:7f787ee8645781f10e8800438ce4357ab38c573ffb191aba380c1e72eba6311c", size = 559423, upload-time = "2026-03-10T17:21:34.766Z" }, ] [[package]]