"""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_.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())