mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-15 05:09:44 +08:00
Evan Lezar
1a9eea5351
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>
107 lines
3.4 KiB
Python
107 lines
3.4 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:
|
|
# 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),
|
|
)
|