inferencekit-cli 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (104) hide show
  1. inferencekit/__init__.py +9 -0
  2. inferencekit/__main__.py +5 -0
  3. inferencekit/analysis/__init__.py +1 -0
  4. inferencekit/analysis/ncu_csv.py +284 -0
  5. inferencekit/analysis/nsys_sqlite.py +518 -0
  6. inferencekit/canonical.py +104 -0
  7. inferencekit/checkpoint.py +396 -0
  8. inferencekit/cli.py +274 -0
  9. inferencekit/cuda/bw_peak.cu +252 -0
  10. inferencekit/cuda/torch_bw.py +60 -0
  11. inferencekit/engines/__init__.py +7 -0
  12. inferencekit/engines/adapters.py +190 -0
  13. inferencekit/engines/clip/__init__.py +13 -0
  14. inferencekit/engines/clip/comfyui.py +749 -0
  15. inferencekit/engines/clip/diffusers_ltx.py +601 -0
  16. inferencekit/engines/clip/noise.py +46 -0
  17. inferencekit/engines/clip/official_ltx.py +316 -0
  18. inferencekit/engines/clip/server.py +542 -0
  19. inferencekit/engines/clip/stepdump.py +182 -0
  20. inferencekit/engines/gates.py +94 -0
  21. inferencekit/engines/sglang_worker.py +86 -0
  22. inferencekit/engines/vllm_worker.py +129 -0
  23. inferencekit/evidence.py +427 -0
  24. inferencekit/gpuinfo.py +93 -0
  25. inferencekit/hardware.py +135 -0
  26. inferencekit/mcp_server.py +178 -0
  27. inferencekit/ops/__init__.py +30 -0
  28. inferencekit/ops/bench.py +184 -0
  29. inferencekit/ops/bench_audio.py +384 -0
  30. inferencekit/ops/bench_clip.py +819 -0
  31. inferencekit/ops/bench_image.py +393 -0
  32. inferencekit/ops/bench_rounds.py +297 -0
  33. inferencekit/ops/bench_serve.py +866 -0
  34. inferencekit/ops/ceiling.py +135 -0
  35. inferencekit/ops/compare.py +378 -0
  36. inferencekit/ops/correctness.py +153 -0
  37. inferencekit/ops/doctor.py +486 -0
  38. inferencekit/ops/gates.py +1772 -0
  39. inferencekit/ops/gates_clip.py +930 -0
  40. inferencekit/ops/intake.py +1321 -0
  41. inferencekit/ops/ledger.py +749 -0
  42. inferencekit/ops/manifest.py +309 -0
  43. inferencekit/ops/ncu.py +183 -0
  44. inferencekit/ops/nsys.py +224 -0
  45. inferencekit/ops/prompts.py +316 -0
  46. inferencekit/ops/quality.py +2161 -0
  47. inferencekit/ops/quality_clip.py +690 -0
  48. inferencekit/ops/quality_clip_steps.py +341 -0
  49. inferencekit/ops/report.py +317 -0
  50. inferencekit/ops/review.py +119 -0
  51. inferencekit/ops/rocprof.py +89 -0
  52. inferencekit/ops/roofline.py +392 -0
  53. inferencekit/ops/store.py +61 -0
  54. inferencekit/proc.py +172 -0
  55. inferencekit/prompts/ATTRIBUTION.md +56 -0
  56. inferencekit/prompts/__init__.py +167 -0
  57. inferencekit/prompts/build.py +518 -0
  58. inferencekit/prompts/dev-load.jsonl.gz +0 -0
  59. inferencekit/prompts/dev-load.manifest.json +212 -0
  60. inferencekit/prompts/dev.jsonl.gz +0 -0
  61. inferencekit/prompts/dev.manifest.json +214 -0
  62. inferencekit/prompts/exclusions.json +54 -0
  63. inferencekit/quality/__init__.py +23 -0
  64. inferencekit/quality/client.py +66 -0
  65. inferencekit/quality/data/gsm8k_LICENSE.txt +21 -0
  66. inferencekit/quality/fake.py +329 -0
  67. inferencekit/quality/gsm8k.py +224 -0
  68. inferencekit/quality/logprobs.py +106 -0
  69. inferencekit/quality/metrics.py +235 -0
  70. inferencekit/quality/niah.py +63 -0
  71. inferencekit/quality/parquet.py +296 -0
  72. inferencekit/quality/sampling.py +262 -0
  73. inferencekit/quality/step_replay.py +217 -0
  74. inferencekit/quality/video_metrics.py +304 -0
  75. inferencekit/registry.py +206 -0
  76. inferencekit/review/__init__.py +6 -0
  77. inferencekit/review/diff.py +163 -0
  78. inferencekit/review/media.py +181 -0
  79. inferencekit/review/page.py +920 -0
  80. inferencekit/review/rundir.py +844 -0
  81. inferencekit/roofline_headers.py +504 -0
  82. inferencekit/schemas/evidence.v1.json +130 -0
  83. inferencekit/schemas/evidence.v2.json +499 -0
  84. inferencekit/schemas/intake.v1.json +276 -0
  85. inferencekit/schemas/run-manifest.v1.json +118 -0
  86. inferencekit/serve/__init__.py +8 -0
  87. inferencekit/serve/audio_client.py +330 -0
  88. inferencekit/serve/client.py +446 -0
  89. inferencekit/serve/clip_client.py +351 -0
  90. inferencekit/serve/clip_rounds.py +429 -0
  91. inferencekit/serve/image_client.py +373 -0
  92. inferencekit/serve/loadgen.py +136 -0
  93. inferencekit/serve/metrics.py +203 -0
  94. inferencekit/serve/profiles.py +105 -0
  95. inferencekit/serve/prompts_io.py +125 -0
  96. inferencekit/serve/rounds.py +1111 -0
  97. inferencekit/serve/servers.py +897 -0
  98. inferencekit/stats.py +249 -0
  99. inferencekit_cli-0.1.0.dist-info/METADATA +233 -0
  100. inferencekit_cli-0.1.0.dist-info/RECORD +104 -0
  101. inferencekit_cli-0.1.0.dist-info/WHEEL +4 -0
  102. inferencekit_cli-0.1.0.dist-info/entry_points.txt +2 -0
  103. inferencekit_cli-0.1.0.dist-info/licenses/LICENSE +202 -0
  104. inferencekit_cli-0.1.0.dist-info/licenses/NOTICE +9 -0
@@ -0,0 +1,309 @@
1
+ """manifest validate: check a run manifest before anyone replays it.
2
+
3
+ The run manifest (``runs/<run-id>/manifest.json``, schema ``schemas/run-manifest.v1.json``)
4
+ is what a re-run of the manifest replays on a clean host: the model at a pinned
5
+ revision, each engine's install and launch commands, the workload and the hash of the
6
+ harness wheel. A manifest that cannot be replayed exactly is useless as evidence, so
7
+ validation has two layers:
8
+
9
+ 1. JSON Schema (types, required fields, allowed values);
10
+ 2. semantic rules the schema cannot express:
11
+ * engine labels are unique (they name directories and result columns);
12
+ * every engine has a non-empty install and launch command;
13
+ * the model revision is a full 40-character commit sha (a branch or tag can move);
14
+ * engine ports are distinct (engines may run side by side);
15
+ * no secret values in an engine's env (pass secrets by name to a re-run);
16
+ * an engine's ``host``, when given, is a loopback address (a re-run runs every
17
+ engine on the host it measures);
18
+ * hashes, when given, are 64 hex characters;
19
+ * ``workload.extra_body`` (merged into every timed request, e.g. sampling settings) sets
20
+ nothing that changes what is timed or how it is read (what the harness sends itself, the
21
+ exact-length and token-counting fields, early stops, extra sequences; EXTRA_BODY_RESERVED),
22
+ keeps chat template settings in ``workload.chat_template_kwargs`` and holds no secret.
23
+
24
+ Warnings flag things a re-run can still run but a reviewer should see, e.g. a
25
+ launch command that names neither the model revision nor a local model path, or a
26
+ local ``models/...`` path without ``model.local_dir`` (a re-run would not put the
27
+ weights there).
28
+
29
+ Paths: ``model.local_dir`` and an engine's ``cwd`` are relative to the re-run's working
30
+ directory, which is the uploaded copy of the agent's run directory; install commands and
31
+ servers run there.
32
+
33
+ An invalid manifest is an error (exit 1, every problem in the message) so that
34
+ ``manifest validate m.json && ...`` stops; ``--no-strict`` returns the result with
35
+ ``valid: false`` instead. Nothing is measured here, so no evidence record is written; the
36
+ manifest's sha256 is returned so a report can cite exactly what was checked.
37
+ """
38
+
39
+ from __future__ import annotations
40
+
41
+ import json
42
+ import re
43
+ from importlib import resources
44
+ from pathlib import Path
45
+
46
+ from inferencekit.canonical import file_sha256
47
+ from inferencekit.evidence import looks_secret_key, looks_secret_name
48
+ from inferencekit.registry import Arg, OpError, operation
49
+
50
+ SCHEMA_FILE = "run-manifest.v1.json"
51
+ FULL_SHA1 = re.compile(r"^[0-9a-f]{40}$")
52
+ SHA256 = re.compile(r"^[0-9a-f]{64}$")
53
+ LOOPBACK = ("127.0.0.1", "localhost", "::1", "[::1]")
54
+
55
+
56
+ def _is_local_path(arg: str, local_dir: str | None) -> bool:
57
+ """A launch argument that names a path on the host rather than a Hub repo id."""
58
+ if arg.startswith(("/", "./", "../", "~", "models/")):
59
+ return True
60
+ if local_dir:
61
+ d = local_dir.rstrip("/")
62
+ return arg == d or arg.startswith(d + "/") or arg in ("./" + d, "./" + d + "/")
63
+ return False
64
+
65
+
66
+ def schema_path() -> Path:
67
+ """The schema shipped in the wheel, or the repo copy in a source checkout."""
68
+ try:
69
+ packaged = resources.files("inferencekit").joinpath(f"schemas/{SCHEMA_FILE}")
70
+ if packaged.is_file():
71
+ return Path(str(packaged))
72
+ except (ModuleNotFoundError, FileNotFoundError):
73
+ pass
74
+ return Path(__file__).resolve().parents[3] / "schemas" / SCHEMA_FILE
75
+
76
+
77
+ _schema_cache: dict | None = None
78
+
79
+
80
+ def load_schema() -> dict:
81
+ global _schema_cache
82
+ if _schema_cache is None:
83
+ _schema_cache = json.loads(schema_path().read_text())
84
+ return _schema_cache
85
+
86
+
87
+ def schema_errors(manifest) -> list[str]:
88
+ import jsonschema
89
+
90
+ validator = jsonschema.Draft202012Validator(load_schema())
91
+ out = []
92
+ for err in sorted(validator.iter_errors(manifest), key=lambda e: [str(p) for p in e.path]):
93
+ where = "/".join(str(p) for p in err.path) or "(root)"
94
+ out.append(f"{where}: {err.message}")
95
+ return out
96
+
97
+
98
+ # Request fields an extra body must not set, because they change what is timed or how it is read: what the harness
99
+ # sends itself (prompt, length, streaming), the exact-length fields bench serve sets (serve/profiles.py), fields that
100
+ # end or steer generation early or add sequences, the token-id and logprob fields bench serve counts tokens with, and
101
+ # fields that change the prompt or the reply text. An extra body is for sampling settings. A re-run refuses the
102
+ # same fields when it re-runs a manifest.
103
+ EXTRA_BODY_RESERVED = {
104
+ "model": "the harness sets model itself (from the manifest)",
105
+ "messages": "the harness sets messages itself (from the prompts)",
106
+ "prompt": "the harness sets prompt itself (from the prompts)",
107
+ "max_tokens": "the harness sets max_tokens itself (from the prompts and workload.max_tokens)",
108
+ "max_completion_tokens": "the harness sets max_completion_tokens itself (as max_tokens)",
109
+ "stream": "the harness sets stream itself (every timed request streams)",
110
+ "stream_options": "the harness sets stream_options itself (usage in the last chunk)",
111
+ "n": "the harness sets n itself (one sequence per request)",
112
+ "best_of": "best_of generates extra sequences per request, so the timing is no longer one sequence's",
113
+ "ignore_eos": "bench serve sets ignore_eos for exact-length decoding (bench serve --no-exact-length turns it off)",
114
+ "min_tokens": "bench serve sets min_tokens for exact-length decoding (bench serve --no-exact-length turns it off)",
115
+ "stop": "stop ends generations early, so engines would not decode the same number of tokens",
116
+ "stop_token_ids": "stop_token_ids ends generations early, so engines would not decode the same number of tokens",
117
+ "logit_bias": "logit_bias steers which tokens are chosen (including the end of the reply)",
118
+ "return_token_ids": "bench serve sets return_token_ids to count streamed tokens",
119
+ "return_tokens_as_token_ids": "bench serve sets return_tokens_as_token_ids to count streamed tokens",
120
+ "logprobs": "bench serve sets logprobs when it counts tokens through them (and logprobs cost decode time)",
121
+ "top_logprobs": "bench serve sets top_logprobs when it counts tokens through logprobs",
122
+ "echo": "echo returns the prompt with the reply, which changes what is streamed and counted",
123
+ "suffix": "suffix changes the prompt",
124
+ }
125
+ # Settings that only matter when sampling: with none of them and no temperature, bench serve stays greedy.
126
+ SAMPLING_ONLY = ("top_k", "top_p", "min_p", "seed")
127
+
128
+
129
+ def _extra_body_schema() -> dict:
130
+ return load_schema()["properties"]["workload"]["properties"]["extra_body"]
131
+
132
+
133
+ def extra_body_problems(body, where: str = "workload/extra_body", typed: bool = True) -> tuple[list[str], list[str]]:
134
+ """(errors, warnings) for an extra request body (``workload.extra_body`` or ``bench rounds --extra-body``).
135
+
136
+ ``typed``: also check the sampling fields' types and ranges against the run-manifest schema (temperature >= 0,
137
+ 0 < top_p <= 1, ...); manifest validate passes False because its schema check already does.
138
+ """
139
+ if not isinstance(body, dict):
140
+ return [f"{where}: must be a JSON object"], []
141
+ errors, warnings = [], []
142
+ for key in body:
143
+ k = str(key)
144
+ if k in EXTRA_BODY_RESERVED:
145
+ errors.append(f"{where}/{k}: {EXTRA_BODY_RESERVED[k]}; an extra body must not set it")
146
+ elif k == "chat_template_kwargs":
147
+ errors.append(f"{where}/chat_template_kwargs: put chat template settings in "
148
+ "workload.chat_template_kwargs (they are merged per prompt there)")
149
+ elif looks_secret_key(k): # the JSON-key rule: skip_special_tokens and min_tokens are request fields
150
+ errors.append(f"{where}/{k}: looks like a secret; an extra body is stored with the evidence")
151
+ if typed:
152
+ import jsonschema
153
+
154
+ validator = jsonschema.Draft202012Validator(_extra_body_schema())
155
+ for err in sorted(validator.iter_errors(body), key=lambda e: [str(p) for p in e.path]):
156
+ errors.append(f"{'/'.join([where, *(str(p) for p in err.path)])}: {err.message}")
157
+ used = [k for k in SAMPLING_ONLY if k in body]
158
+ if used and not body.get("temperature"):
159
+ warnings.append(f"{where}: sets {', '.join(used)} but no temperature above 0; bench serve sends temperature 0 "
160
+ "(greedy decoding), so these settings change nothing on most engines")
161
+ return errors, warnings
162
+
163
+
164
+ def semantic_checks(m: dict) -> tuple[list[str], list[str]]:
165
+ """(errors, warnings) for rules JSON Schema cannot express. Tolerates a schema-invalid manifest."""
166
+ errors, warnings = [], []
167
+ model = m.get("model") if isinstance(m.get("model"), dict) else {}
168
+ rev = model.get("revision")
169
+ if isinstance(rev, str) and not FULL_SHA1.match(rev):
170
+ errors.append(f"model/revision: {rev!r} is not a full 40-character lowercase commit sha "
171
+ "(a branch, tag or short sha can point at different weights later)")
172
+ tok_rev = model.get("tokenizer_revision")
173
+ if isinstance(tok_rev, str) and tok_rev and not FULL_SHA1.match(tok_rev):
174
+ errors.append(f"model/tokenizer_revision: {tok_rev!r} is not a full 40-character commit sha")
175
+
176
+ local_dir = model.get("local_dir") if isinstance(model.get("local_dir"), str) else None
177
+ engines = m.get("engines") if isinstance(m.get("engines"), list) else []
178
+ labels: dict[str, int] = {}
179
+ ports: dict[int, str] = {}
180
+ for i, e in enumerate(engines):
181
+ if not isinstance(e, dict):
182
+ continue
183
+ where = f"engines/{i}"
184
+ label = e.get("label")
185
+ if isinstance(label, str):
186
+ if label in labels:
187
+ errors.append(f"{where}/label: {label!r} is already used by engines/{labels[label]}; "
188
+ "labels must be unique")
189
+ else:
190
+ labels[label] = i
191
+ install = e.get("install")
192
+ if not isinstance(install, list) or not [c for c in install if isinstance(c, str) and c.strip()]:
193
+ errors.append(f"{where}/install: needs at least one non-empty command")
194
+ elif any(isinstance(c, str) and not c.strip() for c in install):
195
+ errors.append(f"{where}/install: contains an empty command")
196
+ launch = e.get("launch")
197
+ if not isinstance(launch, list) or not launch or not (isinstance(launch[0], str) and launch[0].strip()):
198
+ errors.append(f"{where}/launch: needs a non-empty argv (program first)")
199
+ port = e.get("port")
200
+ if isinstance(port, int) and not isinstance(port, bool):
201
+ if port in ports:
202
+ errors.append(f"{where}/port: {port} is already used by {ports[port]}; ports must be distinct")
203
+ else:
204
+ ports[port] = label if isinstance(label, str) else where
205
+ for k, v in (e.get("env") or {}).items() if isinstance(e.get("env"), dict) else ():
206
+ if (looks_secret_name(k) or looks_secret_key(k)) and v:
207
+ errors.append(f"{where}/env/{k}: looks like a secret; never put secret values in a manifest "
208
+ "(a re-run passes secrets by environment variable name)")
209
+ host = e.get("host")
210
+ if isinstance(host, str) and host.strip().lower() not in LOOPBACK and not host.startswith("127."):
211
+ errors.append(f"{where}/host: {host!r} is not a loopback address; a re-run runs every engine on the "
212
+ "host it measures (use 127.0.0.1)")
213
+ ready = e.get("ready_check")
214
+ if isinstance(ready, str) and ready and not ready.startswith("/"):
215
+ errors.append(f"{where}/ready_check: {ready!r} must be an HTTP path starting with '/'")
216
+ if not e.get("ready_check"):
217
+ warnings.append(f"{where}: no ready_check; the harness must guess when the server is up")
218
+ if isinstance(launch, list) and isinstance(rev, str) and FULL_SHA1.match(rev):
219
+ text = " ".join(str(a) for a in launch)
220
+ if rev not in text and not any(_is_local_path(str(a), local_dir) for a in launch[1:]):
221
+ warnings.append(f"{where}/launch: does not mention model/revision {rev[:12]}... or a local model "
222
+ "path; the engine may load a different revision (a re-run runs engines "
223
+ "offline, so an unpinned repo id fails to load there)")
224
+ if isinstance(launch, list) and not local_dir and any(str(a).startswith(("models/", "./models/"))
225
+ for a in launch[1:]):
226
+ warnings.append(f"{where}/launch: uses a models/... path but model.local_dir is not set; a re-run "
227
+ "downloads the model to the HF cache, so that path will not exist unless an install "
228
+ "command creates it")
229
+ if not e.get("resolved_config"):
230
+ warnings.append(f"{where}: no resolved_config (what the engine logged it used); knob claims "
231
+ "cannot be checked against it")
232
+
233
+ harness = m.get("harness") if isinstance(m.get("harness"), dict) else {}
234
+ for key in ("wheel_sha256", "dev_prompt_set_sha256"):
235
+ value = harness.get(key)
236
+ if isinstance(value, str) and value and not SHA256.match(value):
237
+ errors.append(f"harness/{key}: {value!r} is not a 64-character lowercase sha256")
238
+ if harness.get("wheel_sha256") in (None, ""):
239
+ warnings.append("harness/wheel_sha256 is empty: a re-run refuses a manifest without the wheel hash")
240
+
241
+ workload = m.get("workload") if isinstance(m.get("workload"), dict) else {}
242
+ modes = workload.get("modes") if isinstance(workload.get("modes"), list) else []
243
+ slo = workload.get("slo") if isinstance(workload.get("slo"), dict) else {}
244
+ if "load" in modes and not (slo.get("ttft_p95_ms") and slo.get("tpot_p95_ms")):
245
+ warnings.append("workload: load mode without slo.ttft_p95_ms and slo.tpot_p95_ms; N* cannot be computed")
246
+ if isinstance(workload.get("extra_body"), dict):
247
+ body_errors, body_warnings = extra_body_problems(workload["extra_body"], typed=False) # schema: types
248
+ errors += body_errors
249
+ warnings += body_warnings
250
+ return errors, warnings
251
+
252
+
253
+ def pin_errors(manifest, run_dir) -> list[str]:
254
+ """The intake's pinned workload.chat_template_kwargs (``<run_dir>/intake.normalized.json``) against the
255
+ manifest's: every pinned key must be there with the pinned value."""
256
+ from inferencekit.ops.intake import PIN_FILE, intake_pin, pin_conflicts
257
+
258
+ pin = intake_pin(run_dir)
259
+ if not pin or not isinstance(manifest, dict):
260
+ return []
261
+ have = (manifest.get("workload") or {}).get("chat_template_kwargs") if isinstance(
262
+ manifest.get("workload"), dict) else None
263
+ have = have if isinstance(have, dict) else {}
264
+ bad = pin_conflicts(have, pin) + [f"{k}: missing (pinned {json.dumps(v)})" for k, v in sorted(pin.items())
265
+ if k not in have]
266
+ if not bad:
267
+ return []
268
+ return [f"workload/chat_template_kwargs: does not match the intake's pin in {PIN_FILE} ({'; '.join(bad)}); "
269
+ f"set workload.chat_template_kwargs to {json.dumps(pin, sort_keys=True)}"]
270
+
271
+
272
+ def check_manifest(manifest, run_dir=None) -> dict:
273
+ """Validate a parsed manifest: {valid, errors, warnings}. Used by a re-run script too. With
274
+ ``run_dir``, the intake's pinned chat_template_kwargs there (intake.normalized.json) must be honoured."""
275
+ errors = schema_errors(manifest)
276
+ sem_errors, warnings = semantic_checks(manifest) if isinstance(manifest, dict) else ([], [])
277
+ errors += [e for e in sem_errors if e not in errors]
278
+ errors += pin_errors(manifest, run_dir)
279
+ return {"valid": not errors, "errors": errors, "warnings": warnings}
280
+
281
+
282
+ def load_manifest(path: str | Path) -> dict:
283
+ p = Path(path)
284
+ if not p.is_file():
285
+ raise OpError(f"manifest not found: {path}")
286
+ try:
287
+ return json.loads(p.read_text())
288
+ except ValueError as e:
289
+ raise OpError(f"{path} is not valid JSON: {e}") from None
290
+
291
+
292
+ @operation("validate_manifest", ("manifest", "validate"), "Check a run manifest (schema + semantic rules: unique "
293
+ "engine labels, non-empty install/launch, full commit sha revision, distinct ports, no secrets, and the "
294
+ "intake's pinned chat_template_kwargs when intake.normalized.json sits next to it) before anyone "
295
+ "re-runs it on a clean machine.",
296
+ [Arg("file", "string", "Path to runs/<run-id>/manifest.json.", required=True, positional=True),
297
+ Arg("strict", "boolean", "Exit non-zero (an error listing every problem) when the manifest is invalid; "
298
+ "--no-strict returns the result with valid: false instead.", default=True)])
299
+ def validate_manifest(call, file, strict): # noqa: A002
300
+ manifest = load_manifest(file)
301
+ result = check_manifest(manifest, Path(file).resolve().parent) # runs/<run-id>/: the intake pin, if any
302
+ sha, size = file_sha256(file)
303
+ out = {"file": str(file), "sha256": sha, "bytes": size, "schema": "run-manifest.v1", **result,
304
+ "engines": [e.get("label") for e in manifest.get("engines", []) if isinstance(e, dict)]
305
+ if isinstance(manifest, dict) else []}
306
+ if strict is not False and not result["valid"]:
307
+ raise OpError(f"{file}: invalid run manifest ({len(result['errors'])} problem(s)): "
308
+ + "; ".join(result["errors"]))
309
+ return out
@@ -0,0 +1,183 @@
1
+ """ncu capture / ncu summarize: hardware counters for a few chosen kernel launches.
2
+
3
+ Capture command (allowlisted environment):
4
+
5
+ ncu --clock-control none --kernel-name-base function -k regex:^(?:R)$
6
+ --launch-skip N --launch-count C --set detailed -o <out> <argv>
7
+
8
+ --clock-control none keep the GPU's normal clocks. By default ncu locks clocks to
9
+ base, which makes memory-bound kernels look different from
10
+ how they run in production.
11
+ --kernel-name-base function + an anchored regex
12
+ match the plain function name exactly, so "gemv" does not
13
+ also catch "gemv_splitk_reduce" by accident.
14
+ --launch-skip / --launch-count
15
+ skip warm-up launches, then profile C launches. ncu replays
16
+ each kernel many times to collect all counters, so profile
17
+ as few as possible.
18
+
19
+ The .ncu-rep stores the full profiler command line (so the profiled command too), the
20
+ host name and the host's IP addresses: treat it as sensitive. The evidence record keeps
21
+ only a digest of the command and the names of the variables passed.
22
+ """
23
+
24
+ from __future__ import annotations
25
+
26
+ from pathlib import Path
27
+
28
+ from inferencekit import evidence as ev
29
+ from inferencekit.analysis import ncu_csv
30
+ from inferencekit.gpuinfo import gpu_identity
31
+ from inferencekit.proc import clean_env, dropped_config_env, require, run, tool_version
32
+ from inferencekit.registry import Arg, OpError, operation
33
+
34
+
35
+ def anchor(regex: str) -> str:
36
+ """Anchor a kernel-name regex at both ends unless the caller already anchored it."""
37
+ if regex.startswith("^") or regex.endswith("$"):
38
+ return regex
39
+ return f"^(?:{regex})$"
40
+
41
+
42
+ @operation("capture_ncu", ("ncu", "capture"), "Targeted Nsight Compute run (--clock-control none, anchored kernel "
43
+ "regex, launch skip/count) under an allowlisted environment.",
44
+ [Arg("kernel_regex", "string", "Kernel function-name regex (anchored at both ends unless it has ^ or $).",
45
+ required=True),
46
+ Arg("launch_skip", "integer", "Matching launches to skip first (warm-up).", default=0),
47
+ Arg("count", "integer", "Matching launches to profile.", default=1),
48
+ Arg("set", "string", "ncu section set.", default="detailed", choices=("basic", "detailed", "full")),
49
+ Arg("keep_env", "array", "Extra environment variable NAMES to pass through (names only in the record; "
50
+ "secret-looking names are refused).", default=[]),
51
+ Arg("timeout", "integer", "Seconds before the run (and every process it started) is killed.",
52
+ default=3600)],
53
+ passthrough=True)
54
+ def capture_ncu(call, kernel_regex, launch_skip, count, set, keep_env, timeout, argv): # noqa: A002
55
+ ncu = require("ncu", "Install Nsight Compute (it ships with the CUDA toolkit).")
56
+ env, env_names = clean_env(keep_env) # raises on secret-looking names, before anything runs
57
+ dropped = dropped_config_env(env)
58
+ store = ev.Store()
59
+ with store.staging_dir() as staging: # removed again if anything below fails
60
+ base = staging / "kernels"
61
+ cmd = [ncu, "--clock-control", "none", "--kernel-name-base", "function", "-k",
62
+ "regex:" + anchor(kernel_regex), "--launch-skip", str(launch_skip), "--launch-count", str(count),
63
+ "--set", set, "-f", "-o", str(base)]
64
+ r = run(cmd + list(argv), timeout=timeout, env=env, stdout_path=staging / "stdout.txt",
65
+ stderr_path=staging / "stderr.txt")
66
+ text = r.stdout + r.stderr
67
+ for name in ("stdout.txt", "stderr.txt"): # the program's own output may echo secrets
68
+ p = staging / name
69
+ p.write_text(ev.redact_text(p.read_text(errors="replace")))
70
+ rep = base.with_suffix(".ncu-rep")
71
+ limitations, missing = [], []
72
+ if "ERR_NVGPUCTRPERM" in text:
73
+ limitations.append("no permission for GPU performance counters (run as root or set "
74
+ "NVreg_RestrictProfilingToAdminUsers=0)")
75
+ missing.append("counters")
76
+ if "No kernels were profiled" in text or not rep.is_file():
77
+ limitations.append("no kernel matched: check the regex against 'nsys analyze' top_kernels names")
78
+ missing.append("profiled kernels")
79
+ if r.returncode != 0:
80
+ limitations.append(f"profiled command exited with code {r.returncode}")
81
+ if rep.is_file():
82
+ limitations.append("the .ncu-rep stores the full profiler command line (the profiled command "
83
+ "included), the host name and the host's IP addresses: treat it as sensitive")
84
+ if dropped:
85
+ limitations.append("set here but not passed to the profiled run, so the configuration may differ from a "
86
+ "benchmark run with the full environment: " + ", ".join(dropped)
87
+ + " (pass each one that matters with --keep-env NAME)")
88
+ artifacts = [ev.artifact(staging / "stderr.txt", staging, "stderr")]
89
+ if rep.is_file():
90
+ artifacts.insert(0, ev.artifact(rep, staging, "ncu_rep"))
91
+ summary = (f"ncu {set} on /{anchor(kernel_regex)}/ skip {launch_skip} count {count}: "
92
+ + ("report written" if rep.is_file() else "NO REPORT"))
93
+ record = ev.build(call, kind="ncu_capture", claim_class="observed", summary=summary,
94
+ identity_=ev.identity(gpu=gpu_identity()), artifacts=artifacts, missing=missing,
95
+ limitations=limitations,
96
+ data={"ncu_version": tool_version("ncu"), "exit_code": r.returncode,
97
+ "env_names": env_names, "env_dropped": dropped,
98
+ "kernel_regex": anchor(kernel_regex), "launch_skip": launch_skip, "count": count,
99
+ "set": set, "clock_control": "none"})
100
+ return store.save(record, staging)
101
+
102
+
103
+ def _peak(peak_gbs: str | None) -> tuple[float | None, str | None]:
104
+ if not peak_gbs:
105
+ return None, None
106
+ try:
107
+ return float(peak_gbs), "user"
108
+ except ValueError:
109
+ pass
110
+ from inferencekit.ops.roofline import _bw_from_record
111
+
112
+ value, info, _ = _bw_from_record(ev.Store().load(peak_gbs))
113
+ return value, info.get("evidence_id")
114
+
115
+
116
+ @operation("summarize_ncu", ("ncu", "summarize"), "Per kernel: duration, DRAM GB/s and % of measured peak, SM %, "
117
+ "occupancy and its limiter, top stall reasons, bound label.",
118
+ [Arg("source", "string", "A .ncu-rep, an ncu_capture evidence id, or a raw-page CSV file.",
119
+ required=True, positional=True),
120
+ Arg("details_csv", "string", "Details-page CSV (when source is a CSV file)."),
121
+ Arg("peak_gbs", "string", "Measured DRAM peak: GB/s number or a kind=ceiling evidence id.")])
122
+ def summarize_ncu(call, source, details_csv, peak_gbs):
123
+ store = ev.Store()
124
+ src_rec, inputs = None, []
125
+ p = Path(source)
126
+ if p.suffix == ".csv" and p.is_file():
127
+ raw_text = p.read_text(errors="replace")
128
+ details_text = Path(details_csv).read_text(errors="replace") if details_csv else ""
129
+ inputs = [p] + ([Path(details_csv)] if details_csv else [])
130
+ ncu_version = None
131
+ else:
132
+ if p.suffix == ".ncu-rep" and p.is_file():
133
+ rep = p
134
+ else:
135
+ src_rec = store.load(source)
136
+ if src_rec.get("kind") != "ncu_capture":
137
+ raise OpError(f"{source}: expected a .ncu-rep, a raw-page CSV or an ncu_capture evidence id")
138
+ art = next((a for a in src_rec.get("artifacts", []) if a.get("role") == "ncu_rep"), None)
139
+ if art is None:
140
+ raise OpError(f"{source}: no .ncu-rep artifact in that record")
141
+ rep = store.run_dir(src_rec["evidence_id"]) / art["path"]
142
+ ncu = require("ncu", "ncu is needed to read .ncu-rep files")
143
+ raw = run([ncu, "--import", str(rep), "--page", "raw", "--csv"], timeout=600)
144
+ det = run([ncu, "--import", str(rep), "--page", "details", "--csv"], timeout=600)
145
+ if not raw.ok and not det.ok:
146
+ raise OpError(f"ncu --import failed: {raw.stderr[-500:]}")
147
+ raw_text, details_text = raw.stdout, det.stdout
148
+ inputs = [rep]
149
+ ncu_version = tool_version("ncu")
150
+ peak, peak_src = _peak(peak_gbs)
151
+ kernels = ncu_csv.summarize(raw_text, details_text, peak)
152
+ if not kernels:
153
+ raise OpError("no kernel rows found in the ncu output")
154
+ metrics, limitations = [], []
155
+ for k in kernels:
156
+ subject = f"#{k['index']} {k['kernel'][:80]}"
157
+ metrics.append(ev.metric("duration_us", k["duration_us"], "us", "single_launch", n=1, subject=subject))
158
+ metrics.append(ev.metric("dram_gbs", k["dram_gbs"], "GB/s", "single_launch", n=1, subject=subject))
159
+ if k["dram_pct_of_measured_peak"] is not None:
160
+ metrics.append(ev.metric("dram_pct_of_measured_peak", k["dram_pct_of_measured_peak"], "%",
161
+ "single_launch", n=1, subject=subject))
162
+ if peak is None:
163
+ limitations.append("no measured peak given: % of peak is ncu's nominal DRAM peak only (pass --peak-gbs)")
164
+ unparsed = [f"#{k['index']} {p}" for k in kernels for p in k["unparsed_units"]]
165
+ if unparsed:
166
+ limitations.append("values with an unknown or unexpected unit were left empty, not guessed: "
167
+ + "; ".join(unparsed[:8]) + (" ..." if len(unparsed) > 8 else ""))
168
+ if src_rec and ev.integrity_problem(src_rec):
169
+ limitations.append(f"source capture {ev.integrity_problem(src_rec)}")
170
+ limitations.append("bound labels are speed-of-light rules of thumb (>=60% of a unit's peak = bound by it; "
171
+ "both <40% = latency-bound), not a full analysis")
172
+ if any(k["duration_us"] and k["duration_us"] < 5 for k in kernels):
173
+ limitations.append("kernels under 5 us: launch and tail effects dominate, counters are noisy")
174
+ labels = ", ".join(f"#{k['index']} {k['kernel'][:30]}: {k['bound']}" for k in kernels[:4])
175
+ summary = f"{len(kernels)} kernel launch(es): {labels}"
176
+ gpu = src_rec["identity"]["gpu"] if src_rec else {}
177
+ record = ev.build(call, kind="ncu_kernels", claim_class="observed", summary=summary,
178
+ identity_=ev.identity(gpu=gpu), metrics=metrics,
179
+ artifacts=[ev.artifact(i, None, "input") for i in inputs], limitations=limitations,
180
+ data={"kernels": kernels, "peak_gbs": peak, "peak_source": peak_src,
181
+ "ncu_version": ncu_version,
182
+ "source_evidence": src_rec["evidence_id"] if src_rec else None})
183
+ return store.save(record)