From c0f9dcf1bc55e8f56e866dc350ebb6b1beb8e050 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 23 Mar 2026 14:55:57 +0100
Subject: [PATCH] feat(tasks): wire e2e:gpu to bootstrap cluster with GPU
 support

Pass CLUSTER_GPU=1 inline in e2e:python:gpu's depends so that the
cluster is bootstrapped with --gpu when GPU e2e tests are run.

Add --gpu flag handling to cluster-bootstrap.sh and default
OPENSHELL_E2E_GPU_IMAGE to an empty string so the server resolves
the default sandbox image when no override is provided.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 e2e/python/conftest.py             | 8 ++++----
 tasks/scripts/cluster-bootstrap.sh | 4 ++++
 tasks/test.toml                    | 2 +-
 3 files changed, 9 insertions(+), 5 deletions(-)

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"