code/scripts/check_environment.py

#!/usr/bin/env python3
"""Validate the pinned OraRL runtime and an optional local checkpoint."""

from __future__ import annotations

import argparse
import importlib.util
import json
import shutil
import subprocess
import sys
from importlib import metadata
from pathlib import Path
from typing import Sequence

EXPECTED_DISTRIBUTIONS = {
    "flash-attn": "2.8.3",
    "flash-linear-attention": "0.4.2",
    "orarl": "0.1.0",
    "qwen-vl-utils": "0.0.14",
    "ray": "2.54.0",
    "torchcodec": "0.10.0",
    "transformers": "5.5.4",
    "vllm": "0.19.1",
}
TRAINING_MODULES = (
    "verl.trainer.main",
    "verl.trainer.ray_trainer",
    "verl.workers.fsdp_workers",
)
REQUIRED_MODEL_FILES = {
    "config.json",
    "generation_config.json",
    "model.safetensors.index.json",
    "processor_config.json",
    "tokenizer.json",
    "tokenizer_config.json",
}


def _arguments(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--require-gpu",
        action="store_true",
        help="fail when CUDA is unavailable (use this inside a GPU allocation)",
    )
    parser.add_argument(
        "--require-h20",
        action="store_true",
        help="fail unless at least one visible GPU is an NVIDIA H20",
    )
    parser.add_argument(
        "--model",
        type=Path,
        help="optionally validate checkpoint metadata and every indexed weight shard",
    )
    parser.add_argument(
        "--evaluation-only",
        action="store_true",
        help="skip the checks that only training runs need",
    )
    return parser.parse_args(argv)


def _check_versions(errors: list[str]) -> None:
    if sys.version_info < (3, 10):
        errors.append(f"Python 3.10+ is required, found {sys.version.split()[0]}")

    for distribution, expected in EXPECTED_DISTRIBUTIONS.items():
        try:
            actual = metadata.version(distribution)
        except metadata.PackageNotFoundError:
            errors.append(f"missing package: {distribution}=={expected}")
            continue
        print(f"{distribution}=={actual}")
        if actual != expected:
            errors.append(f"{distribution}: expected {expected}, found {actual}")

    qwen_module = "transformers.models.qwen3_5.modeling_qwen3_5"
    try:
        qwen_spec = importlib.util.find_spec(qwen_module)
    except (ImportError, ModuleNotFoundError):
        qwen_spec = None
    if qwen_spec is None:
        errors.append(f"Transformers does not expose {qwen_module}")


def _check_training_runtime(errors: list[str]) -> None:
    """Confirm the bundled trainer resolves from the installed environment."""

    for module in TRAINING_MODULES:
        try:
            spec = importlib.util.find_spec(module)
        except (ImportError, ModuleNotFoundError) as exc:
            errors.append(f"cannot resolve training module {module}: {exc}")
            continue
        if spec is None or spec.origin is None:
            errors.append(f"missing training module: {module}")
            continue
        print(f"{module}: {spec.origin}")


def _check_torch(require_gpu: bool, require_h20: bool, errors: list[str]) -> None:
    try:
        import torch
    except Exception as exc:
        errors.append(f"cannot import torch: {exc}")
        return

    print(f"torch=={torch.__version__}")
    print(f"torch CUDA build: {torch.version.cuda}")
    if not str(torch.__version__).startswith("2.10.0+cu129"):
        errors.append(f"expected torch 2.10.0+cu129, found {torch.__version__}")
    if torch.version.cuda != "12.9":
        errors.append(f"expected a CUDA 12.9 PyTorch build, found {torch.version.cuda}")

    if not torch.cuda.is_available():
        print("CUDA device: unavailable")
        if require_gpu or require_h20:
            errors.append("CUDA is unavailable inside this process")
        return

    names: list[str] = []
    for index in range(torch.cuda.device_count()):
        name = torch.cuda.get_device_name(index)
        capability = torch.cuda.get_device_capability(index)
        names.append(name)
        print(f"CUDA device {index}: {name} (compute capability {capability[0]}.{capability[1]})")

    if require_h20 and not any("H20" in name.upper() for name in names):
        errors.append(f"expected an NVIDIA H20, found: {', '.join(names)}")


def _check_video_runtime(errors: list[str]) -> None:
    executable = shutil.which("ffmpeg")
    if executable is None:
        errors.append("missing FFmpeg executable; install ffmpeg=7 from conda-forge")
    else:
        result = subprocess.run(
            [executable, "-version"],
            check=False,
            capture_output=True,
            text=True,
        )
        version_line = result.stdout.splitlines()[0] if result.stdout else ""
        print(version_line or f"ffmpeg: {executable}")
        if result.returncode != 0:
            errors.append(f"FFmpeg failed with exit code {result.returncode}")

    try:
        import torchcodec
    except Exception as exc:
        errors.append(f"cannot import torchcodec with FFmpeg runtime: {exc}")
    else:
        print(f"torchcodec runtime=={torchcodec.__version__}")


def _check_model(model: Path, errors: list[str]) -> None:
    model = model.expanduser().resolve()
    if not model.is_dir():
        errors.append(f"model directory does not exist: {model}")
        return

    present = {path.name for path in model.iterdir() if path.is_file()}
    missing = sorted(REQUIRED_MODEL_FILES - present)
    if missing:
        errors.append(f"model is missing required files: {', '.join(missing)}")
        return

    index_path = model / "model.safetensors.index.json"
    try:
        index = json.loads(index_path.read_text(encoding="utf-8"))
        weight_map = index["weight_map"]
        shards = {str(name) for name in weight_map.values()}
    except (OSError, ValueError, KeyError, TypeError) as exc:
        errors.append(f"invalid model shard index: {exc}")
        return

    missing_shards = sorted(shard for shard in shards if not (model / shard).is_file())
    if missing_shards:
        errors.append(f"model is missing indexed weight shards: {', '.join(missing_shards)}")
        return

    try:
        from transformers import AutoConfig, AutoProcessor

        config = AutoConfig.from_pretrained(model, trust_remote_code=False)
        AutoProcessor.from_pretrained(model, trust_remote_code=False)
    except Exception as exc:
        errors.append(f"Transformers cannot load checkpoint metadata: {exc}")
        return

    print(f"model: {model}")
    print(f"model type: {config.model_type}")
    print(f"indexed weight shards: {len(shards)}")


def main(argv: Sequence[str] | None = None) -> int:
    args = _arguments(argv)
    errors: list[str] = []

    _check_versions(errors)
    if not args.evaluation_only:
        _check_training_runtime(errors)
    _check_torch(args.require_gpu, args.require_h20, errors)
    _check_video_runtime(errors)
    if args.model is not None:
        _check_model(args.model, errors)

    if errors:
        print("\nEnvironment validation failed:", file=sys.stderr)
        for error in errors:
            print(f"- {error}", file=sys.stderr)
        return 1

    print("\nOraRL environment validation passed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())