File size: 4,897 Bytes
b6a9d87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Shared pytest configuration."""
from __future__ import annotations
import contextlib
import os
from pathlib import Path
import tempfile
from filelock import FileLock
import pytest
def _pin_xdist_worker_to_gpu() -> None:
"""Pin each pytest-xdist worker to a single GPU.
Runs at conftest import time, which is *before* any test module
(and therefore any ``import torch``) executes inside the worker
subprocess. pytest-xdist exposes the worker id as ``PYTEST_XDIST_WORKER``
(e.g. ``gw0``, ``gw1``). We map ``gwN`` to the Nth GPU visible to the
parent process so each worker sees exactly one GPU and they don't
contend for memory.
No-op when running outside xdist (single-process pytest).
"""
worker = os.environ.get("PYTEST_XDIST_WORKER")
if not worker or not worker.startswith("gw"):
return
try:
idx = int(worker[2:])
except ValueError:
return
visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
if visible:
gpus = [g for g in visible.split(",") if g.strip()]
if 0 <= idx < len(gpus):
os.environ["CUDA_VISIBLE_DEVICES"] = gpus[idx]
return
os.environ["CUDA_VISIBLE_DEVICES"] = str(idx)
_pin_xdist_worker_to_gpu()
def _configure_shared_caches() -> None:
"""Set shared cache env vars in os.environ once for the whole test session.
HF cache dirs are content-addressed, so all test groups safely share one
location. UV_PROJECT_ENVIRONMENT forwards the active venv to uv
subprocesses so ``uv run`` uses the same installed packages as the test
runner. Tests that need an isolated venv (e.g. SO100's lerobot_conversion
step) can strip UV_PROJECT_ENVIRONMENT from their local env dict.
"""
from test_support.runtime import build_shared_hf_cache_env, resolve_shared_uv_cache_dir
# Single shared HF cache for all test groups.
hf_env = build_shared_hf_cache_env("shared")
os.environ.update(hf_env)
if hf_env:
print(
f"\n[conftest] shared HF cache: {hf_env.get('HF_HOME', 'default')}",
flush=True,
)
uv_cache = resolve_shared_uv_cache_dir()
if uv_cache is not None:
os.environ["UV_CACHE_DIR"] = str(uv_cache)
print(f"[conftest] UV_CACHE_DIR={uv_cache}", flush=True)
# Forward the active venv to uv subprocesses.
if not os.environ.get("UV_PROJECT_ENVIRONMENT"):
venv = os.environ.get("VIRTUAL_ENV", "")
if venv:
os.environ["UV_PROJECT_ENVIRONMENT"] = venv
print(f"[conftest] UV_PROJECT_ENVIRONMENT={venv}", flush=True)
def pytest_configure(config) -> None: # noqa: ARG001
# Set before any test runs so subprocesses launched via run_bash_blocks /
# uv run inherit it — PYTEST_CURRENT_TEST alone can be cleared by uv.
os.environ["GROOT_PATCH_MISTRAL"] = "1"
os.environ["GROOT_HF_LOCAL_FIRST"] = "1"
os.environ.setdefault("GROOT_SKIP_HF_MODEL_WEIGHTS", "1")
_configure_shared_caches()
@pytest.fixture
def serialize_subprocess_spawns():
"""Let only one subprocess-spawning test run at a time, host-wide.
Under ``-n auto`` every core is already claimed by an xdist worker, so a
test that additionally spawns its own torch-importing processes
oversubscribes the box and can blow its wall-clock timeout. A host-wide
lock serializes such tests across workers; they finish in well under a
second when not starved, so the serialization is effectively free.
"""
lock_path = Path(tempfile.gettempdir()) / "gr00t-test-subprocess-spawn.lock"
with FileLock(str(lock_path)):
yield
@pytest.fixture(scope="session")
def load_hf_model_weights():
"""Temporarily opt a test into normal Hugging Face checkpoint weight loading."""
@contextlib.contextmanager
def _enabled():
previous = os.environ.get("GROOT_SKIP_HF_MODEL_WEIGHTS")
os.environ["GROOT_SKIP_HF_MODEL_WEIGHTS"] = "0"
try:
yield
finally:
if previous is None:
os.environ.pop("GROOT_SKIP_HF_MODEL_WEIGHTS", None)
else:
os.environ["GROOT_SKIP_HF_MODEL_WEIGHTS"] = previous
return _enabled
|