#!/usr/bin/env python3 """No-GPU unit tests for zl_bench backend adapters and plugin policy. FlagGems pytest invocations pass an explicit node, so this file is not collected during operator runs. Run it with `python test_backends.py`. """ from __future__ import annotations import json import os import sys import tempfile import types import unittest from pathlib import Path ROOT = Path(__file__).resolve().parent if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from backends import ( # noqa: E402 UnknownBackendError, get_backend, normalize_backend_name, ) from backends.base import ALL_VISIBILITY_KEYS # noqa: E402 from batch_summary import SUMMARY_FIELDS # noqa: E402 from plugin_plan import plan_plugins # noqa: E402 from providers.flaggems import ( # noqa: E402 VENDOR_BY_BACKEND, emit_run_env, flaggems_vendor, parse_test_node, ) from run_batch import STATUS_FIELDS, resolve_device_arg # noqa: E402 from _run_identity import init_identity, load_identity, update_identity # noqa: E402 class BackendResolveTests(unittest.TestCase): def test_empty_and_alias_resolve_to_cuda(self): self.assertEqual(normalize_backend_name(""), "cuda") self.assertEqual(get_backend("nvidia").name, "cuda") self.assertEqual(get_backend("nv").name, "cuda") self.assertEqual(get_backend("").name, "cuda") def test_other_aliases(self): self.assertEqual(get_backend("amd").name, "rocm") self.assertEqual(get_backend("hip").name, "rocm") self.assertEqual(get_backend("npu").name, "ascend") self.assertEqual(get_backend("huawei").name, "ascend") def test_unknown_backend_fails_closed(self): with self.assertRaises(UnknownBackendError): get_backend("tpu") with self.assertRaises(UnknownBackendError): get_backend("not-a-backend") class VisibilityTests(unittest.TestCase): def test_musa_drops_leftover_cuda(self): env = {"CUDA_VISIBLE_DEVICES": "7", "PATH": "/bin"} out = get_backend("musa").apply_visibility(env, "1") self.assertNotIn("CUDA_VISIBLE_DEVICES", out) self.assertEqual(out["MUSA_VISIBLE_DEVICES"], "1") self.assertEqual(out["PATH"], "/bin") for key in ALL_VISIBILITY_KEYS: if key != "MUSA_VISIBLE_DEVICES": self.assertNotIn(key, out) def test_rocm_writes_hip_and_rocr(self): out = get_backend("rocm").apply_visibility( {"CUDA_VISIBLE_DEVICES": "0"}, "3") self.assertEqual(out["HIP_VISIBLE_DEVICES"], "3") self.assertEqual(out["ROCR_VISIBLE_DEVICES"], "3") self.assertNotIn("CUDA_VISIBLE_DEVICES", out) def test_ascend_writes_both_keys(self): out = get_backend("ascend").apply_visibility({}, "2") self.assertEqual(out["ASCEND_RT_VISIBLE_DEVICES"], "2") self.assertEqual(out["ASCEND_VISIBLE_DEVICES"], "2") self.assertNotIn("CUDA_VISIBLE_DEVICES", out) def test_cuda_only_sets_cuda(self): out = get_backend("cuda").apply_visibility( {"MUSA_VISIBLE_DEVICES": "9"}, "0") self.assertEqual(out["CUDA_VISIBLE_DEVICES"], "0") self.assertNotIn("MUSA_VISIBLE_DEVICES", out) class FlagGemsVendorTests(unittest.TestCase): def test_vendor_map(self): self.assertEqual(VENDOR_BY_BACKEND["cuda"], "nvidia") self.assertEqual(VENDOR_BY_BACKEND["rocm"], "amd") self.assertEqual(flaggems_vendor(get_backend("cuda")), "nvidia") self.assertEqual(flaggems_vendor(get_backend("rocm")), "amd") self.assertEqual(flaggems_vendor(get_backend("musa")), "musa") self.assertEqual(flaggems_vendor(get_backend("ascend")), "ascend") class FlagGemsNodeTests(unittest.TestCase): _NODE = ( "/workspace/FlagGems/benchmark/test_fused_marlin_moe.py" "::test_fused_marlin_moe_mxfp4" ) def test_parse_user_command_node(self): parsed = parse_test_node(self._NODE) self.assertEqual(parsed["op"], "fused_marlin_moe_mxfp4") self.assertEqual(parsed["op_file"], "fused_marlin_moe") self.assertEqual(parsed["flaggems_dir"], "/workspace/FlagGems") self.assertTrue(parsed["test_file"].endswith( "test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4")) def test_test_node_defaults_to_upstream_shapes(self): env = emit_run_env( test_node=self._NODE, op="softmax", op_file="softmax", flaggems_dir="/workspace/dev/FlagGems", ) self.assertEqual(env["OP"], "fused_marlin_moe_mxfp4") self.assertEqual(env["OP_FILE"], "fused_marlin_moe") self.assertEqual(env["FLAGGEMS_DIR"], "/workspace/FlagGems") self.assertEqual(env["SHAPE_MODE"], "upstream") self.assertEqual(env["PASS_SHAPE_FILE"], "0") self.assertEqual(env["TEST_FILE"], self._NODE) def test_plain_run_keeps_inline_shapes(self): env = emit_run_env( op="fused_marlin_moe_mxfp4", op_file="fused_marlin_moe", flaggems_dir="/workspace/dev/FlagGems", ) self.assertEqual(env["SHAPE_MODE"], "inline") self.assertEqual(env["PASS_SHAPE_FILE"], "1") self.assertIn("test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4", env["TEST_FILE"]) def test_explicit_shape_file_wins(self): env = emit_run_env( test_node=self._NODE, shape_file="/tmp/my_shapes.yaml", ) self.assertEqual(env["SHAPE_MODE"], "file") self.assertEqual(env["PASS_SHAPE_FILE"], "1") def test_resolve_uses_existing_local_root(self): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) / "FlagGems" py = root / "benchmark" / "test_fused_marlin_moe.py" py.parent.mkdir(parents=True) py.write_text("# stub\n", encoding="utf-8") env = emit_run_env( op="fused_marlin_moe_mxfp4", op_file="fused_marlin_moe", flaggems_dir=str(root), ) self.assertEqual(env["FLAGGEMS_DIR"], str(root)) self.assertTrue(Path(env["TEST_FILE"].split("::", 1)[0]).is_file()) class PluginPlanTests(unittest.TestCase): def _plan(self, backend, **kwargs): kwargs.setdefault("extra", []) kwargs.setdefault("enable_autotune", True) kwargs.setdefault("enable_ir", True) kwargs.setdefault("enable_mm_cluster", False) return plan_plugins(backend, **kwargs) def test_cuda_does_not_put_cudagraph_on_the_list(self): names = self._plan("cuda") self.assertIn("_timing_plugin", names) self.assertIn("_device_setup_plugin", names) self.assertNotIn("_cudagraph_plugin", names) self.assertNotIn("_mm_cluster_fix_plugin", names) def test_musa_loads_its_compat_plugin_first(self): names = self._plan("musa") self.assertIn("_musa_benchmark_compat_plugin", names) # It rewrites collection-time markers, so nothing may inspect the item # list before it runs. self.assertEqual(names[0], "_musa_benchmark_compat_plugin") def test_compat_plugins_do_not_leak_to_other_backends(self): for name in ("cuda", "rocm", "ascend"): self.assertNotIn("_musa_benchmark_compat_plugin", self._plan(name), name) def test_musa_same_core_without_cluster_fix(self): # _cudagraph_plugin is absent from the plan on every backend: the timing # plugin loads it, so its absence here says nothing about capture support. names = self._plan("musa") self.assertIn("_timing_plugin", names) self.assertNotIn("_cudagraph_plugin", names) self.assertNotIn("_mm_cluster_fix_plugin", names) def test_mm_cluster_ignored_on_non_cuda(self): names = self._plan("musa", enable_mm_cluster=True) self.assertNotIn("_mm_cluster_fix_plugin", names) names = self._plan("cuda", enable_mm_cluster=True) self.assertIn("_mm_cluster_fix_plugin", names) def test_extra_plugins_isolated(self): names = self._plan("cuda", extra=["_cudagraph_plugin"]) self.assertIn("_cudagraph_plugin", names) names = self._plan("cuda", extra=[]) self.assertNotIn("_cudagraph_plugin", names) def test_extra_none_reads_env(self): old = os.environ.get("ZL_EXTRA_PLUGINS") os.environ["ZL_EXTRA_PLUGINS"] = "_cudagraph_plugin" try: names = plan_plugins( "cuda", extra=None, enable_autotune=True, enable_ir=True, enable_mm_cluster=False) self.assertIn("_cudagraph_plugin", names) finally: if old is None: os.environ.pop("ZL_EXTRA_PLUGINS", None) else: os.environ["ZL_EXTRA_PLUGINS"] = old class CaptureClassificationTests(unittest.TestCase): """Graph-capture failures split three ways; the wording is vendor text.""" def test_cuda_and_musa_both_capture(self): self.assertTrue(get_backend("cuda").supports_cuda_graph) # torch_musa supplies Stream/Event/graph/MUSAGraph; measured on S5000. self.assertTrue(get_backend("musa").supports_cuda_graph) def test_recoverable_is_the_default(self): # A kernel refusing capture is ordinary and must reach the fallback. for name in ("cuda", "musa"): v = get_backend(name).classify_capture_error( RuntimeError("operation not permitted when stream is capturing")) self.assertEqual(v, "recoverable", name) def test_cuda_fatal_markers(self): b = get_backend("cuda") self.assertEqual( b.classify_capture_error(RuntimeError( "CUDA error: an illegal memory access was encountered")), "fatal") self.assertEqual( b.classify_capture_error(RuntimeError("CUDA out of memory")), "fatal") def test_torch_dsa_footer_does_not_make_everything_fatal(self): # torch appends this to every AcceleratorError; matching it loosely # would classify each capture refusal as fatal and kill the fallback. exc = RuntimeError("Cannot call current_seed during capture. Compile with " "`TORCH_USE_CUDA_DSA` to enable device-side assertions.") self.assertEqual(get_backend("cuda").classify_capture_error(exc), "recoverable") def test_musa_unknown_error_is_poisoning_not_recoverable(self): b = get_backend("musa") exc = RuntimeError("MUSA error: unknown error") self.assertEqual(b.classify_capture_error(exc), "poisoned") # The same text carries no meaning for CUDA, so it must not leak across. self.assertEqual(get_backend("cuda").classify_capture_error(exc), "recoverable") def test_musa_rng_seed_failure_stays_recoverable(self): exc = RuntimeError("Cannot call MUSAGeneratorImpl::current_seed during " "MUSA graph capture") self.assertEqual(get_backend("musa").classify_capture_error(exc), "recoverable") def test_deny_list_is_per_backend(self): self.assertTrue(get_backend("musa").denies_capture("unique2")) self.assertFalse(get_backend("musa").denies_capture("softmax")) self.assertFalse(get_backend("cuda").denies_capture("unique2")) self.assertFalse(get_backend("musa").denies_capture("")) def test_deny_list_is_keyed_by_pytest_function_name(self): # FLAGGEMS_PERF_CURRENT_OP holds the pytest function name minus `test_`, # so a deny entry written as the benchmark FILE name never matches and # fails open. Observed once: "unique" listed, op=unique2 at runtime. musa = get_backend("musa") self.assertTrue(musa.denies_capture("unique2")) self.assertFalse(musa.denies_capture("unique")) class DirectRepeatTests(unittest.TestCase): """Repetition on the direct timer. _cudagraph_plugin imports triton at module scope, so it is loaded here against a stub: the repeat logic is plain arithmetic and should stay testable on a machine with no GPU and no Triton. """ @classmethod def setUpClass(cls): cls._saved = {k: sys.modules.get(k) for k in ("triton", "triton.testing")} triton = types.ModuleType("triton") testing = types.ModuleType("triton.testing") testing.do_bench = lambda *a, **k: 0.0 testing.do_bench_cudagraph = lambda *a, **k: 0.0 testing._summarize_statistics = lambda *a, **k: 0.0 triton.testing = testing sys.modules["triton"] = triton sys.modules["triton.testing"] = testing import _cudagraph_plugin # noqa: E402 cls.cg = _cudagraph_plugin @classmethod def tearDownClass(cls): for name, mod in cls._saved.items(): if mod is None: sys.modules.pop(name, None) else: sys.modules[name] = mod sys.modules.pop("_cudagraph_plugin", None) def _run(self, value, **kwargs): """-> (result, how many times do_bench was called)""" seq, calls = [3.0, 1.0, 2.0, 9.0, 5.0], [] self.cg._ORIGINAL_DO_BENCH = ( lambda fn, **kw: (calls.append(1), seq[(len(calls) - 1) % len(seq)])[1]) old = os.environ.get("ZL_DIRECT_REPEAT") os.environ["ZL_DIRECT_REPEAT"] = value try: out = self.cg.run_direct(lambda: None, 25, 100, None, kwargs.get("quantiles"), kwargs.get("return_mode", "mean")) finally: if old is None: os.environ.pop("ZL_DIRECT_REPEAT", None) else: os.environ["ZL_DIRECT_REPEAT"] = old return out, len(calls) def test_default_is_one_call(self): self.assertEqual(self._run("1"), (3.0, 1)) def test_unset_is_one_call(self): old = os.environ.pop("ZL_DIRECT_REPEAT", None) try: self.assertEqual(self.cg.direct_repeat(), 1) finally: if old is not None: os.environ["ZL_DIRECT_REPEAT"] = old def test_median_over_calls(self): # 3, 1, 2, 9, 5 -> 3 self.assertEqual(self._run("5"), (3.0, 5)) def test_garbage_and_zero_fall_back_to_one(self): self.assertEqual(self._run("abc"), (3.0, 1)) self.assertEqual(self._run("0"), (3.0, 1)) self.assertEqual(self._run("-4"), (3.0, 1)) def test_sequence_results_are_not_repeated(self): # quantiles / return_mode="all" make do_bench return a sequence, and a # median across calls is undefined for it; take the single call. self.assertEqual(self._run("5", quantiles=[0.5])[1], 1) self.assertEqual(self._run("5", return_mode="all")[1], 1) class DeviceArgTests(unittest.TestCase): def test_devices_preferred(self): self.assertEqual(resolve_device_arg("0,1", None), "0,1") self.assertEqual(resolve_device_arg(None, "2"), "2") self.assertIsNone(resolve_device_arg(None, None)) def test_both_set_is_error(self): with self.assertRaises(ValueError): resolve_device_arg("0", "1") class CsvFieldTests(unittest.TestCase): def test_status_fields(self): for name in ("backend", "device", "gpu", "timer"): self.assertIn(name, STATUS_FIELDS) def test_summary_fields(self): for name in ("backend", "device", "timer"): self.assertIn(name, SUMMARY_FIELDS) self.assertLess(SUMMARY_FIELDS.index("backend"), SUMMARY_FIELDS.index("fg_op_name")) class IdentityMergeTests(unittest.TestCase): def test_init_then_update(self): old_out = os.environ.get("OUT_DIR") try: with tempfile.TemporaryDirectory() as tmp: init_identity(tmp, backend="cuda", device="0", op="softmax") update_identity(tmp, timer="direct") ident = load_identity(tmp) self.assertEqual(ident["schema_version"], 1) self.assertEqual(ident["provider"], "flaggems") self.assertEqual(ident["backend"], "cuda") self.assertEqual(ident["device"], "0") self.assertEqual(ident["timer"], "direct") self.assertEqual(ident["op"], "softmax") finally: if old_out is None: os.environ.pop("OUT_DIR", None) else: os.environ["OUT_DIR"] = old_out class AutotuneReplayBackendTests(unittest.TestCase): def setUp(self): self._old = os.environ.get("ZL_BACKEND") os.environ["ZL_BACKEND"] = "cuda" def tearDown(self): if self._old is None: os.environ.pop("ZL_BACKEND", None) else: os.environ["ZL_BACKEND"] = self._old def test_missing_meta_still_replays(self): from _autotune_record_plugin import _replay_backend_ok with tempfile.TemporaryDirectory() as tmp: self.assertTrue(_replay_backend_ok(tmp)) def test_matching_backend_ok(self): from _autotune_record_plugin import _replay_backend_ok with tempfile.TemporaryDirectory() as tmp: Path(tmp, "meta.json").write_text( json.dumps({"backend": "cuda"}), encoding="utf-8") self.assertTrue(_replay_backend_ok(tmp)) def test_mismatch_refuses_replay(self): from _autotune_record_plugin import _replay_backend_ok with tempfile.TemporaryDirectory() as tmp: Path(tmp, "meta.json").write_text( json.dumps({"backend": "musa"}), encoding="utf-8") self.assertFalse(_replay_backend_ok(tmp)) class CaliberGuardTests(unittest.TestCase): """两次运行口径不同产出的报告，和有效报告长得一模一样。 这些用例钉住的是本仓已经踩过的三次：record 对 record、共享 tuner cache、 timer 不同。它们都不会让任何一步报错，只会让数字变成另一个含义。 """ BASE = { "backend": "musa", "device": "S5000", "timer": "cuda_graph", "bench_level": "core", "shape_mode": "default", "tuner_cache": "isolated", "direct_repeat": 5, "autotune_mode": "replay", } def _run(self, root, name, **over): d = Path(root, name) d.mkdir(parents=True) if over.get("_no_identity"): return d ident = dict(self.BASE) ident.update({k: v for k, v in over.items() if not k.startswith("_")}) Path(d, "run_identity.json").write_text( json.dumps(ident), encoding="utf-8") return d def test_same_caliber_has_no_complaints(self): from batch_summary import caliber_complaints with tempfile.TemporaryDirectory() as tmp: runs = [("A", self._run(tmp, "a")), ("B", self._run(tmp, "b"))] self.assertEqual(caliber_complaints(runs), []) def test_record_side_pinning_the_other_is_normal(self): """A/B 的正常形态：一侧 record 定 config，另一侧 replay 它。""" from batch_summary import caliber_complaints with tempfile.TemporaryDirectory() as tmp: runs = [("r1_A", self._run(tmp, "a", autotune_mode="record")), ("r1_B", self._run(tmp, "b"))] self.assertEqual(caliber_complaints(runs), []) def test_record_against_record_is_flagged(self): from batch_summary import caliber_complaints with tempfile.TemporaryDirectory() as tmp: runs = [("A", self._run(tmp, "a", autotune_mode="record")), ("B", self._run(tmp, "b", autotune_mode="record"))] self.assertTrue(any("record" in c for c in caliber_complaints(runs))) def test_each_caliber_field_is_checked(self): from batch_summary import CALIBER_FIELDS, caliber_complaints others = {"backend": "cuda", "device": "H20", "timer": "direct", "bench_level": "comprehensive", "shape_mode": "yaml", "tuner_cache": "shared", "direct_repeat": 1} self.assertEqual(set(others), set(CALIBER_FIELDS)) for field, other in others.items(): with tempfile.TemporaryDirectory() as tmp: runs = [("A", self._run(tmp, "a")), ("B", self._run(tmp, "b", **{field: other}))] got = caliber_complaints(runs) self.assertTrue(any(c.startswith(field) for c in got), f"{field} 不一致时没被拦下: {got}") def test_missing_identity_is_flagged_not_assumed_equal(self): from batch_summary import caliber_complaints with tempfile.TemporaryDirectory() as tmp: runs = [("A", self._run(tmp, "a")), ("B", self._run(tmp, "b", _no_identity=True))] self.assertTrue(any("run_identity" in c for c in caliber_complaints(runs))) def test_shared_cache_record_warns_but_does_not_block(self): from batch_summary import caliber_complaints, caliber_hazards with tempfile.TemporaryDirectory() as tmp: runs = [("r1_A", self._run(tmp, "a", autotune_mode="record", tuner_cache="shared(~/.flaggems)")), ("r1_B", self._run(tmp, "b", tuner_cache="shared(~/.flaggems)"))] self.assertEqual(caliber_complaints(runs), []) self.assertTrue(caliber_hazards(runs)) def test_isolated_cache_has_no_hazard(self): from batch_summary import caliber_hazards with tempfile.TemporaryDirectory() as tmp: runs = [("A", self._run(tmp, "a", autotune_mode="record")), ("B", self._run(tmp, "b"))] self.assertEqual(caliber_hazards(runs), []) def test_strict_refuses_and_permissive_continues(self): from batch_summary import _report_caliber with tempfile.TemporaryDirectory() as tmp: runs = [("A", self._run(tmp, "a")), ("B", self._run(tmp, "b", timer="direct"))] self.assertFalse(_report_caliber(runs, "op", strict=True)) self.assertTrue(_report_caliber(runs, "op", strict=False)) LOG = ("Operator: op (dtype=torch.float16, mode=kernel, level=core)\n" "SUCCESS 1.000 2.000 0.500 shape\n") def test_report_ab_refuses_mismatched_pair(self): from batch_summary import CaliberMismatch, report_ab with tempfile.TemporaryDirectory() as tmp: a = self._run(tmp, "a") b = self._run(tmp, "b", timer="direct") for d in (a, b): Path(d, "run.log").write_text(self.LOG, encoding="utf-8") with self.assertRaises(CaliberMismatch): report_ab(a / "run.log", b / "run.log", "op") # 放行开关在，但必须是显式的 self.assertEqual( len(report_ab(a / "run.log", b / "run.log", "op", strict=False)), 1) class TelemetryTests(unittest.TestCase): """采样器命令与解析。两边归一到同一组列，campaign 才能并排读。""" # 取自 S5000 实机：值带单位，"Power Draw " 的键名有尾随空格 MUSA_JSON = """ { "Attached GPUs": "8", "GPU": [ { "Index": "7", "Temperature": {"GPU Current Temp": "36C"}, "Clocks": {"Graphics": "1888MHz", "Memory": "2500MHz", "Video": "0MHz"}, "Power Readings": {"Power Draw ": "477.68W", "Current Power Limit": "950.00W"}, "Utilization": {"Gpu": "87%", "Memory": "14%"} } ] } """ CUDA_CSV = "0, 1980, 2619, 41, 122.35, 63\n" def test_musa_parses_units_off_values(self): rows = get_backend("musa").parse_telemetry(self.MUSA_JSON) self.assertEqual(rows, [{ "index": "7", "sm_clock_mhz": "1888", "mem_clock_mhz": "2500", "temp_c": "36", "power_w": "477.68", "util_pct": "87"}]) def test_cuda_parses_csv(self): rows = get_backend("cuda").parse_telemetry(self.CUDA_CSV) self.assertEqual(rows, [{ "index": "0", "sm_clock_mhz": "1980", "mem_clock_mhz": "2619", "temp_c": "41", "power_w": "122.35", "util_pct": "63"}]) def test_both_backends_yield_the_same_columns(self): import telemetry a = get_backend("musa").parse_telemetry(self.MUSA_JSON)[0] b = get_backend("cuda").parse_telemetry(self.CUDA_CSV)[0] self.assertEqual(set(a), set(b)) self.assertEqual(set(a), set(telemetry.FIELDS)) def test_garbage_parses_to_nothing_rather_than_raising(self): for backend in ("musa", "cuda"): for raw in ("", "no such tool", "{broken json"): self.assertEqual(get_backend(backend).parse_telemetry(raw), []) def test_no_device_drops_the_flag_not_just_the_value(self): # 留下裸 "-i" 会让采样器直接失败，比采全部更糟 argv = get_backend("musa").telemetry_command("") self.assertNotIn("-i", argv) self.assertIn("--json", argv) self.assertEqual(get_backend("musa").telemetry_command("7")[-2:], ["-i", "7"]) def test_backend_without_sampler_returns_empty(self): from backends.base import DeviceBackend self.assertEqual(DeviceBackend().telemetry_command("0"), []) def test_device_comes_from_the_visibility_var_when_unpinned(self): import telemetry backend = get_backend("musa") old = os.environ.get("MUSA_VISIBLE_DEVICES") try: os.environ["MUSA_VISIBLE_DEVICES"] = "7" self.assertEqual(telemetry.resolve_device(backend, ""), "7") # 一组卡不挑第一张，交给采样器采全部 os.environ["MUSA_VISIBLE_DEVICES"] = "0,1,2" self.assertEqual(telemetry.resolve_device(backend, ""), "") self.assertEqual(telemetry.resolve_device(backend, "3"), "3") finally: if old is None: os.environ.pop("MUSA_VISIBLE_DEVICES", None) else: os.environ["MUSA_VISIBLE_DEVICES"] = old if __name__ == "__main__": unittest.main()