feat(sandbox): add gpu sandbox scheduling support

Allow sandbox creation to request GPU resources explicitly or infer them from GPU image names. This wires GPU intent through bootstrap, validates gateway support, and adds dedicated GPU E2E coverage for follow-up cluster testing.
This commit is contained in:
Piotr Mlocek
2026-03-13 10:32:58 -07:00
parent 35037465c7
commit c2603f99bd
16 changed files with 779 additions and 17 deletions
+13
View File
@@ -11,6 +11,7 @@ import grpc
import pytest
from openshell import InferenceRouteClient, Sandbox, SandboxClient
from openshell._proto import datamodel_pb2
if TYPE_CHECKING:
from collections.abc import Callable, Iterator
@@ -91,3 +92,15 @@ def run_python() -> Callable[[Sandbox, str], tuple[int, str, str]]:
return result.exit_code, result.stdout, result.stderr
return _run
@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",
)
return datamodel_pb2.SandboxSpec(
gpu=True,
template=datamodel_pb2.SandboxTemplate(image=image),
)
+30
View File
@@ -0,0 +1,30 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from typing import TYPE_CHECKING
import pytest
from openshell._proto import datamodel_pb2
if TYPE_CHECKING:
from collections.abc import Callable
from openshell import Sandbox
@pytest.mark.gpu
def test_gpu_sandbox_reports_available_gpu(
sandbox: Callable[..., Sandbox],
gpu_sandbox_spec: datamodel_pb2.SandboxSpec,
) -> None:
with sandbox(spec=gpu_sandbox_spec, delete_on_exit=True) as sb:
result = sb.exec(
["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
timeout_seconds=30,
)
assert result.exit_code == 0, result.stderr
assert result.stdout.strip()