mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-12 04:39:51 +08:00
Piotr Mlocek
72e0268028
* 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.
107 lines
3.2 KiB
Python
107 lines
3.2 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from typing import TYPE_CHECKING
|
|
|
|
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
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def cluster_name() -> str | None:
|
|
return os.environ.get("OPENSHELL_GATEWAY")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def sandbox_client(cluster_name: str | None) -> Iterator[SandboxClient]:
|
|
with SandboxClient.from_active_cluster(cluster=cluster_name) as client:
|
|
yield client
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def ensure_sandbox_persistence_ready(sandbox_client: SandboxClient) -> None:
|
|
for _ in range(60):
|
|
try:
|
|
sandbox_client.list_ids(limit=1)
|
|
return
|
|
except grpc.RpcError as exc:
|
|
details = exc.details() or ""
|
|
if exc.code() == grpc.StatusCode.UNAVAILABLE:
|
|
time.sleep(2)
|
|
continue
|
|
if (
|
|
exc.code() == grpc.StatusCode.INTERNAL
|
|
and "no such table: objects" in details
|
|
):
|
|
time.sleep(1)
|
|
continue
|
|
raise
|
|
|
|
pytest.fail(
|
|
"openshell-server persistence is not initialized (missing sqlite objects table); "
|
|
"redeploy the active cluster and rerun e2e sandbox tests"
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sandbox(cluster_name: str | None) -> Callable[..., Sandbox]:
|
|
def _create(*, spec: object | None = None, delete_on_exit: bool = True) -> Sandbox:
|
|
return Sandbox(
|
|
cluster=cluster_name,
|
|
spec=spec,
|
|
delete_on_exit=delete_on_exit,
|
|
# The sandbox image is large (Python, Node.js, coding agents) so the
|
|
# first pod in the cluster may need extra time for the image pull.
|
|
ready_timeout_seconds=300.0,
|
|
)
|
|
|
|
return _create
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def inference_client(sandbox_client: SandboxClient) -> InferenceRouteClient:
|
|
return InferenceRouteClient.from_sandbox_client(sandbox_client)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def _worker_suffix(worker_id: str) -> str:
|
|
"""Return a suffix for worker-unique resource names.
|
|
|
|
Uses the built-in ``worker_id`` fixture from pytest-xdist which returns
|
|
``"gw0"``, ``"gw1"``, etc. for workers, or ``"master"`` for non-xdist runs.
|
|
"""
|
|
if worker_id == "master":
|
|
return ""
|
|
return f"-{worker_id}"
|
|
|
|
|
|
@pytest.fixture
|
|
def run_python() -> Callable[[Sandbox, str], tuple[int, str, str]]:
|
|
def _run(sandbox: Sandbox, code: str) -> tuple[int, str, str]:
|
|
result = sandbox.exec(["python", "-c", code], timeout_seconds=20)
|
|
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),
|
|
)
|