zl_bench/providers/flaggems.py

"""FlagGems-specific mapping kept out of the generic scheduler.

GEMS_VENDOR values are FlagGems' vendor names, not zl_bench backend names.
cuda → nvidia is the important one: FlagGems' device_finder otherwise shells
out to nvidia-smi without a timeout.
"""

from __future__ import annotations

import argparse
import os
import shlex
import sys
from pathlib import Path
from typing import Dict, Optional

_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from backends.base import DeviceBackend

# Live checkouts used by the NV FlagGems product path. Order is the fallback
# sequence when the configured root does not contain the requested test file.
KNOWN_FLAGGEMS_ROOTS = (
    "/workspace/FlagGems",
    "/workspace/dev/FlagGems",
)

VENDOR_ENV_KEYS = (
    "GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND",
)

VENDOR_BY_BACKEND: Dict[str, str] = {
    "cuda": "nvidia",
    "musa": "musa",
    "rocm": "amd",
    "ascend": "ascend",
}


def flaggems_vendor(backend: DeviceBackend) -> str:
    return backend.flaggems_vendor or VENDOR_BY_BACKEND[backend.name]


def apply_vendor_env(backend: DeviceBackend) -> Optional[str]:
    """Set GEMS_VENDOR unless the user already chose one. Returns the value."""
    for key in VENDOR_ENV_KEYS:
        if key in os.environ and os.environ[key].strip():
            return os.environ[key].strip()
    vendor = flaggems_vendor(backend)
    if vendor:
        os.environ["GEMS_VENDOR"] = vendor
    return vendor or None


def make_test_node(flaggems_dir: str, op_file: str, op: str) -> str:
    root = flaggems_dir.replace("\\", "/").rstrip("/")
    return f"{root}/benchmark/test_{op_file}.py::test_{op}"


# Back-compat alias. Do not call this from emit_run_env: the parameter of the
# same name would shadow it and raise TypeError.
test_node = make_test_node


def parse_test_node(node: str) -> Dict[str, Optional[str]]:
    """Parse a FlagGems pytest node into op / op_file / optional repo root.

    Example:
      /workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4
    """
    raw = (node or "").strip()
    if "::" not in raw:
        raise ValueError(f"pytest node must contain '::': {node!r}")
    file_part, func = raw.rsplit("::", 1)
    if not func.startswith("test_"):
        raise ValueError(f"test function must start with 'test_': {func!r}")
    stem = Path(file_part.replace("\\", "/").rsplit("/", 1)[-1]).stem
    if not stem.startswith("test_"):
        raise ValueError(f"benchmark filename must start with 'test_': {file_part!r}")
    posix_file = file_part.replace("\\", "/")
    parent = posix_file.rsplit("/", 1)[0] if "/" in posix_file else ""
    parent_name = parent.rsplit("/", 1)[-1] if parent else ""
    flaggems_dir = (
        parent.rsplit("/", 1)[0] if parent_name == "benchmark" and "/" in parent else None
    )
    return {
        "test_file": f"{file_part}::{func}",
        "file": file_part,
        "op": func[len("test_"):],
        "op_file": stem[len("test_"):],
        "flaggems_dir": flaggems_dir,
    }


def _test_py(flaggems_dir: str, op_file: str) -> Path:
    return Path(flaggems_dir) / "benchmark" / f"test_{op_file}.py"


def resolve_flaggems_dir(configured: str, op_file: str) -> str:
    """Return the first root that actually contains benchmark/test_<op_file>.py."""
    candidates = []
    for item in (configured, *KNOWN_FLAGGEMS_ROOTS):
        if item and item not in candidates:
            candidates.append(item)
    for root in candidates:
        if _test_py(root, op_file).is_file():
            return root
    return configured


def resolve_existing_test_file(test_file: str, op_file: str) -> str:
    """If the node path is missing, retry the same relative path under known roots."""
    if "::" not in test_file:
        return test_file
    file_part, func = test_file.rsplit("::", 1)
    path = Path(file_part)
    if path.is_file():
        return test_file
    for root in KNOWN_FLAGGEMS_ROOTS:
        candidate = _test_py(root, op_file)
        if candidate.is_file():
            return f"{candidate.as_posix()}::{func}"
    return test_file


def emit_run_env(
    test_node: str = "",
    op: str = "",
    op_file: str = "",
    flaggems_dir: str = "",
    shape_file: str = "",
    shape_mode: str = "",
) -> Dict[str, str]:
    """Resolve OP / OP_FILE / TEST_FILE / SHAPE_MODE for run_pytest.sh."""
    parsed = parse_test_node(test_node) if test_node.strip() else None
    if parsed:
        op = parsed["op"] or op
        op_file = parsed["op_file"] or op_file
        if parsed["flaggems_dir"]:
            flaggems_dir = parsed["flaggems_dir"]
        test_file = parsed["test_file"] or ""
    else:
        flaggems_dir = resolve_flaggems_dir(flaggems_dir, op_file)
        test_file = make_test_node(flaggems_dir, op_file, op)

    flaggems_dir = resolve_flaggems_dir(flaggems_dir, op_file)
    test_file = resolve_existing_test_file(test_file, op_file)

    shape_file = shape_file.strip()
    shape_mode = shape_mode.strip().lower()
    if shape_file in ("-", "upstream", "none"):
        shape_file = ""
        shape_mode = "upstream"
    if not shape_mode:
        shape_mode = "upstream" if (test_node.strip() and not shape_file) else "inline"
    if shape_file:
        shape_mode = "file"
        pass_shape_file = "1"
    elif shape_mode == "upstream":
        pass_shape_file = "0"
    else:
        pass_shape_file = "1"
        shape_mode = "inline"

    return {
        "OP": op,
        "OP_FILE": op_file,
        "FLAGGEMS_DIR": flaggems_dir,
        "TEST_FILE": test_file,
        "SHAPE_FILE": shape_file,
        "SHAPE_MODE": shape_mode,
        "PASS_SHAPE_FILE": pass_shape_file,
    }


def _print_shell(env: Dict[str, str]) -> None:
    for key, value in env.items():
        print(f"{key}={shlex.quote(str(value))}")


def main(argv: Optional[list] = None) -> int:
    p = argparse.ArgumentParser(description="FlagGems pytest node helpers")
    p.add_argument("--parse-node", default="")
    p.add_argument("--emit-shell", action="store_true")
    p.add_argument("--op", default="")
    p.add_argument("--op-file", default="")
    p.add_argument("--flaggems-dir", default="")
    p.add_argument("--shape-file", default="")
    p.add_argument("--shape-mode", default="")
    p.add_argument("--shell", action="store_true",
                   help="with --parse-node, print OP/OP_FILE/FLAGGEMS_DIR/TEST_FILE")
    args = p.parse_args(argv)
    if args.emit_shell:
        _print_shell(emit_run_env(
            test_node=args.parse_node,
            op=args.op,
            op_file=args.op_file,
            flaggems_dir=args.flaggems_dir,
            shape_file=args.shape_file,
            shape_mode=args.shape_mode,
        ))
        return 0
    if args.parse_node:
        parsed = parse_test_node(args.parse_node)
        if args.shell:
            _print_shell({k: v or "" for k, v in parsed.items()
                          if k in ("op", "op_file", "flaggems_dir", "test_file")})
            return 0
        for key, value in parsed.items():
            print(f"{key}={value or ''}")
        return 0
    p.error("pass --emit-shell or --parse-node")
    return 2


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