File size: 3,305 Bytes
af66a85 | 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 | from __future__ import annotations
import os
import sys
import types
from argparse import ArgumentParser
from collections.abc import Sequence
from pathlib import Path
import pytest
ENV_PREFIX = "SLIME_CONTRACT_"
def install_paths() -> None:
current = Path(__file__).resolve()
sys.path.insert(0, str(current.parent))
sys.path.insert(0, str(current.parents[1]))
sys.path.insert(0, str(current.parents[2]))
def install_stubs(*, with_sglang_router: bool = False, with_transformers: bool = False) -> None:
if "ray" not in sys.modules:
ray_mod = types.ModuleType("ray")
ray_mod._private = types.SimpleNamespace(
services=types.SimpleNamespace(get_node_ip_address=lambda: "127.0.0.1")
)
sys.modules["ray"] = ray_mod
if with_sglang_router and "sglang_router" not in sys.modules:
mod = types.ModuleType("sglang_router")
mod.__version__ = "0.2.3"
sys.modules["sglang_router"] = mod
if with_transformers and "transformers" not in sys.modules:
mod = types.ModuleType("transformers")
mod.AutoTokenizer = type(
"AutoTokenizer", (), {"from_pretrained": staticmethod(lambda *args, **kwargs: object())}
)
mod.AutoProcessor = type(
"AutoProcessor",
(),
{"from_pretrained": staticmethod(lambda *args, **kwargs: (_ for _ in ()).throw(OSError()))},
)
mod.PreTrainedTokenizerBase = type("PreTrainedTokenizerBase", (), {})
mod.ProcessorMixin = type("ProcessorMixin", (), {})
sys.modules["transformers"] = mod
def contract_env_name(key: str) -> str:
return f"{ENV_PREFIX}{key}"
def get_contract_path(key: str, default: str | None = None) -> str | None:
return os.environ.get(contract_env_name(key), default)
def run_contract_test_for_file(
file: str,
path_args: Sequence[str] = (),
extra_args: Sequence[tuple[str, dict]] = (),
extra_setup=None,
) -> None:
"""Parse ``--xxx-path`` CLI arguments, store as ``SLIME_CONTRACT_XXX_PATH``
env vars, then call pytest on *file*.
Args:
file: ``__file__`` of the calling contract test module.
path_args: Argument names **without** the leading ``--``, e.g.
``["rollout-function-path"]``. Each is exposed as a
positional env var ``FOO_BAR_PATH`` automatically.
extra_args: Additional ``(flag, kwargs)`` pairs forwarded verbatim to
``parser.add_argument``. Useful for boolean flags like
``("--group-rm", {"action": "store_true", "default": False})``.
extra_setup: Optional callable ``fn(parsed_args)`` invoked after env
vars are set, for any file-specific side-effects.
"""
parser = ArgumentParser()
for arg in path_args:
parser.add_argument(f"--{arg}", default=None)
for flag, kwargs in extra_args:
parser.add_argument(flag, **kwargs)
parsed, remaining = parser.parse_known_args()
for key, value in vars(parsed).items():
if value and isinstance(value, str):
os.environ[contract_env_name(key.upper())] = value
if extra_setup is not None:
extra_setup(parsed)
raise SystemExit(pytest.main([file, *remaining]))
|