mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-29 07:23:25 +08:00
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:
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user