diff --git a/gitm/importers/_common.py b/gitm/importers/_common.py index 996f5d1..07882ea 100644 --- a/gitm/importers/_common.py +++ b/gitm/importers/_common.py @@ -3,11 +3,11 @@ from __future__ import annotations import os +import sys import tempfile import warnings from collections.abc import Iterable from dataclasses import dataclass, field -import sys from pathlib import Path from typing import Any diff --git a/gitm/importers/analyze.py b/gitm/importers/analyze.py index 84a02ee..178799f 100644 --- a/gitm/importers/analyze.py +++ b/gitm/importers/analyze.py @@ -18,7 +18,7 @@ from gitm.optimizer.qualification import QualificationResult, qualify from gitm.planner.context import peak_for_sku from gitm.tracer.capture import write_trace_jsonl -from gitm.tracer.schema import KernelEvent, Trace +from gitm.tracer.schema import Trace # Default catalogue peak when SKU is unknown (A100 PCIe figures; called out in caveats). _DEFAULT_PEAK = HardwarePeak(name="A100", peak_flops=312e12, peak_bw_bytes_s=1555e9) diff --git a/gitm/importers/node_rollup.py b/gitm/importers/node_rollup.py index d433d07..c30bfd4 100644 --- a/gitm/importers/node_rollup.py +++ b/gitm/importers/node_rollup.py @@ -11,7 +11,7 @@ from typing import Any from gitm.optimizer.metrics import _merge_intervals -from gitm.tracer.schema import KernelEvent, Trace +from gitm.tracer.schema import Trace # Collective kernel name patterns (case-insensitive). One module-level table; # each entry is a substring match against the demangled/exported kernel name. diff --git a/gitm/importers/nsys.py b/gitm/importers/nsys.py index 68035bd..27a1af7 100644 --- a/gitm/importers/nsys.py +++ b/gitm/importers/nsys.py @@ -17,11 +17,8 @@ ImportError, ImportStats, as_int, - device_count_from_events, file_mtime_ns, - filter_device, finish_trace, - per_device_kernel_counts, ) from gitm.tracer.schema import KernelEvent, MemcpyEvent, SyncEvent, Trace, TraceEvent diff --git a/gitm/importers/torch_trace.py b/gitm/importers/torch_trace.py index c2c8be4..efe7413 100644 --- a/gitm/importers/torch_trace.py +++ b/gitm/importers/torch_trace.py @@ -24,7 +24,7 @@ finish_trace, per_device_kernel_counts, ) -from gitm.tracer.schema import KernelEvent, MemcpyEvent, Trace, TraceEvent +from gitm.tracer.schema import MemcpyEvent, Trace, TraceEvent # Max decompressed size for .json.gz (customer dumps can be multi-GB). DEFAULT_MAX_DECOMPRESSED_BYTES = 20 * 1024 ** 3 # 20 GiB @@ -91,7 +91,7 @@ def _arg_get(args: dict[str, Any] | None, keys: tuple[str, ...]) -> Any: def _as_dims(val: Any) -> tuple[int, int, int]: if val is None: return (1, 1, 1) - if isinstance(val, (list, tuple)) and len(val) >= 1: + if isinstance(val, list | tuple) and len(val) >= 1: x = as_int(val[0], 1) if len(val) > 0 else 1 y = as_int(val[1], 1) if len(val) > 1 else 1 z = as_int(val[2], 1) if len(val) > 2 else 1 diff --git a/gitm/optimizer/metrics.py b/gitm/optimizer/metrics.py index 47a0968..8d93542 100644 --- a/gitm/optimizer/metrics.py +++ b/gitm/optimizer/metrics.py @@ -26,7 +26,7 @@ from collections.abc import Callable, Iterable from dataclasses import dataclass -from gitm.tracer.schema import KernelEvent, MemcpyEvent, SyncEvent, Trace +from gitm.tracer.schema import KernelEvent, Trace FlopsModel = Callable[[KernelEvent], float] diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 1f39309..4090596 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -339,15 +339,15 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # delta is measured even on a box without CUPTI. (Runs before the empty-trace # guard for that reason; attribution below is included only if kernels exist.) if workload in _HFT_INTERVENTION_WORKLOADS: - applicator = getattr(runner, "applicator", None) - if applicator is not None: + runner_applicator = getattr(runner, "applicator", None) + if runner_applicator is not None: return _hft_intervention_result( run_dir=run_dir, run_id=run_id, workload=workload, trace=trace, qual=qual, - applicator=applicator, + applicator=runner_applicator, started_ns=started_ns, trace_path=trace_path, ) @@ -357,15 +357,15 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # A/B, gated on plDDT-equivalence). Before the empty-trace guard so the A/B # still runs on a box without CUPTI; attribution is included if kernels exist. if workload in _OPENFOLD_INTERVENTION_WORKLOADS: - applicator = getattr(runner, "applicator", None) - if applicator is not None: + runner_applicator = getattr(runner, "applicator", None) + if runner_applicator is not None: return _openfold_intervention_result( run_dir=run_dir, run_id=run_id, workload=workload, trace=trace, qual=qual, - applicator=applicator, + applicator=runner_applicator, started_ns=started_ns, trace_path=trace_path, ) @@ -375,15 +375,15 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # the fp32-vs-fp16 A/B, gated on detection-equivalence). Before the empty- # trace guard so the A/B still runs on a box without CUPTI. if workload in _EDGE_INTERVENTION_WORKLOADS: - applicator = getattr(runner, "applicator", None) - if applicator is not None: + runner_applicator = getattr(runner, "applicator", None) + if runner_applicator is not None: return _edge_intervention_result( run_dir=run_dir, run_id=run_id, workload=workload, trace=trace, qual=qual, - applicator=applicator, + applicator=runner_applicator, started_ns=started_ns, trace_path=trace_path, ) diff --git a/pyproject.toml b/pyproject.toml index 230b098..c5cfe65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ s3 = ["boto3>=1.34"] dev = [ "pytest>=8.0", "pytest-cov>=4.1", + "hypothesis>=6.0", # property-based fuzz tests for the importers (tests/test_importers_hypothesis.py) "ruff>=0.4", "mypy>=1.10", "build>=1.0", # `python -m build` for the wheel-build verification check diff --git a/tests/test_importers_hypothesis.py b/tests/test_importers_hypothesis.py index 7a343fb..e45eded 100644 --- a/tests/test_importers_hypothesis.py +++ b/tests/test_importers_hypothesis.py @@ -41,7 +41,7 @@ def test_prop_us_to_ns_monotonic(ts: float, dur: float): } ) # Default import path uses LightKernel for scale; --strict uses KernelEvent. - assert isinstance(ev, (KernelEvent, LightKernel)) + assert isinstance(ev, KernelEvent | LightKernel) assert ev.start_ns == int(ts * 1000.0) assert ev.end_ns == int((ts + dur) * 1000.0) assert ev.end_ns >= ev.start_ns or dur == 0.0 @@ -74,7 +74,7 @@ def test_prop_missing_args_never_crash(cat, name, has_grid, has_stream, has_devi ) # May be None for weird name/cat combos; must not raise. if ev is not None: - assert isinstance(ev, (KernelEvent, MemcpyEvent, LightKernel, LightMemcpy)) + assert isinstance(ev, KernelEvent | MemcpyEvent | LightKernel | LightMemcpy) assert getattr(ev, "kind", None) in ("kernel", "memcpy") assert ev.end_ns >= ev.start_ns # Light events are not pydantic; only validate the full Trace path