forked from zhoulin/zl_bench
Aharrypotter
03fdcbab28
Two pieces the parallel campaign needs. partition.py assigns screened operators to cards longest-first. run_batch.py's existing hash partition ignores weight: on this batch it puts the critical path at 120 minutes and leaves a card idle, where greedy gets 70. The floor is set by the single heaviest operator — topk alone is 70.6 minutes against an ideal share of 45.4 — so partitioning cannot beat it, only keep the other cards from waiting. Weights come from screen.tsv's elapsed_before + elapsed_after, which is one pair, i.e. 2 cells; N rounds is 2N cells, so the estimate is N times the weight. Getting that wrong reported topk as 11.8 minutes. RESCORE=<root> re-judges an existing screening root against the current fingerprint and rewrites its screen.tsv, keeping the original alongside. The criterion has already changed twice — index.tsv used to be hashed, and the switch's own ttg.rlc-* annotations had to be stripped — and the second change moved 24 operators from 23 IR_CHANGED to 16. Rescoring the artifacts already on disk takes seconds; re-running that screen takes 79 minutes.
739 lines
31 KiB
Python
739 lines
31 KiB
Python
#!/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()
|