cascade-profiler 0.1.0__py3-none-any.whl

cascade_profiler/__init__.py

@@ -0,0 +1,14 @@
+"""Profiler orchestration; GPU and Modal dependencies are optional."""
+
+from cascade_profiler.pipeline import RunResult, profile
+from cascade_profiler.planning import trace_plan
+from cascade_profiler.workload import Workload
+
+__all__ = ["RunResult", "Workload", "image", "profile", "register", "trace_plan"]
+
+
+def __getattr__(name):
+    if name in {"image", "register"}:
+        from cascade_profiler import modal_support
+        return getattr(modal_support, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

cascade_profiler/align.py

@@ -0,0 +1,116 @@
+"""Ordered, step-local instruction-trace to timing alignment (no GPU dependencies)."""
+
+import csv
+import json
+import re
+from difflib import SequenceMatcher
+from pathlib import Path
+
+
+def kernel_key(kernel: dict) -> tuple:
+    def dims(value):
+        return tuple(map(int, value.split("x") if isinstance(value, str) else value))
+    return (kernel.get("mangled_name") or kernel["name"], dims(kernel["grid"]), dims(kernel["block"]))
+
+
+def read_launches(traces: Path) -> list[dict]:
+    launches = []
+    for path in sorted(traces.glob("stats_ctx_*")):
+        with path.open() as stream:
+            for row in csv.reader(stream):
+                if not row or not re.fullmatch(r"kernel-(\d+)-ctx_.*\.trace(?:\.xz)?", row[0].strip()):
+                    continue
+                if len(row) < 12:
+                    raise ValueError(f"malformed tracer stats row in {path}: {row}")
+                filename = row[0].strip()
+                raw = traces / filename
+                processed = traces / (filename.removesuffix(".xz") + "z")
+                launches.append({"tracer_uid": int(re.match(r"kernel-(\d+)-", filename)[1]),
+                                 "trace_file": filename, "name": row[1].strip(),
+                                 "grid": list(map(int, row[2:5])), "block": list(map(int, row[6:9])),
+                                 "traced": int(row[10]) > 0 or raw.is_file() or processed.is_file()})
+    ids = [r["tracer_uid"] for r in launches]
+    if len(set(ids)) != len(ids):
+        raise ValueError("duplicate tracer UIDs across contexts; cannot identify kernels unambiguously")
+    return sorted(launches, key=lambda r: r["tracer_uid"])
+
+
+def indexed_steps(steps: list[dict]) -> dict:
+    counts, result = {}, {}
+    for step in sorted(steps, key=lambda s: s["step"]):
+        repeat = step["repeat"]
+        local = step.get("step_in_repeat", counts.get(repeat, 0))
+        identity = repeat, local
+        if identity in result:
+            raise ValueError(f"duplicate repeat-local step {identity}")
+        result[identity] = step
+        counts[repeat] = local + 1
+    return result
+
+
+def align_step(tracer: list[dict], timeline: list[dict]) -> list[tuple[int, int]]:
+    """Dimension-exact anchors first; name-only fallback within the remaining ordered gaps."""
+    a, b = [kernel_key(k) for k in tracer], [kernel_key(k) for k in timeline]
+    matches = []
+    for tag, i, end_i, j, end_j in SequenceMatcher(a=a, b=b, autojunk=False).get_opcodes():
+        if tag == "equal":
+            matches.extend(zip(range(i, end_i), range(j, end_j)))
+        elif tag == "replace":
+            fallback = SequenceMatcher(a=[k[0] for k in a[i:end_i]],
+                                       b=[k[0] for k in b[j:end_j]], autojunk=False)
+            for block in fallback.get_matching_blocks():
+                matches.extend((i + block.a + n, j + block.b + n) for n in range(block.size))
+    return matches
+
+
+def build_alignment(timeline: dict, trace_steps: list[dict], launches: list[dict],
+                    provenance: dict | None = None) -> dict:
+    targets = indexed_steps(timeline["steps"])
+    steps, kernels, missing_tracer, missing_timeline = [], [], [], []
+    used = set()
+    for identity, trace_step in indexed_steps(trace_steps).items():
+        if not trace_step.get("traced"):
+            continue
+        first, end = trace_step.get("first_kernel_id"), trace_step.get("end_kernel_id")
+        if first is None or end is None:
+            raise ValueError(f"missing tracer kernel-id range for step {identity}")
+        source = [k for k in launches if k["traced"] and first <= k["tracer_uid"] < end]
+        target = targets.get(identity)
+        expected = target["kernels"] if target else []
+        pairs = align_step(source, expected)
+        matched_a, matched_b = {a for a, _ in pairs}, {b for _, b in pairs}
+        for a, b in pairs:
+            launch, kernel = source[a], expected[b]
+            if launch["tracer_uid"] in used:
+                raise ValueError("overlapping traced step kernel-id ranges")
+            used.add(launch["tracer_uid"])
+            kernels.append({"tracer_uid": launch["tracer_uid"], "trace_file": launch["trace_file"],
+                            "repeat": identity[0], "step": target["step"], "seq": kernel["seq"],
+                            "grid": launch["grid"], "block": launch["block"],
+                            "name": kernel_key(kernel)[0], "dims_match": kernel_key(launch) == kernel_key(kernel)})
+        missing_tracer.extend({**k, "repeat": identity[0], "step_in_repeat": identity[1]}
+                              for n, k in enumerate(source) if n not in matched_a)
+        missing_timeline.extend({"repeat": identity[0], "step": target["step"], "seq": k["seq"], "name": k["name"]}
+                                for n, k in enumerate(expected) if n not in matched_b)
+        steps.append({"repeat": identity[0], "step": target["step"] if target else None,
+                      "step_in_repeat": identity[1], "n_tracer": len(source), "n_timeline": len(expected),
+                      "matched": len(pairs), "unmatched_tracer": len(source) - len(pairs),
+                      "unmatched_timeline": len(expected) - len(pairs),
+                      "exact": target is not None and [kernel_key(k) for k in source] ==
+                      [kernel_key(k) for k in expected]})
+    return {"provenance": {**(provenance or {}), "run_name": timeline["meta"].get("run_name"),
+                           "traced_steps": [[s["repeat"], s["step_in_repeat"]] for s in steps],
+                           "traced_timing_used": False}, "steps": steps, "kernels": kernels,
+            "unmatched": {"n_tracer": len(missing_tracer), "n_timeline": len(missing_timeline),
+                          "tracer": missing_tracer, "timeline": missing_timeline}}
+
+
+def write_alignment(run_dir: str | Path) -> dict:
+    run = Path(run_dir)
+    manifest = json.loads((run / "run.json").read_text())
+    result = build_alignment(json.loads((run / "corrected_timeline.json").read_text()),
+                             json.loads((run / "traced/trace_steps.json").read_text()),
+                             read_launches(run / "traced/traces"),
+                             {"trace_repeat": manifest["trace_repeat"]})
+    (run / "trace_alignment.json").write_text(json.dumps(result, indent=2))
+    return result

cascade_profiler/clocks.py

@@ -0,0 +1,93 @@
+import subprocess
+import threading
+import time
+from pathlib import Path
+
+
+def _write_clock_snapshot(path: Path):
+    """nvidia-smi clock/throttle snapshot (supported, current, max, application clocks)."""
+    out = subprocess.run(["nvidia-smi", "-q", "-d", "CLOCK,PERFORMANCE,POWER"],
+                         capture_output=True, text=True)
+    path.write_text(out.stdout + out.stderr)
+
+
+class ClockSampler(threading.Thread):
+    """Samples SM/mem clocks + throttle reasons via NVML every `period_s` into a CSV.
+
+    Each row carries wall_ns (time.time_ns, same clock as vllm_trace.py's step log) so samples
+    can be mapped onto the nsys timeline through the step NVTX ranges.
+    """
+
+    COLS = ["wall_ns", "sm_mhz", "mem_mhz", "gr_mhz", "power_w", "temp_c", "throttle_reasons_hex",
+            "throttle_reasons"]
+
+    def __init__(self, path: Path, period_s: float = 0.05):
+        super().__init__(daemon=True, name="clock-sampler")
+        self.path, self.period_s = path, period_s
+        self._stop_event = threading.Event()
+        self.rows: list[dict] = []
+        self.error: str | None = None
+
+    def run(self):
+        try:
+            import pynvml as nv  # nvidia-ml-py, a vLLM dependency
+            nv.nvmlInit()
+            h = nv.nvmlDeviceGetHandleByIndex(0)
+            get_reasons = getattr(nv, "nvmlDeviceGetCurrentClocksEventReasons", None) or \
+                nv.nvmlDeviceGetCurrentClocksThrottleReasons
+            names = {getattr(nv, a): a.split("Reason", 1)[1] for a in dir(nv)
+                     if a.startswith(("nvmlClocksThrottleReason", "nvmlClocksEventReason"))
+                     and "All" not in a and "None" not in a and isinstance(getattr(nv, a), int)}
+            while not self._stop_event.is_set():
+                r = get_reasons(h)
+                self.rows.append({
+                    "wall_ns": time.time_ns(),
+                    "sm_mhz": nv.nvmlDeviceGetClockInfo(h, nv.NVML_CLOCK_SM),
+                    "mem_mhz": nv.nvmlDeviceGetClockInfo(h, nv.NVML_CLOCK_MEM),
+                    "gr_mhz": nv.nvmlDeviceGetClockInfo(h, nv.NVML_CLOCK_GRAPHICS),
+                    "power_w": nv.nvmlDeviceGetPowerUsage(h) / 1000,
+                    "temp_c": nv.nvmlDeviceGetTemperature(h, nv.NVML_TEMPERATURE_GPU),
+                    "throttle_reasons_hex": hex(r),
+                    "throttle_reasons": "|".join(sorted({n for b, n in names.items() if r & b})),
+                })
+                self._stop_event.wait(self.period_s)
+        except Exception as e:
+            self.error = repr(e)
+            print(f"[clocks] NVML sampling failed: {e!r}", flush=True)
+
+    def stop(self):
+        self._stop_event.set()
+        self.join(10)
+        import csv
+        with open(self.path, "w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=self.COLS)
+            w.writeheader()
+            w.writerows(self.rows)
+
+
+def _lock_clocks(mhz: int) -> dict:
+    """Try to lock SM clocks (needs root on the host; usually refused in containers)."""
+    res = {"requested_sm_mhz": mhz}
+    for key, cmd in (("lgc", ["nvidia-smi", "-lgc", f"{mhz},{mhz}"]),):
+        r = subprocess.run(cmd, capture_output=True, text=True)
+        res[key] = {"returncode": r.returncode, "output": (r.stdout + r.stderr).strip()}
+    res["locked"] = res["lgc"]["returncode"] == 0
+    return res
+
+
+def _clock_stats(rows: list[dict], t0_ns: int | None, t1_ns: int | None) -> dict:
+    import statistics
+    sel = [r for r in rows if t0_ns is None or t0_ns <= r["wall_ns"] <= t1_ns] or rows
+    if not sel:
+        return {}
+    out = {"samples": len(sel)}
+    for k in ("sm_mhz", "mem_mhz"):
+        v = [r[k] for r in sel]
+        out[k] = {"min": min(v), "median": statistics.median(v), "max": max(v)}
+    out["throttle_reasons_seen"] = sorted({x for r in sel for x in r["throttle_reasons"].split("|") if x})
+    return out
+
+
+def _gpu_name() -> str:
+    return subprocess.run(["nvidia-smi", "--query-gpu=name,driver_version", "--format=csv,noheader"],
+                          capture_output=True, text=True).stdout.strip()