#!/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 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_same_core_without_graph_or_cluster(self): 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 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)) if __name__ == "__main__": unittest.main()