diff --git a/e2e/python/conftest.py b/e2e/python/conftest.py index dc95e91e..71270492 100644 --- a/e2e/python/conftest.py +++ b/e2e/python/conftest.py @@ -96,10 +96,10 @@ def _run(sandbox: Sandbox, code: str) -> tuple[int, str, str]: @pytest.fixture(scope="session") def gpu_sandbox_spec() -> datamodel_pb2.SandboxSpec: - image = os.environ.get( - "OPENSHELL_E2E_GPU_IMAGE", - "ghcr.io/nvidia/openshell-community/sandboxes/nvidia-gpu:latest", - ) + # Empty string defers image resolution to the server, which substitutes + # the configured default sandbox image. Set OPENSHELL_E2E_GPU_IMAGE to + # override (e.g. a locally-built or registry-mirrored image). + image = os.environ.get("OPENSHELL_E2E_GPU_IMAGE", "") return datamodel_pb2.SandboxSpec( gpu=True, template=datamodel_pb2.SandboxTemplate(image=image), diff --git a/tasks/scripts/cluster-bootstrap.sh b/tasks/scripts/cluster-bootstrap.sh index 4eec288e..def2429b 100755 --- a/tasks/scripts/cluster-bootstrap.sh +++ b/tasks/scripts/cluster-bootstrap.sh @@ -243,6 +243,10 @@ fi DEPLOY_CMD=(openshell gateway start --name "${CLUSTER_NAME}" --port "${GATEWAY_PORT}") +if [ "${CLUSTER_GPU:-0}" = "1" ]; then + DEPLOY_CMD+=(--gpu) +fi + if [ -n "${GATEWAY_HOST:-}" ]; then DEPLOY_CMD+=(--gateway-host "${GATEWAY_HOST}") diff --git a/tasks/test.toml b/tasks/test.toml index 78118760..6231c21e 100644 --- a/tasks/test.toml +++ b/tasks/test.toml @@ -43,6 +43,6 @@ run = "uv run pytest -o python_files='test_*.py' -m 'not gpu' -n ${E2E_PARALLEL: ["e2e:python:gpu"] description = "Run Python GPU e2e tests" -depends = ["python:proto", "cluster"] +depends = ["python:proto", "CLUSTER_GPU=1 cluster"] env = { UV_NO_SYNC = "1", PYTHONPATH = "python" } run = "uv run pytest -o python_files='test_*.py' -m gpu -n ${E2E_PARALLEL:-1} e2e/python"