zl_bench/test_backends.py

#!/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


class ForeignVendorModuleTests(unittest.TestCase):
    """benchmark 的可用性判据常写成 `try: import <厂商实现>; AVAILABLE = True`。

    import 只证明那个 .py 在磁盘上——它不碰设备，所以在任何后端上都成功，gate 于是
    形同虚设。router_gemm 就这样让 NVIDIA Hopper 的 kernel 跑在了 MUSA 卡上，最后
    死在 JIT 深处的 "dynamic_func() missing BLOCK_M ..."，报错完全指不向真正的原因。
    """

    P = "flag_gems.runtime.backend._"

    def _mods(self, *names):
        return {self.P + n: object() for n in names} | {"torch": object()}

    def test_only_own_vendor_is_clean(self):
        from _device_setup_plugin import foreign_vendor_modules
        mods = self._mods("mthreads", "mthreads.ops.mm")
        self.assertEqual(foreign_vendor_modules("mthreads", mods), [])

    def test_foreign_vendor_is_reported_once_per_root(self):
        from _device_setup_plugin import foreign_vendor_modules
        # 同一厂商的多个子模块只报一条
        mods = self._mods("mthreads.ops.mm",
                          "nvidia.hopper.ops.mm", "nvidia.ops.mm", "nvidia")
        self.assertEqual(foreign_vendor_modules("mthreads", mods), ["nvidia"])

    def test_several_foreign_vendors(self):
        from _device_setup_plugin import foreign_vendor_modules
        mods = self._mods("mthreads", "nvidia.ops.mm", "cambricon.ops.mm")
        self.assertEqual(foreign_vendor_modules("mthreads", mods),
                         ["cambricon", "nvidia"])

    def test_a_vendor_whose_name_extends_ours_is_still_foreign(self):
        # 按前缀比会把 _mthreads_extra 认成自家
        from _device_setup_plugin import foreign_vendor_modules
        mods = self._mods("mthreads", "mthreads_extra.ops")
        self.assertEqual(foreign_vendor_modules("mthreads", mods),
                         ["mthreads_extra"])

    def test_unrelated_modules_are_ignored(self):
        from _device_setup_plugin import foreign_vendor_modules
        mods = {"torch": object(), "triton": object(),
                "flag_gems.ops.mm": object()}
        self.assertEqual(foreign_vendor_modules("mthreads", mods), [])


class PartitionTests(unittest.TestCase):
    """按耗时把算子分到各卡。hash 分片在这批活上比它慢 50 分钟。"""

    HEADER = "op\top_file\telapsed_before\telapsed_after\tverdict"

    def _load(self, assignment):
        return sorted(sum(w for _, w in v) for v in assignment.values())

    def _screen(self, tmp, rows, header=None):
        p = Path(tmp, "screen.tsv")
        p.write_text((header or self.HEADER) + "\n" + "\n".join(rows) + "\n",
                     encoding="utf-8")
        return p

    def test_heavy_first_keeps_cards_even(self):
        from partition import plan
        items = [("a", 10), ("b", 9), ("c", 8), ("d", 7), ("e", 6), ("f", 5)]
        self.assertEqual(self._load(plan(items, ["0", "1", "2"])), [15, 15, 15])

    def test_one_item_heavier_than_the_ideal_share_sets_the_floor(self):
        # topk 单个 70 分钟、理想均分 45——加卡不解决，分片只保证别人不空等
        from partition import imbalance, plan
        items = [("topk", 70)] + [(f"op{i}", 5) for i in range(10)]
        a = plan(items, [str(i) for i in range(8)])
        self.assertEqual(max(sum(w for _, w in v) for v in a.values()), 70)
        self.assertGreater(imbalance(a), 1.0)

    def test_more_cards_than_operators_leaves_some_empty(self):
        from partition import plan
        a = plan([("a", 1), ("b", 1)], ["0", "1", "2", "3"])
        self.assertEqual(sum(1 for v in a.values() if not v), 2)

    def test_same_input_gives_the_same_plan(self):
        from partition import plan
        items = [("a", 5), ("b", 5), ("c", 5), ("d", 5)]
        first = plan(items, ["0", "1"])
        for _ in range(5):
            self.assertEqual(plan(list(reversed(items)), ["0", "1"]), first)

    def test_no_cards_is_the_callers_error(self):
        from partition import plan
        with self.assertRaises(ValueError):
            plan([("a", 1)], [])

    def test_perfect_balance_is_one(self):
        from partition import imbalance, plan
        self.assertEqual(imbalance(plan([("a", 5), ("b", 5)], ["0", "1"])), 1.0)

    def test_verdict_filters_and_reports_what_it_dropped(self):
        from partition import read_screen
        with tempfile.TemporaryDirectory() as tmp:
            f = self._screen(tmp, ["a\ta\t10\t10\tIR_CHANGED",
                                   "b\tbb\t5\t5\tNO_IR_DIFF",
                                   "c\tc\t1\t1\tNO_TRITON"])
            items, files, skipped = read_screen(f)
            self.assertEqual(items, [("a", 20.0)])
            self.assertEqual(files["a"], "a")
            self.assertEqual(len(skipped), 2)

    def test_op_file_defaults_to_op_name(self):
        from partition import read_screen
        with tempfile.TemporaryDirectory() as tmp:
            f = self._screen(tmp, ["a\t\t1\t1\tIR_CHANGED"])
            _, files, _ = read_screen(f)
            self.assertEqual(files["a"], "a")

    def test_missing_elapsed_falls_back_to_equal_weights(self):
        from partition import read_screen
        with tempfile.TemporaryDirectory() as tmp:
            f = self._screen(tmp, ["a\ta\tIR_CHANGED", "b\tb\tIR_CHANGED"],
                             header="op\top_file\tverdict")
            items, _, _ = read_screen(f)
            self.assertEqual([w for _, w in items], [1.0, 1.0])


if __name__ == "__main__":
    unittest.main()