sglang-diffusion-benchmark-profile

1from __future__ import annotations2 3import argparse4import csv5import os6import subprocess7from pathlib import Path8 9OUTPUT_DIR_NAMES = {10    "benchmarks": Path("outputs/diffusion_benchmarks"),11    "profiles": Path("outputs/diffusion_profiles"),12}13 14 15def get_repo_root() -> Path:16    # Prefer the checkout that owns this skill. A benchmark helper may be17    # loaded directly from a secondary worktree while another SGLang install18    # is first on sys.path; importing that install would point assets and19    # outputs at the wrong repository.20    for parent in Path(__file__).resolve().parents:21        if (parent / "python" / "sglang" / "__init__.py").is_file():22            return parent23 24    import sglang25 26    return Path(sglang.__file__).resolve().parents[2]27 28 29def get_assets_dir(repo_root: Path | None = None) -> Path:30    root = repo_root or get_repo_root()31    return root / "inputs" / "diffusion_benchmark" / "figs"32 33 34def get_output_dir(name: str, repo_root: Path | None = None) -> Path:35    if name not in OUTPUT_DIR_NAMES:36        raise KeyError(f"Unknown output dir name: {name}")37    root = repo_root or get_repo_root()38    return root / OUTPUT_DIR_NAMES[name]39 40 41def ensure_dir(path: Path) -> Path:42    path.mkdir(parents=True, exist_ok=True)43    return path44 45 46def check_write_access(repo_root: Path | None = None) -> Path:47    root = repo_root or get_repo_root()48    probe_dir = ensure_dir(root / ".cache" / "diffusion_skill_write_test")49    probe_file = probe_dir / "probe.txt"50    probe_file.write_text("ok", encoding="utf-8")51    return probe_file52 53 54def _run_nvidia_smi(query: str) -> list[list[str]]:55    command = [56        "nvidia-smi",57        f"--query-{query}",58        "--format=csv,noheader,nounits",59    ]60    result = subprocess.run(command, check=True, capture_output=True, text=True)61    rows: list[list[str]] = []62    for raw_line in result.stdout.splitlines():63        line = raw_line.strip()64        if not line:65            continue66        rows.append([field.strip() for field in csv.reader([line]).__next__()])67    return rows68 69 70def get_gpu_inventory() -> list[dict[str, int | str]]:71    rows = _run_nvidia_smi("gpu=index,uuid,memory.used,memory.total,utilization.gpu")72    inventory = []73    for index, uuid, memory_used, memory_total, utilization_gpu in rows:74        inventory.append(75            {76                "index": int(index),77                "uuid": uuid,78                "memory_used_mib": int(memory_used),79                "memory_total_mib": int(memory_total),80                "utilization_gpu_pct": int(utilization_gpu),81            }82        )83    return inventory84 85 86def get_busy_gpu_uuids() -> set[str]:87    rows = _run_nvidia_smi("compute-apps=gpu_uuid,pid,process_name,used_gpu_memory")88    return {gpu_uuid for gpu_uuid, *_ in rows}89 90 91def pick_idle_gpus(92    required_gpus: int,93    max_memory_used_mib: int = 32,94    max_utilization_gpu_pct: int = 5,95) -> list[int]:96    inventory = get_gpu_inventory()97    busy_uuids = get_busy_gpu_uuids()98 99    idle = [100        int(gpu["index"])101        for gpu in inventory102        if gpu["uuid"] not in busy_uuids103        and int(gpu["memory_used_mib"]) <= max_memory_used_mib104        and int(gpu["utilization_gpu_pct"]) <= max_utilization_gpu_pct105    ]106    if len(idle) < required_gpus:107        raise RuntimeError(108            "Not enough idle GPUs. "109            f"required={required_gpus}, idle={idle}, inventory={inventory}, busy={sorted(busy_uuids)}"110        )111    return idle[:required_gpus]112 113 114def configure_runtime_env(required_gpus: int = 1) -> str | None:115    os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")116    if os.environ.get("CUDA_VISIBLE_DEVICES"):117        return None118    selected = ",".join(str(index) for index in pick_idle_gpus(required_gpus))119    os.environ["CUDA_VISIBLE_DEVICES"] = selected120    return selected121 122 123def main() -> None:124    parser = argparse.ArgumentParser(125        description="Resolve SGLang diffusion skill paths and idle GPUs."126    )127    parser.add_argument(128        "command",129        choices=[130            "print-root",131            "print-assets-dir",132            "print-output-dir",133            "print-idle-gpus",134            "check-write-access",135        ],136    )137    parser.add_argument(138        "--kind",139        choices=sorted(OUTPUT_DIR_NAMES),140        help="Output directory kind for print-output-dir.",141    )142    parser.add_argument(143        "--count",144        type=int,145        default=1,146        help="Number of idle GPUs to print.",147    )148    parser.add_argument(149        "--mkdir",150        action="store_true",151        help="Create the requested directory before printing it.",152    )153    args = parser.parse_args()154 155    if args.command == "print-root":156        print(get_repo_root())157        return158    if args.command == "print-assets-dir":159        path = get_assets_dir()160        if args.mkdir:161            ensure_dir(path)162        print(path)163        return164    if args.command == "print-output-dir":165        if not args.kind:166            raise SystemExit("--kind is required for print-output-dir")167        path = get_output_dir(args.kind)168        if args.mkdir:169            ensure_dir(path)170        print(path)171        return172    if args.command == "print-idle-gpus":173        print(",".join(str(index) for index in pick_idle_gpus(args.count)))174        return175    if args.command == "check-write-access":176        print(check_write_access())177        return178    raise SystemExit(f"Unhandled command: {args.command}")179 180 181if __name__ == "__main__":182    main()183