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..ff16f3a 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/pyproject.toml b/pyproject.toml index 230b098..40bb220 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 fuzzing for the importers (tests/test_importers_hypothesis.py, _fuzz.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..3237c42 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