bharathtelu commited on
Commit
a9aa4ae
·
verified ·
1 Parent(s): 67b7daf

Deploy auto-tune UI + scripts (work-from-91d0cf0)

Browse files

Branch: work-from-91d0cf0

Adds scripts/auto_tune.py (with hardcoded / llm / llm-explore modes, --model flag, merge-and-test step, --events stream), ui/auto_tune_ui.py (Streamlit frontend), and the /auto-tune FastAPI endpoint for the UI to call against a remote MI300X server. The Space's app_file now points at the auto-tune UI.

.gitignore ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ .claude/
2
+
3
+ # Python
4
+ __pycache__/
5
+ *.py[cod]
6
+ *$py.class
7
+ *.egg-info/
8
+ .venv/
9
+ venv/
10
+ .pytest_cache/
11
+ .ruff_cache/
12
+
13
+ # Goblin-specific
14
+ bench_cache/
15
+ # kb/.embeddings_cache_*.npy is intentionally TRACKED — see README "Deploying
16
+ # to Hugging Face Spaces". Shipping the cache keyed on the current YAML's
17
+ # sha256 keeps the Space cold-start fast (no sentence-transformers download).
18
+ *.trace.csv
19
+ *.trace.json
20
+ .env
21
+ .anthropic-key
22
+
23
+ # Editor / OS
24
+ .vscode/
25
+ .idea/
26
+ .DS_Store
README.md CHANGED
@@ -3,18 +3,17 @@ title: GPU Goblin
3
  emoji: 🧌
4
  colorFrom: red
5
  colorTo: red
6
- sdk: docker
7
- app_port: 7860
 
8
  pinned: false
9
  license: mit
10
- short_description: Qwen agent that hunts wasted compute on AMD MI300X.
11
  tags:
12
  - amd
13
- - amd-hackathon-2026
14
  - mi300x
15
  - rocm
16
  - qwen
17
- - vllm
18
  - huggingface
19
  - agent
20
  - fine-tuning
@@ -382,53 +381,36 @@ Application Platform + Application URL" submission fields.
382
  ### One-time setup
383
 
384
  1. Create a Hugging Face account at [huggingface.co](https://huggingface.co/)
385
- and join the **`lablab-ai-amd-developer-hackathon`** HF Organization
386
- ([direct link](https://huggingface.co/lablab-ai-amd-developer-hackathon)).
387
- The Space MUST live under that org for the hackathon Special Prize and
388
- for the lablab submission to validate.
389
  2. Create a token at [Settings → Access Tokens](https://huggingface.co/settings/tokens)
390
- with **`write`** scope. Export it:
391
- ```bash
392
- export HF_TOKEN=hf_yourtokenhere
393
- ```
394
- 3. On the org's page, click **"New Space"**:
395
- - **Owner:** `lablab-ai-amd-developer-hackathon`
396
- - **Space name:** `gpu-goblin` (or your preferred slug)
397
- - **License:** MIT
398
- - **SDK:** **Streamlit**
399
- - **Hardware:** **CPU basic** (free; the Space's offline-replay default
400
- loads no GPU code path)
401
- - **Visibility:** **Public** (required for the hackathon prize)
402
  4. Don't initialize the Space with anything — leave it empty so the first
403
  push lands cleanly.
404
 
405
- ### Deploy — Option A: scripted upload via `huggingface_hub` (recommended)
406
-
407
- ```bash
408
- export HF_TOKEN=hf_...
409
- python scripts/deploy_to_hf_space.py --space-name gpu-goblin
410
- ```
411
-
412
- The script (`scripts/deploy_to_hf_space.py`) uses `HfApi.upload_file` to
413
- push exactly the files the Space needs (README, requirements.txt, the
414
- `agent/` package, `kb/`, `ui/`, the cached audit fixture, etc.) and
415
- deliberately omits build artifacts like `bench_cache/` and
416
- `__pycache__/` that would bloat the Space repo.
417
-
418
- ### Deploy — Option B: git push (works but uploads everything tracked)
419
 
420
- From the project root:
 
421
 
422
  ```bash
423
- # HTTPS remote auth via your HF token (use the token as the password):
424
- git remote add space https://huggingface.co/spaces/lablab-ai-amd-developer-hackathon/gpu-goblin
425
 
426
- # HF Spaces use 'main' as the default branch:
427
- git push space main
428
  ```
429
 
430
- You'll see a build log at
431
- `https://huggingface.co/spaces/lablab-ai-amd-developer-hackathon/gpu-goblin`.
432
  Cold-start takes 30-60 seconds (Streamlit + the pure-pydantic deps); once
433
  up, the canonical demo trajectory replays in ~10 seconds when a judge
434
  clicks **"Use sample workload"**.
@@ -451,64 +433,30 @@ When a judge opens the Space URL:
451
 
452
  ### Updating the Space
453
 
454
- After any change to the repo, redeploy with whichever flow you used initially:
455
 
456
  ```bash
457
- # Option A:
458
- python scripts/deploy_to_hf_space.py --space-name gpu-goblin
459
-
460
- # Option B:
461
- git push space main
462
  ```
463
 
464
- HF rebuilds the Space automatically on every push/upload.
465
 
466
  ### (Stretch) Live agent in the Space
467
 
468
- The shipped Space defaults to offline-replay (no Space secrets needed).
469
- For a live demo where judges drive a real Qwen agent, three options:
470
-
471
- **Option 1 — In-process Qwen via HF Inference Providers (already wired)**
472
-
473
- Set these as Space **Settings → Variables and secrets**:
474
-
475
- | Secret | Value |
476
- |---|---|
477
- | `HF_TOKEN` | Your HF token with `inference` scope |
478
- | `GOBLIN_AGENT_BACKEND` | `qwen-hf` (default) |
479
- | `GOBLIN_QWEN_MODEL` | `Qwen/Qwen2.5-7B-Instruct` (or any HF Qwen model id) |
480
-
481
- After redeploying, the lane caption flips to `🟢 Live: agent runs Qwen
482
- in-process via Hugging Face Inference Providers` and judges audit live.
483
-
484
- **Option 2 — Connect to a self-hosted vLLM on your AMD droplet**
485
-
486
- If you've followed the AMD Developer Cloud tutorial and have vLLM serving
487
- Qwen on your MI300X at `http://YOUR_DROPLET_IP:8000/v1`, point the Space
488
- at it via Space secrets:
489
-
490
- | Secret | Value |
491
- |---|---|
492
- | `GOBLIN_AGENT_BACKEND` | `qwen-vllm` |
493
- | `GOBLIN_QWEN_VLLM_URL` | `http://YOUR_DROPLET_IP:8000/v1` |
494
- | `GOBLIN_QWEN_VLLM_MODEL` | The model ID vLLM advertises at `/v1/models` |
495
-
496
- **Important:** the AMD Developer Cloud droplet blocks port 8000 by default.
497
- SSH into the droplet and run `ufw allow 8000` (per the [lablab tutorial](https://lablab.ai/ai-tutorials/amd-huggingface-deployment-for-ai-hackathons))
498
- before the Space can reach the endpoint. Verify from outside:
499
- ```bash
500
- curl -s http://YOUR_DROPLET_IP:8000/v1/models
501
- ```
502
-
503
- **Option 3 — Connect to a separate FastAPI backend**
504
 
505
- If you've run `uvicorn agent.server:app` somewhere reachable (an MI300X on
506
- AMD Developer Cloud, an HF Inference Endpoint, a small CPU box), set
507
- `GOBLIN_BACKEND_URL` as a Space secret pointing at the `/audit` endpoint.
508
- Streamlit will stream real SSE from that backend.
 
 
 
 
509
 
510
- All three options are optional; the offline-replay Space is what
511
- satisfies the submission requirement.
512
 
513
  ## Configuration Reference
514
 
@@ -524,4 +472,3 @@ satisfies the submission requirement.
524
  | `GOBLIN_BACKEND_URL` | `http://localhost:8000/audit` | UI's backend endpoint. |
525
  | `ROCM_IMAGE_TAG` | `unknown` | Container tag mixed into the benchmark cache key. |
526
  | `GOBLIN_GPU_ID` | `0` | Which `/dev/dri/renderD*` to bind in `goblin_runner.sh`. |
527
- | `GOBLIN_RUNNER_TIMEOUT_SECONDS` | `1800` | LiveRunner subprocess timeout. Bump if cold-cache model downloads or kernel JIT push past 30 min; LiveRunner falls back to FakeRunner once exceeded. |
 
3
  emoji: 🧌
4
  colorFrom: red
5
  colorTo: red
6
+ sdk: streamlit
7
+ sdk_version: "1.32.0"
8
+ app_file: ui/auto_tune_ui.py
9
  pinned: false
10
  license: mit
11
+ short_description: AI auto-tuner for MI300X fine-tuning workloads.
12
  tags:
13
  - amd
 
14
  - mi300x
15
  - rocm
16
  - qwen
 
17
  - huggingface
18
  - agent
19
  - fine-tuning
 
381
  ### One-time setup
382
 
383
  1. Create a Hugging Face account at [huggingface.co](https://huggingface.co/)
384
+ and accept the invite to the **AMD Developer Hackathon HF Organization**
385
+ (link is on the [hackathon page](https://lablab.ai/ai-hackathons/amd-developer)
386
+ under the Hugging Face section).
 
387
  2. Create a token at [Settings → Access Tokens](https://huggingface.co/settings/tokens)
388
+ with **`write`** scope (you need write access to push to the Space repo).
389
+ Save it as `HF_PUSH_TOKEN`.
390
+ 3. On the HF organization's page, click **"New Space"**:
391
+ - Owner: AMD Developer Hackathon org
392
+ - Space name: `gpu-goblin` (or your preferred slug)
393
+ - License: MIT
394
+ - SDK: **Streamlit**
395
+ - Hardware: **CPU basic** (free; the Space loads no GPU code path)
396
+ - Visibility: Public
 
 
 
397
  4. Don't initialize the Space with anything — leave it empty so the first
398
  push lands cleanly.
399
 
400
+ ### Deploy
 
 
 
 
 
 
 
 
 
 
 
 
 
401
 
402
+ From the project root, push the existing `feat/scaffold` branch to the
403
+ Space's git remote:
404
 
405
  ```bash
406
+ # Add the Space remote (use HTTPS with your username + HF_PUSH_TOKEN as password):
407
+ git remote add space https://huggingface.co/spaces/<org-slug>/gpu-goblin
408
 
409
+ # Push (HF Spaces use 'main' as the default branch):
410
+ git push space feat/scaffold:main
411
  ```
412
 
413
+ You'll see a build log at `https://huggingface.co/spaces/<org-slug>/gpu-goblin`.
 
414
  Cold-start takes 30-60 seconds (Streamlit + the pure-pydantic deps); once
415
  up, the canonical demo trajectory replays in ~10 seconds when a judge
416
  clicks **"Use sample workload"**.
 
433
 
434
  ### Updating the Space
435
 
436
+ After any change to the main repo, redeploy:
437
 
438
  ```bash
439
+ git push space feat/scaffold:main
 
 
 
 
440
  ```
441
 
442
+ HF rebuilds the Space automatically on push.
443
 
444
  ### (Stretch) Live agent in the Space
445
 
446
+ The shipped Space is read-only it doesn't reach a real LLM. If you want
447
+ judges to drive the agent live, two paths:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
448
 
449
+ 1. **Stand up the FastAPI backend somewhere reachable** (an MI300X on AMD
450
+ Developer Cloud, an HF Inference Endpoint, a small CPU box) and set the
451
+ Space's `GOBLIN_BACKEND_URL` secret to that URL. The Streamlit app will
452
+ stream real SSE from your backend instead of the cached replay.
453
+ 2. **Embed the agent loop in-process** (refactor `ui/app.py` to call
454
+ `agent.loop.run_audit` directly via `asyncio.run`). This adds
455
+ `huggingface_hub` to `requirements.txt` and requires `HF_TOKEN` as a
456
+ Space secret. Larger cold-start, fully self-contained.
457
 
458
+ Both are post-MVP; the offline-replay Space is what satisfies the
459
+ submission requirement.
460
 
461
  ## Configuration Reference
462
 
 
472
  | `GOBLIN_BACKEND_URL` | `http://localhost:8000/audit` | UI's backend endpoint. |
473
  | `ROCM_IMAGE_TAG` | `unknown` | Container tag mixed into the benchmark cache key. |
474
  | `GOBLIN_GPU_ID` | `0` | Which `/dev/dri/renderD*` to bind in `goblin_runner.sh`. |
 
agent/server.py CHANGED
@@ -12,14 +12,19 @@ We never crash on missing keys.
12
 
13
  from __future__ import annotations
14
 
 
15
  import json
16
  import os
 
 
17
  import tempfile
18
  from collections.abc import AsyncIterator
19
  from pathlib import Path
 
20
 
21
- from fastapi import FastAPI, File, UploadFile
22
  from fastapi.middleware.cors import CORSMiddleware
 
23
  from sse_starlette.sse import EventSourceResponse
24
 
25
  from agent.backends import active_backend_name
@@ -27,6 +32,9 @@ from agent.loop import run_audit
27
  from agent.schemas import SSEEvent
28
  from agent.tools import ALL_TOOLS
29
 
 
 
 
30
  app = FastAPI(title="GPU Goblin Agent", version="0.1.0")
31
 
32
  app.add_middleware(
@@ -82,20 +90,16 @@ async def _stream_audit(file_path: str) -> AsyncIterator[dict]:
82
  sse-starlette expects. Each yielded dict becomes one `data: ...\\n\\n`
83
  SSE message.
84
  """
85
- # Backend-aware credential check. Only `qwen-hf` (HF Inference Providers)
86
- # needs HF_TOKEN; `qwen-vllm` talks to a self-hosted endpoint and ignores
87
- # tokens by default. Without this guard the server unconditionally errored
88
- # when HF_TOKEN was unset, even though the active backend didn't need it.
89
- if active_backend_name() == "qwen-hf" and not _has_hf_token():
90
  yield {
91
  "data": SSEEvent(
92
  type="error",
93
  data={
94
  "message": (
95
- "HF_TOKEN not set on the server — Qwen-HF agent loop is "
96
- "unavailable. Set HF_TOKEN (or HUGGINGFACEHUB_API_TOKEN), "
97
- "switch GOBLIN_AGENT_BACKEND to qwen-vllm, or use the "
98
- "offline-replay UI lane."
99
  )
100
  },
101
  ).model_dump_json()
@@ -143,6 +147,175 @@ async def audit(file: UploadFile = File(...)) -> EventSourceResponse:
143
  return EventSourceResponse(_stream_audit(tmp_path))
144
 
145
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
146
  # Convenience: support `python -m uvicorn agent.server:app --reload`.
147
  __all__ = ["app"]
148
 
 
12
 
13
  from __future__ import annotations
14
 
15
+ import asyncio
16
  import json
17
  import os
18
+ import subprocess
19
+ import sys
20
  import tempfile
21
  from collections.abc import AsyncIterator
22
  from pathlib import Path
23
+ from typing import Any
24
 
25
+ from fastapi import FastAPI, File, HTTPException, UploadFile
26
  from fastapi.middleware.cors import CORSMiddleware
27
+ from pydantic import BaseModel, Field
28
  from sse_starlette.sse import EventSourceResponse
29
 
30
  from agent.backends import active_backend_name
 
32
  from agent.schemas import SSEEvent
33
  from agent.tools import ALL_TOOLS
34
 
35
+ _REPO_ROOT = Path(__file__).resolve().parent.parent
36
+ _AUTO_TUNE_SCRIPT = _REPO_ROOT / "scripts" / "auto_tune.py"
37
+
38
  app = FastAPI(title="GPU Goblin Agent", version="0.1.0")
39
 
40
  app.add_middleware(
 
90
  sse-starlette expects. Each yielded dict becomes one `data: ...\\n\\n`
91
  SSE message.
92
  """
93
+ if not _has_hf_token():
94
+ # Surface a clean error instead of letting the loop crash on missing key.
 
 
 
95
  yield {
96
  "data": SSEEvent(
97
  type="error",
98
  data={
99
  "message": (
100
+ "HF_TOKEN not set on the server — Qwen agent loop is "
101
+ "unavailable. Set HF_TOKEN (or HUGGINGFACEHUB_API_TOKEN) "
102
+ "or use the offline-replay UI lane."
 
103
  )
104
  },
105
  ).model_dump_json()
 
147
  return EventSourceResponse(_stream_audit(tmp_path))
148
 
149
 
150
+ # ---------------------------------------------------------------------------
151
+ # Auto-tune endpoint — lets a UI on a CPU-only host (e.g. an HF Space) drive
152
+ # scripts/auto_tune.py running on a remote MI300X server. The endpoint
153
+ # spawns the CLI, tails its --events NDJSON stream, and re-emits each line
154
+ # as an SSE message. Subprocess output is discarded; everything the UI
155
+ # needs is in the structured events.
156
+ # ---------------------------------------------------------------------------
157
+
158
+
159
+ class AutoTuneRequest(BaseModel):
160
+ """JSON shape the /auto-tune endpoint accepts. Mirrors the auto_tune.py
161
+ CLI surface so the UI just sends what the user picked in the form."""
162
+
163
+ model: str | None = Field(
164
+ default=None,
165
+ description="HuggingFace model id (e.g. Qwen/Qwen2.5-7B-Instruct). "
166
+ "Mutually exclusive with `workload`.",
167
+ )
168
+ workload: str | None = Field(
169
+ default=None,
170
+ description="Path to a workload script ON THE SERVER's filesystem. "
171
+ "Mutually exclusive with `model`.",
172
+ )
173
+ mode: str = Field(default="hardcoded", pattern="^(hardcoded|llm|llm-explore)$")
174
+ candidates_per_iteration: int = Field(default=3, ge=2, le=10)
175
+ steps: int = Field(default=20, ge=1, le=500)
176
+ max_iterations: int = Field(default=10, ge=1, le=50)
177
+ early_stop_after: int = Field(default=3, ge=1, le=20)
178
+ max_crashes: int = Field(default=4, ge=1, le=20)
179
+ improvement_threshold: float = Field(default=0.0, ge=0.0, le=20.0)
180
+
181
+
182
+ def _build_auto_tune_cmd(req: AutoTuneRequest, events_file: Path) -> list[str]:
183
+ cmd: list[str] = [sys.executable, "-u", str(_AUTO_TUNE_SCRIPT)]
184
+ if req.model:
185
+ cmd.extend(["--model", req.model])
186
+ elif req.workload:
187
+ cmd.append(req.workload)
188
+ cmd.extend([
189
+ "--mode", req.mode,
190
+ "--steps", str(req.steps),
191
+ "--max-iterations", str(req.max_iterations),
192
+ "--early-stop-after", str(req.early_stop_after),
193
+ "--max-crashes", str(req.max_crashes),
194
+ "--improvement-threshold", str(req.improvement_threshold),
195
+ "--events", str(events_file),
196
+ ])
197
+ if req.mode == "llm-explore":
198
+ cmd.extend(["--candidates-per-iteration", str(req.candidates_per_iteration)])
199
+ return cmd
200
+
201
+
202
+ async def _stream_auto_tune(req: AutoTuneRequest) -> AsyncIterator[dict]:
203
+ """Spawn auto_tune.py and forward its NDJSON --events stream as SSE.
204
+
205
+ Each event is forwarded verbatim — the UI gets the same structured
206
+ payload it would see when running auto_tune.py locally. We discard
207
+ the subprocess's stdout/stderr; any errors are surfaced via the
208
+ `summary` event's absence at process exit.
209
+ """
210
+ events_file = Path(tempfile.mktemp(prefix="auto_tune_events_", suffix=".ndjson"))
211
+ events_file.write_text("")
212
+
213
+ cmd = _build_auto_tune_cmd(req, events_file)
214
+
215
+ # Validate at least one of model/workload was provided. (Pydantic
216
+ # can't express "exactly one of A or B" cleanly, so we check here.)
217
+ if not req.model and not req.workload:
218
+ yield {"data": json.dumps({
219
+ "type": "error",
220
+ "message": "Pass either `model` or `workload`, not neither."
221
+ })}
222
+ return
223
+ if req.model and req.workload:
224
+ yield {"data": json.dumps({
225
+ "type": "error",
226
+ "message": "Pass either `model` or `workload`, not both."
227
+ })}
228
+ return
229
+
230
+ proc = subprocess.Popen(
231
+ cmd,
232
+ cwd=str(_REPO_ROOT),
233
+ stdout=subprocess.DEVNULL,
234
+ stderr=subprocess.DEVNULL,
235
+ env={**os.environ},
236
+ )
237
+
238
+ seen_bytes = 0
239
+ try:
240
+ while True:
241
+ # Poll the events file for new lines
242
+ try:
243
+ with events_file.open("r") as f:
244
+ f.seek(seen_bytes)
245
+ chunk = f.read()
246
+ new_seen = f.tell()
247
+ except OSError:
248
+ chunk = ""
249
+ new_seen = seen_bytes
250
+
251
+ if chunk:
252
+ # Drop a trailing partial line — re-read it next tick once
253
+ # the writer has flushed the rest.
254
+ lines = chunk.splitlines(keepends=True)
255
+ if lines and not lines[-1].endswith("\n"):
256
+ partial = lines.pop()
257
+ new_seen -= len(partial.encode("utf-8"))
258
+ for line in lines:
259
+ line = line.strip()
260
+ if line:
261
+ yield {"data": line}
262
+ seen_bytes = new_seen
263
+
264
+ if proc.poll() is not None:
265
+ # Subprocess exited. Drain whatever's left on disk.
266
+ try:
267
+ with events_file.open("r") as f:
268
+ f.seek(seen_bytes)
269
+ tail = f.read()
270
+ except OSError:
271
+ tail = ""
272
+ for line in tail.splitlines():
273
+ line = line.strip()
274
+ if line:
275
+ yield {"data": line}
276
+ if proc.returncode != 0:
277
+ yield {"data": json.dumps({
278
+ "type": "process_exit",
279
+ "returncode": proc.returncode,
280
+ "message": (
281
+ f"auto_tune.py exited with code {proc.returncode}. "
282
+ "Check the server's stderr or check `last_runner_failure_*` "
283
+ "in `bench_cache/` for goblin_runner.sh failure logs."
284
+ ),
285
+ })}
286
+ break
287
+
288
+ await asyncio.sleep(0.5)
289
+ finally:
290
+ if proc.poll() is None:
291
+ proc.terminate()
292
+ try:
293
+ proc.wait(timeout=3)
294
+ except subprocess.TimeoutExpired:
295
+ proc.kill()
296
+ try:
297
+ events_file.unlink()
298
+ except OSError:
299
+ pass
300
+
301
+
302
+ @app.post("/auto-tune")
303
+ async def auto_tune_endpoint(req: AutoTuneRequest) -> EventSourceResponse:
304
+ """Stream auto_tune.py events back to the caller as SSE.
305
+
306
+ Run a UI on any host (HF Spaces, local laptop), point it at this
307
+ endpoint, and the actual GPU work happens on the server hosting the
308
+ FastAPI app. Subprocess output is discarded — only the --events
309
+ NDJSON stream crosses the wire, one structured event per SSE message.
310
+ """
311
+ if not _AUTO_TUNE_SCRIPT.exists():
312
+ raise HTTPException(
313
+ status_code=500,
314
+ detail=f"auto_tune.py not found at {_AUTO_TUNE_SCRIPT}",
315
+ )
316
+ return EventSourceResponse(_stream_auto_tune(req))
317
+
318
+
319
  # Convenience: support `python -m uvicorn agent.server:app --reload`.
320
  __all__ = ["app"]
321
 
brainstorming/architecture.md ADDED
@@ -0,0 +1,359 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # GPU Goblin — Architecture
2
+
3
+ ## System Topology
4
+
5
+ ```
6
+ ┌──────────────────────────────────────────────────────────────────┐
7
+ │ Streamlit Chat UI (browser) │
8
+ │ upload script · live tool-call stream · final report │
9
+ └────────────────────────────┬─────────────────────────────────────┘
10
+ │ HTTP + Server-Sent Events
11
+
12
+ ┌──────────────────────────────────────────────────────────────────┐
13
+ │ Goblin Agent — FastAPI (port 8000) │
14
+ │ Qwen2.5-7B (HF Inference Providers) · agent loop · session │
15
+ └─┬────────┬─────────┬──────────┬──────────┬──────────┬────────────┘
16
+ │ │ │ │ │ │
17
+ ▼ ▼ ▼ ▼ ▼ ▼
18
+ ┌─────┐ ┌────────┐ ┌──────┐ ┌────────┐ ┌─────────┐ ┌────────┐
19
+ │parse│ │profile │ │rocm │ │propose │ │benchmark│ │compare │
20
+ │cfg │ │_run │ │_kb │ │_patch │ │ │ │_runs │
21
+ └──┬──┘ └────┬───┘ └──┬───┘ └────┬───┘ └────┬────┘ └────┬───┘
22
+ │ │ │ │ │ │
23
+ ▼ ▼ ▼ ▼ ▼ ▼
24
+ ┌─────┐ ┌────────────────┐ ┌──────────┐ ┌──────────────────┐
25
+ │AST +│ │torch.profiler +│ │YAML KB + │ │MI300X cloud │
26
+ │regex│ │rocprofv3 + amd- │ │sentence- │ │job runner │
27
+ │ │ │smi │ │transform │ │(subprocess + │
28
+ │ │ │ │ │ers │ │ result cache) │
29
+ └─────┘ └────────────────┘ └──────────┘ └──────────────────┘
30
+ ```
31
+
32
+ ## Components
33
+
34
+ ### 1. Streamlit Chat UI
35
+
36
+ - Single-page app: file upload + chat panel + report panel
37
+ - Streams tool calls as cards: *"calling profile_run…"* with live status
38
+ - Final report: side-by-side metrics, diff viewer, kernel waterfall chart
39
+ - Backup mode: replay a cached "golden run" if MI300X unreachable
40
+
41
+ ### 2. Goblin Agent (FastAPI)
42
+
43
+ - Single endpoint: `POST /audit` — takes uploaded script, returns SSE stream of agent steps
44
+ - Session state in-memory only (hackathon, not production)
45
+ - Hard cap: **max 8 tool calls per audit** (prevents loops)
46
+ - LLM: **Qwen2.5-7B-Instruct** via Hugging Face Inference Providers (`huggingface_hub.AsyncInferenceClient.chat_completion`, OpenAI-shape tool calls). Pluggable via `agent/backends/`; a future `LiveQwenBackend` swaps to a self-hosted vLLM-on-MI300X endpoint with no other code changes.
47
+ - System prompt establishes ROCm expert persona + tool-call etiquette + report format
48
+
49
+ ### 3. The Six Tools
50
+
51
+ #### `parse_config(file_path: str) -> ConfigDict`
52
+ Inputs: HF `TrainingArguments` (Python or JSON), raw PyTorch training script, YAML config.
53
+ Implementation: AST parsing for `.py` (extract `TrainingArguments(...)` kwargs), `json.load` / `yaml.safe_load` for configs, regex fallback.
54
+ Output schema:
55
+ ```python
56
+ {
57
+ "model_name": str,
58
+ "batch_size": int,
59
+ "grad_accum_steps": int,
60
+ "seq_len": int,
61
+ "precision": "fp16" | "bf16" | "fp32",
62
+ "optimizer": str,
63
+ "attention_impl": "sdpa" | "flash" | "eager" | "unknown",
64
+ "gradient_checkpointing": bool,
65
+ "lora_rank": int | None,
66
+ "dataloader_workers": int,
67
+ "lr": float,
68
+ "warmup_steps": int,
69
+ "raw_source": str,
70
+ }
71
+ ```
72
+
73
+ #### `profile_run(config: ConfigDict, steps: int = 10) -> RunMetrics`
74
+ Wraps the user's training command in `torch.profiler` + `rocprofv3`. Runs for N steps after a 2-step warmup. Captures and **summarizes** (never returns raw traces — too big for LLM context).
75
+ Output is a `RunMetrics` (the same type returned by `benchmark` — see schemas below):
76
+ ```python
77
+ RunMetrics(
78
+ steps=10,
79
+ tokens_per_sec=...,
80
+ mfu_pct=...,
81
+ hbm_peak_gb=..., hbm_avg_gb=...,
82
+ gpu_util_pct=...,
83
+ top_kernels=[KernelEntry(name=..., pct_time=...), ...], # top 5
84
+ attention_kernel_loaded=...,
85
+ waste_budget=WasteBudget(...), # see below
86
+ warnings=[...],
87
+ )
88
+ ```
89
+
90
+ ##### Waste Budget Decomposition
91
+
92
+ Inside `RunMetrics.waste_budget`, total step time is decomposed into interpretable buckets:
93
+
94
+ ```
95
+ T_total = T_useful_gpu
96
+ + T_data_wait # dataloader / host→device stalls
97
+ + T_host_gap # CPU→GPU launch latency, eager-mode kernel gaps
98
+ + T_comm_excess # collectives, all-reduce, RCCL overhead
99
+ + T_memory_headroom # HBM left idle relative to model needs
100
+ + T_precision_path # throughput lost vs ideal precision (fp32 where bf16 OK)
101
+ + T_kernel_shape # GEMM tile mismatch, untuned hipBLASLt/MIOpen
102
+ ```
103
+
104
+ This decomposition is the basis for the "where time was lost" chart in the final report and for ranking which rules to apply first. Each rule in the KB declares which bucket it targets (`targets_bucket: data_wait`), so `propose_patch` can avoid stacking redundant fixes for the same bucket.
105
+
106
+ #### `query_rocm_kb(symptom: str, top_k: int = 5) -> List[Rule]`
107
+ Semantic search over the YAML KB. Embeds the symptom string with `sentence-transformers/all-MiniLM-L6-v2`, cosine-similarity against pre-embedded rules.
108
+ Output: list of rules sorted by relevance.
109
+
110
+ #### `propose_patch(config: ConfigDict, rules: List[Rule], metrics: RunMetrics) -> Patch`
111
+ Deterministic — no LLM call. Applies rule-to-config transforms (each rule has a `transform` field describing the diff) and computes a per-rule and aggregate uplift estimate from the waste budget.
112
+ Output:
113
+ ```python
114
+ Patch(
115
+ new_config=ConfigDict(...),
116
+ diff=str, # unified diff
117
+ rationale=[RuleApplication(...)], # one entry per applied rule
118
+ expected_speedup_low=float, # conservative end of range
119
+ expected_speedup_high=float, # optimistic end of range
120
+ confidence=float, # 0..1, see formula below
121
+ )
122
+ ```
123
+
124
+ ##### Uplift Estimate
125
+
126
+ For each applied rule with `targets_bucket=B`, predicted recovery is `recovery_fraction × T_B / T_total`. Aggregate predicted recoverable waste = sum across applied rules (capped per bucket). Predicted speedup = `1 / (1 - recoverable_fraction)`. Reported as a range, not a point estimate.
127
+
128
+ ##### Confidence Score
129
+
130
+ ```
131
+ confidence = evidence_coverage × rule_consistency
132
+ ```
133
+
134
+ - `evidence_coverage` — fraction of waste-budget buckets that had real measurement (vs. defaulted-to-zero because the profiler couldn't observe it). Lower if `profile_run` produced partial traces.
135
+ - `rule_consistency` — 1.0 if all applied rules target distinct buckets and don't have conflicting `transform` fields; lower if rules overlap or conflict.
136
+
137
+ (The report's `historical_calibration` term is intentionally dropped — we have no historical data in a hackathon timeframe. Document this honestly in the system prompt.)
138
+
139
+ #### `benchmark(config: ConfigDict, steps: int = 50) -> RunMetrics`
140
+ Same pipeline as `profile_run` but runs longer and at full quality. Returns the same `RunMetrics` type. Result is cached by `sha256((canonical_config_json, workload_script_sha, rocm_image_tag, runner_script_sha))` — re-running the same config is free, and the version-tagged hash prevents stale cache hits when the container or runner changes.
141
+
142
+ #### `compare_runs(before: RunMetrics, after: RunMetrics) -> Report`
143
+ Pure function. Builds the side-by-side report dict the UI renders. Includes the side-by-side waste-budget bar chart that visually shows which buckets shrank.
144
+
145
+ ### 4. ROCm Knowledge Base
146
+
147
+ Single file: `kb/rocm_rules.yaml`. ~25 rules at MVP. Schema:
148
+
149
+ ```yaml
150
+ - id: precision.bf16_over_fp16_on_mi300x
151
+ category: precision
152
+ symptom: "fp16 used on MI300X"
153
+ detect:
154
+ config.precision: fp16
155
+ fix:
156
+ config.precision: bf16
157
+ expected_impact: "Same throughput, +numerical stability. Reduces NaN risk."
158
+ rocm_version_min: "6.0"
159
+ citation: "AMD ROCm Best Practices Guide §3.2"
160
+ ```
161
+
162
+ Categories:
163
+ - `precision` (bf16 over fp16 on CDNA3 matrix cores; FP8 for stretch inference scenarios)
164
+ - `attention` (flash-attn ROCm fork via Optimum-AMD, PyTorch SDPA, packed sequences)
165
+ - `memory` (batch-size sweetspot for 192 GB HBM3, gradient checkpoint thresholds, activation offloading)
166
+ - `kernels` (hipBLASLt hint logging + offline tuning files; MIOpen `MIOPEN_FIND_*` autotune)
167
+ - `env_vars` (`HSA_FORCE_FINE_GRAIN_PCIE`, `MIOPEN_FIND_MODE`, `NCCL_MIN_NCHANNELS=112`, NUMA auto-balancing disable)
168
+ - `optimizer` (8-bit Adam on ROCm — **warn that bitsandbytes is not officially supported on ROCm**; recommend Optimum-AMD-validated alternatives or CPU-offload optimizers)
169
+ - `data` (`num_workers`, `pin_memory=True`, `prefetch_factor`, `persistent_workers=True`, IterableDataset sharding correctness)
170
+ - `compile` (`torch.compile` when graph-break evidence is low; `torch_compile=True` in `TrainingArguments`)
171
+ - `collectives` (RCCL channel count, one-process-per-GPU vs one-process-many-GPUs)
172
+ - `topology` (tensor parallelism within a single XGMI island; prefer TP over EP on single-node)
173
+
174
+ Rules are hand-curated on Day 1 from ROCm docs + AMD blog posts. **This is the moat.** No LLM-generated rules. Each rule entry includes a `citation` field linking back to a ROCm doc page or AMD blog post — every recommendation in the final report carries this citation.
175
+
176
+ **Footgun guardrail — `ROCPROFSYS_*` are NOT tuning vars.** ROCm Systems Profiler env vars (`ROCPROFSYS_MODE`, `ROCPROFSYS_USE_SAMPLING`, etc.) configure how the *profiler* observes a run; they do not affect how work is dispatched to the GPU. The KB **excludes** all `ROCPROFSYS_*` vars from optimization rules. Additionally, `parse_config` flags any user config that sets `ROCPROFSYS_*` as if it were a perf knob and surfaces a warning in the report ("These configure the profiler, not the workload — they will not change throughput").
177
+
178
+ **Workload-validity disclaimer.** Every recommendation is valid for the specific tuple `(workload script, model, GPU=MI300X, ROCm version, framework version, batch/seq pattern)` observed during profiling. The system prompt instructs the agent to state this explicitly when uncertainty is high, and the final report includes a footer line: *"Recommendations validated against MI300X with ROCm <version> and PyTorch <version>. Re-run audit if you change model, hardware, or framework version."*
179
+
180
+ ### 5. Profiling Pipeline
181
+
182
+ A wrapper script `goblin_runner.sh` invokes the user's training command inside the ROCm container, with:
183
+
184
+ ```bash
185
+ # isolate to a single MI300X to keep concurrent benchmark runs sane
186
+ export ROCR_VISIBLE_DEVICES=${GOBLIN_GPU_ID:-0}
187
+
188
+ rocprofv3 --hsa-trace --kernel-trace -o trace.csv -- \
189
+ python -m torch.profiler.scheduler ... \
190
+ python <user_script> --max_steps=10
191
+ ```
192
+
193
+ Post-processing (`profile_parser.py`):
194
+ - parse `trace.csv` for kernel breakdown
195
+ - parse `torch.profiler` JSON for tokens/sec + MFU
196
+ - parse `amd-smi` polling output for HBM + util
197
+ - merge into single `RunMetrics` with populated `WasteBudget`
198
+
199
+ ### 6. Benchmark Result Cache
200
+
201
+ `bench_cache/` directory. Key: `sha256((canonical_config_json, workload_script_sha, rocm_image_tag, runner_script_sha))`. Value: full `RunMetrics` plus the unhashed key tuple for debuggability. Lets us re-run the demo without burning cloud time. The Streamlit UI exposes a `--no-cache` toggle for the Day-3 dry-run pass that confirms cached results haven't gone stale.
202
+
203
+ ### 7. Synthetic Corpus (`workloads/synthetic/`)
204
+
205
+ A directory of pre-generated, deliberately misconfigured fine-tuning runs. Each entry is `{config.py, expected_findings.yaml, cached_metrics.json}`. Generated on Day 1 by ROCm Lead by perturbing the canonical Qwen2.5-7B LoRA workload along single dimensions: `num_workers=0`, FP32 instead of BF16, naive attention, `torch.compile=False`, untuned hipBLASLt, etc.
206
+
207
+ Two purposes:
208
+ 1. **Backend Lead can develop on a laptop** — test the agent loop, KB queries, and patch generation against `cached_metrics.json` without touching the GPU.
209
+ 2. **Demo lane 1 (offline replay)** — judges can audit any synthetic scenario in <30 seconds without live MI300X. If cloud is unreachable on demo day, this is the safety net.
210
+
211
+ Each scenario's `expected_findings.yaml` lists which rules *should* fire — gives us a regression test for KB recall.
212
+
213
+ ### 8. Demo Lanes
214
+
215
+ The system runs in one of two modes:
216
+
217
+ | Lane | Source of metrics | Use |
218
+ |---|---|---|
219
+ | **Offline replay** | Synthetic corpus + cached `RunMetrics` JSON | Dev, regression testing, demo backup if MI300X unreachable |
220
+ | **Live MI300X** | `goblin_runner.sh` → real GPU | Canonical demo, "before/after with real numbers" pitch moment |
221
+
222
+ The agent loop is identical in both lanes — only the `RunnerProtocol` implementation differs. This is the seam introduced to address the audit's testability finding.
223
+
224
+ ### 9. Secrets & Privacy
225
+
226
+ Uploaded scripts and logs may contain API tokens, dataset paths, internal URLs. Before any artefact is persisted to disk or sent to the LLM:
227
+
228
+ 1. **Regex redaction pass** in `parse_config` for common patterns: `Bearer [A-Za-z0-9._-]+`, `sk-[A-Za-z0-9]{20,}`, `hf_[A-Za-z0-9]{20,}`, `/home/<user>/`, `s3://`, `wss?://`.
229
+ 2. **Ephemeral storage** — `bench_cache/` is gitignored; uploaded scripts are deleted at session end.
230
+ 3. **No model weights or training data required** — we audit configs and metrics only. The system prompt explicitly tells the agent it cannot ask for weights or datasets.
231
+
232
+ This is a one-screen redaction pass, not a full PII pipeline. Sufficient for hackathon; would need expansion for production.
233
+
234
+ ## Data Flow — One Audit
235
+
236
+ ```
237
+ [user uploads script] ──► [redaction pass]
238
+
239
+
240
+ parse_config ──► ConfigDict ────────────────────────────┐
241
+ │ │
242
+ ▼ │
243
+ profile_run ──► RunMetrics + WasteBudget ──► query_rocm_kb │
244
+ │ │ │
245
+ │ ▼ │
246
+ │ List[Rule] ──────►│
247
+ │ ▼
248
+ │ propose_patch (uplift + confidence)
249
+ │ │
250
+ │ ▼
251
+ │ Patch (new ConfigDict + diff)
252
+ │ │
253
+ ▼ ▼
254
+ benchmark(original) ◄──── cache ─────► benchmark(patched)
255
+ │ │
256
+ └────────► compare_runs ◄────────────────────────┘
257
+
258
+
259
+ Report ──► UI
260
+ (waste-budget bar chart + metrics + diff)
261
+ ```
262
+
263
+ ## Tech Stack (frozen)
264
+
265
+ | Layer | Choice | Why |
266
+ |---|---|---|
267
+ | Agent LLM | Qwen2.5-7B-Instruct via HF Inference Providers | Tool calling via OpenAI-compatible API; routes through Together / Fireworks-AI / Nebius (auto). Stretch: self-host on MI300X via vLLM. |
268
+ | Backend | Python 3.11 + FastAPI + anthropic SDK | Standard, async-friendly |
269
+ | Frontend | Streamlit | Ships in hours, not days |
270
+ | Container | `rocm/pytorch:rocm6.1_ubuntu22.04_py3.10_pytorch_2.3` | Official AMD image |
271
+ | Profiling | torch.profiler + rocprofv3 + amd-smi | Native ROCm tools |
272
+ | KB index | sentence-transformers + numpy cosine | No vector DB needed for 25 rules |
273
+ | Workload | Qwen2.5-7B + LoRA + alpaca-cleaned | Canonical, reproducible |
274
+ | Hardware | MI300X cloud (single GPU, 192 GB HBM) | Hackathon constraint |
275
+
276
+ ## Repo Layout
277
+
278
+ ```
279
+ amd-hackathon/
280
+ ├── brainstorming/ # this folder
281
+ ├── kb/
282
+ │ └── rocm_rules.yaml # the moat
283
+ ├── agent/
284
+ │ ├── server.py # FastAPI app
285
+ │ ├── loop.py # agent loop driver
286
+ │ ├── prompts.py # system prompt
287
+ │ └── tools/
288
+ │ ├── parse_config.py
289
+ │ ├── profile_run.py
290
+ │ ├── query_rocm_kb.py
291
+ │ ├── propose_patch.py
292
+ │ ├── benchmark.py
293
+ │ └── compare_runs.py
294
+ ├── runner/
295
+ │ ├── goblin_runner.sh # rocprofv3 wrapper
296
+ │ └── profile_parser.py
297
+ ├── ui/
298
+ │ └── app.py # Streamlit
299
+ ├── workloads/
300
+ │ ├── train_qwen_lora.py # canonical demo workload
301
+ │ └── synthetic/ # pre-generated misconfigured runs + cached metrics
302
+ │ ├── 01_no_workers/
303
+ │ ├── 02_fp32_default/
304
+ │ ├── 03_naive_attention/
305
+ │ └── ...
306
+ ├── bench_cache/ # gitignored
307
+ └── docker/
308
+ └── Dockerfile # rocm/pytorch + our deps
309
+ ```
310
+
311
+ ## Agent Loop — Pseudocode
312
+
313
+ The loop is provider-agnostic. It talks to a `Backend` (see `agent/backends/`); today's only concrete is `QwenHFBackend`.
314
+
315
+ ```python
316
+ def run_audit(user_script_path: str) -> SSEStream:
317
+ backend = make_backend(system_prompt=SYSTEM_PROMPT) # QwenHFBackend
318
+ backend.add_user_message(f"Audit this fine-tuning workload: {user_script_path}")
319
+
320
+ for step in range(MAX_STEPS): # MAX_STEPS = 8
321
+ turn = await backend.next_turn(tool_schemas())
322
+ for text in turn.text_blocks:
323
+ yield {"type": "thought", "text": text}
324
+ for tc in turn.tool_calls:
325
+ yield {"type": "tool_call", "name": tc.name, "input": tc.input}
326
+ result = call_tool(tc.name, **tc.input)
327
+ yield {"type": "tool_result", "name": tc.name, "result": result}
328
+ backend.add_tool_result(tc.id, tc.name, result.content, is_error=not result.ok)
329
+ if turn.stop_reason == "end_turn":
330
+ break
331
+
332
+ yield {"type": "final_report", "report": extract_final_report(tool_results)}
333
+ ```
334
+
335
+ The Backend protocol (`agent/backends/base.py`) is a 3-method contract:
336
+ `add_user_message`, `next_turn`, `add_tool_result`. `QwenHFBackend` translates
337
+ between this neutral shape and OpenAI-compatible chat-completion calls
338
+ through HF Inference Providers.
339
+
340
+ ## Boundaries & Interfaces
341
+
342
+ - **Agent ↔ Tools:** Each tool is a pure function with typed input/output. No shared state. Tools never call other tools — only the agent orchestrates.
343
+ - **Tools ↔ MI300X:** Only `profile_run` and `benchmark` touch the GPU. Both go through `goblin_runner.sh` for consistency.
344
+ - **Backend ↔ UI:** SSE stream of typed events (`thought`, `tool_call`, `tool_result`, `final_report`). UI just renders events; no agent logic on frontend.
345
+ - **KB ↔ Code:** YAML is the single source of truth. Code consumes; never generates rules at runtime.
346
+
347
+ ## What's Deliberately Excluded (YAGNI)
348
+
349
+ - ❌ Multi-GPU / distributed training analysis (single MI300X is enough; collective rules in KB still apply when user uploads multi-GPU configs, we just don't benchmark them)
350
+ - ❌ Pretraining workloads (fine-tuning only)
351
+ - ❌ vLLM / inference serving as a primary mode (Day-4 stretch only; vLLM-specific rules NOT in MVP KB)
352
+ - ❌ Live job watching with intervention (offline audit only — much simpler)
353
+ - ❌ User accounts / persistence / DB (in-memory session, hackathon)
354
+ - ❌ Vector DB (25 rules + numpy is fine)
355
+ - ❌ Polars/DuckDB for run storage (Pydantic + JSON is enough at this scale)
356
+ - ❌ ML model for uplift prediction (deterministic waste-budget calculus is enough)
357
+ - ❌ Auto-applying patches to the user's repo (we output a diff, user applies)
358
+ - ❌ Sliders/what-if interactive panel (chat layer answers what-if conversationally; sliders are stretch only)
359
+ - ❌ Cost calculator beyond a static "$ per training run" line (stretch goal only)
brainstorming/goals.md ADDED
@@ -0,0 +1,233 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # GPU Goblin — Goals & Implementation Plan
2
+
3
+ ## North Star
4
+
5
+ Win the AMD hackathon (**Track 1: AI Agents & Agentic Workflows**) by demonstrating a **real, reproducible 2×+ throughput improvement on a Qwen2.5-7B LoRA fine-tune on MI300X**, driven end-to-end by a Qwen-powered tool-using agent. Hits the Qwen Technology Partner challenge with end-to-end Qwen-on-AMD; deploys as a Hugging Face Space within the hackathon HF Organization.
6
+
7
+ Everything else is in service of that demo.
8
+
9
+ ## Success Criteria
10
+
11
+ | # | Criterion | Measurable As |
12
+ |---|---|---|
13
+ | 1 | Real MI300X speedup | ≥ 2.0× tokens/sec, before vs. after, on the canonical workload |
14
+ | 2 | Agent is actually agentic | ≥ 4 tool calls per audit, visible in UI |
15
+ | 3 | AMD differentiation | 100% of recommendations cite a ROCm-specific KB rule |
16
+ | 4 | Uplift estimate honesty | Predicted speedup range (low–high) brackets the measured speedup on ≥ 8 of 10 synthetic corpus scenarios |
17
+ | 5 | Offline parse coverage | ≥ 90% of synthetic corpus parsed without manual edits |
18
+ | 6 | End-to-end reliability | Demo runs cleanly 5× in a row without manual intervention |
19
+ | 7 | Pitch quality | 3-min demo + 1-min architecture + 1-min impact, rehearsed |
20
+ | 8 | Backup safety | Offline-replay lane works fully when MI300X disconnected |
21
+
22
+ ## Mapping to Judging Criteria
23
+
24
+ | Hackathon criterion | Where we score |
25
+ |---|---|
26
+ | **Application of Technology** | Tool-using agent loop on MI300X; rocprofv3 + torch.profiler integration; deterministic uplift via waste-budget calculus; ROCm-specific KB |
27
+ | **Presentation** | Live before/after benchmark on real MI300X; waste-budget bar chart; chat-based "why?" interaction; 90-sec backup video |
28
+ | **Business Value** | Each audit reports `$ saved per training run` and `time saved per epoch`, framed against $1.99/hr public MI300X pricing reference |
29
+ | **Originality** | "AI for AI builders" — meta-positioning. Real benchmarks, not just LLM advice. Waste-budget decomposition is novel framing |
30
+
31
+ ## Team Roles (3 people × 4 days)
32
+
33
+ ### ROCm / ML Lead
34
+ - Day 1: MI300X cloud env, ROCm container, baseline workload (Qwen2.5-7B LoRA on alpaca)
35
+ - Day 1-2: Hand-curate 20-25 KB rules from ROCm docs + AMD blog
36
+ - Day 2: `profile_run` + `benchmark` tools (rocprofv3 wrapper, parser)
37
+ - Day 3: Validate end-to-end speedup on canonical demo, generate cached results
38
+ - **Owns:** anything that touches the GPU
39
+
40
+ ### Agent / Backend Lead
41
+ - Day 1: FastAPI skeleton, tool schemas, Qwen-via-HF tool-use plumbing
42
+ - Day 1-2: `parse_config`, `propose_patch`, `query_rocm_kb`, `compare_runs` tools
43
+ - Day 2: Full agent loop, system prompt, SSE streaming
44
+ - Day 3: Hardening, error handling, max-steps cap, fallback behaviors
45
+ - **Owns:** agent reasoning quality + tool wiring
46
+
47
+ ### Frontend / Demo Lead
48
+ - Day 1: Streamlit skeleton, file upload, message log
49
+ - Day 2: Tool-call cards (live status), report renderer, diff viewer
50
+ - Day 3: Demo polish, charts, golden-run replay mode
51
+ - Day 4: Pitch deck, 90-second backup video, dry runs
52
+ - **Owns:** what the judges actually see
53
+
54
+ Roles overlap on integration days — pair-program when blocked.
55
+
56
+ ## Day-by-Day Plan
57
+
58
+ ### Day 1 (May 5–6) — Foundations
59
+
60
+ **ROCm Lead**
61
+ - [ ] Provision MI300X cloud instance via AMD Developer Cloud ($100 credits), SSH access, persistent storage
62
+ - [ ] Pull `rocm/pytorch:rocm6.1_*` image, verify GPU visible inside container
63
+ - [ ] Run baseline `train_qwen_lora.py` (batch=4, fp16, naive attention, alpaca, 100 steps)
64
+ - [ ] Capture baseline tokens/sec, MFU, HBM peak — *this is our "before"*
65
+ - [ ] Generate **synthetic corpus** — 5-8 misconfigured variants of the canonical workload (FP32, num_workers=0, naive attention, etc.) with cached `RunMetrics` JSON for each
66
+ - [ ] Start drafting KB rules (target 10 rules by EOD), each tagged with `targets_bucket` matching the waste-budget decomposition
67
+
68
+ **Backend Lead**
69
+ - [ ] Repo scaffold per architecture.md layout
70
+ - [ ] `pip install fastapi anthropic sentence-transformers pyyaml pydantic`
71
+ - [ ] Define `agent/schemas.py` with `RunMetrics`, `WasteBudget`, `ConfigDict`, `Patch`, `Rule`, `Report` as pydantic models — **Day-1 priority** (blocks all tools)
72
+ - [ ] Define `RunnerProtocol` interface; build `FakeRunner` that loads cached metrics from `workloads/synthetic/` (lets backend dev without MI300X)
73
+ - [ ] FastAPI `POST /audit` skeleton with SSE
74
+ - [ ] `parse_config` tool — handle HF `TrainingArguments` first; include regex redaction pass for tokens/paths
75
+ - [ ] Qwen tool-use hello-world (one tool, one round-trip via HF Inference Providers)
76
+
77
+ **Frontend Lead**
78
+ - [ ] Streamlit skeleton with file upload + chat panel
79
+ - [ ] Hardcoded "fake audit" — render canned tool calls + report
80
+ - [ ] Pick chart library (Altair recommended for Streamlit)
81
+
82
+ **Day 1 Exit Criteria**
83
+ - Baseline benchmark numbers in hand
84
+ - Synthetic corpus has ≥ 3 cached scenarios (Backend Lead can now dev without GPU)
85
+ - Schemas (`RunMetrics`, `WasteBudget`, etc.) frozen
86
+ - `RunnerProtocol` + `FakeRunner` working end-to-end
87
+ - Backend can call Qwen with one tool via HF Inference Providers
88
+ - UI renders a fake audit
89
+ - 10 KB rules drafted
90
+
91
+ ### Day 2 (May 6–7) — Core Build
92
+
93
+ **ROCm Lead**
94
+ - [ ] Finish KB to 20-25 rules; hand-tag categories + `targets_bucket`; pre-embed with sentence-transformers
95
+ - [ ] Include the high-impact MI300X rules: BF16-over-FP16, AITER-flash-attn-via-Optimum-AMD, `NCCL_MIN_NCHANNELS=112`, NUMA disable, one-process-per-GPU, hipBLASLt hint logging, MIOpen `MIOPEN_FIND_*`, **bitsandbytes-not-supported-on-ROCm warning**, `num_workers`/`pin_memory`/`prefetch_factor`/`persistent_workers`
96
+ - [ ] `profile_run` tool: rocprofv3 wrapper + torch.profiler + amd-smi → `RunMetrics` with `WasteBudget`
97
+ - [ ] `benchmark` tool: same pipeline, longer run, with version-tagged cache
98
+ - [ ] Validate on baseline workload: profile output makes sense, MFU is plausible, waste budget sums to ~T_total
99
+
100
+ **Backend Lead**
101
+ - [ ] All 6 tools wired and individually tested with fixtures from synthetic corpus
102
+ - [ ] Full agent loop with max-steps cap, SSE event types finalized, error envelope per tool (`{ok, result, error}`)
103
+ - [ ] System prompt iterated against test workloads — include MI300X hardware specs (304 CUs, 192 GB HBM3, ~5.3 TB/s, FP8 native) so the agent reasons quantitatively
104
+ - [ ] `propose_patch` deterministic transformer + uplift estimator (waste-budget × bucket recovery) + confidence formula
105
+
106
+ **Frontend Lead**
107
+ - [ ] Live tool-call cards consuming real SSE stream
108
+ - [ ] Final report layout: side-by-side metrics + diff + kernel chart + **waste-budget bar chart** (where time was lost, before vs after)
109
+ - [ ] Lane toggle: `Live MI300X` vs `Offline replay (synthetic corpus)` — judges can pick either
110
+ - [ ] First end-to-end run through the actual backend
111
+
112
+ **Day 2 Exit Criteria**
113
+ - Real audit runs end-to-end on a toy workload
114
+ - Profile + benchmark return real MI300X numbers
115
+ - KB has 20+ rules, embeddings pre-computed
116
+ - UI streams real agent activity
117
+
118
+ ### Day 3 (May 7–8) — Demo Day Prep
119
+
120
+ **All hands**
121
+ - [ ] Run canonical demo (Qwen2.5-7B LoRA) end-to-end → confirm ≥ 2× speedup
122
+ - [ ] Cache the demo benchmark results — don't burn cloud time on every rehearsal
123
+ - [ ] Build "golden run" replay mode (read cached SSE events, replay timing)
124
+ - [ ] Validate uplift accuracy on synthetic corpus: predicted range should bracket measured speedup on ≥ 8 of 10 scenarios
125
+ - [ ] Polish system prompt for demo-friendly narration ("I'll start by…")
126
+ - [ ] Tighten error handling — agent should never panic in front of judges; tool failures degrade gracefully with `{ok:false}` envelopes
127
+ - [ ] Run with `--no-cache` once to verify cached results aren't masking real bugs
128
+ - [ ] 5× clean dry runs, fix anything flaky
129
+
130
+ **Day 3 Exit Criteria**
131
+ - Canonical demo: cleanly runs in ≤ 4 minutes, ≥ 2× speedup, no manual fixes
132
+ - Cached results enable offline replay
133
+ - Offline-replay lane works fully when MI300X is disconnected (proven by unplugging cloud)
134
+ - Golden-run video recorded as backup
135
+
136
+ ### Day 4 (May 8–9) — Pitch & Stretch
137
+
138
+ **Frontend Lead**
139
+ - [ ] Pitch deck (5 slides): problem, agent loop diagram, demo (live), KB rules sample, impact + ask
140
+ - [ ] Cover image for the submission listing
141
+ - [ ] Final 90-second backup video
142
+ - [ ] Submission form filled (title, short + long description, tags), repo public, README crisp
143
+ - [ ] **Build-in-Public bonus track:** at least one public update post (X/LinkedIn/Discord) showing the agent's first audit; one ROCm/Optimum-AMD feedback note based on what was rough during build
144
+
145
+ **ROCm + Backend Leads (in parallel, optional stretch)**
146
+ - [ ] vLLM inference workload as second demo (only if rock-solid on Day 3)
147
+ - [ ] Cost calculator: `$ saved per training run` line, anchored on $1.99/hr public MI300X reference
148
+ - [ ] What-if slider panel for batch / precision / attention (chat already does this conversationally, sliders are visual icing)
149
+ - [ ] Stretch dream: self-host Qwen on MI300X via vLLM (replacing the HF-Inference-Providers path) — closes the loop entirely on AMD silicon. Mention in pitch even if not running live.
150
+
151
+ **Day 4 Exit Criteria**
152
+ - Submission complete by deadline
153
+ - 5+ rehearsed dry runs of the pitch
154
+ - Cover image, video, slides, repo all linked from submission form
155
+ - Backup video on USB stick, in cloud, on phone
156
+
157
+ ## Scope Discipline (YAGNI)
158
+
159
+ If we're behind schedule, **cut in this order**:
160
+
161
+ 1. ❌ All stretch goals (vLLM, cost calc, self-hosted agent)
162
+ 2. ❌ Live tool-call UI animations — static cards work
163
+ 3. ❌ Some KB categories (keep precision/attention/memory; drop env_vars/collectives)
164
+ 4. ❌ Multi-file script parsing — single-file only
165
+ 5. ❌ Dynamic batch-size search — hardcode the recommended batch for demo workload
166
+
167
+ **Never cut:**
168
+ - Real MI300X benchmark (that's the entire pitch)
169
+ - Tool-using agent loop (that's the track fit)
170
+ - ROCm-specific KB citations (that's the differentiation)
171
+
172
+ ## Risk Register
173
+
174
+ | Risk | Likelihood | Impact | Mitigation |
175
+ |---|---|---|---|
176
+ | MI300X cloud quota burns out | Medium | High | Cache every benchmark, dev with 10-step traces, full benchmarks only for canonical demo |
177
+ | `rocprofv3` flaky / version mismatch | Medium | High | Always run `torch.profiler` in parallel as fallback. Pre-record golden trace. Use `rocprofv3` not deprecated `rocprof`/`rocprofv2` |
178
+ | MI300X cloud unreachable on demo day | Low | Critical | **Offline-replay lane** (synthetic corpus + cached metrics) provides full demo without cloud — judges can't tell the difference for the agent-reasoning portion |
179
+ | Uploaded scripts contain user secrets / tokens | Medium | Medium | Regex redaction pass in `parse_config` before any persistence or LLM call; ephemeral storage; explicit "we never store weights or datasets" line in UI |
180
+ | Agent loops infinitely | Low | Medium | Hard cap of 8 tool calls, fall back to "best effort" report after cap |
181
+ | Recommendations are generic, not ROCm-specific | Medium | Critical | Hand-curate KB on Day 1 *before* wiring agent — KB is the moat |
182
+ | Demo crashes during pitch | Low | Critical | Pre-recorded video backup. Golden-run replay mode. USB + cloud + phone copies |
183
+ | Qwen2.5-7B doesn't fit a 12-batch on MI300X | Low | Medium | Have a fallback config in hand (batch=8 + grad_accum=2) |
184
+ | LoRA on alpaca too easy — speedup looks staged | Low | High | Measure on a non-trivial seq_len (1024+), include MFU not just tokens/sec, show kernel breakdown to prove it's real |
185
+ | HF Inference Provider rate limit / outage during demo | Low | Medium | Offline-replay UI lane plays cached_audit.json without any backend; pre-cache a full recorded session; have a backup HF token; `provider="auto"` already routes around individual provider outages |
186
+ | Team member unavailable (illness, etc.) | Low | High | Pair on critical path (agent loop, KB) so no single point of failure |
187
+
188
+ ## Definition of Done — MVP
189
+
190
+ GPU Goblin is "done enough to ship" when **all** of these are true:
191
+
192
+ - ✅ A judge can upload `train_qwen_lora.py` (we provide it) and get a real audit
193
+ - ✅ Agent makes ≥ 4 visible tool calls
194
+ - ✅ Final report shows ≥ 2× tokens/sec, real numbers from MI300X
195
+ - ✅ Every recommendation in the report cites a ROCm KB rule by ID
196
+ - ✅ User can ask follow-up questions in chat ("why bf16?") and get cited answers
197
+ - ✅ Full audit completes in ≤ 4 minutes
198
+ - ✅ 5 consecutive dry runs succeed without manual intervention
199
+ - ✅ Backup video and cached replay both work
200
+
201
+ ## Stretch Definition of Done
202
+
203
+ - 🎯 Second canonical workload (vLLM inference) audited end-to-end
204
+ - 🎯 Cost calculator: "you save $X per training run, $Y per epoch"
205
+ - 🎯 Agent backed by self-hosted Qwen via vLLM on the same MI300X (the ultimate AMD story — replaces today's HF Inference Providers path with on-cluster serving)
206
+
207
+ ## Compute Budget — AMD Developer Cloud Credits
208
+
209
+ Eligible participants get **$100 in AMD Developer Cloud credits**. Public reference price for MI300X is around $1.99/hr (single VM) or $3.39/hr (8× bare metal). Plan accordingly:
210
+
211
+ | Activity | GPU-hours | Notes |
212
+ |---|---|---|
213
+ | Day-1 baseline + synthetic corpus generation | 4–6 | One-shot, results cached |
214
+ | Day-2 KB validation runs | 2–3 | Sanity-check the rules fire on synthetic scenarios |
215
+ | Day-3 canonical demo dry runs (cached) | 2–4 | Cache hits after the first run |
216
+ | Day-3 `--no-cache` cold validation | 1 | Confirms nothing's stale |
217
+ | Day-4 final dry runs + record video | 2–3 | Lock the demo |
218
+ | **Total estimate** | **~12–17 hrs** | Well within $100 even at bare-metal rates |
219
+
220
+ Backend Lead spends zero MI300X time after Day 1 — develops against synthetic corpus + `FakeRunner`.
221
+
222
+ ## Submission Checklist
223
+
224
+ - [ ] Public GitHub repo with clear README + architecture diagram + setup instructions
225
+ - [ ] 90-second demo video (live agent run, real MI300X numbers)
226
+ - [ ] Pitch deck (PDF or slides URL)
227
+ - [ ] Cover image (project listing visual)
228
+ - [ ] Architecture diagram (PNG)
229
+ - [ ] Sample audit report PDF (one canonical run, before/after, including waste-budget chart)
230
+ - [ ] Short description (1-2 sentences) + long description (paragraph) + technology/category tags
231
+ - [ ] At least one Build-in-Public post + one ROCm/Optimum-AMD feedback note (bonus track)
232
+ - [ ] Hackathon submission form filled
233
+ - [ ] Team member credits + contact
brainstorming/idea.md ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # GPU Goblin — Idea
2
+
3
+ > **An AI agent that hunts wasted compute on AMD MI300X.**
4
+ > Upload a fine-tuning config; the agent profiles a real run, diagnoses inefficiency, recommends ROCm-specific fixes, then re-runs and proves the speedup with hard numbers.
5
+
6
+ ---
7
+
8
+ ## The Problem
9
+
10
+ Most teams fine-tuning LLMs on AMD hardware leave **2–3× of throughput on the floor** without realizing it. The waste hides in plain sight:
11
+
12
+ - **HBM under-fed** — batch sizes copied from NVIDIA tutorials don't exploit MI300X's 192 GB.
13
+ - **Wrong precision** — `fp16` is the reflexive default; `bf16` matches throughput on CDNA3 with better numerics.
14
+ - **Naive attention** — teams ship without flash-attn ROCm fork or PyTorch SDPA enabled.
15
+ - **Generic kernels** — hipBLASLt tuning, MIOpen autotune, RCCL collective tweaks all left at defaults.
16
+ - **Wrong libraries** — `bitsandbytes` is the reflexive choice for 8-bit Adam, but it's **not officially supported on ROCm**. Optimum-AMD validates Flash Attention 2 + GPTQ/AWQ paths instead.
17
+ - **Distributed footguns** — one-process-driving-many-GPUs serializes kernel launches on ROCm; NUMA auto-balancing causes variance; `NCCL_MIN_NCHANNELS` left at default.
18
+ - **Knobs nobody knows about** — `HSA_FORCE_FINE_GRAIN_PCIE`, `MIOPEN_FIND_MODE`, hipBLASLt cache paths.
19
+
20
+ Engineers don't fix this not because they're lazy — it's because the knowledge is **scattered across ROCm docs, AMD blog posts, GitHub issues, and Discord threads**. Nobody has time to read it all before launching a run.
21
+
22
+ ## The Solution — GPU Goblin
23
+
24
+ A tool-using agent that does what an experienced AMD performance engineer would do, in one minute:
25
+
26
+ 1. **Read** the user's fine-tuning script or HF `TrainingArguments`.
27
+ 2. **Profile** a 10-step warm run on MI300X (torch.profiler + `rocprofv3` + amd-smi).
28
+ 3. **Diagnose** against a curated knowledge base of ROCm-specific optimizations.
29
+ 4. **Patch** the config — concrete diff, not vague advice.
30
+ 5. **Benchmark** the patched config on the same MI300X.
31
+ 6. **Report** before/after side-by-side with real numbers and citations.
32
+
33
+ The user keeps interacting in chat: *"why bf16?"*, *"what if I can't change the optimizer?"*, *"how much $ does this save per epoch?"*
34
+
35
+ ## Why This Wins
36
+
37
+ **Most teams build user-facing AI. We build AI for AI builders.** That alone is judge-bait.
38
+
39
+ But the deeper hook is that GPU Goblin is **provably useful** in a way most hackathon projects aren't. We don't have to convince judges the recommendations are good — we *show* them, on the same MI300X, with the same model, in the same demo. Tokens/sec: `142 → 318`. End of debate.
40
+
41
+ ### Track Fit (Track 1: AI Agents & Agentic Workflows)
42
+
43
+ - **Primary track:** AI Agents & Agentic Workflows. Real tool-using loop. The agent observes (profile), hypothesizes (KB query), tests (benchmark), refines (patch). Every step visible in chat. Not a one-shot LLM call dressed up as an agent.
44
+ - **Fine-tuning is what we *audit*, not the track we enter.** Canonical workload is **Qwen2.5-7B-Instruct + LoRA on MI300X**. Recommendations are *fine-tuning specific* — batch sizing for LoRA, gradient checkpointing thresholds, optimizer choice (8-bit Adam on ROCm), seq-length packing.
45
+
46
+ ### Qwen Technology Partner Challenge
47
+
48
+ GPU Goblin satisfies the Qwen partner challenge two ways:
49
+
50
+ 1. **Qwen as the agent brain.** The audit loop runs on `Qwen/Qwen2.5-7B-Instruct` served via Hugging Face Inference Providers. Every tool call, every recommendation, every chat answer is generated by Qwen.
51
+ 2. **Qwen as the audit target.** The canonical demo workload is a Qwen2.5-7B-Instruct LoRA fine-tune. The agent audits *the same model family it runs on*.
52
+
53
+ Result: end-to-end Qwen on AMD silicon. Directly answers the judging criterion *"How effectively the chosen model(s) are integrated into the solution."*
54
+
55
+ ### Hugging Face Integration
56
+
57
+ Hugging Face is the named Technology Partner — both **model hub** and **deployment layer**:
58
+
59
+ - **Model hub:** Qwen models pulled from HF Hub at audit time + at agent-spin-up time.
60
+ - **Optimum-AMD:** several KB rules cite Optimum-AMD validated paths (Flash Attention 2 on MI300, GPTQ/AWQ ROCm path).
61
+ - **Inference Providers:** the agent's Qwen calls go through HF's Inference Providers router (`provider="auto"` selects Together / Fireworks-AI / Nebius based on availability).
62
+ - **Deployment as HF Space:** the Streamlit UI ships as a Space within the AMD Developer Hackathon HF Organization, satisfying the "Demo Application Platform" + "Application URL" submission fields. The Space runs in offline-replay mode by default so it works without any live backend.
63
+
64
+ ### AMD Differentiation
65
+
66
+ Every single recommendation cites a **ROCm-specific rule**, not generic PyTorch advice. The knowledge base is the moat:
67
+
68
+ - ROCm env-var tuning (`HSA_*`, `MIOPEN_*`, `NCCL_MIN_NCHANNELS=112`)
69
+ - hipBLASLt hint logging + offline tuning files; MIOpen `MIOPEN_FIND_*` autotune
70
+ - Flash-attn ROCm fork (validated on MI300 via Optimum-AMD)
71
+ - RCCL topology + tensor-parallel placement within a single XGMI island
72
+ - bitsandbytes-on-ROCm gotchas (not officially supported — recommend alternatives)
73
+ - bf16 vs fp16 on CDNA3 matrix cores
74
+ - One-process-per-GPU vs one-process-many-GPUs (ROCm serializes launches in the latter)
75
+ - NUMA auto-balancing disable for stable benchmarks
76
+ - MI300X-specific: 304 CUs, 192 GB HBM3, ~5.3 TB/s peak bandwidth, native FP8
77
+
78
+ This is the kind of insight you only get from someone who has actually shipped on MI300X. We bottle it.
79
+
80
+ ## Demo Narrative (3 minutes)
81
+
82
+ **Setup (15s):** "Most teams waste 50%+ of their MI300X. Watch GPU Goblin find that waste live."
83
+
84
+ **Live demo (2 min):**
85
+
86
+ 1. Drop in `train_qwen_lora.py` — batch=4, fp16, naive attention. (Looks normal.)
87
+ 2. Agent: *"Parsing config…"* → extracts hyperparams.
88
+ 3. Agent: *"Running 10-step profile on MI300X…"* → shows real metrics: HBM 38%, MFU 24%.
89
+ 4. Agent: *"Querying ROCm playbook…"* → 4 issues found, each with citation.
90
+ 5. Agent: *"Generating optimized config…"* → diff appears: bf16, batch=12, flash-attn ROCm, hipBLASLt env, packed sequences.
91
+ 6. Agent: *"Benchmarking new config — 50 steps on MI300X…"* → live progress.
92
+ 7. Final report: **142 → 318 tokens/sec (2.24×). MFU 24% → 51%. $X saved per epoch.**
93
+
94
+ **Why it works (45s):** Show the agent loop. Show the KB. Land the line: *"the agent runs Qwen on AMD silicon, the audit target is Qwen on AMD silicon, the optimizations are AMD-specific. End-to-end AMD."*
95
+
96
+ ## What This Is *Not*
97
+
98
+ - Not a generic PyTorch profiler GUI (`rocprofv3` + tensorboard already exist)
99
+ - Not a chatbot wrapping a static FAQ (the agent runs real benchmarks)
100
+ - Not a multi-cloud cost tool (focused on MI300X compute waste, not infra cost)
101
+ - Not a fine-tuning trainer itself (we *audit* others' runs)
102
+
103
+ ## One-Line Pitch
104
+
105
+ > **GPU Goblin is the AMD performance engineer you wish was on your team — except it costs five minutes and audits any fine-tuning run on MI300X.**
pyproject.toml ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [project]
2
+ name = "gpu-goblin"
3
+ version = "0.1.0"
4
+ description = "An AI agent that hunts wasted compute on AMD MI300X"
5
+ requires-python = ">=3.10"
6
+ dependencies = [
7
+ "fastapi>=0.110",
8
+ "uvicorn[standard]>=0.27",
9
+ "python-multipart>=0.0.9",
10
+ "huggingface_hub>=0.28",
11
+ "pydantic>=2.6",
12
+ "pyyaml>=6.0",
13
+ "sentence-transformers>=2.7",
14
+ "numpy>=1.26",
15
+ "sse-starlette>=2.0",
16
+ "streamlit>=1.32",
17
+ "altair>=5.2",
18
+ "pandas>=2.2",
19
+ "requests>=2.31",
20
+ "openai>=1.30",
21
+ ]
22
+
23
+ [project.optional-dependencies]
24
+ dev = [
25
+ "pytest>=8.0",
26
+ "pytest-asyncio>=0.23",
27
+ "ruff>=0.4",
28
+ ]
29
+
30
+ [build-system]
31
+ requires = ["setuptools>=68"]
32
+ build-backend = "setuptools.build_meta"
33
+
34
+ [tool.setuptools.packages.find]
35
+ include = ["agent*", "runner*", "ui*"]
36
+
37
+ [tool.ruff]
38
+ line-length = 100
39
+ target-version = "py310"
40
+
41
+ [tool.ruff.lint]
42
+ select = ["E", "F", "W", "I", "B", "UP"]
43
+ ignore = ["E501"]
runner/goblin_runner.sh ADDED
@@ -0,0 +1,169 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # goblin_runner.sh — wrap the user's training command with rocprofv3 + amd-smi.
3
+ #
4
+ # Architecture: brainstorming/architecture.md §5 (Profiling Pipeline).
5
+ #
6
+ # Inputs (env vars):
7
+ # USER_SCRIPT Path to the workload python file. Required.
8
+ # OUT_DIR Directory to write trace.csv / torch_profile.json /
9
+ # amd_smi.csv. Required. Created if missing.
10
+ # STEPS --max_steps argument forwarded to the user script. Default 10.
11
+ # GOBLIN_GPU_ID ROCR_VISIBLE_DEVICES value. Default 0.
12
+ #
13
+ # Outputs (in $OUT_DIR):
14
+ # trace.csv rocprofv3 kernel trace
15
+ # torch_profile.json torch.profiler chrome trace (the user script writes this)
16
+ # amd_smi.csv amd-smi telemetry sampled at 200 ms
17
+ # stdout.log user-script stdout
18
+ # stderr.log user-script stderr
19
+ #
20
+ # Failure mode: any non-zero rocprofv3 exit short-circuits the script. On
21
+ # failure we dump the captured stdout/stderr logs to THIS script's own stderr
22
+ # so `subprocess.run(capture_output=True)` in LiveRunner sees the real error
23
+ # (not just an empty `[]` tail). LiveRunner then archives the whole OUT_DIR
24
+ # under bench_cache/last_runner_failure_<ts>/ so you can inspect after-the-fact.
25
+
26
+ set -uo pipefail
27
+
28
+ : "${USER_SCRIPT:?USER_SCRIPT env var is required}"
29
+ : "${OUT_DIR:?OUT_DIR env var is required}"
30
+ STEPS="${STEPS:-10}"
31
+
32
+ mkdir -p "$OUT_DIR"
33
+
34
+ # Pin to a single MI300X so concurrent benchmark runs don't fight.
35
+ export ROCR_VISIBLE_DEVICES="${GOBLIN_GPU_ID:-0}"
36
+
37
+ # Background HBM/power telemetry. Different amd-smi versions ship slightly
38
+ # different flag names (--interval / --watch / no flag at all in older
39
+ # builds), so try a few variants and gracefully degrade. Telemetry is
40
+ # optional — profile_parser tolerates a missing amd_smi.csv.
41
+ start_amd_smi() {
42
+ if ! command -v amd-smi >/dev/null 2>&1; then
43
+ echo "amd-smi not on PATH; skipping telemetry sidecar" \
44
+ > "$OUT_DIR/amd_smi.err"
45
+ return 1
46
+ fi
47
+ # amd-smi telemetry surface drifted hard across rocm versions. We try
48
+ # subcommands in this order, newest-first; first one that survives 0.5s
49
+ # is kept:
50
+ #
51
+ # 1. `amd-smi metric --watch <s> --mem-usage --usage --csv` (ROCm 7.x).
52
+ # Different code path from `monitor`, so it dodges the known
53
+ # `AttributeError: 'Namespace' object has no attribute 'violation'`
54
+ # crash that some ROCm 7.x point releases ship inside the `monitor`
55
+ # subcommand. Produces VRAM_USED_MB / GFX_ACTIVITY / MEM_ACTIVITY
56
+ # columns that profile_parser._pick_column already recognises.
57
+ #
58
+ # 2. `amd-smi monitor --watch <s> ...` (ROCm 6.x and the 7.x builds
59
+ # where `monitor` actually works). Kept as a fallback for older
60
+ # installs that may not have the `metric --watch` form.
61
+ #
62
+ # 3. `amd-smi monitor --interval <s>` (pre-6.0) and finally bare
63
+ # `amd-smi monitor` (very old / implicit-all).
64
+ local variants=(
65
+ # ROCm 7.x: prefer the metric subcommand — bypasses the monitor bug.
66
+ "amd-smi metric --watch 1 --mem-usage --usage --csv"
67
+ "amd-smi metric -w 1 -m -u --csv"
68
+ # ROCm 7.x monitor (works on builds without the violation-attribute bug)
69
+ "amd-smi monitor --watch 1 --power-usage --gfx --mem --vram-usage --csv"
70
+ "amd-smi monitor -w 1 -p -u -m -v --csv"
71
+ # ROCm 6.x intermediate forms
72
+ "amd-smi monitor --watch 1 --csv"
73
+ "amd-smi monitor --watch 1"
74
+ # Older / fallback
75
+ "amd-smi monitor --interval 1 --csv"
76
+ "amd-smi monitor --interval 1"
77
+ "amd-smi monitor --csv"
78
+ "amd-smi monitor"
79
+ )
80
+ for cmd in "${variants[@]}"; do
81
+ # shellcheck disable=SC2086
82
+ $cmd > "$OUT_DIR/amd_smi.csv" 2> "$OUT_DIR/amd_smi.err" &
83
+ local pid=$!
84
+ sleep 0.5
85
+ if kill -0 "$pid" 2>/dev/null; then
86
+ AMD_SMI_PID=$pid
87
+ return 0
88
+ fi
89
+ done
90
+ # All variants failed. Telemetry is optional; mark it as deliberately
91
+ # skipped and let the main run proceed — profile_parser tolerates a
92
+ # missing/empty amd_smi.csv.
93
+ echo "amd-smi monitor: no compatible flag set in this build; telemetry skipped" \
94
+ > "$OUT_DIR/amd_smi.skipped"
95
+ rm -f "$OUT_DIR/amd_smi.csv"
96
+ return 1
97
+ }
98
+
99
+ AMD_SMI_PID=
100
+ start_amd_smi || true # never block the main run on telemetry
101
+
102
+ cleanup() {
103
+ if [[ -n "${AMD_SMI_PID:-}" ]] && kill -0 "$AMD_SMI_PID" 2>/dev/null; then
104
+ kill "$AMD_SMI_PID" 2>/dev/null || true
105
+ wait "$AMD_SMI_PID" 2>/dev/null || true
106
+ fi
107
+ }
108
+ trap cleanup EXIT
109
+
110
+ # Dump the captured logs to *our* stderr on a non-zero rocprofv3 exit so the
111
+ # Python subprocess that spawned us actually sees the real error message.
112
+ # Without this the redirected stdout.log / stderr.log live inside the tempdir
113
+ # only and LiveRunner's stderr-tail check sees nothing.
114
+ dump_failure_logs() {
115
+ local code=$?
116
+ if [[ $code -ne 0 ]]; then
117
+ {
118
+ echo "=== goblin_runner.sh failed with exit code $code ==="
119
+ echo "=== USER_SCRIPT: $USER_SCRIPT ==="
120
+ echo "=== OUT_DIR: $OUT_DIR ==="
121
+ echo "=== ROCR_VISIBLE_DEVICES: ${ROCR_VISIBLE_DEVICES:-unset} ==="
122
+ echo
123
+ echo "=== last 50 lines of $OUT_DIR/stdout.log ==="
124
+ tail -n 50 "$OUT_DIR/stdout.log" 2>/dev/null || echo "(stdout.log missing)"
125
+ echo
126
+ echo "=== last 50 lines of $OUT_DIR/stderr.log ==="
127
+ tail -n 50 "$OUT_DIR/stderr.log" 2>/dev/null || echo "(stderr.log missing)"
128
+ echo
129
+ echo "=== last 20 lines of $OUT_DIR/amd_smi.err ==="
130
+ tail -n 20 "$OUT_DIR/amd_smi.err" 2>/dev/null || echo "(amd_smi.err missing)"
131
+ } 1>&2
132
+ fi
133
+ return $code
134
+ }
135
+
136
+ # rocprofv3 collects HSA + kernel traces. The user script is responsible for
137
+ # writing torch_profile.json (the agent injects torch.profiler around the
138
+ # training loop in Phase 3). --output-format csv keeps parsing simple.
139
+ set +e
140
+ rocprofv3 \
141
+ --hsa-trace --kernel-trace \
142
+ --output-directory "$OUT_DIR" \
143
+ --output-file trace \
144
+ --output-format csv \
145
+ -- \
146
+ python "$USER_SCRIPT" \
147
+ --max_steps="$STEPS" \
148
+ --torch_profile_out="$OUT_DIR/torch_profile.json" \
149
+ > "$OUT_DIR/stdout.log" 2> "$OUT_DIR/stderr.log"
150
+ ROCPROF_EXIT=$?
151
+ set -e
152
+
153
+ if [[ $ROCPROF_EXIT -ne 0 ]]; then
154
+ (exit $ROCPROF_EXIT) || dump_failure_logs
155
+ exit $ROCPROF_EXIT
156
+ fi
157
+
158
+ # rocprofv3 may write trace_kernel_trace.csv etc. — normalize to trace.csv so
159
+ # profile_parser has one stable filename to look for.
160
+ if [[ ! -f "$OUT_DIR/trace.csv" ]]; then
161
+ for candidate in "$OUT_DIR"/trace*kernel*.csv "$OUT_DIR"/*kernel_trace.csv; do
162
+ if [[ -f "$candidate" ]]; then
163
+ cp "$candidate" "$OUT_DIR/trace.csv"
164
+ break
165
+ fi
166
+ done
167
+ fi
168
+
169
+ exit 0
runner/profile_parser.py ADDED
@@ -0,0 +1,693 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """profile_parser — turn goblin_runner.sh artefacts into RunMetrics.
2
+
3
+ Reads three files written by goblin_runner.sh (architecture.md §5):
4
+
5
+ <out_dir>/trace.csv rocprofv3 kernel trace
6
+ <out_dir>/torch_profile.json torch.profiler chrome trace
7
+ <out_dir>/amd_smi.csv amd-smi telemetry, ~200 ms cadence
8
+
9
+ and produces one `RunMetrics` with a populated `WasteBudget`.
10
+
11
+ Each waste-budget bucket is computed with a deliberately simple, documented
12
+ heuristic — these are best-effort signals for the agent, NOT measured
13
+ ground-truth. See the docstring on `_waste_budget()` for the per-bucket
14
+ formula and its known failure modes.
15
+
16
+ This module is import-tolerant on machines without the rocprofv3 stack —
17
+ it only reads files. Missing or unparseable files degrade individual
18
+ metrics to zero and append a warning to RunMetrics.warnings rather than
19
+ raising. LiveRunner ultimately decides what to do with parse failures.
20
+ """
21
+
22
+ from __future__ import annotations
23
+
24
+ import csv
25
+ import json
26
+ import logging
27
+ import re
28
+ from dataclasses import dataclass
29
+ from pathlib import Path
30
+ from typing import Any
31
+
32
+ from agent.schemas import KernelEntry, RunMetrics, WasteBudget, WorkloadConfig
33
+
34
+ _LOG = logging.getLogger(__name__)
35
+
36
+
37
+ # MI300X has 192 GB HBM3. We treat sustained <70% utilisation as headroom.
38
+ _HBM_TOTAL_GB = 192.0
39
+ _HBM_HEALTHY_TARGET = 0.70
40
+
41
+ # Kernel-name patterns used by the heuristics.
42
+ #
43
+ # We use `(?<![A-Za-z0-9])` / `(?![A-Za-z0-9])` instead of `\b` so that
44
+ # underscores act as token separators — kernel names like
45
+ # `rccl_AllReduce` and `hipBLASLt_generic_gemm` should match.
46
+ _BOUND_L = r"(?<![A-Za-z0-9])"
47
+ _BOUND_R = r"(?![A-Za-z0-9])"
48
+ _RCCL_PATTERN = re.compile(
49
+ _BOUND_L + r"(rccl|nccl|all[_-]?reduce|broadcast|reduce[_-]?scatter)" + _BOUND_R, re.I
50
+ )
51
+ _GEMM_PATTERN = re.compile(_BOUND_L + r"(gemm|matmul|hgemm|sgemm|hipblaslt)" + _BOUND_R, re.I)
52
+ _GENERIC_GEMM_PATTERN = re.compile(
53
+ _BOUND_L + r"(generic|fallback|naive|reference)" + _BOUND_R, re.I
54
+ )
55
+ _FP16_PATTERN = re.compile(_BOUND_L + r"(fp16|half|f16)" + _BOUND_R, re.I)
56
+ _BF16_PATTERN = re.compile(_BOUND_L + r"(bf16|bfloat16)" + _BOUND_R, re.I)
57
+
58
+
59
+ # ---------------------------------------------------------------------------
60
+ # Public entry point
61
+ # ---------------------------------------------------------------------------
62
+
63
+
64
+ def parse(
65
+ out_dir: Path,
66
+ config: WorkloadConfig | None = None,
67
+ steps: int = 10,
68
+ ) -> RunMetrics:
69
+ """Build a `RunMetrics` from goblin_runner.sh artefacts in `out_dir`.
70
+
71
+ `config` is used by some heuristics (e.g. precision_path skips the
72
+ estimate when the config is already bf16). It can be None at the cost
73
+ of a more conservative waste-budget estimate.
74
+
75
+ Always returns a RunMetrics object — individual metrics degrade to
76
+ zero with a warning rather than raising, because LiveRunner ultimately
77
+ decides whether parse failures should trigger fallback.
78
+ """
79
+ warnings: list[str] = []
80
+
81
+ kernels = _read_kernels(out_dir / "trace.csv", warnings)
82
+ torch_summary = _read_torch_profile(out_dir / "torch_profile.json", warnings)
83
+ smi = _read_amd_smi(out_dir / "amd_smi.csv", warnings)
84
+
85
+ top_kernels = _top_kernels(kernels, top_n=5)
86
+ gpu_util_pct = smi.gpu_util_pct if smi.gpu_util_pct is not None else _gpu_util_from_kernels(
87
+ kernels, torch_summary
88
+ )
89
+ # Write the resolved gpu_util back into smi so _waste_budget sees the
90
+ # same number RunMetrics reports. Without this, a build where amd-smi's
91
+ # `usage` column returns N/A leaves smi.gpu_util_pct=None — and
92
+ # _waste_budget reads it as 0%, which forces host_gap to consume the
93
+ # full step time and collapses kernel_shape / precision_path to zero
94
+ # (both multiplied by gpu_util).
95
+ smi.gpu_util_pct = gpu_util_pct
96
+
97
+ waste_budget = _waste_budget(
98
+ kernels=kernels,
99
+ torch_summary=torch_summary,
100
+ smi=smi,
101
+ config=config,
102
+ )
103
+
104
+ attention_kernel_loaded = _detect_attention_kernel(kernels)
105
+
106
+ return RunMetrics(
107
+ steps=steps,
108
+ tokens_per_sec=torch_summary.tokens_per_sec or 0.0,
109
+ mfu_pct=torch_summary.mfu_pct or 0.0,
110
+ hbm_peak_gb=smi.hbm_peak_gb or 0.0,
111
+ hbm_avg_gb=smi.hbm_avg_gb or 0.0,
112
+ gpu_util_pct=gpu_util_pct,
113
+ top_kernels=top_kernels,
114
+ attention_kernel_loaded=attention_kernel_loaded,
115
+ waste_budget=waste_budget,
116
+ warnings=warnings,
117
+ rocm_version=smi.rocm_version or "unknown",
118
+ pytorch_version=torch_summary.pytorch_version or "unknown",
119
+ runner_kind="live",
120
+ )
121
+
122
+
123
+ # ---------------------------------------------------------------------------
124
+ # Internal data carriers
125
+ # ---------------------------------------------------------------------------
126
+
127
+
128
+ @dataclass
129
+ class _Kernel:
130
+ name: str
131
+ duration_ns: int
132
+ """Kernel duration in nanoseconds."""
133
+
134
+ is_collective: bool = False
135
+ is_gemm: bool = False
136
+ is_generic_gemm: bool = False
137
+
138
+
139
+ @dataclass
140
+ class _TorchSummary:
141
+ tokens_per_sec: float | None = None
142
+ mfu_pct: float | None = None
143
+ pytorch_version: str | None = None
144
+ step_time_seconds: float | None = None
145
+ """Average wall-clock seconds per training step (used by waste budget)."""
146
+
147
+ host_busy_fraction: float | None = None
148
+ """Fraction of step time the host (CPU) spent doing non-launch work.
149
+ Heuristic: `cpu_self_time / step_time` reported by torch.profiler."""
150
+
151
+
152
+ @dataclass
153
+ class _SmiSummary:
154
+ hbm_peak_gb: float | None = None
155
+ hbm_avg_gb: float | None = None
156
+ gpu_util_pct: float | None = None
157
+ rocm_version: str | None = None
158
+
159
+
160
+ # ---------------------------------------------------------------------------
161
+ # trace.csv (rocprofv3) → list[_Kernel]
162
+ # ---------------------------------------------------------------------------
163
+
164
+
165
+ def _read_kernels(path: Path, warnings: list[str]) -> list[_Kernel]:
166
+ """Parse rocprofv3 kernel-trace CSV into a list of `_Kernel` records.
167
+
168
+ rocprofv3 column names vary slightly by version. We look up by header and
169
+ accept the common aliases; if the file is missing or unparseable we
170
+ return an empty list and append a warning so the caller can decide.
171
+ """
172
+ if not path.exists():
173
+ warnings.append(f"profile_parser: kernel trace not found at {path}")
174
+ return []
175
+
176
+ try:
177
+ with path.open(newline="") as f:
178
+ reader = csv.DictReader(f)
179
+ if reader.fieldnames is None:
180
+ warnings.append(f"profile_parser: empty kernel trace at {path}")
181
+ return []
182
+ name_col = _pick_column(reader.fieldnames, ["KernelName", "Kernel_Name", "kernel_name", "Name"])
183
+ start_col = _pick_column(reader.fieldnames, ["BeginNs", "start_ns", "BeginNS", "Start"])
184
+ end_col = _pick_column(reader.fieldnames, ["EndNs", "end_ns", "EndNS", "End"])
185
+ duration_col = _pick_column(reader.fieldnames, ["DurationNs", "duration_ns", "Duration"])
186
+
187
+ kernels: list[_Kernel] = []
188
+ for row in reader:
189
+ name = (row.get(name_col) or "").strip() if name_col else ""
190
+ if not name:
191
+ continue
192
+ duration = _row_duration_ns(row, duration_col, start_col, end_col)
193
+ if duration <= 0:
194
+ continue
195
+ kernels.append(
196
+ _Kernel(
197
+ name=name,
198
+ duration_ns=duration,
199
+ is_collective=bool(_RCCL_PATTERN.search(name)),
200
+ is_gemm=bool(_GEMM_PATTERN.search(name)),
201
+ is_generic_gemm=bool(
202
+ _GEMM_PATTERN.search(name) and _GENERIC_GEMM_PATTERN.search(name)
203
+ ),
204
+ )
205
+ )
206
+ return kernels
207
+ except (OSError, csv.Error) as exc:
208
+ warnings.append(f"profile_parser: failed to read kernel trace ({exc})")
209
+ return []
210
+
211
+
212
+ def _pick_column(fieldnames: list[str], candidates: list[str]) -> str | None:
213
+ """Pick the first matching column name, with three fallback tiers:
214
+ 1. Exact match.
215
+ 2. Case-insensitive exact match.
216
+ 3. Substring match (case-insensitive) — every token of any candidate
217
+ must appear somewhere in the field name. Tolerates the column-name
218
+ drift between rocprofv3 / amd-smi versions (e.g. `VRAM_USED` vs
219
+ `vram_used_mb` vs `VRAM USED MB`).
220
+ """
221
+ for c in candidates:
222
+ if c in fieldnames:
223
+ return c
224
+ lower = {f.lower(): f for f in fieldnames}
225
+ for c in candidates:
226
+ if c.lower() in lower:
227
+ return lower[c.lower()]
228
+ # Substring tier: split each candidate on _/space, require all tokens
229
+ # appear in the (lowercased) field name. Avoids matching too eagerly
230
+ # by requiring every token of the candidate.
231
+ for c in candidates:
232
+ tokens = [t for t in c.lower().replace("_", " ").split() if t]
233
+ if not tokens:
234
+ continue
235
+ for fname in fieldnames:
236
+ fl = fname.lower()
237
+ if all(t in fl for t in tokens):
238
+ return fname
239
+ return None
240
+
241
+
242
+ def _row_duration_ns(
243
+ row: dict[str, str], duration_col: str | None, start_col: str | None, end_col: str | None
244
+ ) -> int:
245
+ if duration_col and row.get(duration_col):
246
+ try:
247
+ return int(float(row[duration_col]))
248
+ except ValueError:
249
+ return 0
250
+ if start_col and end_col and row.get(start_col) and row.get(end_col):
251
+ try:
252
+ return int(float(row[end_col])) - int(float(row[start_col]))
253
+ except ValueError:
254
+ return 0
255
+ return 0
256
+
257
+
258
+ def _top_kernels(kernels: list[_Kernel], top_n: int) -> list[KernelEntry]:
259
+ if not kernels:
260
+ return []
261
+ total = sum(k.duration_ns for k in kernels) or 1
262
+ by_name: dict[str, int] = {}
263
+ for k in kernels:
264
+ by_name[k.name] = by_name.get(k.name, 0) + k.duration_ns
265
+ ranked = sorted(by_name.items(), key=lambda kv: kv[1], reverse=True)[:top_n]
266
+ return [KernelEntry(name=name, pct_time=ns / total * 100.0) for name, ns in ranked]
267
+
268
+
269
+ def _detect_attention_kernel(kernels: list[_Kernel]) -> str:
270
+ for k in kernels:
271
+ n = k.name.lower()
272
+ if "flash_attn" in n and "rocm" in n:
273
+ return "flash_rocm"
274
+ if "flash" in n and "attn" in n:
275
+ return "flash"
276
+ if "scaled_dot_product_attention" in n:
277
+ return "sdpa"
278
+ if kernels:
279
+ return "eager" # nothing flash-shaped; default conservative label
280
+ return "unknown"
281
+
282
+
283
+ def _gpu_util_from_kernels(kernels: list[_Kernel], torch_summary: _TorchSummary) -> float:
284
+ """Fallback GPU util when amd-smi is missing.
285
+
286
+ util ≈ sum(kernel duration) / total wall-clock step time.
287
+ """
288
+ if not kernels or torch_summary.step_time_seconds in (None, 0):
289
+ return 0.0
290
+ total_kernel_ns = sum(k.duration_ns for k in kernels)
291
+ wall_ns = torch_summary.step_time_seconds * 1e9
292
+ if wall_ns <= 0:
293
+ return 0.0
294
+ return min(100.0, total_kernel_ns / wall_ns * 100.0)
295
+
296
+
297
+ # ---------------------------------------------------------------------------
298
+ # torch_profile.json (torch.profiler chrome trace) → _TorchSummary
299
+ # ---------------------------------------------------------------------------
300
+
301
+
302
+ def _read_torch_profile(path: Path, warnings: list[str]) -> _TorchSummary:
303
+ """Pull tokens/sec, MFU, and step timing from a torch.profiler artefact.
304
+
305
+ The user script (workloads/train_qwen_lora.py in Phase 3) is
306
+ responsible for embedding `tokens_per_sec`, `mfu_pct`, `pytorch_version`
307
+ and `step_time_seconds` in the trace as `metadata` events. If those are
308
+ missing, we estimate `step_time_seconds` from the total trace duration.
309
+ """
310
+ summary = _TorchSummary()
311
+ if not path.exists():
312
+ warnings.append(f"profile_parser: torch profile not found at {path}")
313
+ return summary
314
+ try:
315
+ data = json.loads(path.read_text())
316
+ except (OSError, json.JSONDecodeError) as exc:
317
+ warnings.append(f"profile_parser: failed to read torch profile ({exc})")
318
+ return summary
319
+
320
+ # torch.profiler chrome trace is `{"traceEvents": [...], "metadata": {...}}`
321
+ metadata = data.get("metadata") if isinstance(data, dict) else None
322
+ if isinstance(metadata, dict):
323
+ summary.tokens_per_sec = _coerce_float(metadata.get("tokens_per_sec"))
324
+ summary.mfu_pct = _coerce_float(metadata.get("mfu_pct"))
325
+ summary.pytorch_version = metadata.get("pytorch_version") or metadata.get("torch_version")
326
+ summary.step_time_seconds = _coerce_float(metadata.get("step_time_seconds"))
327
+ summary.host_busy_fraction = _coerce_float(metadata.get("host_busy_fraction"))
328
+
329
+ events = data.get("traceEvents") if isinstance(data, dict) else None
330
+ if isinstance(events, list):
331
+ if summary.step_time_seconds is None:
332
+ summary.step_time_seconds = _step_time_from_events(events)
333
+ if summary.host_busy_fraction is None:
334
+ summary.host_busy_fraction = _host_busy_from_events(events)
335
+
336
+ return summary
337
+
338
+
339
+ def _coerce_float(v: Any) -> float | None:
340
+ if v is None:
341
+ return None
342
+ try:
343
+ return float(v)
344
+ except (TypeError, ValueError):
345
+ return None
346
+
347
+
348
+ def _step_time_from_events(events: list[dict]) -> float | None:
349
+ """Estimate per-step wall-clock seconds from chrome-trace duration events.
350
+
351
+ Looks for `name == "ProfilerStep#*"` complete events; falls back to the
352
+ overall trace span if those aren't present.
353
+ """
354
+ durations: list[float] = []
355
+ overall_start: float | None = None
356
+ overall_end: float | None = None
357
+ for ev in events:
358
+ if not isinstance(ev, dict):
359
+ continue
360
+ name = ev.get("name", "")
361
+ ts = ev.get("ts")
362
+ dur = ev.get("dur")
363
+ if isinstance(name, str) and name.startswith("ProfilerStep") and isinstance(dur, (int, float)):
364
+ durations.append(float(dur) / 1e6) # us → seconds
365
+ if isinstance(ts, (int, float)) and isinstance(dur, (int, float)):
366
+ start = float(ts)
367
+ end = start + float(dur)
368
+ overall_start = start if overall_start is None else min(overall_start, start)
369
+ overall_end = end if overall_end is None else max(overall_end, end)
370
+ if durations:
371
+ return sum(durations) / len(durations)
372
+ if overall_start is not None and overall_end is not None and overall_end > overall_start:
373
+ return (overall_end - overall_start) / 1e6
374
+ return None
375
+
376
+
377
+ def _host_busy_from_events(events: list[dict]) -> float | None:
378
+ """Heuristic: cpu_op event time / total event span.
379
+
380
+ Used by the data_wait waste bucket to disambiguate "GPU idle because the
381
+ host is busy preparing the next batch" from "GPU idle because nothing
382
+ is running anywhere".
383
+ """
384
+ cpu_op_us = 0.0
385
+ span_min: float | None = None
386
+ span_max: float | None = None
387
+ for ev in events:
388
+ if not isinstance(ev, dict):
389
+ continue
390
+ cat = ev.get("cat", "")
391
+ ts = ev.get("ts")
392
+ dur = ev.get("dur")
393
+ if not (isinstance(ts, (int, float)) and isinstance(dur, (int, float))):
394
+ continue
395
+ ts_f = float(ts)
396
+ dur_f = float(dur)
397
+ end = ts_f + dur_f
398
+ span_min = ts_f if span_min is None else min(span_min, ts_f)
399
+ span_max = end if span_max is None else max(span_max, end)
400
+ if isinstance(cat, str) and "cpu_op" in cat.lower():
401
+ cpu_op_us += dur_f
402
+ if span_min is None or span_max is None or span_max <= span_min:
403
+ return None
404
+ span = span_max - span_min
405
+ if span <= 0:
406
+ return None
407
+ return min(1.0, cpu_op_us / span)
408
+
409
+
410
+ # ---------------------------------------------------------------------------
411
+ # amd_smi.csv → _SmiSummary
412
+ # ---------------------------------------------------------------------------
413
+
414
+
415
+ def _read_amd_smi(path: Path, warnings: list[str]) -> _SmiSummary:
416
+ """Aggregate amd-smi polling output into HBM peak/avg + GPU util."""
417
+ summary = _SmiSummary()
418
+ if not path.exists():
419
+ warnings.append(f"profile_parser: amd-smi telemetry not found at {path}")
420
+ return summary
421
+
422
+ try:
423
+ raw = path.read_text()
424
+ except OSError as exc:
425
+ warnings.append(f"profile_parser: failed to read amd-smi telemetry ({exc})")
426
+ return summary
427
+
428
+ csv_text = _strip_amd_smi_preamble(raw)
429
+ if csv_text is None:
430
+ warnings.append(f"profile_parser: no parseable header in amd-smi telemetry at {path}")
431
+ return summary
432
+
433
+ try:
434
+ import io as _io
435
+
436
+ reader = csv.DictReader(_io.StringIO(csv_text))
437
+ if reader.fieldnames is None:
438
+ warnings.append(f"profile_parser: empty amd-smi telemetry at {path}")
439
+ return summary
440
+ hbm_col = _pick_column(
441
+ reader.fieldnames,
442
+ [
443
+ # `amd-smi metric --mem-usage --csv` (ROCm 7.x) emits
444
+ # "used_vram" in MB. Older `amd-smi monitor --vram-usage`
445
+ # emits "VRAM_USED" / "vram_used_mb" depending on minor
446
+ # version.
447
+ "used_vram",
448
+ "USED_VRAM",
449
+ "VRAM_USED_MB",
450
+ "vram_used_mb",
451
+ "VRAM_USED",
452
+ "VRAM_USED_GB",
453
+ "vram_used",
454
+ "VRAM Used",
455
+ # Older rocm 6.x naming
456
+ "MEM_USED",
457
+ "mem_used",
458
+ ],
459
+ )
460
+ util_col = _pick_column(
461
+ reader.fieldnames,
462
+ [
463
+ # `amd-smi metric --usage --csv` (ROCm 7.x) emits a single
464
+ # consolidated "usage" column. Some builds report N/A here
465
+ # — _coerce_float drops those silently and we fall back to
466
+ # the kernel-trace gpu_util estimate downstream.
467
+ "usage",
468
+ "USAGE",
469
+ # `amd-smi monitor --gfx` (ROCm 7.x) → "gfx_util"
470
+ "GFX_UTIL",
471
+ "gfx_util",
472
+ "GFX_UTILIZATION",
473
+ "gfx_utilization",
474
+ # rocm 6.x and older
475
+ "GFX_ACTIVITY",
476
+ "gfx_activity",
477
+ "GPU_USE",
478
+ "GFX %",
479
+ "Util",
480
+ ],
481
+ )
482
+ rocm_col = _pick_column(reader.fieldnames, ["ROCM_VERSION", "rocm_version"])
483
+
484
+ hbm_samples: list[float] = []
485
+ util_samples: list[float] = []
486
+ for row in reader:
487
+ if hbm_col:
488
+ hbm_gb = _hbm_to_gb(row.get(hbm_col), hbm_col)
489
+ if hbm_gb is not None:
490
+ hbm_samples.append(hbm_gb)
491
+ if util_col:
492
+ util = _coerce_float(row.get(util_col))
493
+ if util is not None:
494
+ util_samples.append(min(100.0, util))
495
+ if rocm_col and summary.rocm_version is None:
496
+ val = (row.get(rocm_col) or "").strip()
497
+ if val:
498
+ summary.rocm_version = val
499
+ if hbm_samples:
500
+ summary.hbm_peak_gb = max(hbm_samples)
501
+ summary.hbm_avg_gb = sum(hbm_samples) / len(hbm_samples)
502
+ if util_samples:
503
+ summary.gpu_util_pct = sum(util_samples) / len(util_samples)
504
+ return summary
505
+ except csv.Error as exc:
506
+ warnings.append(f"profile_parser: failed to parse amd-smi telemetry ({exc})")
507
+ return summary
508
+
509
+
510
+ def _strip_amd_smi_preamble(raw: str) -> str | None:
511
+ """Drop everything before the first real CSV header and dedupe repeated
512
+ header lines — both noise produced by `amd-smi <subcmd> --watch`.
513
+
514
+ --watch prints a "'CTRL' + 'C' to stop watching output:" banner once at
515
+ the top, then re-emits the CSV header on every iteration. csv.DictReader
516
+ naively reads the banner as fieldnames and treats every subsequent
517
+ header as a misshapen data row. Pre-strip both before handing it off.
518
+
519
+ Returns a CSV string ready for DictReader, or None if no header line
520
+ is recognisable.
521
+ """
522
+ lines = raw.splitlines()
523
+ header_idx: int | None = None
524
+ for i, line in enumerate(lines):
525
+ if "," not in line:
526
+ continue
527
+ lower = line.lower()
528
+ # Recognised tokens come from the columns amd-smi metric / monitor
529
+ # actually emit. Match conservatively — banners or other noise can
530
+ # contain commas too.
531
+ if any(tok in lower for tok in ("vram", "gfx_", "timestamp,", "gpu_use", "gpu,")):
532
+ header_idx = i
533
+ break
534
+ if header_idx is None:
535
+ return None
536
+
537
+ header = lines[header_idx]
538
+ data = [
539
+ line
540
+ for line in lines[header_idx + 1 :]
541
+ if line.strip() and line != header
542
+ ]
543
+ return header + "\n" + "\n".join(data) + ("\n" if data else "")
544
+
545
+
546
+ def _hbm_to_gb(raw: str | None, column_name: str | None = None) -> float | None:
547
+ """amd-smi sometimes reports VRAM in MB, sometimes in GB.
548
+
549
+ First check the column name — `*vram*` / `*_mb` columns are MB-typed
550
+ in every amd-smi build we've seen; `*_gb` is GB. Without a column-name
551
+ hint, fall back to a value heuristic. The old "v > 1024 ⇒ MB" heuristic
552
+ misclassified small idle samples (e.g. 285 MB at GPU idle) as GB and
553
+ inflated the peak across the run, which then forced memory_headroom to
554
+ zero downstream.
555
+ """
556
+ if not raw:
557
+ return None
558
+ try:
559
+ v = float(str(raw).strip().replace("MB", "").replace("GB", "").replace(",", ""))
560
+ except ValueError:
561
+ return None
562
+ if column_name:
563
+ lower = column_name.lower()
564
+ if "_mb" in lower or lower.endswith("mb"):
565
+ return v / 1024.0
566
+ if "_gb" in lower or lower.endswith("gb"):
567
+ return v
568
+ if "vram" in lower or "mem" in lower:
569
+ # amd-smi metric / monitor default to MB for the bare
570
+ # `used_vram` / `mem_used` columns on every ROCm 6.x+ build.
571
+ return v / 1024.0
572
+ # No column hint — fall back to value heuristic.
573
+ if v > 1024.0:
574
+ return v / 1024.0
575
+ return v
576
+
577
+
578
+ # ---------------------------------------------------------------------------
579
+ # Waste-budget heuristics
580
+ # ---------------------------------------------------------------------------
581
+
582
+
583
+ def _waste_budget(
584
+ *,
585
+ kernels: list[_Kernel],
586
+ torch_summary: _TorchSummary,
587
+ smi: _SmiSummary,
588
+ config: WorkloadConfig | None,
589
+ ) -> WasteBudget:
590
+ """Decompose step time into the seven WasteBudget buckets (architecture.md §3).
591
+
592
+ These are HEURISTICS, not measurements. Each bucket is in seconds-per-step
593
+ so they can be summed against `step_time_seconds`. If we can't observe a
594
+ bucket we set it to 0 — `evidence_coverage` in `propose_patch` then
595
+ discounts confidence accordingly.
596
+
597
+ Per-bucket logic:
598
+
599
+ data_wait
600
+ Fraction of step time where GPU util was below 30% AND the host was
601
+ busy (host_busy_fraction > 0.5). Maps to dataloader / H2D copy stalls.
602
+
603
+ precision_path
604
+ If the user is already on bf16/fp8 we skip this bucket (recovery is 0).
605
+ Otherwise we estimate from kernel names: time spent in fp16-tagged
606
+ GEMMs is the recoverable surface; bf16-tagged kernels are not.
607
+
608
+ kernel_shape
609
+ Fraction of total GEMM kernel time spent on kernels whose names
610
+ match `generic|fallback|naive|reference` — i.e. hipBLASLt/MIOpen
611
+ couldn't pick a tuned tile size and fell back to a slow path.
612
+
613
+ host_gap
614
+ Time the GPU was idle while the host was NOT busy either — pure
615
+ launch latency / eager-mode kernel gaps. We approximate as
616
+ `(1 - gpu_util) * (1 - host_busy)` × step_time.
617
+
618
+ comm_excess
619
+ Sum of all collective-kernel duration (anything matching the rccl
620
+ pattern). Treated as 100% recoverable elsewhere — the rule decides
621
+ recovery_fraction.
622
+
623
+ memory_headroom
624
+ `(192 - hbm_peak) / 192 × small_constant`. Only counts the headroom
625
+ that exceeded a healthy 70% target — running at 60% of HBM costs us
626
+ roughly 0.07 of step time worth of optimisation surface.
627
+
628
+ useful_gpu
629
+ Whatever step time is left. Will be roughly the busy GPU time minus
630
+ comm_excess and the kernel-shape penalty.
631
+ """
632
+ step_t = torch_summary.step_time_seconds or 0.0
633
+ gpu_util = (smi.gpu_util_pct or 0.0) / 100.0 # 0..1
634
+ host_busy = torch_summary.host_busy_fraction or 0.0
635
+ if step_t <= 0:
636
+ return WasteBudget()
637
+
638
+ # data_wait: GPU under-utilised AND host busy → dataloader bottleneck.
639
+ if gpu_util < 0.30 and host_busy > 0.5:
640
+ data_wait = step_t * (1.0 - gpu_util) * host_busy
641
+ else:
642
+ data_wait = 0.0
643
+
644
+ # host_gap: GPU idle while host idle too → launch latency / kernel gaps.
645
+ host_gap = step_t * (1.0 - gpu_util) * (1.0 - host_busy)
646
+
647
+ # comm_excess: total time in collective kernels (in seconds).
648
+ comm_excess_ns = sum(k.duration_ns for k in kernels if k.is_collective)
649
+ comm_excess = comm_excess_ns / 1e9
650
+
651
+ # kernel_shape: fraction of GEMM time spent on un-tuned/generic kernels.
652
+ gemm_total_ns = sum(k.duration_ns for k in kernels if k.is_gemm) or 1
653
+ generic_gemm_ns = sum(k.duration_ns for k in kernels if k.is_generic_gemm)
654
+ kernel_shape = (generic_gemm_ns / gemm_total_ns) * step_t * gpu_util # cap at "real" GPU time
655
+
656
+ # precision_path: only meaningful if config is fp16/fp32. Estimate from
657
+ # kernel name tags. If config is bf16+ we leave as 0 — already optimal.
658
+ precision_path = 0.0
659
+ if config is None or config.precision in {"fp16", "fp32"}:
660
+ fp16_ns = sum(k.duration_ns for k in kernels if _FP16_PATTERN.search(k.name))
661
+ bf16_ns = sum(k.duration_ns for k in kernels if _BF16_PATTERN.search(k.name))
662
+ denom = fp16_ns + bf16_ns
663
+ if denom > 0:
664
+ # Fraction of compute still on fp16. On MI300X bf16 is faster +
665
+ # more numerically stable, so this fraction is the recoverable
666
+ # precision_path surface.
667
+ precision_path = (fp16_ns / denom) * step_t * gpu_util * 0.10
668
+
669
+ # memory_headroom: headroom past the 70% healthy target × small constant.
670
+ memory_headroom = 0.0
671
+ hbm_peak = smi.hbm_peak_gb
672
+ if hbm_peak is not None and hbm_peak > 0:
673
+ utilisation = hbm_peak / _HBM_TOTAL_GB
674
+ if utilisation < _HBM_HEALTHY_TARGET:
675
+ slack = (_HBM_HEALTHY_TARGET - utilisation) / _HBM_HEALTHY_TARGET
676
+ # Small constant: HBM slack is real but only enables a fraction of
677
+ # potential gain (you still need a larger batch to use it). 0.05
678
+ # of step time per "unit of slack" is a conservative anchor.
679
+ memory_headroom = slack * step_t * 0.05
680
+
681
+ # useful_gpu: everything else. Clamp to >= 0.
682
+ spent = data_wait + host_gap + comm_excess + kernel_shape + precision_path + memory_headroom
683
+ useful_gpu = max(0.0, step_t - spent)
684
+
685
+ return WasteBudget(
686
+ useful_gpu=useful_gpu,
687
+ data_wait=data_wait,
688
+ host_gap=host_gap,
689
+ comm_excess=comm_excess,
690
+ memory_headroom=memory_headroom,
691
+ precision_path=precision_path,
692
+ kernel_shape=kernel_shape,
693
+ )
runner/protocol.py CHANGED
@@ -222,41 +222,6 @@ def _archive_failure(out_dir: Path, proc: subprocess.CompletedProcess) -> Path:
222
  return dest
223
 
224
 
225
- def _default_runner_timeout_seconds() -> int:
226
- """Resolve the LiveRunner subprocess timeout from env, with safe defaults.
227
-
228
- Reads ``GOBLIN_RUNNER_TIMEOUT_SECONDS`` and validates it as a positive
229
- int. Anything missing, empty, non-numeric, or non-positive falls back to
230
- 1800 seconds (30 minutes).
231
-
232
- Why this helper exists: live MI300X audits sometimes overshoot the old
233
- 600s default — model download on a cold cache, ROCm kernel JIT on the
234
- first step, or torch silently running on CPU after a botched pip
235
- install. Operators need a knob to extend the budget without editing
236
- code; this turns it into a single env var.
237
- """
238
- raw = os.environ.get("GOBLIN_RUNNER_TIMEOUT_SECONDS", "").strip()
239
- if not raw:
240
- return 1800
241
- try:
242
- val = int(raw)
243
- except ValueError:
244
- _LOG.warning(
245
- "LiveRunner: GOBLIN_RUNNER_TIMEOUT_SECONDS=%r is not an int; "
246
- "falling back to 1800s.",
247
- raw,
248
- )
249
- return 1800
250
- if val <= 0:
251
- _LOG.warning(
252
- "LiveRunner: GOBLIN_RUNNER_TIMEOUT_SECONDS=%d is not positive; "
253
- "falling back to 1800s.",
254
- val,
255
- )
256
- return 1800
257
- return val
258
-
259
-
260
  class LiveRunner:
261
  """Real-MI300X path: shells out to goblin_runner.sh and parses artefacts.
262
 
@@ -271,26 +236,17 @@ class LiveRunner:
271
  self,
272
  runner_script: Path | str = _DEFAULT_RUNNER_SCRIPT,
273
  user_script: Path | str = _DEFAULT_USER_SCRIPT,
274
- timeout_seconds: int | None = None,
275
  fake_fallback: FakeRunner | None = None,
276
  ) -> None:
277
- # `timeout_seconds=None` (the default) means "consult
278
- # GOBLIN_RUNNER_TIMEOUT_SECONDS, then fall back to 1800s (30 min)".
279
- # Explicit ints from callers/tests still win.
280
- #
281
- # 30 min is generous on a healthy MI300X — a 50-step benchmark of
282
- # Qwen2.5-7B LoRA at bs=1/seq=512 finishes in 2–5 min when torch is
283
- # correctly using ROCm. The extra headroom absorbs cold model
284
- # downloads, kernel JIT on first run, and slow CPU-fallback hiccups
285
- # if torch ends up CPU-only. If your workload needs longer, bump
286
- # the env var:
287
- # export GOBLIN_RUNNER_TIMEOUT_SECONDS=3600
288
  self.runner_script = Path(runner_script)
289
  self.user_script = Path(user_script)
290
- self.timeout_seconds = (
291
- timeout_seconds if timeout_seconds is not None
292
- else _default_runner_timeout_seconds()
293
- )
294
  self._fake = fake_fallback or FakeRunner()
295
 
296
  # ------------------------------------------------------------------
 
222
  return dest
223
 
224
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
225
  class LiveRunner:
226
  """Real-MI300X path: shells out to goblin_runner.sh and parses artefacts.
227
 
 
236
  self,
237
  runner_script: Path | str = _DEFAULT_RUNNER_SCRIPT,
238
  user_script: Path | str = _DEFAULT_USER_SCRIPT,
239
+ timeout_seconds: int = 600,
240
  fake_fallback: FakeRunner | None = None,
241
  ) -> None:
242
+ # Default 600s (10 min). Profile runs (10 steps) finish in seconds
243
+ # on a healthy MI300X; benchmarks (50 steps) in a couple of minutes.
244
+ # 30 minutes was a leftover from a workload that wasn't honoring
245
+ # --max_steps and silently trained for hours. With max_steps wired
246
+ # correctly, 600s is generous.
 
 
 
 
 
 
247
  self.runner_script = Path(runner_script)
248
  self.user_script = Path(user_script)
249
+ self.timeout_seconds = timeout_seconds
 
 
 
250
  self._fake = fake_fallback or FakeRunner()
251
 
252
  # ------------------------------------------------------------------
scripts/auto_tune.py ADDED
@@ -0,0 +1,1918 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ """Iterative auto-tuner for AMD MI300X / ROCm 7.0 workloads.
3
+
4
+ Three modes, picked with `--mode`:
5
+
6
+ hardcoded (default)
7
+ Walks through a curated list of MI300X-specific tuning changes one
8
+ at a time. Deterministic, no LLM required — experiments are
9
+ derived from the rules in kb/rocm_rules.yaml.
10
+
11
+ llm
12
+ On each iteration, asks the LLM backend (qwen-hf via HF_TOKEN, or
13
+ qwen-vllm via GOBLIN_QWEN_VLLM_URL) for ONE next experiment given
14
+ the live waste_budget, history, and KB rules. Greedy coordinate
15
+ descent — accept changes that beat the current best by the
16
+ improvement threshold, otherwise revert.
17
+
18
+ llm-explore
19
+ On each iteration, asks the LLM for K candidate experiments at
20
+ once (--candidates-per-iteration, default 3). Benchmarks all K,
21
+ picks the one with the highest tokens/sec, and accepts only if it
22
+ beats the current best. Higher GPU cost (~Kx benchmarks per
23
+ iteration) but better at finding interaction effects that greedy
24
+ one-at-a-time can miss.
25
+
26
+ After each change, runs a real benchmark via goblin_runner.sh and keeps
27
+ the change only if tokens/sec improved meaningfully (>1% by default —
28
+ the threshold cuts measurement noise). Stops when N consecutive
29
+ experiments produce no improvement, or when the source of experiments
30
+ is exhausted.
31
+
32
+ Usage:
33
+ # hardcoded mode (default):
34
+ python scripts/auto_tune.py workloads/train_qwen_lora.py --steps 20
35
+
36
+ # LLM-driven greedy mode:
37
+ python scripts/auto_tune.py workloads/train_qwen_lora.py \\
38
+ --mode llm --steps 20
39
+
40
+ # LLM-driven multi-candidate exploration:
41
+ python scripts/auto_tune.py workloads/train_qwen_lora.py \\
42
+ --mode llm-explore --candidates-per-iteration 3 --steps 20
43
+
44
+ Output:
45
+ - A row-by-row log of each experiment attempted, accepted or rejected
46
+ - A final summary with cumulative speedup
47
+ - A pointer to a temp file containing the best workload script for
48
+ diff-against-baseline inspection
49
+
50
+ Extending hardcoded mode: add an Experiment to EXPERIMENTS. The
51
+ substitutions field is a list of (regex_pattern, replacement) tuples
52
+ applied with re.subn against the workload source. env_vars are exported
53
+ into the goblin_runner.sh subprocess and persist on every accepted
54
+ iteration.
55
+ """
56
+
57
+ from __future__ import annotations
58
+
59
+ import argparse
60
+ import asyncio
61
+ import json
62
+ import os
63
+ import re
64
+ import subprocess
65
+ import sys
66
+ import tempfile
67
+ from dataclasses import dataclass, field
68
+ from pathlib import Path
69
+
70
+ REPO_ROOT = Path(__file__).resolve().parent.parent
71
+ GOBLIN_RUNNER = REPO_ROOT / "runner" / "goblin_runner.sh"
72
+ sys.path.insert(0, str(REPO_ROOT))
73
+
74
+ # Optional structured-events output. When `--events FILE` is passed, the
75
+ # script appends one JSON object per line at key milestones (baseline,
76
+ # iteration start, candidate done, iteration done, summary). Used by the
77
+ # Streamlit UI to render progress live; CLI users typically don't need it.
78
+ _EVENTS_PATH: Path | None = None
79
+
80
+
81
+ def _emit(event: dict) -> None:
82
+ """Append one NDJSON event to the events file if one was configured."""
83
+ if _EVENTS_PATH is None:
84
+ return
85
+ try:
86
+ with _EVENTS_PATH.open("a") as f:
87
+ f.write(json.dumps(event, default=str) + "\n")
88
+ f.flush() # so a UI tailing the file sees events promptly
89
+ except OSError:
90
+ pass # never crash the run on an event-write failure
91
+
92
+ # Default workload template — used when the user passes --model instead
93
+ # of an explicit workload path. We just substitute MODEL_ID and reuse all
94
+ # the other defaults (fp16, batch=4, eager attention, LoRA r=16, …).
95
+ _DEFAULT_WORKLOAD_TEMPLATE = REPO_ROOT / "workloads" / "train_qwen_lora.py"
96
+
97
+
98
+ def _generate_workload_from_model(model_id: str, dest: Path) -> Path:
99
+ """Build a baseline workload by substituting MODEL_ID into the demo
100
+ template (`workloads/train_qwen_lora.py`). Writes to `dest`, returns
101
+ the path.
102
+
103
+ Caveats:
104
+ - Uses the demo's LoRA target_modules (`q_proj`, `v_proj`) which work
105
+ for the major decoder-only LLM families (Qwen, Llama, Mistral,
106
+ Gemma). MoE / GPT-2-style architectures will need a custom workload.
107
+ - The template overwrites HF_TOKEN with a redactable fake. Public
108
+ models load fine; gated models (Llama, etc.) need the user to edit
109
+ the generated workload or use a custom one.
110
+ """
111
+ if not _DEFAULT_WORKLOAD_TEMPLATE.exists():
112
+ raise SystemExit(
113
+ f"--model needs the template at {_DEFAULT_WORKLOAD_TEMPLATE}, but it's missing"
114
+ )
115
+ template_src = _DEFAULT_WORKLOAD_TEMPLATE.read_text()
116
+ new_src, n = re.subn(
117
+ r'MODEL_ID = "[^"]*"',
118
+ f'MODEL_ID = "{model_id}"',
119
+ template_src,
120
+ )
121
+ if n == 0:
122
+ raise SystemExit(
123
+ f"Couldn't find `MODEL_ID = \"...\"` in {_DEFAULT_WORKLOAD_TEMPLATE} "
124
+ "to substitute. Has the template format changed?"
125
+ )
126
+ dest.write_text(new_src)
127
+ return dest
128
+
129
+
130
+ # POSIX env var name: leading letter or underscore, then alnum/underscore.
131
+ # subprocess.run() raises ValueError if any key in the env dict violates
132
+ # this. We validate up-front rather than letting the subprocess crash.
133
+ _VALID_ENV_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
134
+
135
+
136
+ def _sanitize_env_vars(envs: dict, context: str = "") -> dict[str, str]:
137
+ """Clean an env_vars dict from the LLM:
138
+ 1. Strip dotted prefixes (`env_vars.X` → `X`) the LLM mimics from the
139
+ KB transform notation.
140
+ 2. Drop any key that still isn't a valid POSIX env var name. Warns
141
+ instead of crashing — the LLM occasionally embeds shell syntax
142
+ (e.g. `'NUMACTL_INTERLEAVE=1'` as a key) which would make
143
+ subprocess.run raise ValueError.
144
+ """
145
+ cleaned: dict[str, str] = {}
146
+ for k, v in envs.items():
147
+ key = str(k)
148
+ if "." in key:
149
+ stripped = key.rsplit(".", 1)[-1]
150
+ tag = f" [{context}]" if context else ""
151
+ print(f" [warn]{tag} dotted env key {key!r}; using {stripped!r}")
152
+ key = stripped
153
+ if not _VALID_ENV_NAME.match(key):
154
+ tag = f" [{context}]" if context else ""
155
+ print(
156
+ f" [warn]{tag} dropping invalid env var name {key!r} "
157
+ "(must match [A-Za-z_][A-Za-z0-9_]*)"
158
+ )
159
+ continue
160
+ cleaned[key] = str(v)
161
+ return cleaned
162
+
163
+
164
+ @dataclass
165
+ class Experiment:
166
+ name: str
167
+ description: str
168
+ rationale: str
169
+ substitutions: list[tuple[str, str]] = field(default_factory=list)
170
+ env_vars: dict[str, str] = field(default_factory=dict)
171
+
172
+
173
+ # Curated for ROCm 7.0 + MI300X (CDNA3, 192 GB HBM3). Ordered roughly by
174
+ # typical impact on Qwen-shaped LoRA fine-tuning workloads. Each
175
+ # experiment stacks on top of any previously accepted ones.
176
+ EXPERIMENTS: list[Experiment] = [
177
+ Experiment(
178
+ name="bf16_over_fp16",
179
+ description="Switch precision from fp16 to bf16",
180
+ rationale=(
181
+ "MI300X (CDNA3) prefers bf16: same throughput, larger numeric "
182
+ "range, no loss-scaler needed. fp16 underutilizes the matrix "
183
+ "engine on this arch."
184
+ ),
185
+ substitutions=[
186
+ (r"torch_dtype=torch\.float16", "torch_dtype=torch.bfloat16"),
187
+ (r"\bfp16=True\b", "bf16=True"),
188
+ ],
189
+ ),
190
+ Experiment(
191
+ name="batch_size_8",
192
+ description="Increase per_device_train_batch_size 4 → 8",
193
+ rationale="MI300X has 192 GB HBM; batch=4 leaves it on the floor.",
194
+ substitutions=[
195
+ (r"per_device_train_batch_size=4\b", "per_device_train_batch_size=8"),
196
+ ],
197
+ ),
198
+ Experiment(
199
+ name="batch_size_16",
200
+ description="Further increase per_device_train_batch_size to 16",
201
+ rationale="If batch=8 fit and improved, try doubling again.",
202
+ substitutions=[
203
+ (r"per_device_train_batch_size=\d+", "per_device_train_batch_size=16"),
204
+ ],
205
+ ),
206
+ Experiment(
207
+ name="batch_size_32",
208
+ description="Push per_device_train_batch_size to 32",
209
+ rationale=(
210
+ "MI300X has 192 GB HBM3 — batch 16 typically peaks ~130 GB. "
211
+ "If 16 fit, 32 likely fits too and reduces step overhead per "
212
+ "token. Reverts cleanly via OOM-as-crash if not."
213
+ ),
214
+ substitutions=[
215
+ (r"per_device_train_batch_size=\d+", "per_device_train_batch_size=32"),
216
+ ],
217
+ ),
218
+ Experiment(
219
+ name="sdpa_attention",
220
+ description="Switch attention from eager to SDPA",
221
+ rationale=(
222
+ "Eager attention is the slowest path. SDPA dispatches to the "
223
+ "best available kernel (flash on ROCm 7.x where supported, "
224
+ "memory-efficient elsewhere)."
225
+ ),
226
+ substitutions=[
227
+ (r'attn_implementation="eager"', 'attn_implementation="sdpa"'),
228
+ ],
229
+ ),
230
+ Experiment(
231
+ name="dataloader_workers_4",
232
+ description="Bump dataloader_num_workers 0 → 4",
233
+ rationale=(
234
+ "0 workers means the GPU sits idle while the host loads the "
235
+ "next batch. 4 is a safe value across most CPU configs."
236
+ ),
237
+ substitutions=[
238
+ (r"dataloader_num_workers=0", "dataloader_num_workers=4"),
239
+ (r"num_workers=0", "num_workers=4"),
240
+ ],
241
+ ),
242
+ Experiment(
243
+ name="pin_memory",
244
+ description="Enable dataloader_pin_memory",
245
+ rationale=(
246
+ "Pinned host buffers make H2D copies async and overlap with "
247
+ "the GPU. Worth it once you have >0 dataloader workers."
248
+ ),
249
+ substitutions=[
250
+ (r"dataloader_pin_memory=False", "dataloader_pin_memory=True"),
251
+ (r"\bpin_memory=False\b", "pin_memory=True"),
252
+ ],
253
+ ),
254
+ Experiment(
255
+ name="env_hipblaslt",
256
+ description="Set TORCH_BLAS_PREFER_HIPBLASLT=1",
257
+ rationale=(
258
+ "hipBLASLt is significantly faster than rocBLAS for the GEMM "
259
+ "shapes Qwen produces (LoRA-projected attention)."
260
+ ),
261
+ env_vars={"TORCH_BLAS_PREFER_HIPBLASLT": "1"},
262
+ ),
263
+ Experiment(
264
+ name="env_tunable_op",
265
+ description="Set PYTORCH_TUNABLEOP_ENABLED=1",
266
+ rationale=(
267
+ "Enables runtime kernel auto-tuning. Pays a first-run "
268
+ "warmup cost in exchange for a steady-state win on every "
269
+ "subsequent step."
270
+ ),
271
+ env_vars={"PYTORCH_TUNABLEOP_ENABLED": "1"},
272
+ ),
273
+ Experiment(
274
+ name="env_miopen_find",
275
+ description="Set MIOPEN_FIND_MODE=3",
276
+ rationale=(
277
+ "MIOpen FAST mode picks already-tuned kernels without on-the-"
278
+ "fly search. Reduces per-step variance."
279
+ ),
280
+ env_vars={"MIOPEN_FIND_MODE": "3"},
281
+ ),
282
+ ]
283
+
284
+
285
+ # ---------------------------------------------------------------------------
286
+ # Helpers
287
+ # ---------------------------------------------------------------------------
288
+
289
+
290
+ def apply_substitutions(source: str, subs: list[tuple[str, str]]) -> str | None:
291
+ """Apply each (pattern, replacement) in order. Returns the new source,
292
+ or None if any pattern matched zero times (already applied or N/A for
293
+ this workload)."""
294
+ out = source
295
+ for pattern, replacement in subs:
296
+ new, n = re.subn(pattern, replacement, out)
297
+ if n == 0:
298
+ return None
299
+ out = new
300
+ return out
301
+
302
+
303
+ def benchmark(
304
+ workload_path: Path,
305
+ steps: int,
306
+ env_overrides: dict[str, str],
307
+ timeout: int = 600,
308
+ ) -> dict | None:
309
+ """Run goblin_runner.sh on the workload, return parsed RunMetrics dict
310
+ or None on failure."""
311
+ with tempfile.TemporaryDirectory(prefix="auto_tune_run_") as out_dir_str:
312
+ out_dir = Path(out_dir_str)
313
+ env = os.environ.copy()
314
+ env["USER_SCRIPT"] = str(workload_path)
315
+ env["OUT_DIR"] = str(out_dir)
316
+ env["STEPS"] = str(steps)
317
+ # Candidate workload lives in /tmp, so its self-bootstrap line
318
+ # `sys.path.insert(0, dirname(dirname(__file__)))` resolves to /tmp
319
+ # — which has no `workloads/` package. Inject the real repo root via
320
+ # PYTHONPATH so `from workloads._runtime import ...` succeeds.
321
+ existing_pp = env.get("PYTHONPATH", "")
322
+ env["PYTHONPATH"] = (
323
+ str(REPO_ROOT) + (os.pathsep + existing_pp if existing_pp else "")
324
+ )
325
+ env.update(env_overrides)
326
+
327
+ try:
328
+ proc = subprocess.run(
329
+ [str(GOBLIN_RUNNER)],
330
+ env=env,
331
+ capture_output=True,
332
+ text=True,
333
+ timeout=timeout,
334
+ )
335
+ except subprocess.TimeoutExpired:
336
+ print(f" TIMEOUT after {timeout}s")
337
+ return None
338
+ except ValueError as exc:
339
+ # subprocess.run validates env var names and raises ValueError
340
+ # for malformed keys (e.g. names containing '=' or spaces). The
341
+ # LLM has occasionally emitted those; we sanitize earlier but
342
+ # this is the last-resort backstop so a single bad candidate
343
+ # doesn't crash the whole tuning run.
344
+ print(f" REJECTED — illegal env var name(s): {exc}")
345
+ print(f" env keys offered: {list(env_overrides.keys())}")
346
+ return None
347
+ except OSError as exc:
348
+ print(f" REJECTED — could not spawn goblin_runner.sh: {exc}")
349
+ return None
350
+
351
+ if proc.returncode != 0:
352
+ print(f" goblin_runner.sh failed (exit {proc.returncode})")
353
+ tail = (proc.stderr or "").strip().splitlines()[-8:]
354
+ for line in tail:
355
+ print(f" | {line}")
356
+ return None
357
+
358
+ try:
359
+ from runner import profile_parser
360
+
361
+ metrics = profile_parser.parse(out_dir, steps=steps)
362
+ return metrics.model_dump()
363
+ except Exception as exc: # parser is defensive but be safe
364
+ print(f" profile_parser raised: {type(exc).__name__}: {exc}")
365
+ return None
366
+
367
+
368
+ def _delta_pct(new: float, baseline: float) -> float:
369
+ if baseline <= 0:
370
+ return 0.0
371
+ return (new - baseline) / baseline * 100.0
372
+
373
+
374
+ # ---------------------------------------------------------------------------
375
+ # LLM-driven experiment generator
376
+ # ---------------------------------------------------------------------------
377
+
378
+
379
+ _LLM_SYSTEM_PROMPT = """\
380
+ You are an expert at tuning AMD MI300X (ROCm 7.0, CDNA3 arch, 192 GB
381
+ HBM3) training workloads. The user is iteratively benchmarking changes
382
+ to a transformers/peft fine-tuning script. On each turn you suggest ONE
383
+ specific parameter change to try next, targeting the largest non-useful
384
+ waste bucket in the most recent benchmark.
385
+
386
+ Your output MUST be a single JSON object with this exact shape (no
387
+ prose, no markdown fences, just the object):
388
+
389
+ {
390
+ "name": "short_snake_case_name",
391
+ "rationale": "1-3 sentences on why this change addresses the worst waste bucket",
392
+ "substitutions": [["regex_pattern", "replacement"]],
393
+ "env_vars": {"VAR_NAME": "value"}
394
+ }
395
+
396
+ CRITICAL output rules — read carefully:
397
+
398
+ 1. env_vars keys are LITERAL POSIX shell environment variable names.
399
+ They MUST match the regex [A-Za-z_][A-Za-z0-9_]* — letters, digits,
400
+ underscores only, starting with a letter or underscore.
401
+ - NEVER prefix them with "env_vars." or any other dotted path.
402
+ - NEVER include "=" or shell syntax in the key — env var names are
403
+ identifiers, NOT assignments and NOT commands.
404
+ - If you want to invoke a command-line tool like `numactl` or
405
+ `taskset`, that CANNOT be expressed as an env_var. Don't try.
406
+ Either propose a `substitutions` change to the script, or skip.
407
+ Wrong: {"env_vars.MIOPEN_FIND_MODE": "3"}
408
+ Wrong: {"NUMACTL_INTERLEAVE=1": "numactl --interleave=all"}
409
+ Wrong: {"export FOO": "bar"}
410
+ Right: {"MIOPEN_FIND_MODE": "3"}
411
+ Right: {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}
412
+
413
+ 2. substitutions are (regex_pattern, replacement) pairs applied with
414
+ re.subn against the current workload source. Patterns must match at
415
+ least one occurrence in the source — if zero matches, the experiment
416
+ is auto-skipped (counted as no improvement).
417
+
418
+ 3. When the previous change for a parameter improved tokens/sec, push
419
+ that parameter further in the same direction next time. E.g. if
420
+ batch_size 4 → 8 won, try 8 → 16. If 16 won and HBM is still under
421
+ ~150 GB, try 32. Don't be timid — MI300X has 192 GB HBM3.
422
+
423
+ 4. Don't repeat any (name OR substitution OR env_var combo) from
424
+ history. If a change was rejected, don't propose the same numerical
425
+ value again — try a different one.
426
+
427
+ 5. If you cannot think of a productive next change, output:
428
+ {"name": "STOP", "rationale": "<why>", "substitutions": [], "env_vars": {}}
429
+
430
+ CONCRETE OUTPUT EXAMPLES — match this shape exactly:
431
+
432
+ Switch fp16 → bf16 (precision_path bucket):
433
+ {"name": "bf16_over_fp16",
434
+ "rationale": "MI300X CDNA3 matrix cores prefer bf16: same throughput, larger numeric range, no loss-scaler.",
435
+ "substitutions": [["fp16=True", "bf16=True"], ["torch_dtype=torch\\\\.float16", "torch_dtype=torch.bfloat16"]],
436
+ "env_vars": {}}
437
+
438
+ Increase batch size to 16 (memory_headroom bucket):
439
+ {"name": "batch_size_16",
440
+ "rationale": "Current HBM peak is well under 192 GB; bigger batch saturates the GPU.",
441
+ "substitutions": [["per_device_train_batch_size=\\\\d+", "per_device_train_batch_size=16"]],
442
+ "env_vars": {}}
443
+
444
+ Switch attention to SDPA (kernel_shape bucket):
445
+ {"name": "sdpa_attention",
446
+ "rationale": "Eager attention is the slowest path; SDPA dispatches to a tuned kernel.",
447
+ "substitutions": [["attn_implementation=\\"eager\\"", "attn_implementation=\\"sdpa\\""]],
448
+ "env_vars": {}}
449
+
450
+ Bump dataloader workers (data_wait bucket):
451
+ {"name": "dataloader_workers_4",
452
+ "rationale": "0 workers starves the GPU between batches.",
453
+ "substitutions": [["dataloader_num_workers=0", "dataloader_num_workers=4"]],
454
+ "env_vars": {}}
455
+
456
+ Set MIOpen FAST mode (kernel_shape bucket, env-only):
457
+ {"name": "miopen_find_fast",
458
+ "rationale": "FAST mode picks already-tuned kernels without on-the-fly search.",
459
+ "substitutions": [],
460
+ "env_vars": {"MIOPEN_FIND_MODE": "3"}}
461
+
462
+ Prefer hipBLASLt (kernel_shape bucket, env-only):
463
+ {"name": "prefer_hipblaslt",
464
+ "rationale": "hipBLASLt is faster than rocBLAS for Qwen GEMM shapes on MI300X.",
465
+ "substitutions": [],
466
+ "env_vars": {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}}
467
+ """
468
+
469
+
470
+ _LLM_USER_TEMPLATE = """\
471
+ Hardware facts (use these — do not contradict):
472
+ - AMD MI300X, CDNA3 architecture, 192 GB HBM3
473
+ - bf16 throughput on CDNA3 ≈ same as fp16, > fp32 (matrix engine is fp16/bf16/fp8 native)
474
+ - fp32 is the SLOWEST option on this arch — never suggest it as an improvement
475
+
476
+ Known incompatibilities for THIS workload (peft + LoRA on transformers Trainer):
477
+ {incompatibilities}
478
+
479
+ KB rules (one-liner per rule, for grounding):
480
+ {kb_summary}
481
+
482
+ Current accepted workload state — these are the literal values in the
483
+ script after every change accepted so far. The next change you propose
484
+ should mutate one of these (or set an env var). DO NOT propose a value
485
+ that's already present here.
486
+ {tunables}
487
+
488
+ Latest benchmark (this is the result of the most recent ACCEPTED state):
489
+ - tokens_per_sec: {tps:.1f}
490
+ - mfu_pct: {mfu:.2f} (% of MI300X dense bf16 peak; healthy LoRA ranges 30-50%)
491
+ - gpu_util_pct: {util:.1f}
492
+ - hbm_peak_gb: {hbm:.2f}
493
+ - waste_budget (seconds/step):
494
+ {waste_lines}
495
+
496
+ Sorted recoverable waste (largest first — go after these):
497
+ {recoverable_sorted}
498
+
499
+ History of changes already tried this run (newest first; outcomes are
500
+ "accepted" / "rejected" / "crashed" / "skipped"):
501
+ {history_lines}
502
+
503
+ If the latest entry is "crashed", the change you propose next must be
504
+ STRUCTURALLY different (different parameter, not just a different value
505
+ of the same one).
506
+
507
+ Suggest ONE next change targeting the largest recoverable bucket. JSON only.
508
+ """
509
+
510
+
511
+ # Workload-specific incompatibilities the LLM otherwise wastes iterations on.
512
+ # Keep this list short and concrete — it goes into every prompt.
513
+ _KNOWN_INCOMPATIBILITIES = [
514
+ "gradient_checkpointing=True requires `model.enable_input_require_grads()`"
515
+ " before peft wrapping for LoRA models. Setting it via a single substitution"
516
+ " WILL CRASH the workload. Don't propose it.",
517
+ "bitsandbytes-based optimizers (`adamw_8bit`, `paged_adamw_8bit`) and"
518
+ " `load_in_8bit=True` are NOT supported on ROCm 7.x. Don't propose them.",
519
+ "torch_compile=True with peft/LoRA on ROCm 7.x triggers compile-time"
520
+ " errors with the current PyTorch nightly (2.9.x). Don't propose it"
521
+ " unless you have specific evidence it works on this version.",
522
+ "flash_attention_2 may not be installed (try `attn_implementation=\"sdpa\"`"
523
+ " before `\"flash_attention_2\"`).",
524
+ "persistent_workers=True requires num_workers > 0. PyTorch raises"
525
+ " `ValueError: persistent_workers option needs num_workers > 0` if you"
526
+ " enable it while num_workers=0. If the current workload has"
527
+ " dataloader_num_workers=0, do NOT propose persistent_workers=True"
528
+ " alone — pair it with `dataloader_num_workers=4` (or higher) in the"
529
+ " SAME experiment via two substitutions, or wait until a previous"
530
+ " experiment has bumped num_workers above 0.",
531
+ "dataloader_prefetch_factor only works when num_workers > 0 (same"
532
+ " constraint as persistent_workers). Same rule: bump num_workers in"
533
+ " the same experiment, or skip.",
534
+ ]
535
+
536
+
537
+ def _kb_summary(rules_yaml_path: Path, max_chars: int = 6000) -> str:
538
+ """Return a compact one-line-per-rule summary of kb/rocm_rules.yaml.
539
+
540
+ Notably we DO NOT show the raw `transform` field — earlier versions
541
+ did and the LLM ended up copying its dotted-path notation literally
542
+ (`env_vars.MIOPEN_FIND_MODE` as the env var name, not as a dict
543
+ accessor). The system prompt's CONCRETE EXAMPLES section is the
544
+ canonical source of truth for output shape; this summary just
545
+ grounds the LLM's reasoning in the catalog of known issues.
546
+ """
547
+ if not rules_yaml_path.exists():
548
+ return "(KB rules file not found)"
549
+ try:
550
+ import yaml
551
+
552
+ rules = yaml.safe_load(rules_yaml_path.read_text()) or []
553
+ except Exception as exc:
554
+ return f"(failed to parse KB: {exc})"
555
+
556
+ lines = []
557
+ for r in rules:
558
+ if not isinstance(r, dict):
559
+ continue
560
+ rid = r.get("id", "?")
561
+ bucket = r.get("targets_bucket", "?")
562
+ sym = (r.get("symptom") or "").strip().replace("\n", " ")
563
+ if len(sym) > 110:
564
+ sym = sym[:107] + "..."
565
+ lines.append(f"- {rid:55s} [{bucket}] {sym}")
566
+ text = "\n".join(lines)
567
+ if len(text) > max_chars:
568
+ text = text[:max_chars] + "\n... (truncated)"
569
+ return text
570
+
571
+
572
+ # Map of (substring-in-source) → (parameter description, example regex
573
+ # pattern, example replacement template). Each entry is a hint shown to
574
+ # the LLM so it has a concrete target to point its substitutions at —
575
+ # instead of guessing what the workload's literal config text looks like.
576
+ _TUNABLE_HINTS: list[tuple[str, str, str, str]] = [
577
+ # (token to detect, description, regex_for_substitution, replacement_template)
578
+ ("torch_dtype=torch.float16",
579
+ "model precision (matches `torch_dtype=torch.float16`)",
580
+ r"torch_dtype=torch\.float16",
581
+ "torch_dtype=torch.bfloat16"),
582
+ ("torch_dtype=torch.bfloat16",
583
+ "model precision (already bf16)",
584
+ r"torch_dtype=torch\.bfloat16",
585
+ "torch_dtype=torch.float16"),
586
+ ("fp16=True",
587
+ "TrainingArguments fp16 (matches `fp16=True`)",
588
+ r"\bfp16=True\b",
589
+ "bf16=True"),
590
+ ("bf16=True",
591
+ "TrainingArguments bf16 (already bf16)",
592
+ r"\bbf16=True\b",
593
+ "fp16=True"),
594
+ ("attn_implementation=\"eager\"",
595
+ "attention impl (matches `attn_implementation=\"eager\"`)",
596
+ r'attn_implementation="eager"',
597
+ 'attn_implementation="sdpa"'),
598
+ ("attn_implementation=\"sdpa\"",
599
+ "attention impl (currently sdpa; could try flash_attention_2)",
600
+ r'attn_implementation="sdpa"',
601
+ 'attn_implementation="flash_attention_2"'),
602
+ ("per_device_train_batch_size=",
603
+ "per-device batch size (matches `per_device_train_batch_size=<N>`)",
604
+ r"per_device_train_batch_size=\d+",
605
+ "per_device_train_batch_size=<NEW_VALUE>"),
606
+ ("dataloader_num_workers=",
607
+ "dataloader workers (matches `dataloader_num_workers=<N>`)",
608
+ r"dataloader_num_workers=\d+",
609
+ "dataloader_num_workers=<NEW_VALUE>"),
610
+ ("dataloader_pin_memory=",
611
+ "dataloader pin_memory (matches `dataloader_pin_memory=<bool>`)",
612
+ r"dataloader_pin_memory=(True|False)",
613
+ "dataloader_pin_memory=True"),
614
+ ("gradient_checkpointing=",
615
+ "gradient checkpointing toggle",
616
+ r"gradient_checkpointing=(True|False)",
617
+ "gradient_checkpointing=True"),
618
+ ("torch_compile=",
619
+ "torch.compile toggle (use cautiously on ROCm 7.x)",
620
+ r"torch_compile=(True|False)",
621
+ "torch_compile=True"),
622
+ ("optim=\"adamw_torch\"",
623
+ "optimizer choice (currently adamw_torch)",
624
+ r'optim="adamw_torch"',
625
+ 'optim="adamw_torch_fused"'),
626
+ ]
627
+
628
+
629
+ def _tunables_summary(source: str) -> str:
630
+ """Detect which tunable parameters are present in the workload source
631
+ and surface their current literal values + ready-to-use regex patterns
632
+ so the LLM has concrete substitution targets.
633
+
634
+ Skips comment lines when reporting the "current" value — many workloads
635
+ document expected findings in a top-of-file comment block, and we want
636
+ the LLM to see the live config line, not the doc string.
637
+ """
638
+ lines: list[str] = []
639
+ source_lines = source.splitlines()
640
+ for token, desc, pattern, replacement in _TUNABLE_HINTS:
641
+ live_line: str | None = None
642
+ for raw in source_lines:
643
+ stripped = raw.lstrip()
644
+ if stripped.startswith("#"):
645
+ continue
646
+ if token in raw:
647
+ live_line = raw.strip()
648
+ break
649
+ if live_line is None:
650
+ continue
651
+ lines.append(
652
+ f" • {desc}\n"
653
+ f" current: {live_line}\n"
654
+ f" pattern: {pattern!r} replacement template: {replacement!r}"
655
+ )
656
+ if not lines:
657
+ return " (no recognized tunables — substitutions will need to match other text)"
658
+ return "\n".join(lines)
659
+
660
+
661
+ def _recoverable_sorted(waste: dict) -> str:
662
+ """List the non-useful_gpu buckets sorted by size, so the LLM can
663
+ explicitly target the biggest one first."""
664
+ if not waste:
665
+ return " (no waste_budget available)"
666
+ items = [
667
+ (name, value)
668
+ for name, value in waste.items()
669
+ if name != "useful_gpu" and isinstance(value, (int, float))
670
+ ]
671
+ items.sort(key=lambda kv: kv[1], reverse=True)
672
+ if not items:
673
+ return " (no recoverable buckets)"
674
+ return "\n".join(f" {i + 1}. {name:18s} = {value:.4f}" for i, (name, value) in enumerate(items))
675
+
676
+
677
+ def _config_snippet(source: str, max_lines: int = 80) -> str:
678
+ """Return the lines around `TrainingArguments(` and `from_pretrained(` so
679
+ the LLM sees the actual config it's modifying without us shipping the
680
+ whole script. Gives ~max_lines of context.
681
+ """
682
+ lines = source.splitlines()
683
+ keep: list[tuple[int, str]] = []
684
+ for i, line in enumerate(lines):
685
+ lower = line.lower()
686
+ if any(
687
+ tok in lower
688
+ for tok in (
689
+ "trainingarguments(",
690
+ "from_pretrained(",
691
+ "loraconfig(",
692
+ "dataloader(",
693
+ "torch_dtype",
694
+ "attn_implementation",
695
+ "fp16=",
696
+ "bf16=",
697
+ "per_device_train_batch_size",
698
+ "dataloader_num_workers",
699
+ "dataloader_pin_memory",
700
+ "gradient_checkpointing",
701
+ "torch_compile",
702
+ "optim=",
703
+ )
704
+ ):
705
+ keep.append((i, line))
706
+ if not keep:
707
+ return source[:2000]
708
+ # Coalesce nearby line indices into windows for readability
709
+ windows: list[list[str]] = []
710
+ last_idx = -10
711
+ cur: list[str] = []
712
+ for i, line in keep:
713
+ if i - last_idx > 3:
714
+ if cur:
715
+ windows.append(cur)
716
+ cur = []
717
+ cur.append(f"{i + 1:4d}: {line}")
718
+ last_idx = i
719
+ if cur:
720
+ windows.append(cur)
721
+ out = "\n\n".join("\n".join(w) for w in windows)
722
+ if out.count("\n") > max_lines:
723
+ out_lines = out.splitlines()[:max_lines]
724
+ out = "\n".join(out_lines) + "\n... (truncated)"
725
+ return out
726
+
727
+
728
+ def _format_history(history: list[dict]) -> str:
729
+ if not history:
730
+ return "(none yet — this is the first iteration)"
731
+ lines = []
732
+ for h in reversed(history[-12:]): # last 12 newest-first
733
+ outcome = h.get("outcome", "?")
734
+ delta = h.get("delta_pct")
735
+ delta_s = f"{delta:+.2f}%" if delta is not None else "n/a"
736
+ subs = h.get("substitutions") or []
737
+ envs = h.get("env_vars") or {}
738
+ change_repr = f"subs={subs} env={envs}"
739
+ lines.append(f"- {h['name']:25s} {outcome:9s} Δ {delta_s:8s} {change_repr}")
740
+ return "\n".join(lines)
741
+
742
+
743
+ def _format_waste(waste: dict) -> str:
744
+ keys = (
745
+ "useful_gpu",
746
+ "data_wait",
747
+ "host_gap",
748
+ "comm_excess",
749
+ "memory_headroom",
750
+ "precision_path",
751
+ "kernel_shape",
752
+ )
753
+ return "\n".join(f" {k:18s} = {waste.get(k, 0.0):.4f}" for k in keys)
754
+
755
+
756
+ def _build_llm_backend(system_prompt: str = _LLM_SYSTEM_PROMPT, max_tokens: int = 1024):
757
+ """Construct the same backend the agent loop uses. Surfaces a clear
758
+ message if neither HF_TOKEN nor a vLLM URL is configured."""
759
+ has_hf = bool(os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN"))
760
+ has_vllm = bool(os.environ.get("GOBLIN_QWEN_VLLM_URL"))
761
+ backend_kind = os.environ.get("GOBLIN_AGENT_BACKEND", "qwen-hf").lower()
762
+ if backend_kind in ("qwen-hf", "qwen", "hf", "") and not has_hf:
763
+ raise SystemExit(
764
+ "LLM mode requires HF_TOKEN (qwen-hf backend) or "
765
+ "GOBLIN_AGENT_BACKEND=qwen-vllm + GOBLIN_QWEN_VLLM_URL."
766
+ )
767
+ if backend_kind in ("qwen-vllm", "qwen_vllm", "vllm", "local") and not has_vllm:
768
+ raise SystemExit(
769
+ "LLM mode with qwen-vllm backend requires GOBLIN_QWEN_VLLM_URL."
770
+ )
771
+ from agent.backends import make_backend
772
+
773
+ return make_backend(system_prompt=system_prompt, max_tokens=max_tokens)
774
+
775
+
776
+ async def _ask_llm_for_experiment(
777
+ backend,
778
+ *,
779
+ kb_summary: str,
780
+ source: str,
781
+ metrics: dict,
782
+ history: list[dict],
783
+ ) -> Experiment | None:
784
+ """One LLM turn → one Experiment (or None for STOP / parse failure)."""
785
+ waste = metrics.get("waste_budget") or {}
786
+ prompt = _LLM_USER_TEMPLATE.format(
787
+ incompatibilities="\n".join(f"- {line}" for line in _KNOWN_INCOMPATIBILITIES),
788
+ kb_summary=kb_summary,
789
+ tunables=_tunables_summary(source),
790
+ tps=metrics.get("tokens_per_sec", 0.0),
791
+ mfu=metrics.get("mfu_pct", 0.0),
792
+ util=metrics.get("gpu_util_pct", 0.0),
793
+ hbm=metrics.get("hbm_peak_gb", 0.0),
794
+ waste_lines=_format_waste(waste),
795
+ recoverable_sorted=_recoverable_sorted(waste),
796
+ history_lines=_format_history(history),
797
+ )
798
+ backend.add_user_message(prompt)
799
+ turn = await backend.next_turn(tool_schemas=[])
800
+ raw = " ".join(turn.text_blocks).strip()
801
+
802
+ obj = _extract_json_object(raw)
803
+ if obj is None:
804
+ print(f" LLM response was not parseable JSON. Raw: {raw[:300]!r}")
805
+ return None
806
+
807
+ name = (obj.get("name") or "").strip()
808
+ if not name or name.upper() == "STOP":
809
+ print(f" LLM signaled STOP: {obj.get('rationale', '(no rationale)')}")
810
+ return None
811
+
812
+ subs_raw = obj.get("substitutions") or []
813
+ envs = obj.get("env_vars") or {}
814
+ if not subs_raw and not envs:
815
+ print(f" LLM returned an empty experiment ({name}); skipping")
816
+ return None
817
+
818
+ subs: list[tuple[str, str]] = []
819
+ for entry in subs_raw:
820
+ if isinstance(entry, list) and len(entry) == 2:
821
+ subs.append((str(entry[0]), str(entry[1])))
822
+ elif isinstance(entry, dict) and "pattern" in entry and "replacement" in entry:
823
+ subs.append((str(entry["pattern"]), str(entry["replacement"])))
824
+
825
+ cleaned_envs = _sanitize_env_vars(envs, context=name)
826
+ if not subs and not cleaned_envs:
827
+ # Everything got dropped during sanitization (bad env names + no
828
+ # valid substitutions). Treat as a no-op rather than benchmarking
829
+ # an unchanged workload.
830
+ print(f" LLM experiment {name!r} had nothing valid after sanitization; skipping")
831
+ return None
832
+
833
+ return Experiment(
834
+ name=name,
835
+ description=obj.get("description") or name,
836
+ rationale=str(obj.get("rationale") or ""),
837
+ substitutions=subs,
838
+ env_vars=cleaned_envs,
839
+ )
840
+
841
+
842
+ # ---------------------------------------------------------------------------
843
+ # llm-explore mode: ask for K candidates per iteration
844
+ # ---------------------------------------------------------------------------
845
+
846
+
847
+ _LLM_EXPLORE_SYSTEM_PROMPT = """\
848
+ You are an expert at tuning AMD MI300X (ROCm 7.0, CDNA3 arch, 192 GB
849
+ HBM3) training workloads. The user is running a multi-candidate
850
+ exploration: on every iteration you suggest K STRUCTURALLY-DIFFERENT
851
+ candidate changes, the user benchmarks all of them, and the best one
852
+ is accepted (if it beats the current best by the threshold).
853
+
854
+ Your output MUST be a JSON ARRAY of K objects, no prose, no markdown
855
+ fences, just the array:
856
+
857
+ [
858
+ {"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}},
859
+ {"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}},
860
+ {"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}}
861
+ ]
862
+
863
+ CRITICAL output rules:
864
+
865
+ 1. Each candidate must target a DIFFERENT waste bucket or parameter
866
+ category than the others. Diversity beats redundancy — don't propose
867
+ three batch-size bumps; propose one batch bump, one env var, one
868
+ precision/attention/dataloader change.
869
+
870
+ 2. env_vars keys are LITERAL POSIX shell environment variable names —
871
+ they MUST match the regex [A-Za-z_][A-Za-z0-9_]*. NEVER prefix them
872
+ with "env_vars." or any other dotted path. NEVER include "=" or
873
+ shell syntax in the key. If you want to invoke a CLI tool like
874
+ `numactl`, that's NOT an env var — skip the candidate entirely.
875
+ Wrong: {"env_vars.MIOPEN_FIND_MODE": "3"}
876
+ Wrong: {"NUMACTL_INTERLEAVE=1": "numactl --interleave=all"}
877
+ Right: {"MIOPEN_FIND_MODE": "3"}
878
+
879
+ 3. substitutions are (regex_pattern, replacement) pairs applied with
880
+ re.subn. Patterns must match at least one occurrence — if zero
881
+ matches, that candidate is skipped.
882
+
883
+ 4. NEVER propose a (substitutions, env_vars) combination that already
884
+ appears in history with outcome rejected/crashed. Diversify within
885
+ the array AND across the run.
886
+
887
+ 5. If you genuinely cannot find K productive candidates, output fewer
888
+ (e.g. 2 if K=3). The user will benchmark whatever you provide. If
889
+ you have zero productive candidates, output:
890
+ [{"name": "STOP", "rationale": "<why>", "substitutions": [], "env_vars": {}}]
891
+
892
+ CONCRETE OUTPUT EXAMPLES (for K=3):
893
+
894
+ [
895
+ {"name": "bf16_over_fp16",
896
+ "rationale": "Largest recoverable bucket is precision_path; CDNA3 prefers bf16.",
897
+ "substitutions": [["fp16=True", "bf16=True"], ["torch_dtype=torch\\\\.float16", "torch_dtype=torch.bfloat16"]],
898
+ "env_vars": {}},
899
+ {"name": "batch_size_16",
900
+ "rationale": "HBM peak well under 192 GB; bigger batch saturates the GPU.",
901
+ "substitutions": [["per_device_train_batch_size=\\\\d+", "per_device_train_batch_size=16"]],
902
+ "env_vars": {}},
903
+ {"name": "prefer_hipblaslt",
904
+ "rationale": "hipBLASLt outperforms rocBLAS on Qwen GEMM shapes.",
905
+ "substitutions": [],
906
+ "env_vars": {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}}
907
+ ]
908
+ """
909
+
910
+
911
+ _LLM_EXPLORE_USER_TEMPLATE = """\
912
+ Hardware facts (use these — do not contradict):
913
+ - AMD MI300X, CDNA3 architecture, 192 GB HBM3
914
+ - bf16 throughput on CDNA3 ≈ same as fp16, > fp32 (matrix engine is fp16/bf16/fp8 native)
915
+ - fp32 is the SLOWEST option on this arch — never suggest it as an improvement
916
+
917
+ Known incompatibilities for THIS workload (peft + LoRA on transformers Trainer):
918
+ {incompatibilities}
919
+
920
+ KB rules (one-liner per rule, for grounding):
921
+ {kb_summary}
922
+
923
+ Current accepted workload state — the literal values in the script
924
+ after every change accepted so far. Each candidate you propose should
925
+ mutate one of these (or set an env var). DO NOT propose a value that's
926
+ already present here.
927
+ {tunables}
928
+
929
+ Latest benchmark (this is the result of the most recent ACCEPTED state):
930
+ - tokens_per_sec: {tps:.1f}
931
+ - mfu_pct: {mfu:.2f} (% of MI300X dense bf16 peak; healthy LoRA ranges 30-50%)
932
+ - gpu_util_pct: {util:.1f}
933
+ - hbm_peak_gb: {hbm:.2f}
934
+ - waste_budget (seconds/step):
935
+ {waste_lines}
936
+
937
+ Sorted recoverable waste (largest first — go after these):
938
+ {recoverable_sorted}
939
+
940
+ Previously rejected (full fingerprint — DO NOT repropose any of these):
941
+ {rejected_fingerprints}
942
+
943
+ History of changes already tried this run (newest first; outcomes are
944
+ "accepted" / "rejected" / "crashed" / "skipped"):
945
+ {history_lines}
946
+
947
+ Suggest {num_candidates} STRUCTURALLY-DIFFERENT candidate changes.
948
+ Each must target a different waste bucket or parameter category. JSON
949
+ array only.
950
+ """
951
+
952
+
953
+ async def _ask_llm_for_experiments(
954
+ backend,
955
+ *,
956
+ kb_summary: str,
957
+ source: str,
958
+ metrics: dict,
959
+ history: list[dict],
960
+ num_candidates: int,
961
+ ) -> list[Experiment]:
962
+ """One LLM turn → up to `num_candidates` Experiments.
963
+
964
+ Returns an empty list on parse failure or STOP signal.
965
+ """
966
+ waste = metrics.get("waste_budget") or {}
967
+ prompt = _LLM_EXPLORE_USER_TEMPLATE.format(
968
+ num_candidates=num_candidates,
969
+ incompatibilities="\n".join(f"- {line}" for line in _KNOWN_INCOMPATIBILITIES),
970
+ kb_summary=kb_summary,
971
+ tunables=_tunables_summary(source),
972
+ tps=metrics.get("tokens_per_sec", 0.0),
973
+ mfu=metrics.get("mfu_pct", 0.0),
974
+ util=metrics.get("gpu_util_pct", 0.0),
975
+ hbm=metrics.get("hbm_peak_gb", 0.0),
976
+ waste_lines=_format_waste(waste),
977
+ recoverable_sorted=_recoverable_sorted(waste),
978
+ rejected_fingerprints=_format_rejected_fingerprints(history),
979
+ history_lines=_format_history(history),
980
+ )
981
+ backend.add_user_message(prompt)
982
+ turn = await backend.next_turn(tool_schemas=[])
983
+ raw = " ".join(turn.text_blocks).strip()
984
+
985
+ arr = _extract_json_array(raw)
986
+ if not arr:
987
+ print(f" LLM response was not parseable JSON array. Raw: {raw[:300]!r}")
988
+ return []
989
+
990
+ experiments: list[Experiment] = []
991
+ for obj in arr:
992
+ if not isinstance(obj, dict):
993
+ continue
994
+ name = (obj.get("name") or "").strip()
995
+ if not name:
996
+ continue
997
+ if name.upper() == "STOP":
998
+ print(f" LLM signaled STOP: {obj.get('rationale', '(no rationale)')}")
999
+ return []
1000
+ subs_raw = obj.get("substitutions") or []
1001
+ envs_raw = obj.get("env_vars") or {}
1002
+ if not subs_raw and not envs_raw:
1003
+ continue
1004
+ subs = []
1005
+ for entry in subs_raw:
1006
+ if isinstance(entry, list) and len(entry) == 2:
1007
+ subs.append((str(entry[0]), str(entry[1])))
1008
+ elif isinstance(entry, dict) and "pattern" in entry and "replacement" in entry:
1009
+ subs.append((str(entry["pattern"]), str(entry["replacement"])))
1010
+ cleaned_envs = _sanitize_env_vars(envs_raw, context=name)
1011
+ if not subs and not cleaned_envs:
1012
+ print(f" candidate {name!r} had nothing valid after sanitization; dropping")
1013
+ continue
1014
+ experiments.append(
1015
+ Experiment(
1016
+ name=name,
1017
+ description=obj.get("description") or name,
1018
+ rationale=str(obj.get("rationale") or ""),
1019
+ substitutions=subs,
1020
+ env_vars=cleaned_envs,
1021
+ )
1022
+ )
1023
+ return experiments
1024
+
1025
+
1026
+ def _extract_json_array(text: str) -> list | None:
1027
+ """Pull the first JSON array out of an LLM response, tolerating
1028
+ markdown fences and leading prose. Returns None if nothing parseable."""
1029
+ if not text:
1030
+ return None
1031
+ fence_match = re.search(r"```(?:json)?\s*(\[.*?\])\s*```", text, re.DOTALL)
1032
+ if fence_match:
1033
+ try:
1034
+ obj = json.loads(fence_match.group(1))
1035
+ if isinstance(obj, list):
1036
+ return obj
1037
+ except json.JSONDecodeError:
1038
+ pass
1039
+ depth = 0
1040
+ start = -1
1041
+ for i, ch in enumerate(text):
1042
+ if ch == "[":
1043
+ if depth == 0:
1044
+ start = i
1045
+ depth += 1
1046
+ elif ch == "]":
1047
+ depth -= 1
1048
+ if depth == 0 and start >= 0:
1049
+ blob = text[start : i + 1]
1050
+ try:
1051
+ obj = json.loads(blob)
1052
+ if isinstance(obj, list):
1053
+ return obj
1054
+ except json.JSONDecodeError:
1055
+ start = -1
1056
+ continue
1057
+ return None
1058
+
1059
+
1060
+ # ---------------------------------------------------------------------------
1061
+ # Dedup + history utilities (used by all LLM modes)
1062
+ # ---------------------------------------------------------------------------
1063
+
1064
+
1065
+ def _experiment_fingerprint(exp: Experiment) -> tuple:
1066
+ """Hashable identity for an experiment — substitutions + env_vars,
1067
+ NOT name (the LLM tends to give the same change different names)."""
1068
+ subs = tuple(sorted(tuple(s) for s in exp.substitutions))
1069
+ envs = tuple(sorted(exp.env_vars.items()))
1070
+ return (subs, envs)
1071
+
1072
+
1073
+ def _build_merged_experiment(
1074
+ exps: list[Experiment], base_source: str
1075
+ ) -> tuple[Experiment | None, str]:
1076
+ """Try to combine 2+ experiments into one. The merged experiment
1077
+ applies all of their substitutions in sequence and unions their
1078
+ env_vars. Returns (merged, "") on success, (None, reason) when the
1079
+ merge is structurally unsafe — caller should fall back to using just
1080
+ the individual winner.
1081
+
1082
+ Conflict detection:
1083
+ - A later substitution's pattern must still match after earlier
1084
+ substitutions have been applied (zero matches → conflict, e.g.
1085
+ cand A rewrote `fp16=True` and cand B was also targeting it).
1086
+ - Env var keys with conflicting values (same name, different value)
1087
+ → conflict.
1088
+ - Bad regex anywhere → conflict.
1089
+ """
1090
+ if len(exps) < 2:
1091
+ return None, "need at least 2 experiments"
1092
+
1093
+ merged_subs: list[tuple[str, str]] = []
1094
+ merged_envs: dict[str, str] = {}
1095
+ test_source = base_source
1096
+
1097
+ for exp in exps:
1098
+ for pattern, replacement in exp.substitutions:
1099
+ try:
1100
+ new_source, n = re.subn(pattern, replacement, test_source)
1101
+ except re.error as e:
1102
+ return None, f"bad regex in '{exp.name}': {e}"
1103
+ if n == 0:
1104
+ return None, (
1105
+ f"'{exp.name}' substitution {pattern!r} no longer matches "
1106
+ "after prior merges (likely overwrites an earlier change)"
1107
+ )
1108
+ test_source = new_source
1109
+ merged_subs.append((pattern, replacement))
1110
+ for k, v in exp.env_vars.items():
1111
+ if k in merged_envs and merged_envs[k] != v:
1112
+ return None, (
1113
+ f"env var conflict on {k!r}: {merged_envs[k]!r} vs {v!r}"
1114
+ )
1115
+ merged_envs[k] = v
1116
+
1117
+ short_names = "+".join(e.name[:14] for e in exps)
1118
+ full_names = " + ".join(e.name for e in exps)
1119
+ return (
1120
+ Experiment(
1121
+ name=f"merge[{short_names}]"[:60],
1122
+ description=f"Merged: {full_names}",
1123
+ rationale=(
1124
+ f"Combined {len(exps)} candidates that each had positive delta "
1125
+ "against the current best this iteration. Tests the compound "
1126
+ "effect; falls back to the individual winner if it doesn't help."
1127
+ ),
1128
+ substitutions=merged_subs,
1129
+ env_vars=merged_envs,
1130
+ ),
1131
+ "",
1132
+ )
1133
+
1134
+
1135
+ def _is_duplicate_of_history(exp: Experiment, history: list[dict]) -> dict | None:
1136
+ """If `exp` matches a prior history entry by fingerprint, return that
1137
+ entry. Otherwise None."""
1138
+ fp = _experiment_fingerprint(exp)
1139
+ for h in history:
1140
+ h_subs = tuple(
1141
+ sorted(
1142
+ (str(s[0]), str(s[1]))
1143
+ for s in (h.get("substitutions") or [])
1144
+ if isinstance(s, (list, tuple)) and len(s) == 2
1145
+ )
1146
+ )
1147
+ h_envs = tuple(sorted((h.get("env_vars") or {}).items()))
1148
+ if fp == (h_subs, h_envs):
1149
+ return h
1150
+ return None
1151
+
1152
+
1153
+ def _format_rejected_fingerprints(history: list[dict]) -> str:
1154
+ """Compact list of every (substitutions, env_vars) the LLM has already
1155
+ tried with outcome rejected/crashed/skipped — so it can't propose them
1156
+ again under a different name."""
1157
+ seen: set[tuple] = set()
1158
+ lines: list[str] = []
1159
+ for h in history:
1160
+ outcome = h.get("outcome", "")
1161
+ if outcome not in ("rejected", "crashed", "skipped"):
1162
+ continue
1163
+ subs = tuple(
1164
+ sorted(
1165
+ (str(s[0]), str(s[1]))
1166
+ for s in (h.get("substitutions") or [])
1167
+ if isinstance(s, (list, tuple)) and len(s) == 2
1168
+ )
1169
+ )
1170
+ envs = tuple(sorted((h.get("env_vars") or {}).items()))
1171
+ fp = (subs, envs)
1172
+ if fp in seen:
1173
+ continue
1174
+ seen.add(fp)
1175
+ lines.append(f" - {outcome:9s} subs={list(subs)} env={dict(envs)}")
1176
+ if not lines:
1177
+ return " (none yet)"
1178
+ return "\n".join(lines)
1179
+
1180
+
1181
+ def _print_waste(metrics: dict, prefix: str = " waste: ") -> None:
1182
+ """Print a one-line summary of waste_budget — useful is highlighted
1183
+ first, then non-zero recoverable buckets sorted by size."""
1184
+ wb = metrics.get("waste_budget") or {}
1185
+ if not wb:
1186
+ return
1187
+ parts = [f"useful_gpu={wb.get('useful_gpu', 0.0):.3f}"]
1188
+ others = [(k, v) for k, v in wb.items() if k != "useful_gpu" and isinstance(v, (int, float)) and v > 0]
1189
+ others.sort(key=lambda kv: kv[1], reverse=True)
1190
+ parts.extend(f"{k}={v:.3f}" for k, v in others)
1191
+ print(prefix + ", ".join(parts))
1192
+
1193
+
1194
+ # ---------------------------------------------------------------------------
1195
+ # JSON object extractor (used by single-experiment llm mode)
1196
+ # ---------------------------------------------------------------------------
1197
+
1198
+
1199
+ def _extract_json_object(text: str) -> dict | None:
1200
+ """Pull the first JSON object out of an LLM response, tolerating
1201
+ markdown fences / leading prose."""
1202
+ if not text:
1203
+ return None
1204
+ # strip ```json ... ``` fences if present
1205
+ fence_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
1206
+ if fence_match:
1207
+ try:
1208
+ return json.loads(fence_match.group(1))
1209
+ except json.JSONDecodeError:
1210
+ pass
1211
+ # otherwise grab the first balanced { ... }
1212
+ depth = 0
1213
+ start = -1
1214
+ for i, ch in enumerate(text):
1215
+ if ch == "{":
1216
+ if depth == 0:
1217
+ start = i
1218
+ depth += 1
1219
+ elif ch == "}":
1220
+ depth -= 1
1221
+ if depth == 0 and start >= 0:
1222
+ blob = text[start : i + 1]
1223
+ try:
1224
+ return json.loads(blob)
1225
+ except json.JSONDecodeError:
1226
+ start = -1
1227
+ continue
1228
+ return None
1229
+
1230
+
1231
+ # ---------------------------------------------------------------------------
1232
+ # Main
1233
+ # ---------------------------------------------------------------------------
1234
+
1235
+
1236
+ def main() -> int:
1237
+ p = argparse.ArgumentParser(
1238
+ description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
1239
+ )
1240
+ p.add_argument(
1241
+ "workload",
1242
+ type=Path,
1243
+ nargs="?",
1244
+ default=None,
1245
+ help=(
1246
+ "Path to a workload script (omit if using --model). When given, "
1247
+ "the script is used as-is for the baseline benchmark."
1248
+ ),
1249
+ )
1250
+ p.add_argument(
1251
+ "--model",
1252
+ type=str,
1253
+ default=None,
1254
+ help=(
1255
+ "HuggingFace model id (e.g. Qwen/Qwen2.5-7B-Instruct, "
1256
+ "meta-llama/Llama-3.2-3B). Generates a baseline workload from "
1257
+ "workloads/train_qwen_lora.py with this MODEL_ID substituted in. "
1258
+ "Use this OR a workload path, not both. For gated models, "
1259
+ "ensure HF_TOKEN is set in your shell."
1260
+ ),
1261
+ )
1262
+ p.add_argument(
1263
+ "--mode",
1264
+ choices=("hardcoded", "llm", "llm-explore"),
1265
+ default="hardcoded",
1266
+ help=(
1267
+ "hardcoded (default): walk through the priority-ordered EXPERIMENTS list. "
1268
+ "llm: ask the LLM for ONE next experiment per iteration (greedy). "
1269
+ "llm-explore: ask for K candidates per iteration, benchmark all, keep "
1270
+ "the best (slower but better at finding interaction effects)."
1271
+ ),
1272
+ )
1273
+ p.add_argument(
1274
+ "--candidates-per-iteration",
1275
+ type=int,
1276
+ default=3,
1277
+ help="Only used when --mode llm-explore. Default 3.",
1278
+ )
1279
+ p.add_argument("--steps", type=int, default=20, help="Steps per benchmark")
1280
+ p.add_argument(
1281
+ "--max-iterations",
1282
+ type=int,
1283
+ default=0,
1284
+ help=(
1285
+ "Cap on experiments to try. Default: len(EXPERIMENTS) for hardcoded mode, "
1286
+ "10 for llm mode."
1287
+ ),
1288
+ )
1289
+ p.add_argument(
1290
+ "--early-stop-after",
1291
+ type=int,
1292
+ default=3,
1293
+ help=(
1294
+ "Stop after N consecutive non-improvements. Crashes do NOT count "
1295
+ "toward this — crashes mean the change was structurally bad, not "
1296
+ "that we've exhausted ideas."
1297
+ ),
1298
+ )
1299
+ p.add_argument(
1300
+ "--max-crashes",
1301
+ type=int,
1302
+ default=4,
1303
+ help=(
1304
+ "Stop after N total subprocess crashes (separate from "
1305
+ "--early-stop-after). Default 4 leaves room for the LLM to try "
1306
+ "structurally different changes after a bad one."
1307
+ ),
1308
+ )
1309
+ p.add_argument(
1310
+ "--improvement-threshold",
1311
+ type=float,
1312
+ default=0.0,
1313
+ help=(
1314
+ "Min %% improvement over current best to accept. Default 0.0 "
1315
+ "(any positive delta wins). Bump to 1.0 if your benchmarks are "
1316
+ "noisy and you want to ignore sub-1%% deltas."
1317
+ ),
1318
+ )
1319
+ p.add_argument(
1320
+ "--events",
1321
+ type=Path,
1322
+ default=None,
1323
+ help=(
1324
+ "Optional NDJSON event stream output. If set, the script appends "
1325
+ "one JSON event per line at baseline / iter / candidate / summary "
1326
+ "milestones. Used by the Streamlit UI; CLI users don't need this."
1327
+ ),
1328
+ )
1329
+ args = p.parse_args()
1330
+ if args.events is not None:
1331
+ global _EVENTS_PATH
1332
+ _EVENTS_PATH = args.events
1333
+ try:
1334
+ args.events.write_text("") # truncate any prior contents
1335
+ except OSError as exc:
1336
+ sys.stderr.write(f"--events: cannot open {args.events} for writing ({exc})\n")
1337
+ return 1
1338
+ if args.max_iterations <= 0:
1339
+ if args.mode == "hardcoded":
1340
+ args.max_iterations = len(EXPERIMENTS)
1341
+ elif args.mode == "llm-explore":
1342
+ args.max_iterations = 5 # K candidates per iter so 5 iters = 5K benchmarks
1343
+ else:
1344
+ args.max_iterations = 10
1345
+
1346
+ # Validate that exactly one workload source was provided
1347
+ if args.workload is None and args.model is None:
1348
+ sys.stderr.write(
1349
+ "Pass either a workload path or --model MODEL_ID. "
1350
+ "Examples:\n"
1351
+ " python scripts/auto_tune.py workloads/train_qwen_lora.py\n"
1352
+ " python scripts/auto_tune.py --model Qwen/Qwen2.5-7B-Instruct\n"
1353
+ )
1354
+ return 1
1355
+ if args.workload is not None and args.model is not None:
1356
+ sys.stderr.write(
1357
+ "Pass EITHER a workload path OR --model, not both.\n"
1358
+ )
1359
+ return 1
1360
+ if not GOBLIN_RUNNER.exists():
1361
+ sys.stderr.write(f"goblin_runner.sh not found at {GOBLIN_RUNNER}\n")
1362
+ return 1
1363
+
1364
+ workspace = Path(tempfile.mkdtemp(prefix="auto_tune_workloads_"))
1365
+
1366
+ if args.workload is not None:
1367
+ workload = args.workload.resolve()
1368
+ if not workload.exists():
1369
+ sys.stderr.write(f"workload not found: {workload}\n")
1370
+ return 1
1371
+ workload_label = str(workload)
1372
+ else:
1373
+ # Generate baseline workload from --model
1374
+ generated = workspace / "_generated_baseline.py"
1375
+ workload = _generate_workload_from_model(args.model, generated)
1376
+ workload_label = f"(generated from --model {args.model})\n "
1377
+ workload_label += f" {workload}\n "
1378
+ workload_label += f" template: {_DEFAULT_WORKLOAD_TEMPLATE}"
1379
+
1380
+ _emit({
1381
+ "type": "started",
1382
+ "mode": args.mode,
1383
+ "workload": str(workload),
1384
+ "model": args.model,
1385
+ "steps": args.steps,
1386
+ "max_iterations": args.max_iterations,
1387
+ "early_stop_after": args.early_stop_after,
1388
+ "max_crashes": args.max_crashes,
1389
+ "improvement_threshold": args.improvement_threshold,
1390
+ "candidates_per_iteration": (
1391
+ args.candidates_per_iteration if args.mode == "llm-explore" else 1
1392
+ ),
1393
+ "workspace": str(workspace),
1394
+ })
1395
+ print(f"Auto-tune workspace: {workspace}")
1396
+ print(f"Mode: {args.mode}")
1397
+ print(f"Workload: {workload_label}")
1398
+ print(f"Steps per benchmark: {args.steps}")
1399
+ print(f"Max iterations: {args.max_iterations}")
1400
+ print(f"Early stop after: {args.early_stop_after} non-improvements")
1401
+ print(f"Max crashes: {args.max_crashes} total")
1402
+ print(f"Accept threshold: {args.improvement_threshold:.1f}%\n")
1403
+
1404
+ # LLM mode setup happens before the baseline so we fail fast on missing
1405
+ # credentials rather than after burning a baseline benchmark. Each LLM
1406
+ # mode gets its own system prompt — the explore mode needs a much
1407
+ # larger token budget to emit K JSON objects.
1408
+ llm_backend = None
1409
+ kb_summary = ""
1410
+ if args.mode == "llm":
1411
+ llm_backend = _build_llm_backend(_LLM_SYSTEM_PROMPT, max_tokens=1024)
1412
+ kb_summary = _kb_summary(REPO_ROOT / "kb" / "rocm_rules.yaml")
1413
+ print("LLM backend ready (single-candidate). KB summary loaded.\n")
1414
+ elif args.mode == "llm-explore":
1415
+ llm_backend = _build_llm_backend(_LLM_EXPLORE_SYSTEM_PROMPT, max_tokens=2048)
1416
+ kb_summary = _kb_summary(REPO_ROOT / "kb" / "rocm_rules.yaml")
1417
+ print(
1418
+ f"LLM backend ready (multi-candidate, K={args.candidates_per_iteration}). "
1419
+ "KB summary loaded.\n"
1420
+ )
1421
+
1422
+ baseline_source = workload.read_text()
1423
+ baseline_path = workspace / "00_baseline.py"
1424
+ baseline_path.write_text(baseline_source)
1425
+
1426
+ print("=" * 60)
1427
+ print("Baseline benchmark")
1428
+ print("=" * 60)
1429
+ baseline = benchmark(baseline_path, args.steps, {})
1430
+ if baseline is None:
1431
+ sys.stderr.write("Baseline benchmark failed; cannot continue.\n")
1432
+ return 1
1433
+
1434
+ baseline_tps = baseline["tokens_per_sec"]
1435
+ print(f" tokens/sec: {baseline_tps:.1f}")
1436
+ print(f" mfu_pct: {baseline.get('mfu_pct', 0.0):.2f}")
1437
+ print(f" hbm_peak_gb: {baseline['hbm_peak_gb']:.2f}")
1438
+ print(f" gpu_util_pct: {baseline['gpu_util_pct']:.1f}")
1439
+ print(
1440
+ " waste_budget: "
1441
+ + ", ".join(f"{k}={v:.3f}" for k, v in baseline["waste_budget"].items() if v > 0)
1442
+ )
1443
+ _emit({"type": "baseline", "metrics": baseline})
1444
+
1445
+ best_source = baseline_source
1446
+ best_tps = baseline_tps
1447
+ best_env: dict[str, str] = {}
1448
+ last_metrics = baseline
1449
+ accepted: list[tuple[str, float, float]] = [] # (name, tps, delta_pct)
1450
+ rejected: list[tuple[str, str]] = [] # (name, reason)
1451
+ history: list[dict] = [] # for LLM context
1452
+ consecutive_no_improvement = 0
1453
+ total_crashes = 0
1454
+ file_counter = 0 # monotonically increases across all candidates
1455
+
1456
+ for i in range(args.max_iterations):
1457
+ # ---- Get candidates list (1 for hardcoded/llm, K for llm-explore) ----
1458
+ if args.mode == "hardcoded":
1459
+ if i >= len(EXPERIMENTS):
1460
+ print("\nReached end of EXPERIMENTS list.")
1461
+ break
1462
+ candidates = [EXPERIMENTS[i]]
1463
+ elif args.mode == "llm":
1464
+ print(f"\n[asking LLM for next experiment, iteration {i + 1}...]")
1465
+ try:
1466
+ exp = asyncio.run(
1467
+ _ask_llm_for_experiment(
1468
+ llm_backend,
1469
+ kb_summary=kb_summary,
1470
+ source=best_source,
1471
+ metrics=last_metrics,
1472
+ history=history,
1473
+ )
1474
+ )
1475
+ except Exception as exc:
1476
+ print(f" LLM call failed: {type(exc).__name__}: {exc}")
1477
+ exp = None
1478
+ if exp is None:
1479
+ print("LLM produced no experiment — stopping.")
1480
+ break
1481
+ candidates = [exp]
1482
+ else: # llm-explore
1483
+ K = args.candidates_per_iteration
1484
+ print(f"\n[asking LLM for {K} candidates, iteration {i + 1}...]")
1485
+ try:
1486
+ candidates = asyncio.run(
1487
+ _ask_llm_for_experiments(
1488
+ llm_backend,
1489
+ kb_summary=kb_summary,
1490
+ source=best_source,
1491
+ metrics=last_metrics,
1492
+ history=history,
1493
+ num_candidates=K,
1494
+ )
1495
+ )
1496
+ except Exception as exc:
1497
+ print(f" LLM call failed: {type(exc).__name__}: {exc}")
1498
+ candidates = []
1499
+ if not candidates:
1500
+ print("LLM produced no candidates — stopping.")
1501
+ break
1502
+ print(f" LLM proposed {len(candidates)} candidate(s): "
1503
+ + ", ".join(c.name for c in candidates))
1504
+
1505
+ print()
1506
+ print("=" * 60)
1507
+ n_label = f" ({len(candidates)} candidates)" if len(candidates) > 1 else ""
1508
+ print(f"Iteration {i + 1}{n_label}")
1509
+ print("=" * 60)
1510
+ _emit({
1511
+ "type": "iter_start",
1512
+ "iteration": i + 1,
1513
+ "candidates": [
1514
+ {
1515
+ "name": c.name,
1516
+ "rationale": c.rationale,
1517
+ "substitutions": c.substitutions,
1518
+ "env_vars": c.env_vars,
1519
+ }
1520
+ for c in candidates
1521
+ ],
1522
+ })
1523
+
1524
+ # ---- Evaluate each candidate against the CURRENT best ----
1525
+ # Crucial for llm-explore: every candidate is benchmarked against
1526
+ # the same best_source / best_env baseline, so the comparison is
1527
+ # apples-to-apples. State updates only happen after the iteration's
1528
+ # winner is chosen.
1529
+ eval_results: list[dict] = [] # candidates that produced metrics
1530
+ seen_this_iter: set[tuple] = set() # within-batch dedup
1531
+ crashed_this_iter = False
1532
+ max_crashes_hit = False
1533
+
1534
+ for j, exp in enumerate(candidates):
1535
+ cand_label = f" Candidate {j + 1}/{len(candidates)}" if len(candidates) > 1 else " Candidate"
1536
+ print(f"\n{cand_label}: {exp.name}")
1537
+ print(f" description: {exp.description}")
1538
+ print(f" rationale: {exp.rationale}")
1539
+
1540
+ # Helper to emit a per-candidate event with the consistent shape
1541
+ # the UI expects. Called at every terminus below.
1542
+ def _cand_event(outcome: str, metrics: dict | None = None,
1543
+ delta_vs_best: float | None = None,
1544
+ reason: str = "") -> None:
1545
+ _emit({
1546
+ "type": "candidate",
1547
+ "iteration": i + 1,
1548
+ "candidate_index": j + 1,
1549
+ "n_candidates": len(candidates),
1550
+ "name": exp.name,
1551
+ "rationale": exp.rationale,
1552
+ "substitutions": exp.substitutions,
1553
+ "env_vars": exp.env_vars,
1554
+ "outcome": outcome,
1555
+ "metrics": metrics,
1556
+ "delta_vs_best": delta_vs_best,
1557
+ "reason": reason,
1558
+ })
1559
+
1560
+ # Dedup: against prior iterations' history
1561
+ dup = _is_duplicate_of_history(exp, history)
1562
+ if dup is not None:
1563
+ print(f" SKIPPED — already tried as '{dup.get('name', '?')}' "
1564
+ f"(outcome '{dup.get('outcome', '?')}')")
1565
+ history.append({
1566
+ "name": exp.name, "outcome": "skipped",
1567
+ "delta_pct": None,
1568
+ "substitutions": exp.substitutions, "env_vars": exp.env_vars,
1569
+ })
1570
+ _cand_event("skipped", reason=f"duplicate of '{dup.get('name', '?')}'")
1571
+ continue
1572
+
1573
+ # Dedup: within the current batch (llm-explore can collide)
1574
+ fp = _experiment_fingerprint(exp)
1575
+ if fp in seen_this_iter:
1576
+ print(" SKIPPED — duplicate of an earlier candidate in this iteration")
1577
+ history.append({
1578
+ "name": exp.name, "outcome": "skipped",
1579
+ "delta_pct": None,
1580
+ "substitutions": exp.substitutions, "env_vars": exp.env_vars,
1581
+ })
1582
+ _cand_event("skipped", reason="duplicate of an earlier candidate this iteration")
1583
+ continue
1584
+ seen_this_iter.add(fp)
1585
+
1586
+ # Apply substitutions
1587
+ if exp.substitutions:
1588
+ try:
1589
+ candidate_source = apply_substitutions(best_source, exp.substitutions)
1590
+ except re.error as exc:
1591
+ print(f" SKIPPED — invalid regex from LLM: {exc}")
1592
+ rejected.append((exp.name, f"bad regex: {exc}"))
1593
+ history.append({
1594
+ "name": exp.name, "outcome": "rejected",
1595
+ "delta_pct": None,
1596
+ "substitutions": exp.substitutions, "env_vars": exp.env_vars,
1597
+ })
1598
+ _cand_event("rejected", reason=f"bad regex: {exc}")
1599
+ continue
1600
+ if candidate_source is None:
1601
+ print(" SKIPPED — substitution patterns didn't match")
1602
+ rejected.append((exp.name, "patterns didn't match"))
1603
+ history.append({
1604
+ "name": exp.name, "outcome": "skipped",
1605
+ "delta_pct": None,
1606
+ "substitutions": exp.substitutions, "env_vars": exp.env_vars,
1607
+ })
1608
+ _cand_event("skipped", reason="substitution patterns didn't match")
1609
+ continue
1610
+ else:
1611
+ candidate_source = best_source
1612
+
1613
+ file_counter += 1
1614
+ safe_name = re.sub(r"[^A-Za-z0-9_]+", "_", exp.name)[:40] or "exp"
1615
+ candidate_path = workspace / f"{file_counter:03d}_iter{i + 1:02d}_{safe_name}.py"
1616
+ candidate_path.write_text(candidate_source)
1617
+
1618
+ candidate_env = {**best_env, **exp.env_vars}
1619
+ if exp.env_vars:
1620
+ print(f" env vars: {exp.env_vars}")
1621
+
1622
+ m = benchmark(candidate_path, args.steps, candidate_env)
1623
+ if m is None:
1624
+ rejected.append((exp.name, "benchmark crashed"))
1625
+ history.append({
1626
+ "name": exp.name, "outcome": "crashed",
1627
+ "delta_pct": None,
1628
+ "substitutions": exp.substitutions, "env_vars": exp.env_vars,
1629
+ })
1630
+ total_crashes += 1
1631
+ crashed_this_iter = True
1632
+ print(
1633
+ f" CRASHED — counted toward max-crashes "
1634
+ f"({total_crashes}/{args.max_crashes})"
1635
+ )
1636
+ _cand_event("crashed", reason="benchmark subprocess failed")
1637
+ if total_crashes >= args.max_crashes:
1638
+ max_crashes_hit = True
1639
+ break
1640
+ continue
1641
+
1642
+ tps = m["tokens_per_sec"]
1643
+ delta_vs_best = _delta_pct(tps, best_tps)
1644
+ print(f" tokens/sec: {tps:.1f} (Δ {delta_vs_best:+.2f}% vs current best)")
1645
+ print(f" mfu_pct: {m.get('mfu_pct', 0.0):.2f}")
1646
+ print(f" hbm_peak_gb: {m['hbm_peak_gb']:.2f}")
1647
+ print(f" gpu_util_pct:{m['gpu_util_pct']:.1f}")
1648
+ _print_waste(m, prefix=" waste: ")
1649
+
1650
+ # Emit "evaluated" — outcome (accepted/rejected) is decided
1651
+ # later when the iteration's winner is picked across all
1652
+ # candidates. For UI display purposes the per-candidate metrics
1653
+ # are already useful.
1654
+ _cand_event("evaluated", metrics=m, delta_vs_best=delta_vs_best)
1655
+
1656
+ eval_results.append({
1657
+ "exp": exp,
1658
+ "candidate_source": candidate_source,
1659
+ "candidate_env": candidate_env,
1660
+ "metrics": m,
1661
+ "delta_vs_best": delta_vs_best,
1662
+ })
1663
+
1664
+ if max_crashes_hit:
1665
+ print(
1666
+ f"\nReached max-crashes ({args.max_crashes}) — stopping to "
1667
+ "avoid burning more GPU on structurally bad changes."
1668
+ )
1669
+ break
1670
+
1671
+ # ---- Pick the iteration's winner from eval_results ----
1672
+ if not eval_results:
1673
+ # Every candidate was skipped or crashed
1674
+ if crashed_this_iter:
1675
+ print("\n All candidates crashed or were skipped this iteration.")
1676
+ else:
1677
+ print("\n All candidates were skipped this iteration.")
1678
+ consecutive_no_improvement += 1
1679
+ else:
1680
+ winner = max(eval_results, key=lambda r: r["metrics"]["tokens_per_sec"])
1681
+ winner_delta = winner["delta_vs_best"]
1682
+
1683
+ # ---- Optional merge step (llm-explore only) ----
1684
+ # If 2+ candidates this iteration each beat the baseline, try
1685
+ # combining them into one experiment and benchmark the merge.
1686
+ # The merge replaces `winner` only if it strictly exceeds the
1687
+ # individual winner's tokens/sec.
1688
+ if args.mode == "llm-explore":
1689
+ positives = [r for r in eval_results if r["delta_vs_best"] > 0]
1690
+ if len(positives) >= 2:
1691
+ merged_exp, merge_reason = _build_merged_experiment(
1692
+ [r["exp"] for r in positives], best_source
1693
+ )
1694
+ if merged_exp is None:
1695
+ print(f"\n MERGE SKIPPED — {merge_reason}")
1696
+ _emit({
1697
+ "type": "merge_attempt",
1698
+ "iteration": i + 1,
1699
+ "outcome": "skipped",
1700
+ "reason": merge_reason,
1701
+ "candidate_names": [r["exp"].name for r in positives],
1702
+ })
1703
+ else:
1704
+ print(
1705
+ f"\n Merging {len(positives)} positive candidates: "
1706
+ f"{merged_exp.description}"
1707
+ )
1708
+ # Apply substitutions to get the merged source
1709
+ merged_source = best_source
1710
+ for pattern, replacement in merged_exp.substitutions:
1711
+ merged_source = re.sub(pattern, replacement, merged_source)
1712
+ merged_env = {**best_env, **merged_exp.env_vars}
1713
+
1714
+ file_counter += 1
1715
+ merged_path = workspace / f"{file_counter:03d}_iter{i + 1:02d}_merge.py"
1716
+ merged_path.write_text(merged_source)
1717
+ if merged_exp.env_vars:
1718
+ print(f" env vars: {merged_exp.env_vars}")
1719
+
1720
+ m = benchmark(merged_path, args.steps, merged_env)
1721
+ if m is None:
1722
+ total_crashes += 1
1723
+ crashed_this_iter = True
1724
+ print(
1725
+ f" MERGE CRASHED — counted toward max-crashes "
1726
+ f"({total_crashes}/{args.max_crashes})"
1727
+ )
1728
+ _emit({
1729
+ "type": "merge_attempt",
1730
+ "iteration": i + 1,
1731
+ "outcome": "crashed",
1732
+ "candidate_names": [r["exp"].name for r in positives],
1733
+ "merged_name": merged_exp.name,
1734
+ })
1735
+ if total_crashes >= args.max_crashes:
1736
+ max_crashes_hit = True
1737
+ else:
1738
+ tps = m["tokens_per_sec"]
1739
+ delta_vs_best = _delta_pct(tps, best_tps)
1740
+ print(
1741
+ f" Merged tokens/sec: {tps:.1f} "
1742
+ f"(Δ {delta_vs_best:+.2f}% vs baseline)"
1743
+ )
1744
+ print(f" mfu_pct: {m.get('mfu_pct', 0.0):.2f}")
1745
+ print(f" hbm_peak_gb: {m['hbm_peak_gb']:.2f}")
1746
+
1747
+ individual_best_tps = winner["metrics"]["tokens_per_sec"]
1748
+ if tps > individual_best_tps:
1749
+ print(
1750
+ f" MERGE WINS — exceeds individual winner "
1751
+ f"'{winner['exp'].name}' "
1752
+ f"({tps:.1f} > {individual_best_tps:.1f})"
1753
+ )
1754
+ _emit({
1755
+ "type": "merge_attempt",
1756
+ "iteration": i + 1,
1757
+ "outcome": "wins",
1758
+ "candidate_names": [r["exp"].name for r in positives],
1759
+ "merged_name": merged_exp.name,
1760
+ "metrics": m,
1761
+ "delta_vs_best": delta_vs_best,
1762
+ "individual_best_name": winner["exp"].name,
1763
+ "individual_best_tps": individual_best_tps,
1764
+ })
1765
+ # Promote merged to be the new winner
1766
+ winner = {
1767
+ "exp": merged_exp,
1768
+ "candidate_source": merged_source,
1769
+ "candidate_env": merged_env,
1770
+ "metrics": m,
1771
+ "delta_vs_best": delta_vs_best,
1772
+ }
1773
+ winner_delta = delta_vs_best
1774
+ else:
1775
+ print(
1776
+ f" Merge didn't beat individual winner; "
1777
+ f"keeping '{winner['exp'].name}'"
1778
+ )
1779
+ _emit({
1780
+ "type": "merge_attempt",
1781
+ "iteration": i + 1,
1782
+ "outcome": "lost",
1783
+ "candidate_names": [r["exp"].name for r in positives],
1784
+ "merged_name": merged_exp.name,
1785
+ "metrics": m,
1786
+ "delta_vs_best": delta_vs_best,
1787
+ "individual_best_name": winner["exp"].name,
1788
+ "individual_best_tps": individual_best_tps,
1789
+ })
1790
+
1791
+ if winner_delta >= args.improvement_threshold:
1792
+ print(
1793
+ f"\n ACCEPTED — '{winner['exp'].name}' wins "
1794
+ f"(Δ {winner_delta:+.2f}% vs current best)"
1795
+ )
1796
+ best_source = winner["candidate_source"]
1797
+ best_tps = winner["metrics"]["tokens_per_sec"]
1798
+ best_env = winner["candidate_env"]
1799
+ last_metrics = winner["metrics"]
1800
+ accepted.append((winner["exp"].name, best_tps, winner_delta))
1801
+ history.append({
1802
+ "name": winner["exp"].name, "outcome": "accepted",
1803
+ "delta_pct": winner_delta,
1804
+ "substitutions": winner["exp"].substitutions,
1805
+ "env_vars": winner["exp"].env_vars,
1806
+ })
1807
+ # Other candidates of this iteration get marked rejected
1808
+ for r in eval_results:
1809
+ if r is winner:
1810
+ continue
1811
+ rejected.append((r["exp"].name, f"{r['delta_vs_best']:+.2f}%"))
1812
+ history.append({
1813
+ "name": r["exp"].name, "outcome": "rejected",
1814
+ "delta_pct": r["delta_vs_best"],
1815
+ "substitutions": r["exp"].substitutions,
1816
+ "env_vars": r["exp"].env_vars,
1817
+ })
1818
+ consecutive_no_improvement = 0
1819
+ _emit({
1820
+ "type": "iter_done",
1821
+ "iteration": i + 1,
1822
+ "outcome": "accepted",
1823
+ "winner_name": winner["exp"].name,
1824
+ "winner_delta": winner_delta,
1825
+ "best_tps": best_tps,
1826
+ "best_metrics": winner["metrics"],
1827
+ "best_env_vars": best_env,
1828
+ })
1829
+ else:
1830
+ print(
1831
+ f"\n ALL REJECTED — best candidate '{winner['exp'].name}' "
1832
+ f"only Δ {winner_delta:+.2f}% (threshold {args.improvement_threshold:.1f}%)"
1833
+ )
1834
+ for r in eval_results:
1835
+ rejected.append((r["exp"].name, f"{r['delta_vs_best']:+.2f}%"))
1836
+ history.append({
1837
+ "name": r["exp"].name, "outcome": "rejected",
1838
+ "delta_pct": r["delta_vs_best"],
1839
+ "substitutions": r["exp"].substitutions,
1840
+ "env_vars": r["exp"].env_vars,
1841
+ })
1842
+ # Update last_metrics with the winner anyway so the LLM sees
1843
+ # the latest waste_budget on the next turn.
1844
+ if args.mode in ("llm", "llm-explore"):
1845
+ last_metrics = winner["metrics"]
1846
+ consecutive_no_improvement += 1
1847
+ _emit({
1848
+ "type": "iter_done",
1849
+ "iteration": i + 1,
1850
+ "outcome": "all_rejected",
1851
+ "winner_name": winner["exp"].name,
1852
+ "winner_delta": winner_delta,
1853
+ "best_tps": best_tps,
1854
+ })
1855
+
1856
+ if consecutive_no_improvement >= args.early_stop_after:
1857
+ print(
1858
+ f"\nNo improvement for {args.early_stop_after} consecutive iterations — early stopping."
1859
+ )
1860
+ break
1861
+
1862
+ # Save best
1863
+ best_path = workspace / "best.py"
1864
+ best_path.write_text(best_source)
1865
+
1866
+ # Summary
1867
+ print()
1868
+ print("=" * 60)
1869
+ print("AUTO-TUNE SUMMARY")
1870
+ print("=" * 60)
1871
+ print(f"Baseline tokens/sec: {baseline_tps:.1f}")
1872
+ print(
1873
+ f"Best tokens/sec: {best_tps:.1f} "
1874
+ f"({_delta_pct(best_tps, baseline_tps):+.2f}% vs baseline)"
1875
+ )
1876
+ print()
1877
+ print(f"Accepted ({len(accepted)}):")
1878
+ for name, tps, delta in accepted:
1879
+ print(f" + {name:25s} {tps:8.1f} tok/s (Δ {delta:+.2f}%)")
1880
+ print()
1881
+ print(f"Rejected ({len(rejected)}):")
1882
+ for name, reason in rejected:
1883
+ print(f" - {name:25s} {reason}")
1884
+ print()
1885
+
1886
+ if best_env:
1887
+ print("Required env vars for best config:")
1888
+ for k, v in best_env.items():
1889
+ print(f" export {k}={v}")
1890
+ print()
1891
+
1892
+ print(f"Best workload script: {best_path}")
1893
+ print(f"Diff vs baseline: diff {workload} {best_path}")
1894
+
1895
+ _emit({
1896
+ "type": "summary",
1897
+ "baseline_metrics": baseline,
1898
+ "best_metrics": last_metrics,
1899
+ "baseline_tps": baseline_tps,
1900
+ "best_tps": best_tps,
1901
+ "improvement_pct": _delta_pct(best_tps, baseline_tps),
1902
+ "accepted": [
1903
+ {"name": name, "tps": tps, "delta_pct": delta}
1904
+ for name, tps, delta in accepted
1905
+ ],
1906
+ "rejected": [
1907
+ {"name": name, "reason": reason}
1908
+ for name, reason in rejected
1909
+ ],
1910
+ "best_env_vars": best_env,
1911
+ "best_workload_path": str(best_path),
1912
+ "baseline_workload_path": str(workload),
1913
+ })
1914
+ return 0
1915
+
1916
+
1917
+ if __name__ == "__main__":
1918
+ raise SystemExit(main())
tests/fixtures/sample_train.json ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
3
+ "per_device_train_batch_size": 8,
4
+ "gradient_accumulation_steps": 4,
5
+ "max_seq_length": 4096,
6
+ "learning_rate": 0.0003,
7
+ "warmup_steps": 200,
8
+ "bf16": true,
9
+ "optim": "adamw_torch_fused",
10
+ "gradient_checkpointing": true,
11
+ "torch_compile": true,
12
+ "dataloader_num_workers": 4,
13
+ "dataloader_pin_memory": true,
14
+ "dataloader_prefetch_factor": 4,
15
+ "dataloader_persistent_workers": true,
16
+ "attn_implementation": "flash",
17
+ "num_train_epochs": 3,
18
+ "save_steps": 500,
19
+ "logging_steps": 25,
20
+ "output_dir": "./out",
21
+ "hub_token": "hf_jsonsamplehfabcdefghijklmnopqrs",
22
+ "checkpoint_uri": "s3://team-bucket/runs/qwen-lora-001/",
23
+ "env_vars": {
24
+ "HSA_FORCE_FINE_GRAIN_PCIE": "1",
25
+ "NCCL_MIN_NCHANNELS": "112"
26
+ }
27
+ }
tests/fixtures/sample_train.py ADDED
@@ -0,0 +1,101 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Realistic-looking HF Trainer fine-tuning script with secrets sprinkled in.
2
+
3
+ Used as a fixture for parse_config — exercises every code path we care about:
4
+ TrainingArguments kwargs, DataLoader kwargs, torch.compile, gradient
5
+ checkpointing, os.environ assignments, LoRA config, and from_pretrained.
6
+ """
7
+
8
+ import os
9
+
10
+ import torch
11
+ from torch.utils.data import DataLoader
12
+ from transformers import (
13
+ AutoModelForCausalLM,
14
+ AutoTokenizer,
15
+ Trainer,
16
+ TrainingArguments,
17
+ )
18
+ from peft import LoraConfig, get_peft_model
19
+ from datasets import load_dataset
20
+
21
+ # Secrets we expect parse_config to redact before storing raw_source.
22
+ HF_TOKEN = "hf_abcdefghijklmnopqrstuvwxyz123456"
23
+ OPENAI_KEY = "sk-abcdefghijklmnopqrstuvwxyz1234567890"
24
+ GH_TOKEN = "gho_abcdefghijklmnopqrstuvwxyz123456"
25
+ AUTH_HEADER = "Authorization: Bearer eyJhbGciOi.JIUzI1NiJ9.signature123"
26
+ DATA_ROOT = "/home/researcher/datasets/alpaca"
27
+ S3_BUCKET = "s3://my-team/checkpoints/qwen-lora/"
28
+ WS_LOG = "wss://logs.internal.example.com/stream"
29
+
30
+ # Environment variables the agent should capture into env_vars.
31
+ os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
32
+ os.environ["MIOPEN_FIND_MODE"] = "3"
33
+ os.environ["NCCL_MIN_NCHANNELS"] = "112"
34
+
35
+ MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
36
+
37
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
38
+ model = AutoModelForCausalLM.from_pretrained(
39
+ MODEL_ID,
40
+ torch_dtype=torch.bfloat16,
41
+ attn_implementation="eager",
42
+ token=HF_TOKEN,
43
+ )
44
+
45
+ # LoRA — rank should land in WorkloadConfig.lora_rank.
46
+ lora_config = LoraConfig(
47
+ r=16,
48
+ lora_alpha=32,
49
+ target_modules=["q_proj", "v_proj"],
50
+ lora_dropout=0.05,
51
+ bias="none",
52
+ task_type="CAUSAL_LM",
53
+ )
54
+ model = get_peft_model(model, lora_config)
55
+
56
+ # Should set gradient_checkpointing=True via the explicit enable() call.
57
+ model.gradient_checkpointing_enable()
58
+
59
+ # Should flip torch_compile=True.
60
+ model = torch.compile(model, mode="reduce-overhead")
61
+
62
+ dataset = load_dataset("yahma/alpaca-cleaned", split="train")
63
+
64
+ train_loader = DataLoader(
65
+ dataset,
66
+ batch_size=4,
67
+ num_workers=0,
68
+ pin_memory=False,
69
+ prefetch_factor=2,
70
+ persistent_workers=False,
71
+ )
72
+
73
+ training_args = TrainingArguments(
74
+ output_dir="./out",
75
+ per_device_train_batch_size=4,
76
+ gradient_accumulation_steps=8,
77
+ num_train_epochs=3,
78
+ learning_rate=2e-4,
79
+ warmup_steps=100,
80
+ fp16=True,
81
+ optim="adamw_torch",
82
+ logging_steps=10,
83
+ save_steps=500,
84
+ dataloader_num_workers=0,
85
+ dataloader_pin_memory=False,
86
+ gradient_checkpointing=True,
87
+ torch_compile=False,
88
+ report_to="none",
89
+ push_to_hub=False,
90
+ hub_token=HF_TOKEN,
91
+ )
92
+
93
+ trainer = Trainer(
94
+ model=model,
95
+ args=training_args,
96
+ train_dataset=dataset,
97
+ tokenizer=tokenizer,
98
+ )
99
+
100
+ if __name__ == "__main__":
101
+ trainer.train()
tests/fixtures/sample_train.yaml ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # YAML training config — same schema as JSON, exercised separately so we know
2
+ # yaml.safe_load + dict extraction agree with the JSON path.
3
+ model_name_or_path: Qwen/Qwen2.5-7B-Instruct
4
+ per_device_train_batch_size: 2
5
+ gradient_accumulation_steps: 16
6
+ max_seq_length: 8192
7
+ learning_rate: 1.0e-4
8
+ warmup_steps: 50
9
+ bf16: true
10
+ optim: adamw_torch
11
+ gradient_checkpointing: true
12
+ torch_compile: false
13
+ dataloader_num_workers: 8
14
+ dataloader_pin_memory: true
15
+ dataloader_prefetch_factor: 2
16
+ dataloader_persistent_workers: true
17
+ attn_implementation: sdpa
18
+ num_train_epochs: 1
19
+ output_dir: ./out
20
+ # Secrets that should be scrubbed:
21
+ hub_token: "hf_yamlsamplehfabcdefghijklmnopqrs1"
22
+ auth_header: "Bearer eyJ.payload.signaturetoken"
23
+ data_path: "/home/teamuser/datasets/alpaca-cleaned"
24
+ env_vars:
25
+ HSA_FORCE_FINE_GRAIN_PCIE: "1"
26
+ MIOPEN_FIND_MODE: "3"
tests/test_compare_runs_normalize.py ADDED
@@ -0,0 +1,155 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for compare_runs's flat-Patch recovery path.
2
+
3
+ Live AMD-GPU lesson: Qwen models routinely forward a flat WorkloadConfig (or
4
+ even just the changed-fields subset) as the ``patch=`` argument to
5
+ compare_runs, instead of the full Patch envelope. ``_normalize_patch`` is
6
+ the safety net — it must:
7
+
8
+ 1. Pass real Patch dicts through unchanged.
9
+ 2. Detect any flat-config shape (full WorkloadConfig, just dataloader
10
+ fields, just env_vars, etc.) — NOT just dicts with model_name.
11
+ 3. Recover by substituting the cached propose_patch result when one
12
+ exists, or wrapping the flat config minimally as a last resort.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import pytest
18
+
19
+ from agent.tools import compare_runs as cr_mod
20
+ from agent.tools import propose_patch as pp_mod
21
+
22
+
23
+ # ---------------------------------------------------------------------------
24
+ # _looks_like_flat_config detection
25
+ # ---------------------------------------------------------------------------
26
+
27
+
28
+ class TestFlatConfigDetection:
29
+ def test_real_patch_is_not_flat(self) -> None:
30
+ real = {
31
+ "new_config": {"model_name": "x"},
32
+ "diff": "(no changes)",
33
+ "rationale": [],
34
+ "expected_speedup_low": 1.0,
35
+ "expected_speedup_high": 1.0,
36
+ "confidence": 0.0,
37
+ }
38
+ assert cr_mod._looks_like_flat_config(real) is False
39
+
40
+ def test_full_workload_config_is_flat(self) -> None:
41
+ flat = {
42
+ "model_name": "Qwen/Qwen2.5-7B-Instruct",
43
+ "precision": "bf16",
44
+ "attention_impl": "flash_rocm",
45
+ "batch_size": 12,
46
+ }
47
+ assert cr_mod._looks_like_flat_config(flat) is True
48
+
49
+ def test_dataloader_only_diff_is_flat(self) -> None:
50
+ # The exact failure mode from the live MI300X audit: model only
51
+ # passed the *changed* dataloader fields, no model_name in sight.
52
+ flat = {
53
+ "dataloader_persistent_workers": True,
54
+ "dataloader_pin_memory": True,
55
+ "dataloader_workers": 8,
56
+ }
57
+ assert cr_mod._looks_like_flat_config(flat) is True
58
+
59
+ def test_env_vars_only_diff_is_flat(self) -> None:
60
+ flat = {"env_vars": {"NCCL_MIN_NCHANNELS": "112"}}
61
+ assert cr_mod._looks_like_flat_config(flat) is True
62
+
63
+ def test_precision_only_diff_is_flat(self) -> None:
64
+ flat = {"precision": "bf16"}
65
+ assert cr_mod._looks_like_flat_config(flat) is True
66
+
67
+ def test_unrelated_dict_not_flat(self) -> None:
68
+ # Garbage dict with no WorkloadConfig fields → don't claim it's flat.
69
+ assert cr_mod._looks_like_flat_config({"foo": 1, "bar": 2}) is False
70
+
71
+ def test_non_dict_not_flat(self) -> None:
72
+ assert cr_mod._looks_like_flat_config(None) is False # type: ignore[arg-type]
73
+ assert cr_mod._looks_like_flat_config("a string") is False # type: ignore[arg-type]
74
+ assert cr_mod._looks_like_flat_config([1, 2, 3]) is False # type: ignore[arg-type]
75
+
76
+
77
+ # ---------------------------------------------------------------------------
78
+ # _normalize_patch + cached-patch recovery
79
+ # ---------------------------------------------------------------------------
80
+
81
+
82
+ @pytest.fixture
83
+ def cached_patch(monkeypatch):
84
+ """Plant a fake `latest_patch()` so the recovery path picks it up."""
85
+ fake = {
86
+ "new_config": {"model_name": "Qwen/Qwen2.5-7B-Instruct", "precision": "bf16"},
87
+ "diff": "- precision: fp16\n+ precision: bf16",
88
+ "rationale": [
89
+ {
90
+ "rule_id": "precision.bf16_over_fp16_on_mi300x",
91
+ "rationale": "r",
92
+ "citation": "c",
93
+ "targets_bucket": "precision_path",
94
+ "estimated_recovery_seconds": 0.09,
95
+ }
96
+ ],
97
+ "expected_speedup_low": 1.05,
98
+ "expected_speedup_high": 1.30,
99
+ "confidence": 0.85,
100
+ }
101
+ monkeypatch.setattr(pp_mod, "_LAST_PATCH", fake)
102
+ yield fake
103
+
104
+
105
+ class TestNormalizePatch:
106
+ def test_real_patch_passes_through(self) -> None:
107
+ real = {
108
+ "new_config": {"model_name": "x"},
109
+ "diff": "...",
110
+ "rationale": [],
111
+ "expected_speedup_low": 1.0,
112
+ "expected_speedup_high": 1.0,
113
+ "confidence": 0.0,
114
+ }
115
+ out, notes = cr_mod._normalize_patch(real)
116
+ assert out is real
117
+ assert notes == []
118
+
119
+ def test_dataloader_only_diff_recovers_via_cached(self, cached_patch) -> None:
120
+ # The exact live-AMD-GPU failure: model forwarded only the changed
121
+ # dataloader fields. Old code's narrow sentinel set (model_name etc.)
122
+ # would miss this. New behavior: detected, cached patch substituted.
123
+ flat = {
124
+ "dataloader_persistent_workers": True,
125
+ "dataloader_pin_memory": True,
126
+ "dataloader_workers": 8,
127
+ }
128
+ out, notes = cr_mod._normalize_patch(flat)
129
+ assert out is cached_patch # full fidelity restored
130
+ assert any("substituted the cached" in n for n in notes)
131
+
132
+ def test_flat_config_falls_back_to_minimal_wrap_when_no_cache(
133
+ self, monkeypatch
134
+ ) -> None:
135
+ # No cached patch — must still produce a Patch-shape dict so
136
+ # compare_runs doesn't crash on Pydantic validation.
137
+ monkeypatch.setattr(pp_mod, "_LAST_PATCH", None)
138
+ flat = {"precision": "bf16"}
139
+ out, notes = cr_mod._normalize_patch(flat)
140
+ assert "new_config" in out
141
+ assert "diff" in out
142
+ assert out["expected_speedup_low"] == 1.0
143
+ assert out["confidence"] == 0.0
144
+ assert any("synthesized a minimal Patch" in n for n in notes)
145
+
146
+ def test_non_flat_garbage_passes_through_for_pydantic_to_reject(
147
+ self, monkeypatch
148
+ ) -> None:
149
+ # If it's neither a real Patch nor a recognizable flat config, let
150
+ # pydantic produce the clear ValidationError — don't silently mangle.
151
+ monkeypatch.setattr(pp_mod, "_LAST_PATCH", None)
152
+ garbage = {"foo": 1, "bar": [2]}
153
+ out, notes = cr_mod._normalize_patch(garbage)
154
+ assert out is garbage
155
+ assert notes == []
tests/test_loop.py ADDED
@@ -0,0 +1,322 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for the agent loop driver.
2
+
3
+ We never hit a real LLM. The loop talks to a `Backend` (see
4
+ `agent/backends/`); each test injects a `FakeBackend` whose `next_turn`
5
+ returns a queued sequence of scripted `AgentTurn` objects. Tools are
6
+ stubbed so we can drive specific control-flow paths.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from typing import Any
12
+
13
+ import pytest
14
+
15
+ from agent import loop as loop_module
16
+ from agent.backends.base import AgentTurn, Backend, ToolCall
17
+ from agent.schemas import SSEEvent, ToolResult
18
+
19
+
20
+ # ---------------------------------------------------------------------------
21
+ # Fake backend
22
+ # ---------------------------------------------------------------------------
23
+
24
+
25
+ class FakeBackend(Backend):
26
+ """A scripted Backend for testing the loop in isolation.
27
+
28
+ Each test queues a list of `AgentTurn`s; calling `next_turn` pops the
29
+ next one. We also record every tool result the loop hands back so tests
30
+ can assert that error / id / content were threaded through correctly.
31
+ """
32
+
33
+ name = "fake"
34
+
35
+ def __init__(
36
+ self,
37
+ scripted_turns: list[AgentTurn] | None = None,
38
+ next_turn_raises: BaseException | None = None,
39
+ ) -> None:
40
+ self._scripted = list(scripted_turns or [])
41
+ self._raise_on_next = next_turn_raises
42
+ self.user_messages: list[str] = []
43
+ self.tool_results: list[dict[str, Any]] = []
44
+ self.turn_count = 0
45
+
46
+ def add_user_message(self, content: str) -> None:
47
+ self.user_messages.append(content)
48
+
49
+ def add_tool_result(
50
+ self,
51
+ tool_call_id: str,
52
+ name: str,
53
+ content: str,
54
+ is_error: bool,
55
+ ) -> None:
56
+ self.tool_results.append(
57
+ {
58
+ "id": tool_call_id,
59
+ "name": name,
60
+ "content": content,
61
+ "is_error": is_error,
62
+ }
63
+ )
64
+
65
+ async def next_turn(self, tool_schemas: list[dict[str, Any]]) -> AgentTurn:
66
+ self.turn_count += 1
67
+ if self._raise_on_next is not None:
68
+ exc = self._raise_on_next
69
+ self._raise_on_next = None
70
+ raise exc
71
+ if not self._scripted:
72
+ raise AssertionError(
73
+ "FakeBackend exhausted — loop made more turns than expected"
74
+ )
75
+ return self._scripted.pop(0)
76
+
77
+
78
+ # ---------------------------------------------------------------------------
79
+ # Helpers
80
+ # ---------------------------------------------------------------------------
81
+
82
+
83
+ def _install_backend(monkeypatch, backend: Backend) -> Backend:
84
+ """Replace `make_backend` so the loop sees our fake."""
85
+ monkeypatch.setattr(loop_module, "make_backend", lambda **_kwargs: backend)
86
+ return backend
87
+
88
+
89
+ def _install_make_backend_raises(monkeypatch, exc: BaseException) -> None:
90
+ def boom(**_kwargs: Any) -> Backend:
91
+ raise exc
92
+
93
+ monkeypatch.setattr(loop_module, "make_backend", boom)
94
+
95
+
96
+ def _install_fake_tools(
97
+ monkeypatch, tool_responses: dict[str, ToolResult]
98
+ ) -> list[str]:
99
+ """Replace `tools_module.call` and `tool_schemas`. Returns a list
100
+ that records the order tools were invoked.
101
+ """
102
+ invoked: list[str] = []
103
+
104
+ def fake_call(name: str, **_kwargs: Any) -> ToolResult:
105
+ invoked.append(name)
106
+ return tool_responses.get(
107
+ name, ToolResult(ok=False, error=f"no fake registered for {name}")
108
+ )
109
+
110
+ monkeypatch.setattr(loop_module.tools_module, "call", fake_call)
111
+ monkeypatch.setattr(loop_module.tools_module, "tool_schemas", lambda: [])
112
+ return invoked
113
+
114
+
115
+ async def _collect(stream) -> list[SSEEvent]:
116
+ out: list[SSEEvent] = []
117
+ async for event in stream:
118
+ out.append(event)
119
+ return out
120
+
121
+
122
+ # ---------------------------------------------------------------------------
123
+ # Tests
124
+ # ---------------------------------------------------------------------------
125
+
126
+
127
+ @pytest.mark.asyncio
128
+ async def test_emits_thought_then_tool_call_then_tool_result(monkeypatch) -> None:
129
+ backend = FakeBackend(
130
+ scripted_turns=[
131
+ AgentTurn(
132
+ text_blocks=["I'll start by parsing the config."],
133
+ tool_calls=[
134
+ ToolCall(id="tu_1", name="parse_config", input={"file_path": "/x.py"})
135
+ ],
136
+ stop_reason="tool_use",
137
+ ),
138
+ AgentTurn(text_blocks=["Done."], tool_calls=[], stop_reason="end_turn"),
139
+ ]
140
+ )
141
+ _install_backend(monkeypatch, backend)
142
+ invoked = _install_fake_tools(
143
+ monkeypatch,
144
+ {"parse_config": ToolResult(ok=True, result={"model_name": "x"})},
145
+ )
146
+
147
+ events = await _collect(loop_module.run_audit("/x.py"))
148
+ types = [e.type for e in events]
149
+
150
+ assert types[0] == "thought"
151
+ assert types[1] == "tool_call"
152
+ assert types[2] == "tool_result"
153
+ # No compare_runs ⇒ final event is the "no final report" error.
154
+ assert types[-1] == "error"
155
+ assert "without producing a final report" in events[-1].data["message"]
156
+ assert invoked == ["parse_config"]
157
+
158
+ # tool_call carries id/name/input; tool_result mirrors that plus ok/result/error.
159
+ assert events[1].data == {
160
+ "id": "tu_1",
161
+ "name": "parse_config",
162
+ "input": {"file_path": "/x.py"},
163
+ }
164
+ assert events[2].data["ok"] is True
165
+ assert events[2].data["result"] == {"model_name": "x"}
166
+ assert events[2].data["error"] is None
167
+
168
+ # The user message and tool result were threaded into the backend.
169
+ assert backend.user_messages == ["Audit this fine-tuning workload: /x.py"]
170
+ assert backend.tool_results == [
171
+ {
172
+ "id": "tu_1",
173
+ "name": "parse_config",
174
+ "content": '{"model_name": "x"}',
175
+ "is_error": False,
176
+ }
177
+ ]
178
+
179
+
180
+ @pytest.mark.asyncio
181
+ async def test_final_report_extracted_from_compare_runs(monkeypatch) -> None:
182
+ fake_report = {"workload_name": "test", "speedup_actual": 2.0}
183
+ backend = FakeBackend(
184
+ scripted_turns=[
185
+ AgentTurn(
186
+ text_blocks=["Wrapping up."],
187
+ tool_calls=[
188
+ ToolCall(
189
+ id="tu_compare",
190
+ name="compare_runs",
191
+ input={
192
+ "workload_name": "t",
193
+ "before": {},
194
+ "after": {},
195
+ "patch": {},
196
+ },
197
+ )
198
+ ],
199
+ stop_reason="end_turn",
200
+ ),
201
+ ]
202
+ )
203
+ _install_backend(monkeypatch, backend)
204
+ _install_fake_tools(
205
+ monkeypatch, {"compare_runs": ToolResult(ok=True, result=fake_report)}
206
+ )
207
+
208
+ events = await _collect(loop_module.run_audit("/x.py"))
209
+
210
+ assert events[-1].type == "final_report"
211
+ assert events[-1].data["report"] == fake_report
212
+
213
+
214
+ @pytest.mark.asyncio
215
+ async def test_tool_error_passes_through_does_not_crash(monkeypatch) -> None:
216
+ backend = FakeBackend(
217
+ scripted_turns=[
218
+ AgentTurn(
219
+ text_blocks=["Trying parse."],
220
+ tool_calls=[
221
+ ToolCall(id="tu_1", name="parse_config", input={"file_path": "/bogus"})
222
+ ],
223
+ stop_reason="tool_use",
224
+ ),
225
+ AgentTurn(text_blocks=["Giving up."], tool_calls=[], stop_reason="end_turn"),
226
+ ]
227
+ )
228
+ _install_backend(monkeypatch, backend)
229
+ _install_fake_tools(
230
+ monkeypatch,
231
+ {"parse_config": ToolResult(ok=False, error="file not found")},
232
+ )
233
+
234
+ events = await _collect(loop_module.run_audit("/bogus"))
235
+ tool_result_events = [e for e in events if e.type == "tool_result"]
236
+ assert len(tool_result_events) == 1
237
+ assert tool_result_events[0].data["ok"] is False
238
+ assert tool_result_events[0].data["error"] == "file not found"
239
+ # The loop kept iterating rather than bailing.
240
+ assert events[-1].type == "error" # no compare_runs ⇒ "no final report"
241
+
242
+ # Backend received an is_error=True tool result with the error message.
243
+ assert backend.tool_results[-1]["is_error"] is True
244
+ assert backend.tool_results[-1]["content"] == "file not found"
245
+
246
+
247
+ @pytest.mark.asyncio
248
+ async def test_backend_construction_failure_yields_error_event(monkeypatch) -> None:
249
+ _install_make_backend_raises(
250
+ monkeypatch, RuntimeError("HF_TOKEN is not set; Qwen backend cannot run.")
251
+ )
252
+ events = await _collect(loop_module.run_audit("/x.py"))
253
+ assert len(events) == 1
254
+ assert events[0].type == "error"
255
+ assert "HF_TOKEN" in events[0].data["message"]
256
+
257
+
258
+ @pytest.mark.asyncio
259
+ async def test_mid_loop_exception_is_caught(monkeypatch) -> None:
260
+ backend = FakeBackend(next_turn_raises=RuntimeError("boom"))
261
+ _install_backend(monkeypatch, backend)
262
+ monkeypatch.setattr(loop_module.tools_module, "tool_schemas", lambda: [])
263
+
264
+ events = await _collect(loop_module.run_audit("/x.py"))
265
+ assert events[-1].type == "error"
266
+ assert "boom" in events[-1].data["message"]
267
+
268
+
269
+ @pytest.mark.asyncio
270
+ async def test_loop_caps_at_max_steps(monkeypatch) -> None:
271
+ """Even if the model never says end_turn, we bail after MAX_STEPS."""
272
+ backend = FakeBackend(
273
+ scripted_turns=[
274
+ AgentTurn(
275
+ text_blocks=[f"step {i}"],
276
+ tool_calls=[
277
+ ToolCall(id=f"tu_{i}", name="parse_config", input={"file_path": "/x.py"})
278
+ ],
279
+ stop_reason="tool_use",
280
+ )
281
+ for i in range(loop_module.MAX_STEPS + 2) # extra so we'd overrun
282
+ ]
283
+ )
284
+ _install_backend(monkeypatch, backend)
285
+ _install_fake_tools(monkeypatch, {"parse_config": ToolResult(ok=True, result={})})
286
+
287
+ events = await _collect(loop_module.run_audit("/x.py"))
288
+ # Backend's next_turn was called exactly MAX_STEPS times.
289
+ assert backend.turn_count == loop_module.MAX_STEPS
290
+ # Last event is the "no final report" error (not a crash).
291
+ assert events[-1].type == "error"
292
+
293
+
294
+ @pytest.mark.asyncio
295
+ async def test_tool_call_id_is_threaded_to_backend(monkeypatch) -> None:
296
+ """The loop must hand the tool_call id back to the backend so the next
297
+ turn's request can correlate the tool_result with the originating call.
298
+ """
299
+ backend = FakeBackend(
300
+ scripted_turns=[
301
+ AgentTurn(
302
+ text_blocks=["parse"],
303
+ tool_calls=[
304
+ ToolCall(id="tu_abc", name="parse_config", input={"file_path": "/x"})
305
+ ],
306
+ stop_reason="tool_use",
307
+ ),
308
+ AgentTurn(text_blocks=["done"], tool_calls=[], stop_reason="end_turn"),
309
+ ]
310
+ )
311
+ _install_backend(monkeypatch, backend)
312
+ _install_fake_tools(
313
+ monkeypatch, {"parse_config": ToolResult(ok=True, result={"a": 1})}
314
+ )
315
+
316
+ await _collect(loop_module.run_audit("/x"))
317
+
318
+ # Backend got exactly one tool_result with id=tu_abc.
319
+ assert len(backend.tool_results) == 1
320
+ assert backend.tool_results[0]["id"] == "tu_abc"
321
+ assert backend.tool_results[0]["name"] == "parse_config"
322
+ assert backend.tool_results[0]["is_error"] is False
tests/test_misnested_args.py ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for the defensive misnested-arg extraction in benchmark + profile_run.
2
+
3
+ Live AMD-GPU lesson: Qwen2.5-7B (and probably others) occasionally JSON-nests
4
+ ``steps`` / ``cache`` *inside* the ``config`` dict instead of at the top level
5
+ alongside it. WorkloadConfig strict-validates extras, so without this defense
6
+ the call errors out and a tool slot is wasted. The well-tuned scenario run
7
+ on 2026-05-07 burned two of the eight available slots on this exact mistake;
8
+ fixing it costs nothing and saves the audit.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import shutil
14
+
15
+ from agent.tools import call
16
+
17
+
18
+ def _baseline_config() -> dict:
19
+ return {
20
+ "model_name": "Qwen/Qwen2.5-7B-Instruct",
21
+ "batch_size": 4,
22
+ "precision": "fp16",
23
+ "attention_impl": "eager",
24
+ "dataloader_workers": 0,
25
+ }
26
+
27
+
28
+ class TestBenchmarkMisnestedArgs:
29
+ def setup_method(self) -> None:
30
+ # Each test starts with an empty cache so cache-hit doesn't mask the
31
+ # behavior under test.
32
+ shutil.rmtree("bench_cache", ignore_errors=True)
33
+
34
+ def test_steps_nested_in_config_is_extracted(self) -> None:
35
+ """Old behavior: ``WorkloadConfig`` validation explodes with
36
+ 'Extra inputs are not permitted [steps]'. New behavior: defensive
37
+ extraction pulls ``steps`` back to the top-level arg, call succeeds.
38
+ """
39
+ cfg = {**_baseline_config(), "steps": 25}
40
+ result = call("benchmark", config=cfg)
41
+ assert result.ok, result.error
42
+ assert result.result["steps"] == 25
43
+
44
+ def test_cache_nested_in_config_is_extracted(self) -> None:
45
+ cfg = {**_baseline_config(), "cache": False}
46
+ result = call("benchmark", config=cfg)
47
+ assert result.ok, result.error
48
+
49
+ def test_force_rerun_nested_in_config_is_extracted(self) -> None:
50
+ cfg = {**_baseline_config(), "force_rerun": True}
51
+ result = call("benchmark", config=cfg)
52
+ assert result.ok, result.error
53
+
54
+ def test_explicit_top_level_wins_over_nested(self) -> None:
55
+ """If caller passes BOTH (config has steps + top-level steps), the
56
+ explicit non-default top-level wins. Defensive code is for the
57
+ accident case, not for letting nesting silently override."""
58
+ cfg = {**_baseline_config(), "steps": 25}
59
+ result = call("benchmark", config=cfg, steps=37)
60
+ assert result.ok, result.error
61
+ assert result.result["steps"] == 37
62
+
63
+ def test_all_three_nested_at_once(self) -> None:
64
+ """The exact failure mode from the live run: model nested three
65
+ runtime args inside config. All three should get pulled out.
66
+ """
67
+ cfg = {
68
+ **_baseline_config(),
69
+ "steps": 30,
70
+ "cache": False,
71
+ "force_rerun": True,
72
+ }
73
+ result = call("benchmark", config=cfg)
74
+ assert result.ok, result.error
75
+ assert result.result["steps"] == 30
76
+
77
+
78
+ class TestProfileRunMisnestedArgs:
79
+ def test_steps_nested_in_config_is_extracted(self) -> None:
80
+ cfg = {**_baseline_config(), "steps": 7}
81
+ result = call("profile_run", config=cfg)
82
+ assert result.ok, result.error
83
+ assert result.result["steps"] == 7
84
+
85
+ def test_explicit_top_level_wins(self) -> None:
86
+ cfg = {**_baseline_config(), "steps": 7}
87
+ result = call("profile_run", config=cfg, steps=15)
88
+ assert result.ok, result.error
89
+ assert result.result["steps"] == 15
90
+
91
+ def test_clean_config_unaffected(self) -> None:
92
+ """Sanity: when nothing is misnested, behavior is unchanged."""
93
+ result = call("profile_run", config=_baseline_config())
94
+ assert result.ok, result.error
95
+ assert result.result["steps"] == 10 # default
tests/test_parse_config.py ADDED
@@ -0,0 +1,377 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for ``agent.tools.parse_config._parse_config``.
2
+
3
+ Covers all three input shapes (Python AST, JSON, YAML), redaction of every
4
+ secret pattern we ship, error paths for missing/malformed inputs, and the
5
+ WorkloadConfig field mapping for HF TrainingArguments + DataLoader kwargs.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import json
11
+ from pathlib import Path
12
+
13
+ import pytest
14
+
15
+ from agent.schemas import WorkloadConfig
16
+ from agent.tools.parse_config import PARSE_CONFIG, _parse_config, _parse_config_full
17
+
18
+ FIXTURES = Path(__file__).parent / "fixtures"
19
+
20
+
21
+ # ---------------------------------------------------------------------------
22
+ # Python script path
23
+ # ---------------------------------------------------------------------------
24
+
25
+
26
+ class TestPythonScript:
27
+ def test_returns_ok(self) -> None:
28
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
29
+ assert result.ok, result.error
30
+ assert result.result is not None
31
+
32
+ def test_extracts_training_arguments_kwargs(self) -> None:
33
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
34
+ cfg = result.result
35
+ assert cfg["batch_size"] == 4
36
+ assert cfg["grad_accum_steps"] == 8
37
+ assert cfg["lr"] == pytest.approx(2e-4)
38
+ assert cfg["warmup_steps"] == 100
39
+ assert cfg["optimizer"] == "adamw_torch"
40
+ # fp16=True in TrainingArguments → precision should resolve to fp16.
41
+ assert cfg["precision"] == "fp16"
42
+
43
+ def test_dataloader_kwargs_captured(self) -> None:
44
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
45
+ cfg = result.result
46
+ assert cfg["dataloader_workers"] == 0
47
+ assert cfg["dataloader_pin_memory"] is False
48
+ assert cfg["dataloader_prefetch_factor"] == 2
49
+ assert cfg["dataloader_persistent_workers"] is False
50
+
51
+ def test_torch_compile_call_flips_flag(self) -> None:
52
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
53
+ # The script calls torch.compile(model, ...), even though the
54
+ # TrainingArguments has torch_compile=False — the explicit call wins.
55
+ assert result.result["torch_compile"] is True
56
+
57
+ def test_gradient_checkpointing_enable_call(self) -> None:
58
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
59
+ assert result.result["gradient_checkpointing"] is True
60
+
61
+ def test_env_vars_captured(self) -> None:
62
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
63
+ env = result.result["env_vars"]
64
+ assert env["HSA_FORCE_FINE_GRAIN_PCIE"] == "1"
65
+ assert env["MIOPEN_FIND_MODE"] == "3"
66
+ assert env["NCCL_MIN_NCHANNELS"] == "112"
67
+
68
+ def test_lora_rank_extracted(self) -> None:
69
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
70
+ assert result.result["lora_rank"] == 16
71
+
72
+ def test_attention_impl_from_from_pretrained(self) -> None:
73
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
74
+ assert result.result["attention_impl"] == "eager"
75
+
76
+ def test_model_name_resolved(self) -> None:
77
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
78
+ assert result.result["model_name"] == "Qwen/Qwen2.5-7B-Instruct"
79
+
80
+
81
+ # ---------------------------------------------------------------------------
82
+ # JSON path
83
+ # ---------------------------------------------------------------------------
84
+
85
+
86
+ class TestJsonConfig:
87
+ def test_returns_ok(self) -> None:
88
+ result = _parse_config(str(FIXTURES / "sample_train.json"))
89
+ assert result.ok, result.error
90
+
91
+ def test_field_mapping(self) -> None:
92
+ cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
93
+ assert cfg["model_name"] == "Qwen/Qwen2.5-7B-Instruct"
94
+ assert cfg["batch_size"] == 8
95
+ assert cfg["grad_accum_steps"] == 4
96
+ assert cfg["seq_len"] == 4096
97
+ assert cfg["precision"] == "bf16"
98
+ assert cfg["optimizer"] == "adamw_torch_fused"
99
+ assert cfg["torch_compile"] is True
100
+ assert cfg["gradient_checkpointing"] is True
101
+ assert cfg["dataloader_workers"] == 4
102
+ assert cfg["dataloader_pin_memory"] is True
103
+ assert cfg["dataloader_persistent_workers"] is True
104
+ assert cfg["attention_impl"] == "flash"
105
+
106
+ def test_env_vars_dict(self) -> None:
107
+ cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
108
+ assert cfg["env_vars"]["HSA_FORCE_FINE_GRAIN_PCIE"] == "1"
109
+
110
+ def test_extras_collects_unmapped_fields(self) -> None:
111
+ cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
112
+ # num_train_epochs has no slot in WorkloadConfig — must land in extras.
113
+ assert cfg["extras"]["num_train_epochs"] == 3
114
+ assert cfg["extras"]["save_steps"] == 500
115
+
116
+
117
+ # ---------------------------------------------------------------------------
118
+ # YAML path
119
+ # ---------------------------------------------------------------------------
120
+
121
+
122
+ class TestYamlConfig:
123
+ def test_returns_ok(self) -> None:
124
+ result = _parse_config(str(FIXTURES / "sample_train.yaml"))
125
+ assert result.ok, result.error
126
+
127
+ def test_field_mapping(self) -> None:
128
+ cfg = _parse_config(str(FIXTURES / "sample_train.yaml")).result
129
+ assert cfg["batch_size"] == 2
130
+ assert cfg["grad_accum_steps"] == 16
131
+ assert cfg["seq_len"] == 8192
132
+ assert cfg["precision"] == "bf16"
133
+ assert cfg["torch_compile"] is False
134
+ assert cfg["gradient_checkpointing"] is True
135
+ assert cfg["attention_impl"] == "sdpa"
136
+ assert cfg["dataloader_workers"] == 8
137
+ assert cfg["dataloader_persistent_workers"] is True
138
+
139
+
140
+ # ---------------------------------------------------------------------------
141
+ # Redaction
142
+ # ---------------------------------------------------------------------------
143
+
144
+
145
+ class TestRedaction:
146
+ def test_python_script_redactions(self) -> None:
147
+ cfg = _parse_config(str(FIXTURES / "sample_train.py")).result
148
+ labels = set(cfg["redactions"])
149
+ # Every secret pattern in sample_train.py should fire.
150
+ assert "hf_token" in labels
151
+ assert "openai_key" in labels
152
+ assert "github_token" in labels
153
+ assert "bearer_token" in labels
154
+ assert "home_path" in labels
155
+ assert "s3_uri" in labels
156
+ assert "ws_uri" in labels
157
+
158
+ def test_raw_source_is_scrubbed(self) -> None:
159
+ # `raw_source` is intentionally stripped from the tool result envelope
160
+ # (keeps the LLM conversation small) — use the `_full` helper to read
161
+ # it. The redaction labels list still proves which patterns fired.
162
+ cfg = _parse_config_full(str(FIXTURES / "sample_train.py"))
163
+ assert isinstance(cfg, WorkloadConfig)
164
+ raw = cfg.raw_source
165
+ assert "hf_abcdefghijklmnopqrstuvwxyz123456" not in raw
166
+ assert "sk-abcdefghijklmnopqrstuvwxyz1234567890" not in raw
167
+ assert "gho_abcdefghijklmnopqrstuvwxyz123456" not in raw
168
+ assert "/home/researcher/datasets/alpaca" not in raw
169
+ assert "s3://my-team/checkpoints/qwen-lora/" not in raw
170
+ assert "wss://logs.internal.example.com/stream" not in raw
171
+ assert "<REDACTED:hf_token>" in raw
172
+ assert "<REDACTED:openai_key>" in raw
173
+
174
+ def test_raw_source_excluded_from_tool_result(self) -> None:
175
+ # The tool result MUST NOT carry raw_source — it bloated the audit
176
+ # conversation past 8K on Qwen2.5-7B during the live AMD GPU run.
177
+ cfg = _parse_config(str(FIXTURES / "sample_train.py")).result
178
+ assert "raw_source" not in cfg
179
+
180
+ def test_json_redactions(self) -> None:
181
+ cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
182
+ labels = set(cfg["redactions"])
183
+ assert "hf_token" in labels
184
+ assert "s3_uri" in labels
185
+ # raw_source is no longer in the result; verify scrubbing via the
186
+ # full-config helper.
187
+ full = _parse_config_full(str(FIXTURES / "sample_train.json"))
188
+ assert isinstance(full, WorkloadConfig)
189
+ assert "hf_jsonsamplehfabcdefghijklmnopqrs" not in full.raw_source
190
+
191
+ def test_extras_values_are_scrubbed(self) -> None:
192
+ # Secret-shaped values that landed in extras must also be redacted —
193
+ # otherwise the leak just moves from raw_source into extras.
194
+ cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
195
+ extras = cfg["extras"]
196
+ assert "hf_jsonsamplehfabcdefghijklmnopqrs" not in extras.get("hub_token", "")
197
+ assert extras.get("hub_token", "").startswith("<REDACTED:")
198
+ assert extras.get("checkpoint_uri", "").startswith("<REDACTED:")
199
+
200
+ def test_yaml_redactions(self) -> None:
201
+ cfg = _parse_config(str(FIXTURES / "sample_train.yaml")).result
202
+ labels = set(cfg["redactions"])
203
+ assert "hf_token" in labels
204
+ assert "bearer_token" in labels
205
+ assert "home_path" in labels
206
+
207
+
208
+ # ---------------------------------------------------------------------------
209
+ # Failure modes
210
+ # ---------------------------------------------------------------------------
211
+
212
+
213
+ class TestErrors:
214
+ def test_missing_file(self) -> None:
215
+ result = _parse_config("/tmp/definitely-does-not-exist-xyz.py")
216
+ assert result.ok is False
217
+ assert "not found" in (result.error or "").lower()
218
+
219
+ def test_unsupported_extension(self, tmp_path: Path) -> None:
220
+ bad = tmp_path / "config.toml"
221
+ bad.write_text("model_name = 'foo'\n")
222
+ result = _parse_config(str(bad))
223
+ assert result.ok is False
224
+ assert "unsupported" in (result.error or "").lower()
225
+
226
+ def test_malformed_python(self, tmp_path: Path) -> None:
227
+ bad = tmp_path / "broken.py"
228
+ bad.write_text("def oops(:\n pass\n")
229
+ result = _parse_config(str(bad))
230
+ assert result.ok is False
231
+ assert "parse error" in (result.error or "").lower()
232
+
233
+ def test_malformed_json(self, tmp_path: Path) -> None:
234
+ bad = tmp_path / "broken.json"
235
+ bad.write_text("{not really json")
236
+ result = _parse_config(str(bad))
237
+ assert result.ok is False
238
+ assert "json" in (result.error or "").lower()
239
+
240
+ def test_json_top_level_must_be_dict(self, tmp_path: Path) -> None:
241
+ bad = tmp_path / "list.json"
242
+ bad.write_text(json.dumps([{"foo": 1}]))
243
+ result = _parse_config(str(bad))
244
+ assert result.ok is False
245
+
246
+ def test_yaml_top_level_must_be_mapping(self, tmp_path: Path) -> None:
247
+ bad = tmp_path / "scalar.yaml"
248
+ bad.write_text("- 1\n- 2\n")
249
+ result = _parse_config(str(bad))
250
+ assert result.ok is False
251
+
252
+
253
+ # ---------------------------------------------------------------------------
254
+ # Schema invariants
255
+ # ---------------------------------------------------------------------------
256
+
257
+
258
+ class TestSchema:
259
+ def test_result_round_trips_through_workload_config(self) -> None:
260
+ result = _parse_config(str(FIXTURES / "sample_train.py"))
261
+ # Must be reconstructible — guards against extras-vs-fields collisions.
262
+ cfg = WorkloadConfig(**result.result)
263
+ assert cfg.model_name == "Qwen/Qwen2.5-7B-Instruct"
264
+
265
+ def test_defaults_when_field_absent(self, tmp_path: Path) -> None:
266
+ # Minimal config — only model_name. Everything else should fall back to schema defaults.
267
+ path = tmp_path / "tiny.json"
268
+ path.write_text(json.dumps({"model_name": "test/tiny"}))
269
+ result = _parse_config(str(path))
270
+ assert result.ok
271
+ cfg = result.result
272
+ assert cfg["batch_size"] == 1
273
+ assert cfg["precision"] == "fp16"
274
+ assert cfg["optimizer"] == "adamw_torch"
275
+ assert cfg["gradient_checkpointing"] is False
276
+ assert cfg["redactions"] == []
277
+
278
+ def test_tool_definition_unchanged_in_shape(self) -> None:
279
+ # The Tool definition should still expose name/description/input_schema/fn.
280
+ assert PARSE_CONFIG.name == "parse_config"
281
+ assert PARSE_CONFIG.fn is _parse_config
282
+ assert "file_path" in PARSE_CONFIG.input_schema["properties"]
283
+
284
+
285
+ # ---------------------------------------------------------------------------
286
+ # Regression: canonical + scenario workloads must parse with all the right
287
+ # audit-relevant fields. These are what the live agent actually sees, so a
288
+ # regression here directly degrades audit quality (the agent reasons over
289
+ # HF defaults instead of the script's settings).
290
+ # ---------------------------------------------------------------------------
291
+
292
+
293
+ REPO_ROOT = Path(__file__).resolve().parent.parent
294
+
295
+
296
+ class TestCanonicalWorkload:
297
+ def test_canonical_workload_extracts_full_config(self) -> None:
298
+ """The canonical demo workload must yield batch_size=4, lr=2e-4, etc.
299
+ — not HF defaults. Catches the `**dict_var` splat regression where
300
+ every TrainingArguments kwarg disappears.
301
+ """
302
+ result = _parse_config(str(REPO_ROOT / "workloads" / "train_qwen_lora.py"))
303
+ assert result.ok, result.error
304
+ cfg = result.result
305
+ assert cfg["model_name"] == "Qwen/Qwen2.5-7B-Instruct"
306
+ assert cfg["batch_size"] == 4, (
307
+ "expected batch_size=4 from per_device_train_batch_size; "
308
+ "did `**_ta_kwargs` splat hide the kwargs?"
309
+ )
310
+ assert cfg["grad_accum_steps"] == 8
311
+ assert cfg["lr"] == 2e-4
312
+ assert cfg["warmup_steps"] == 100
313
+ assert cfg["precision"] == "fp16"
314
+ assert cfg["attention_impl"] == "eager"
315
+ assert cfg["dataloader_workers"] == 0
316
+ assert cfg["dataloader_pin_memory"] is False
317
+ assert cfg["lora_rank"] == 16
318
+ assert cfg["torch_compile"] is False
319
+ assert cfg["env_vars"]["HSA_FORCE_FINE_GRAIN_PCIE"] == "1"
320
+
321
+
322
+ class TestSplatKwargsResolution:
323
+ """`_ta = dict(k=v); Foo(**_ta)` must resolve back through the dict
324
+ constant. Defensive — the canonical workload no longer uses this
325
+ pattern, but third-party scripts often do.
326
+ """
327
+
328
+ def test_dict_function_call_splat(self, tmp_path) -> None:
329
+ src = """
330
+ from transformers import TrainingArguments
331
+
332
+ _ta = dict(
333
+ per_device_train_batch_size=8,
334
+ gradient_accumulation_steps=2,
335
+ fp16=True,
336
+ optim=\"adamw_torch_fused\",
337
+ )
338
+ training_args = TrainingArguments(output_dir=\"./out\", **_ta)
339
+ """
340
+ p = tmp_path / "splat.py"
341
+ p.write_text(src)
342
+ cfg = _parse_config(str(p)).result
343
+ assert cfg["batch_size"] == 8
344
+ assert cfg["grad_accum_steps"] == 2
345
+ assert cfg["precision"] == "fp16"
346
+ assert cfg["optimizer"] == "adamw_torch_fused"
347
+
348
+ def test_dict_literal_splat(self, tmp_path) -> None:
349
+ src = """
350
+ from transformers import TrainingArguments
351
+
352
+ _ta = {
353
+ "per_device_train_batch_size": 16,
354
+ "bf16": True,
355
+ }
356
+ training_args = TrainingArguments(output_dir=\"./out\", **_ta)
357
+ """
358
+ p = tmp_path / "splat_literal.py"
359
+ p.write_text(src)
360
+ cfg = _parse_config(str(p)).result
361
+ assert cfg["batch_size"] == 16
362
+ assert cfg["precision"] == "bf16"
363
+
364
+ def test_explicit_kwarg_overrides_splat(self, tmp_path) -> None:
365
+ src = """
366
+ from transformers import TrainingArguments
367
+
368
+ _ta = dict(per_device_train_batch_size=8)
369
+ training_args = TrainingArguments(per_device_train_batch_size=32, **_ta)
370
+ """
371
+ p = tmp_path / "splat_override.py"
372
+ p.write_text(src)
373
+ cfg = _parse_config(str(p)).result
374
+ # Explicit kwarg wins over splat (both occur in the kwargs list,
375
+ # explicit comes first in the AST → setdefault keeps it).
376
+ assert cfg["batch_size"] == 32
377
+
tests/test_query_rocm_kb.py ADDED
@@ -0,0 +1,348 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for ``agent.tools.query_rocm_kb``.
2
+
3
+ Coverage:
4
+ * The shipped ``kb/rocm_rules.yaml`` is loadable, validates against
5
+ :class:`Rule`, and contains every required-by-spec rule id.
6
+ * Each rule's ``targets_bucket`` is a valid :class:`WasteBucket` and its
7
+ ``category`` is a valid :class:`RuleCategory` (caught for free by
8
+ pydantic, but assert here for a clearer failure when the YAML drifts).
9
+ * Semantic search returns the bf16 rule first for an "fp16 on MI300X"
10
+ query and the flash-attn / sdpa rules first for an attention query.
11
+ * Bad inputs (empty symptom, ``top_k <= 0``, ``top_k`` larger than the
12
+ rule count) are handled gracefully via ``ToolResult.ok=False`` or
13
+ clamped to the rule count.
14
+ * Embeddings cache: the cache file is created on first import, and a
15
+ second ``_embed_rules`` call with the same YAML bytes hits the cache
16
+ without re-encoding.
17
+ * The frozen ``Tool`` definition retains its ``name``, ``description``,
18
+ and ``input_schema`` shape — the agent registry depends on these.
19
+ """
20
+
21
+ from __future__ import annotations
22
+
23
+ from pathlib import Path
24
+ from typing import get_args
25
+
26
+ import pytest
27
+ import yaml
28
+
29
+ from agent.schemas import Rule, RuleCategory, ToolResult, WasteBucket
30
+ from agent.tools.query_rocm_kb import (
31
+ _KB_YAML,
32
+ _RULES,
33
+ QUERY_ROCM_KB,
34
+ _cache_path,
35
+ _embed_rules,
36
+ _load_rules,
37
+ _query_rocm_kb,
38
+ )
39
+
40
+
41
+ REQUIRED_RULE_IDS = {
42
+ "precision.bf16_over_fp16_on_mi300x",
43
+ "attention.flash_rocm_over_eager",
44
+ "attention.sdpa_over_eager",
45
+ "memory.batch_too_small_for_192gb",
46
+ "memory.gradient_checkpointing_for_long_seq",
47
+ "data.dataloader_workers_zero",
48
+ "data.pin_memory_false",
49
+ "data.prefetch_factor_default",
50
+ "data.persistent_workers_false",
51
+ "compile.torch_compile_off",
52
+ "env.nccl_min_nchannels",
53
+ "env.numa_auto_balancing_disable",
54
+ "env.hsa_force_fine_grain_pcie",
55
+ "kernels.hipblaslt_hint_logging",
56
+ "kernels.miopen_find_mode_2",
57
+ "optimizer.bitsandbytes_not_supported_warning",
58
+ "collectives.one_process_per_gpu",
59
+ "topology.tp_within_xgmi_island",
60
+ }
61
+
62
+
63
+ # ---------------------------------------------------------------------------
64
+ # YAML KB invariants
65
+ # ---------------------------------------------------------------------------
66
+
67
+
68
+ class TestYamlKB:
69
+ def test_kb_yaml_exists(self) -> None:
70
+ assert _KB_YAML.exists(), f"Missing {_KB_YAML}"
71
+
72
+ def test_kb_loads_and_validates(self) -> None:
73
+ rules, _raw = _load_rules(_KB_YAML)
74
+ # Every rule pydantic-validated.
75
+ for r in rules:
76
+ assert isinstance(r, Rule)
77
+ # Spec calls for 20-25 rules; allow a small wiggle room either way.
78
+ assert 18 <= len(rules) <= 30, f"Expected 18-30 rules, got {len(rules)}"
79
+
80
+ def test_required_rule_ids_present(self) -> None:
81
+ ids = {r.id for r in _RULES}
82
+ missing = REQUIRED_RULE_IDS - ids
83
+ assert not missing, f"Required rule ids missing from KB: {sorted(missing)}"
84
+
85
+ def test_rule_ids_unique(self) -> None:
86
+ ids = [r.id for r in _RULES]
87
+ dupes = {i for i in ids if ids.count(i) > 1}
88
+ assert not dupes, f"Duplicate rule ids: {sorted(dupes)}"
89
+
90
+ def test_categories_are_valid(self) -> None:
91
+ valid = set(get_args(RuleCategory))
92
+ for r in _RULES:
93
+ assert r.category in valid, f"{r.id}: invalid category {r.category!r}"
94
+
95
+ def test_targets_bucket_valid(self) -> None:
96
+ valid = set(get_args(WasteBucket))
97
+ for r in _RULES:
98
+ assert r.targets_bucket in valid, (
99
+ f"{r.id}: invalid targets_bucket {r.targets_bucket!r}"
100
+ )
101
+
102
+ def test_recovery_fraction_in_range(self) -> None:
103
+ for r in _RULES:
104
+ assert 0.0 <= r.expected_recovery_fraction <= 1.0, (
105
+ f"{r.id}: expected_recovery_fraction out of [0,1] "
106
+ f"({r.expected_recovery_fraction})"
107
+ )
108
+
109
+ def test_citations_non_empty(self) -> None:
110
+ for r in _RULES:
111
+ assert r.citation and r.citation.strip(), f"{r.id}: empty citation"
112
+
113
+ def test_bitsandbytes_is_warning_only(self) -> None:
114
+ rule = next(r for r in _RULES if r.id == "optimizer.bitsandbytes_not_supported_warning")
115
+ # Warning rule has empty transform — propose_patch must not auto-fix.
116
+ assert rule.transform == {}
117
+
118
+ def test_categories_cover_spec(self) -> None:
119
+ # Architecture §4 lists 10 categories. We expect rules in at least
120
+ # the high-impact ones the spec calls out as required.
121
+ cats = {r.category for r in _RULES}
122
+ for required in (
123
+ "precision",
124
+ "attention",
125
+ "memory",
126
+ "data",
127
+ "compile",
128
+ "env_vars",
129
+ "kernels",
130
+ "optimizer",
131
+ "collectives",
132
+ "topology",
133
+ ):
134
+ assert required in cats, f"No rule for category {required!r}"
135
+
136
+
137
+ # ---------------------------------------------------------------------------
138
+ # Skip-and-warn behaviour for invalid entries
139
+ # ---------------------------------------------------------------------------
140
+
141
+
142
+ class TestLoadRulesResilience:
143
+ def test_invalid_entry_is_skipped_with_warning(self, tmp_path: Path) -> None:
144
+ bad = tmp_path / "rules.yaml"
145
+ bad.write_text(
146
+ yaml.safe_dump(
147
+ [
148
+ {
149
+ "id": "good.rule",
150
+ "category": "precision",
151
+ "targets_bucket": "precision_path",
152
+ "symptom": "fp16 used",
153
+ "expected_impact": "switch to bf16",
154
+ "citation": "ROCm guide",
155
+ },
156
+ {"id": "bad.rule", "category": "not_a_real_category"},
157
+ ]
158
+ )
159
+ )
160
+ with pytest.warns(UserWarning):
161
+ rules, _raw = _load_rules(bad)
162
+ assert [r.id for r in rules] == ["good.rule"]
163
+
164
+ def test_top_level_must_be_list(self, tmp_path: Path) -> None:
165
+ bad = tmp_path / "rules.yaml"
166
+ bad.write_text("not_a_list: 1\n")
167
+ with pytest.raises(ValueError, match="top-level"):
168
+ _load_rules(bad)
169
+
170
+
171
+ # ---------------------------------------------------------------------------
172
+ # Semantic search behaviour
173
+ # ---------------------------------------------------------------------------
174
+
175
+
176
+ class TestQuery:
177
+ def test_returns_ok_for_real_query(self) -> None:
178
+ result = _query_rocm_kb("fp16 used on MI300X with eager attention", top_k=5)
179
+ assert isinstance(result, ToolResult)
180
+ assert result.ok, result.error
181
+ rules = result.result["rules"]
182
+ assert 1 <= len(rules) <= 5
183
+
184
+ def test_fp16_query_returns_bf16_rule_in_top_results(self) -> None:
185
+ result = _query_rocm_kb("fp16 used on MI300X / CDNA3", top_k=3)
186
+ ids = [r["id"] for r in result.result["rules"]]
187
+ assert "precision.bf16_over_fp16_on_mi300x" in ids
188
+
189
+ def test_eager_attention_query_returns_attention_rules(self) -> None:
190
+ result = _query_rocm_kb(
191
+ "eager attention with no flash kernel loaded on MI300X", top_k=3
192
+ )
193
+ ids = [r["id"] for r in result.result["rules"]]
194
+ attention_ids = {
195
+ "attention.flash_rocm_over_eager",
196
+ "attention.sdpa_over_eager",
197
+ }
198
+ assert attention_ids & set(ids), (
199
+ f"Expected at least one of {attention_ids} in top 3, got {ids}"
200
+ )
201
+
202
+ def test_dataloader_query_returns_data_rules(self) -> None:
203
+ result = _query_rocm_kb(
204
+ "DataLoader num_workers is zero, GPU starves waiting for batches",
205
+ top_k=3,
206
+ )
207
+ ids = [r["id"] for r in result.result["rules"]]
208
+ data_ids = {
209
+ "data.dataloader_workers_zero",
210
+ "data.pin_memory_false",
211
+ "data.prefetch_factor_default",
212
+ "data.persistent_workers_false",
213
+ }
214
+ assert data_ids & set(ids), (
215
+ f"Expected at least one data.* rule in top 3, got {ids}"
216
+ )
217
+
218
+ def test_top_k_bounds_returned(self) -> None:
219
+ result = _query_rocm_kb("anything", top_k=2)
220
+ assert len(result.result["rules"]) == 2
221
+
222
+ def test_top_k_clamped_to_rule_count(self) -> None:
223
+ # Asking for more than we have should not crash; we should get every
224
+ # rule back, ordered by score.
225
+ result = _query_rocm_kb("anything", top_k=100)
226
+ assert len(result.result["rules"]) == len(_RULES)
227
+
228
+ def test_results_sorted_by_descending_score(self) -> None:
229
+ # If two queries with different focus return different top rules,
230
+ # ordering is real — top_1 differs by query.
231
+ a = _query_rocm_kb("fp16 numerical instability", top_k=1).result["rules"][0]
232
+ b = _query_rocm_kb("eager attention slow on long sequences", top_k=1).result[
233
+ "rules"
234
+ ][0]
235
+ assert a["id"] != b["id"], (
236
+ "Top rule should depend on query, but both queries returned "
237
+ f"{a['id']} — semantic search is degenerate."
238
+ )
239
+
240
+ def test_rule_payload_is_lite_shape(self) -> None:
241
+ # The LLM-facing rule payload is intentionally trimmed from the full
242
+ # Rule schema — only id / symptom / transform / expected_impact /
243
+ # citation make it through. This shrinks the audit conversation enough
244
+ # to fit Qwen2.5-7B's 8K window. Full Rule lookup happens server-side
245
+ # in propose_patch via the loaded KB.
246
+ result = _query_rocm_kb("any query", top_k=1)
247
+ payload = result.result["rules"][0]
248
+ assert set(payload.keys()) == {
249
+ "id",
250
+ "symptom",
251
+ "transform",
252
+ "expected_impact",
253
+ "citation",
254
+ }
255
+ # And the id resolves against the loaded KB so propose_patch can
256
+ # reconstruct the full Rule.
257
+ from agent.tools.query_rocm_kb import _RULES
258
+
259
+ kb_ids = {r.id for r in _RULES}
260
+ assert payload["id"] in kb_ids
261
+
262
+
263
+ # ---------------------------------------------------------------------------
264
+ # Failure modes
265
+ # ---------------------------------------------------------------------------
266
+
267
+
268
+ class TestErrors:
269
+ def test_empty_symptom(self) -> None:
270
+ result = _query_rocm_kb("", top_k=3)
271
+ assert result.ok is False
272
+ assert "symptom" in (result.error or "")
273
+
274
+ def test_whitespace_only_symptom(self) -> None:
275
+ result = _query_rocm_kb(" \t\n ", top_k=3)
276
+ assert result.ok is False
277
+
278
+ def test_top_k_zero(self) -> None:
279
+ result = _query_rocm_kb("fp16", top_k=0)
280
+ assert result.ok is False
281
+ assert "top_k" in (result.error or "")
282
+
283
+ def test_top_k_negative(self) -> None:
284
+ result = _query_rocm_kb("fp16", top_k=-3)
285
+ assert result.ok is False
286
+
287
+
288
+ # ---------------------------------------------------------------------------
289
+ # Embeddings cache
290
+ # ---------------------------------------------------------------------------
291
+
292
+
293
+ class TestEmbeddingsCache:
294
+ def test_cache_file_was_created_on_import(self) -> None:
295
+ raw = _KB_YAML.read_bytes()
296
+ cache = _cache_path(raw)
297
+ assert cache.exists(), (
298
+ f"Expected embeddings cache at {cache}; cache write failed silently."
299
+ )
300
+
301
+ def test_second_embed_call_uses_cache_without_recoding(
302
+ self, monkeypatch: pytest.MonkeyPatch
303
+ ) -> None:
304
+ rules, raw = _load_rules(_KB_YAML)
305
+ # Sentinel: replace the lazy model getter so any encode() call would
306
+ # blow up. If the cache is hit, the model is never consulted.
307
+ from agent.tools import query_rocm_kb as kb_module
308
+
309
+ def explode() -> None:
310
+ raise AssertionError(
311
+ "_get_model() called on a cache-hit path; cache is not being used."
312
+ )
313
+
314
+ monkeypatch.setattr(kb_module, "_get_model", explode)
315
+ embeddings = _embed_rules(rules, raw)
316
+ assert embeddings.shape[0] == len(rules)
317
+ assert embeddings.dtype.kind == "f"
318
+
319
+
320
+ # ---------------------------------------------------------------------------
321
+ # Tool registry shape — the agent loop depends on these fields being stable.
322
+ # ---------------------------------------------------------------------------
323
+
324
+
325
+ class TestToolDefinition:
326
+ def test_name_is_query_rocm_kb(self) -> None:
327
+ assert QUERY_ROCM_KB.name == "query_rocm_kb"
328
+
329
+ def test_description_unchanged_keywords(self) -> None:
330
+ # Must still describe the search-by-symptom semantics; the system
331
+ # prompt references this language.
332
+ desc = QUERY_ROCM_KB.description
333
+ assert "ROCm" in desc and "symptom" in desc
334
+
335
+ def test_input_schema_shape(self) -> None:
336
+ schema = QUERY_ROCM_KB.input_schema
337
+ assert schema["type"] == "object"
338
+ # Both single-query (`symptom`) and batched (`symptoms`) shapes are
339
+ # advertised — either works at runtime, neither is strictly required
340
+ # in the schema because the impl validates "at least one" itself.
341
+ assert "symptom" in schema["properties"]
342
+ assert "symptoms" in schema["properties"]
343
+ assert "top_k" in schema["properties"]
344
+ assert schema["properties"]["symptoms"]["type"] == "array"
345
+ assert schema["properties"]["symptoms"]["items"] == {"type": "string"}
346
+
347
+ def test_fn_is_module_query(self) -> None:
348
+ assert QUERY_ROCM_KB.fn is _query_rocm_kb
tests/test_qwen_vllm_backend.py ADDED
@@ -0,0 +1,280 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for QwenVLLMBackend.
2
+
3
+ Mocks the openai AsyncClient at the chat.completions level so no real
4
+ network call ever happens. Verifies the OpenAI-shape conversation
5
+ threading, tool-call extraction, and finish_reason mapping.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import os
11
+ from types import SimpleNamespace
12
+ from typing import Any
13
+
14
+ import pytest
15
+
16
+ from agent.backends import active_backend_name, make_backend
17
+ from agent.backends.base import AgentTurn
18
+ from agent.backends.qwen_vllm import (
19
+ QwenVLLMBackend,
20
+ _normalize_finish_reason,
21
+ _to_openai_tools,
22
+ )
23
+
24
+
25
+ # ---------------------------------------------------------------------------
26
+ # Helpers
27
+ # ---------------------------------------------------------------------------
28
+
29
+
30
+ class _FakeChatCompletions:
31
+ """Stands in for ``client.chat.completions``. Records every call and
32
+ returns scripted responses one-by-one."""
33
+
34
+ def __init__(self, responses: list[Any]) -> None:
35
+ self.responses = list(responses)
36
+ self.calls: list[dict[str, Any]] = []
37
+
38
+ async def create(self, **kwargs: Any) -> Any:
39
+ self.calls.append(kwargs)
40
+ if not self.responses:
41
+ raise AssertionError(
42
+ "FakeChatCompletions exhausted — backend made more calls than expected"
43
+ )
44
+ return self.responses.pop(0)
45
+
46
+
47
+ class _FakeClient:
48
+ def __init__(self, responses: list[Any]) -> None:
49
+ self.chat = SimpleNamespace(completions=_FakeChatCompletions(responses))
50
+
51
+
52
+ def _scripted_response(
53
+ *,
54
+ content: str | None = None,
55
+ tool_calls: list[dict[str, Any]] | None = None,
56
+ finish_reason: str = "stop",
57
+ ) -> Any:
58
+ """Build the SimpleNamespace shape the openai SDK returns from
59
+ chat.completions.create(...).
60
+ """
61
+ tcs = []
62
+ if tool_calls:
63
+ for tc in tool_calls:
64
+ tcs.append(
65
+ SimpleNamespace(
66
+ id=tc["id"],
67
+ function=SimpleNamespace(
68
+ name=tc["name"],
69
+ arguments=tc.get("arguments", "{}"),
70
+ ),
71
+ )
72
+ )
73
+ msg = SimpleNamespace(content=content, tool_calls=tcs or None)
74
+ choice = SimpleNamespace(message=msg, finish_reason=finish_reason)
75
+ return SimpleNamespace(choices=[choice])
76
+
77
+
78
+ def _backend_with(responses: list[Any]) -> QwenVLLMBackend:
79
+ """Construct a QwenVLLMBackend with a scripted client."""
80
+ backend = QwenVLLMBackend.__new__(QwenVLLMBackend)
81
+ backend._system = "you are a test agent"
82
+ backend._model = "Qwen/Qwen2.5-7B-Instruct"
83
+ backend._base_url = "http://fake-vllm:8000/v1"
84
+ backend._api_key = "EMPTY"
85
+ backend._max_tokens = 1024
86
+ backend._client = _FakeClient(responses)
87
+ backend._conversation = [{"role": "system", "content": backend._system}]
88
+ return backend
89
+
90
+
91
+ # ---------------------------------------------------------------------------
92
+ # Tool-schema translation
93
+ # ---------------------------------------------------------------------------
94
+
95
+
96
+ def test_to_openai_tools_translates_neutral_shape() -> None:
97
+ neutral = [
98
+ {
99
+ "name": "parse_config",
100
+ "description": "Parse the file.",
101
+ "input_schema": {"type": "object", "properties": {"file_path": {"type": "string"}}},
102
+ },
103
+ ]
104
+ out = _to_openai_tools(neutral)
105
+ assert out == [
106
+ {
107
+ "type": "function",
108
+ "function": {
109
+ "name": "parse_config",
110
+ "description": "Parse the file.",
111
+ "parameters": {
112
+ "type": "object",
113
+ "properties": {"file_path": {"type": "string"}},
114
+ },
115
+ },
116
+ }
117
+ ]
118
+
119
+
120
+ def test_to_openai_tools_handles_missing_input_schema() -> None:
121
+ out = _to_openai_tools([{"name": "x", "description": "y"}])
122
+ assert out[0]["function"]["parameters"] == {"type": "object", "properties": {}}
123
+
124
+
125
+ def test_finish_reason_normalization() -> None:
126
+ assert _normalize_finish_reason("stop") == "end_turn"
127
+ assert _normalize_finish_reason("tool_calls") == "tool_use"
128
+ assert _normalize_finish_reason("length") == "max_tokens"
129
+ assert _normalize_finish_reason(None) == "other"
130
+ assert _normalize_finish_reason("weird") == "weird"
131
+
132
+
133
+ # ---------------------------------------------------------------------------
134
+ # next_turn behavior
135
+ # ---------------------------------------------------------------------------
136
+
137
+
138
+ @pytest.mark.asyncio
139
+ async def test_next_turn_emits_text_block_and_end_turn() -> None:
140
+ backend = _backend_with([_scripted_response(content="hello there", finish_reason="stop")])
141
+ backend.add_user_message("audit /tmp/x.py")
142
+ turn = await backend.next_turn(tool_schemas=[])
143
+ assert isinstance(turn, AgentTurn)
144
+ assert turn.text_blocks == ["hello there"]
145
+ assert turn.tool_calls == []
146
+ assert turn.stop_reason == "end_turn"
147
+
148
+
149
+ @pytest.mark.asyncio
150
+ async def test_next_turn_emits_tool_calls_with_parsed_args() -> None:
151
+ backend = _backend_with(
152
+ [
153
+ _scripted_response(
154
+ content="calling parse_config",
155
+ tool_calls=[
156
+ {
157
+ "id": "tc-1",
158
+ "name": "parse_config",
159
+ "arguments": '{"file_path": "/tmp/x.py"}',
160
+ }
161
+ ],
162
+ finish_reason="tool_calls",
163
+ )
164
+ ]
165
+ )
166
+ backend.add_user_message("audit /tmp/x.py")
167
+ turn = await backend.next_turn(tool_schemas=[])
168
+ assert turn.stop_reason == "tool_use"
169
+ assert len(turn.tool_calls) == 1
170
+ tc = turn.tool_calls[0]
171
+ assert tc.id == "tc-1"
172
+ assert tc.name == "parse_config"
173
+ assert tc.input == {"file_path": "/tmp/x.py"}
174
+
175
+
176
+ @pytest.mark.asyncio
177
+ async def test_next_turn_handles_malformed_tool_arguments() -> None:
178
+ """vLLM occasionally emits unparseable JSON in arguments — don't crash."""
179
+ backend = _backend_with(
180
+ [
181
+ _scripted_response(
182
+ tool_calls=[
183
+ {"id": "tc-1", "name": "parse_config", "arguments": "{not-json"}
184
+ ],
185
+ finish_reason="tool_calls",
186
+ )
187
+ ]
188
+ )
189
+ backend.add_user_message("x")
190
+ turn = await backend.next_turn(tool_schemas=[])
191
+ # We get the call but with empty args rather than raising.
192
+ assert turn.tool_calls[0].input == {}
193
+
194
+
195
+ @pytest.mark.asyncio
196
+ async def test_tool_result_is_threaded_into_next_request() -> None:
197
+ backend = _backend_with(
198
+ [
199
+ _scripted_response(
200
+ tool_calls=[
201
+ {"id": "tc-1", "name": "parse_config", "arguments": "{}"}
202
+ ],
203
+ finish_reason="tool_calls",
204
+ ),
205
+ _scripted_response(content="done", finish_reason="stop"),
206
+ ]
207
+ )
208
+ backend.add_user_message("audit")
209
+ await backend.next_turn(tool_schemas=[])
210
+ backend.add_tool_result("tc-1", "parse_config", '{"ok": true}', is_error=False)
211
+ await backend.next_turn(tool_schemas=[])
212
+
213
+ # The second create() call should include role="tool" referencing tc-1.
214
+ second_call = backend._client.chat.completions.calls[1]
215
+ msgs = second_call["messages"]
216
+ tool_msgs = [m for m in msgs if m["role"] == "tool"]
217
+ assert len(tool_msgs) == 1
218
+ assert tool_msgs[0]["tool_call_id"] == "tc-1"
219
+ assert tool_msgs[0]["content"] == '{"ok": true}'
220
+
221
+
222
+ @pytest.mark.asyncio
223
+ async def test_is_error_prefix_added_to_failed_tool_results() -> None:
224
+ backend = _backend_with([_scripted_response(content="adapting", finish_reason="stop")])
225
+ backend.add_tool_result("tc-1", "parse_config", "file not found", is_error=True)
226
+ await backend.next_turn(tool_schemas=[])
227
+ msgs = backend._client.chat.completions.calls[0]["messages"]
228
+ tool_msg = next(m for m in msgs if m["role"] == "tool")
229
+ assert tool_msg["content"].startswith("ERROR:")
230
+
231
+
232
+ # ---------------------------------------------------------------------------
233
+ # Factory selection via env var
234
+ # ---------------------------------------------------------------------------
235
+
236
+
237
+ def test_make_backend_picks_vllm_when_env_var_set(monkeypatch) -> None:
238
+ monkeypatch.setenv("GOBLIN_AGENT_BACKEND", "qwen-vllm")
239
+ monkeypatch.setenv("GOBLIN_QWEN_VLLM_URL", "http://test:8000/v1")
240
+ assert active_backend_name() == "qwen-vllm"
241
+ backend = make_backend(system_prompt="x")
242
+ assert isinstance(backend, QwenVLLMBackend)
243
+ assert backend._base_url == "http://test:8000/v1"
244
+
245
+
246
+ def test_active_backend_name_aliases(monkeypatch) -> None:
247
+ for alias in ("qwen-vllm", "qwen_vllm", "vllm", "local", "QWEN-VLLM", "Vllm"):
248
+ monkeypatch.setenv("GOBLIN_AGENT_BACKEND", alias)
249
+ assert active_backend_name() == "qwen-vllm", alias
250
+ for alias in ("qwen-hf", "qwen", "hf", ""):
251
+ monkeypatch.setenv("GOBLIN_AGENT_BACKEND", alias)
252
+ assert active_backend_name() == "qwen-hf", alias
253
+
254
+
255
+ def test_make_backend_default_is_hf(monkeypatch) -> None:
256
+ monkeypatch.delenv("GOBLIN_AGENT_BACKEND", raising=False)
257
+ monkeypatch.setenv("HF_TOKEN", "fake-test-token")
258
+ assert active_backend_name() == "qwen-hf"
259
+
260
+
261
+ # ---------------------------------------------------------------------------
262
+ # Construction-time error: openai SDK missing
263
+ # ---------------------------------------------------------------------------
264
+
265
+
266
+ def test_constructor_raises_when_openai_not_installed(monkeypatch) -> None:
267
+ """If openai isn't on PYTHONPATH, the backend's _build_client raises a
268
+ clear RuntimeError instead of an opaque ImportError."""
269
+ import sys
270
+
271
+ saved = sys.modules.pop("openai", None)
272
+ sys.modules["openai"] = None # block re-import
273
+ try:
274
+ with pytest.raises(RuntimeError, match="openai"):
275
+ QwenVLLMBackend(system_prompt="x")
276
+ finally:
277
+ if saved is not None:
278
+ sys.modules["openai"] = saved
279
+ else:
280
+ sys.modules.pop("openai", None)
tests/test_runner.py ADDED
@@ -0,0 +1,311 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tests for runner/protocol.py and runner/profile_parser.py.
2
+
3
+ Two laptop-only invariants:
4
+ 1. FakeRunner still works exactly as before (the Phase 1 contract).
5
+ 2. LiveRunner gracefully falls back to FakeRunner whenever GPU/profiler
6
+ tools are missing — this dev box has no AMD GPU, so every test here
7
+ should exercise the fallback path.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import csv
13
+ import json
14
+ from pathlib import Path
15
+ from unittest import mock
16
+
17
+ import pytest
18
+
19
+ from agent.schemas import RunMetrics, WorkloadConfig
20
+ from runner import profile_parser
21
+ from runner.protocol import FakeRunner, LiveRunner, _default_runner, gpu_available
22
+
23
+
24
+ # ---------------------------------------------------------------------------
25
+ # Helpers
26
+ # ---------------------------------------------------------------------------
27
+
28
+
29
+ def _baseline_config() -> WorkloadConfig:
30
+ return WorkloadConfig(
31
+ model_name="Qwen/Qwen2.5-7B-Instruct",
32
+ precision="fp16",
33
+ batch_size=4,
34
+ attention_impl="eager",
35
+ dataloader_workers=0,
36
+ )
37
+
38
+
39
+ # ---------------------------------------------------------------------------
40
+ # FakeRunner — unchanged contract
41
+ # ---------------------------------------------------------------------------
42
+
43
+
44
+ class TestFakeRunner:
45
+ def test_matches_baseline_scenario(self):
46
+ runner = FakeRunner()
47
+ metrics = runner.run(_baseline_config(), steps=10)
48
+ assert isinstance(metrics, RunMetrics)
49
+ assert metrics.runner_kind == "fake"
50
+ assert metrics.steps == 10
51
+ # 01_baseline_bad fixture
52
+ assert metrics.tokens_per_sec == pytest.approx(142.0)
53
+
54
+ def test_steps_override_takes_precedence(self):
55
+ runner = FakeRunner()
56
+ metrics = runner.run(_baseline_config(), steps=99)
57
+ assert metrics.steps == 99
58
+
59
+ def test_default_metrics_when_no_match(self):
60
+ runner = FakeRunner()
61
+ # An unknown model_name forces the no-match path.
62
+ cfg = _baseline_config().model_copy(update={"model_name": "unknown/model"})
63
+ metrics = runner.run(cfg, steps=7)
64
+ assert metrics.runner_kind == "fake"
65
+ assert metrics.steps == 7
66
+ assert any("FakeRunner" in w for w in metrics.warnings)
67
+
68
+ def test_corpus_dir_missing_returns_default(self, tmp_path):
69
+ runner = FakeRunner(corpus_dir=tmp_path / "nope")
70
+ metrics = runner.run(_baseline_config(), steps=10)
71
+ assert metrics.runner_kind == "fake"
72
+
73
+
74
+ # ---------------------------------------------------------------------------
75
+ # gpu_available — pure detection
76
+ # ---------------------------------------------------------------------------
77
+
78
+
79
+ class TestGpuAvailable:
80
+ def test_no_rocprofv3(self):
81
+ with mock.patch("runner.protocol.shutil.which", return_value=None):
82
+ ok, reason = gpu_available()
83
+ assert ok is False
84
+ assert reason and "rocprofv3" in reason
85
+
86
+ def test_no_amd_smi(self):
87
+ def which(name):
88
+ return "/usr/bin/rocprofv3" if name == "rocprofv3" else None
89
+
90
+ with mock.patch("runner.protocol.shutil.which", side_effect=which):
91
+ ok, reason = gpu_available()
92
+ assert ok is False
93
+ assert reason and "amd-smi" in reason
94
+
95
+ def test_no_render_device(self):
96
+ with mock.patch(
97
+ "runner.protocol.shutil.which",
98
+ side_effect=lambda name: f"/usr/bin/{name}",
99
+ ), mock.patch("runner.protocol._has_render_device", return_value=False):
100
+ ok, reason = gpu_available()
101
+ assert ok is False
102
+ assert reason and "renderD" in reason
103
+
104
+ def test_all_present(self):
105
+ with mock.patch(
106
+ "runner.protocol.shutil.which",
107
+ side_effect=lambda name: f"/usr/bin/{name}",
108
+ ), mock.patch("runner.protocol._has_render_device", return_value=True):
109
+ ok, reason = gpu_available()
110
+ assert ok is True
111
+ assert reason is None
112
+
113
+
114
+ # ---------------------------------------------------------------------------
115
+ # LiveRunner — must fall back on this no-GPU dev machine
116
+ # ---------------------------------------------------------------------------
117
+
118
+
119
+ class TestLiveRunnerFallback:
120
+ def test_falls_back_when_gpu_unavailable(self):
121
+ runner = LiveRunner()
122
+ metrics = runner.run(_baseline_config(), steps=10)
123
+ # On a laptop, gpu_available() returns False → FakeRunner path.
124
+ assert metrics.runner_kind == "fake"
125
+ # The warning must be the FIRST entry (LiveRunner prepends it).
126
+ assert metrics.warnings, "LiveRunner must surface a fallback warning"
127
+ assert "LiveRunner" in metrics.warnings[0]
128
+
129
+ def test_falls_back_when_runner_script_missing(self, tmp_path):
130
+ with mock.patch("runner.protocol.gpu_available", return_value=(True, None)):
131
+ runner = LiveRunner(runner_script=tmp_path / "does_not_exist.sh")
132
+ metrics = runner.run(_baseline_config(), steps=10)
133
+ assert metrics.runner_kind == "fake"
134
+ assert any("runner script not found" in w for w in metrics.warnings)
135
+
136
+ def test_falls_back_when_runner_script_not_executable(self, tmp_path):
137
+ script = tmp_path / "goblin_runner.sh"
138
+ script.write_text("#!/bin/sh\nexit 0\n")
139
+ # Deliberately don't chmod +x
140
+ with mock.patch("runner.protocol.gpu_available", return_value=(True, None)):
141
+ runner = LiveRunner(runner_script=script)
142
+ metrics = runner.run(_baseline_config(), steps=10)
143
+ assert metrics.runner_kind == "fake"
144
+ assert any("not executable" in w for w in metrics.warnings)
145
+
146
+ def test_falls_back_when_subprocess_returns_nonzero(self, tmp_path):
147
+ script = tmp_path / "goblin_runner.sh"
148
+ script.write_text("#!/usr/bin/env bash\nexit 7\n")
149
+ script.chmod(0o755)
150
+ with mock.patch("runner.protocol.gpu_available", return_value=(True, None)):
151
+ runner = LiveRunner(runner_script=script)
152
+ metrics = runner.run(_baseline_config(), steps=10)
153
+ assert metrics.runner_kind == "fake"
154
+ assert any("exited with code 7" in w for w in metrics.warnings)
155
+
156
+
157
+ # ---------------------------------------------------------------------------
158
+ # _default_runner — module-level factory
159
+ # ---------------------------------------------------------------------------
160
+
161
+
162
+ def test_default_runner_is_live_runner():
163
+ runner = _default_runner()
164
+ assert isinstance(runner, LiveRunner)
165
+
166
+
167
+ # ---------------------------------------------------------------------------
168
+ # profile_parser — graceful degradation when artefacts missing
169
+ # ---------------------------------------------------------------------------
170
+
171
+
172
+ class TestProfileParser:
173
+ def test_empty_dir_returns_zero_metrics_with_warnings(self, tmp_path):
174
+ metrics = profile_parser.parse(tmp_path, config=_baseline_config(), steps=10)
175
+ assert metrics.tokens_per_sec == 0.0
176
+ assert metrics.mfu_pct == 0.0
177
+ assert metrics.gpu_util_pct == 0.0
178
+ assert len(metrics.warnings) >= 3 # one warning per missing artefact
179
+
180
+ def test_parses_synthetic_artefacts(self, tmp_path):
181
+ # Minimal rocprofv3-shaped CSV
182
+ trace = tmp_path / "trace.csv"
183
+ with trace.open("w", newline="") as f:
184
+ w = csv.writer(f)
185
+ w.writerow(["KernelName", "DurationNs"])
186
+ w.writerow(["aten::matmul (fp16)", 5_000_000])
187
+ w.writerow(["aten::scaled_dot_product_attention", 3_000_000])
188
+ w.writerow(["rccl_AllReduce", 1_000_000])
189
+ w.writerow(["hipBLASLt_generic_gemm", 2_000_000])
190
+
191
+ # Minimal torch.profiler chrome trace with embedded metadata
192
+ torch_profile = {
193
+ "metadata": {
194
+ "tokens_per_sec": 142.0,
195
+ "mfu_pct": 24.0,
196
+ "pytorch_version": "2.3.0+rocm6.1",
197
+ "step_time_seconds": 0.5,
198
+ "host_busy_fraction": 0.6,
199
+ },
200
+ "traceEvents": [],
201
+ }
202
+ (tmp_path / "torch_profile.json").write_text(json.dumps(torch_profile))
203
+
204
+ # Minimal amd-smi telemetry
205
+ smi = tmp_path / "amd_smi.csv"
206
+ with smi.open("w", newline="") as f:
207
+ w = csv.writer(f)
208
+ w.writerow(["VRAM_USED_GB", "GFX_ACTIVITY"])
209
+ w.writerow(["72.0", "20.0"]) # < 30% util → triggers data_wait
210
+ w.writerow(["75.0", "22.0"])
211
+
212
+ metrics = profile_parser.parse(tmp_path, config=_baseline_config(), steps=10)
213
+ assert metrics.tokens_per_sec == pytest.approx(142.0)
214
+ assert metrics.mfu_pct == pytest.approx(24.0)
215
+ assert metrics.hbm_peak_gb == pytest.approx(75.0)
216
+ assert metrics.hbm_avg_gb == pytest.approx(73.5)
217
+ # comm_excess detected (rccl kernel, 1 ms)
218
+ assert metrics.waste_budget.comm_excess == pytest.approx(0.001)
219
+ # data_wait triggered (gpu util < 30, host_busy > 0.5)
220
+ assert metrics.waste_budget.data_wait > 0.0
221
+ # precision_path triggered (config.precision='fp16' AND fp16 kernels present)
222
+ assert metrics.waste_budget.precision_path > 0.0
223
+ # kernel_shape: generic GEMM detected
224
+ assert metrics.waste_budget.kernel_shape > 0.0
225
+ # memory_headroom: 75 GB used << 70% × 192 GB = 134.4 GB → slack
226
+ assert metrics.waste_budget.memory_headroom > 0.0
227
+
228
+ def test_bf16_config_skips_precision_path(self, tmp_path):
229
+ # Even with fp16-tagged kernels, a bf16 config means precision_path = 0
230
+ # because the user is already on the optimal precision.
231
+ trace = tmp_path / "trace.csv"
232
+ with trace.open("w", newline="") as f:
233
+ w = csv.writer(f)
234
+ w.writerow(["KernelName", "DurationNs"])
235
+ w.writerow(["matmul_fp16_kernel", 5_000_000])
236
+ torch_profile = {
237
+ "metadata": {
238
+ "tokens_per_sec": 318.0,
239
+ "mfu_pct": 51.0,
240
+ "step_time_seconds": 0.3,
241
+ "host_busy_fraction": 0.2,
242
+ },
243
+ "traceEvents": [],
244
+ }
245
+ (tmp_path / "torch_profile.json").write_text(json.dumps(torch_profile))
246
+ smi = tmp_path / "amd_smi.csv"
247
+ smi.write_text("VRAM_USED_GB,GFX_ACTIVITY\n168.0,86.0\n")
248
+
249
+ bf16_config = _baseline_config().model_copy(update={"precision": "bf16"})
250
+ metrics = profile_parser.parse(tmp_path, config=bf16_config, steps=50)
251
+ assert metrics.waste_budget.precision_path == 0.0
252
+
253
+
254
+ # ---------------------------------------------------------------------------
255
+ # Caching — exercise the benchmark tool's cache layer
256
+ # ---------------------------------------------------------------------------
257
+
258
+
259
+ class TestBenchmarkCache:
260
+ """The benchmark tool writes to the real bench_cache/ directory; isolate it."""
261
+
262
+ @pytest.fixture(autouse=True)
263
+ def _isolate_cache(self, tmp_path, monkeypatch):
264
+ monkeypatch.setattr("agent.tools.benchmark._CACHE_DIR", tmp_path / "bench_cache")
265
+ # Force ROCM_IMAGE_TAG to a known value so the key is reproducible.
266
+ monkeypatch.setenv("ROCM_IMAGE_TAG", "test-tag")
267
+ yield
268
+
269
+ def test_cache_hit_on_second_call(self):
270
+ from agent.tools.benchmark import _benchmark
271
+
272
+ cfg = _baseline_config().model_dump()
273
+ r1 = _benchmark(cfg, steps=50)
274
+ assert r1.ok
275
+ # Second call should HIT the cache and warn about it.
276
+ r2 = _benchmark(cfg, steps=50)
277
+ assert r2.ok
278
+ assert any("cache hit" in w for w in r2.result["warnings"])
279
+
280
+ def test_force_rerun_bypasses_cache(self):
281
+ from agent.tools.benchmark import _benchmark
282
+
283
+ cfg = _baseline_config().model_dump()
284
+ _benchmark(cfg, steps=50)
285
+ r2 = _benchmark(cfg, steps=50, force_rerun=True)
286
+ assert r2.ok
287
+ assert not any("cache hit" in w for w in r2.result["warnings"])
288
+
289
+ def test_different_steps_invalidate_cache(self):
290
+ from agent.tools.benchmark import _benchmark
291
+
292
+ cfg = _baseline_config().model_dump()
293
+ _benchmark(cfg, steps=50)
294
+ r2 = _benchmark(cfg, steps=100)
295
+ # Same config, different steps → different cache key → cold call.
296
+ assert not any("cache hit" in w for w in r2.result["warnings"])
297
+
298
+ def test_runner_script_change_invalidates_cache(self, tmp_path, monkeypatch):
299
+ from agent.tools.benchmark import _benchmark
300
+
301
+ cfg = _baseline_config().model_dump()
302
+ _benchmark(cfg, steps=50)
303
+
304
+ # Pretend the runner script changed by swapping the path the cache
305
+ # key reads from. (Simulates "container/runner version bump".)
306
+ fake_script = tmp_path / "different_runner.sh"
307
+ fake_script.write_text("# different content\n")
308
+ monkeypatch.setattr("agent.tools.benchmark._RUNNER_SCRIPT", fake_script)
309
+
310
+ r2 = _benchmark(cfg, steps=50)
311
+ assert not any("cache hit" in w for w in r2.result["warnings"])
ui/auto_tune_ui.py ADDED
@@ -0,0 +1,799 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """GPU Goblin auto-tune UI — Streamlit frontend for scripts/auto_tune.py.
2
+
3
+ Run:
4
+ streamlit run ui/auto_tune_ui.py
5
+
6
+ The UI:
7
+ 1. Form: pick model OR workload path, mode, steps, etc.
8
+ 2. Run: launches scripts/auto_tune.py as a subprocess with --events FILE
9
+ 3. Live progress: tails the events file, renders iteration cards as they
10
+ arrive, updates the best-tokens/sec metric on every accepted change
11
+ 4. Final report: improvement vs baseline, accepted vs rejected
12
+ experiments, waste-budget reduction chart, and a copyable diff.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import json
18
+ import os
19
+ import subprocess
20
+ import sys
21
+ import tempfile
22
+ import time
23
+ from pathlib import Path
24
+ from typing import Any
25
+
26
+ # Streamlit only puts ui/ on sys.path, so add the repo root for any helper
27
+ # imports that might land in shared modules later.
28
+ _REPO_ROOT = Path(__file__).resolve().parent.parent
29
+ if str(_REPO_ROOT) not in sys.path:
30
+ sys.path.insert(0, str(_REPO_ROOT))
31
+
32
+ import altair as alt
33
+ import pandas as pd
34
+ import requests
35
+ import streamlit as st
36
+
37
+ REPO_ROOT = _REPO_ROOT
38
+ AUTO_TUNE_SCRIPT = REPO_ROOT / "scripts" / "auto_tune.py"
39
+
40
+ WASTE_BUCKETS = [
41
+ "useful_gpu",
42
+ "data_wait",
43
+ "host_gap",
44
+ "comm_excess",
45
+ "memory_headroom",
46
+ "precision_path",
47
+ "kernel_shape",
48
+ ]
49
+
50
+ # ---------------------------------------------------------------------------
51
+ # Page setup
52
+ # ---------------------------------------------------------------------------
53
+
54
+ st.set_page_config(
55
+ page_title="GPU Goblin — Auto-Tune",
56
+ page_icon="🔧",
57
+ layout="wide",
58
+ )
59
+
60
+ st.title("🔧 GPU Goblin — Auto-Tune")
61
+ st.caption(
62
+ "Iteratively find the fastest configuration for a fine-tuning workload "
63
+ "on AMD MI300X. Pick a model, hit Run, watch tokens/sec climb."
64
+ )
65
+
66
+
67
+ # ---------------------------------------------------------------------------
68
+ # Form: inputs
69
+ # ---------------------------------------------------------------------------
70
+
71
+ with st.sidebar:
72
+ st.header("Run configuration")
73
+
74
+ # ---- Backend mode (Local subprocess vs remote GPU server) ----
75
+ default_backend = os.environ.get("GOBLIN_AUTO_TUNE_URL", "")
76
+ backend_mode = st.radio(
77
+ "Backend",
78
+ options=("Local subprocess", "Remote GPU server"),
79
+ index=1 if default_backend else 0,
80
+ help=(
81
+ "Local: launches scripts/auto_tune.py here (this host needs an MI300X). "
82
+ "Remote: POSTs to /auto-tune on a FastAPI server you've stood up on "
83
+ "the GPU host. Use Remote when running the UI on HF Spaces."
84
+ ),
85
+ )
86
+ backend_url = ""
87
+ if backend_mode == "Remote GPU server":
88
+ backend_url = st.text_input(
89
+ "Backend URL",
90
+ value=default_backend or "http://localhost:8000",
91
+ help="Base URL of the FastAPI server (no trailing slash). "
92
+ "/auto-tune is appended automatically.",
93
+ )
94
+
95
+ st.divider()
96
+ workload_source = st.radio(
97
+ "Workload source",
98
+ options=("Model id", "Custom workload script"),
99
+ index=0,
100
+ help=(
101
+ "Model id: auto-generates a baseline workload from the demo "
102
+ "template (Qwen-style LoRA fine-tune). Custom: point at any "
103
+ "Python training script."
104
+ ),
105
+ )
106
+
107
+ model_id = ""
108
+ workload_path = ""
109
+ if workload_source == "Model id":
110
+ model_id = st.text_input(
111
+ "Model id",
112
+ value="Qwen/Qwen2.5-7B-Instruct",
113
+ help=(
114
+ "Any HuggingFace causal-LM model id. For gated models "
115
+ "(Llama, etc.) ensure HF_TOKEN is set in the environment."
116
+ ),
117
+ )
118
+ else:
119
+ workload_path = st.text_input(
120
+ "Path to workload script (relative to repo root)",
121
+ value="workloads/train_qwen_lora.py",
122
+ )
123
+
124
+ st.divider()
125
+ st.subheader("Tuning strategy")
126
+
127
+ mode = st.selectbox(
128
+ "Mode",
129
+ options=("hardcoded", "llm", "llm-explore"),
130
+ index=0,
131
+ help=(
132
+ "hardcoded: priority-ordered playbook (no API key). "
133
+ "llm: LLM picks one experiment per iteration (greedy). "
134
+ "llm-explore: LLM proposes K candidates per iteration; best wins."
135
+ ),
136
+ )
137
+ candidates_per_iteration = 3
138
+ if mode == "llm-explore":
139
+ candidates_per_iteration = st.slider(
140
+ "Candidates per iteration (K)",
141
+ min_value=2,
142
+ max_value=6,
143
+ value=3,
144
+ help="Each iteration runs K benchmarks. Higher = broader search, more GPU time.",
145
+ )
146
+
147
+ steps = st.slider(
148
+ "Steps per benchmark",
149
+ min_value=10,
150
+ max_value=100,
151
+ value=20,
152
+ step=5,
153
+ help="More steps = lower variance but longer benchmarks.",
154
+ )
155
+
156
+ max_iterations = st.slider(
157
+ "Max iterations",
158
+ min_value=1,
159
+ max_value=20,
160
+ value=10,
161
+ help="Cap on tuning iterations. Default 10 for llm modes, 5 for llm-explore.",
162
+ )
163
+
164
+ early_stop_after = st.slider(
165
+ "Early stop after N non-improvements",
166
+ min_value=1,
167
+ max_value=10,
168
+ value=3,
169
+ )
170
+
171
+ max_crashes = st.slider(
172
+ "Max total crashes",
173
+ min_value=1,
174
+ max_value=10,
175
+ value=4,
176
+ )
177
+
178
+ improvement_threshold = st.number_input(
179
+ "Improvement threshold (%)",
180
+ min_value=0.0,
181
+ max_value=10.0,
182
+ value=0.0,
183
+ step=0.1,
184
+ help="Min % gain to accept. 0.0 = any positive delta wins.",
185
+ )
186
+
187
+ st.divider()
188
+ run_pressed = st.button("🚀 Run auto-tune", type="primary", use_container_width=True)
189
+
190
+
191
+ # ---------------------------------------------------------------------------
192
+ # Render helpers
193
+ # ---------------------------------------------------------------------------
194
+
195
+
196
+ def _format_tps(v: float) -> str:
197
+ return f"{v:,.0f}" if v else "—"
198
+
199
+
200
+ def _format_pct(v: float | None, suffix: str = "%") -> str:
201
+ if v is None:
202
+ return "—"
203
+ return f"{v:+.2f}{suffix}"
204
+
205
+
206
+ def _waste_chart(baseline: dict, current: dict | None) -> alt.Chart:
207
+ """Two-row stacked bar: baseline vs current waste_budget breakdown."""
208
+ rows = []
209
+ sources = [("baseline", baseline)]
210
+ if current is not None and current is not baseline:
211
+ sources.append(("current best", current))
212
+ for label, m in sources:
213
+ wb = (m or {}).get("waste_budget") or {}
214
+ for bucket in WASTE_BUCKETS:
215
+ v = float(wb.get(bucket, 0.0))
216
+ if v <= 0:
217
+ continue
218
+ rows.append({"run": label, "bucket": bucket, "seconds": v})
219
+ if not rows:
220
+ return None
221
+ df = pd.DataFrame(rows)
222
+ return (
223
+ alt.Chart(df)
224
+ .mark_bar()
225
+ .encode(
226
+ x=alt.X("seconds:Q", title="seconds / step"),
227
+ y=alt.Y("run:N", title=None, sort=["baseline", "current best"]),
228
+ color=alt.Color(
229
+ "bucket:N",
230
+ scale=alt.Scale(scheme="tableau10"),
231
+ sort=WASTE_BUCKETS,
232
+ ),
233
+ order=alt.Order("bucket:N", sort="ascending"),
234
+ tooltip=["run", "bucket", alt.Tooltip("seconds:Q", format=".3f")],
235
+ )
236
+ .properties(height=120)
237
+ )
238
+
239
+
240
+ def _render_metrics_row(baseline: dict | None, current: dict | None) -> None:
241
+ """Top-of-page metrics: tokens/sec, mfu_pct, hbm, with deltas vs baseline."""
242
+ cols = st.columns(4)
243
+ if baseline is None:
244
+ for c, label in zip(cols, ["tokens/sec", "mfu_pct", "hbm_peak (GB)", "iterations"]):
245
+ c.metric(label, "—")
246
+ return
247
+
248
+ cur = current if current is not None else baseline
249
+ base_tps = float(baseline.get("tokens_per_sec") or 0)
250
+ cur_tps = float(cur.get("tokens_per_sec") or 0)
251
+ base_mfu = float(baseline.get("mfu_pct") or 0)
252
+ cur_mfu = float(cur.get("mfu_pct") or 0)
253
+ base_hbm = float(baseline.get("hbm_peak_gb") or 0)
254
+ cur_hbm = float(cur.get("hbm_peak_gb") or 0)
255
+
256
+ cols[0].metric(
257
+ "tokens/sec",
258
+ _format_tps(cur_tps),
259
+ delta=f"{cur_tps - base_tps:+,.0f}" if base_tps else None,
260
+ )
261
+ cols[1].metric(
262
+ "MFU %",
263
+ f"{cur_mfu:.2f}",
264
+ delta=f"{cur_mfu - base_mfu:+.2f} pts" if base_mfu else None,
265
+ )
266
+ cols[2].metric(
267
+ "HBM peak (GB)",
268
+ f"{cur_hbm:.1f}",
269
+ delta=f"{cur_hbm - base_hbm:+.1f}" if base_hbm else None,
270
+ )
271
+
272
+
273
+ def _render_candidate_card(event: dict) -> None:
274
+ """One candidate's outcome inside an iteration."""
275
+ name = event.get("name", "?")
276
+ rationale = event.get("rationale", "")
277
+ outcome = event.get("outcome", "?")
278
+ metrics = event.get("metrics") or {}
279
+ delta = event.get("delta_vs_best")
280
+ reason = event.get("reason", "")
281
+
282
+ icon = {
283
+ "evaluated": "📊",
284
+ "skipped": "⏭️",
285
+ "rejected": "❌",
286
+ "crashed": "💥",
287
+ }.get(outcome, "•")
288
+
289
+ title = f"{icon} **{name}**"
290
+ if outcome == "evaluated" and delta is not None:
291
+ title += f" — Δ {delta:+.2f}%"
292
+ elif outcome != "evaluated":
293
+ title += f" — *{outcome}*"
294
+
295
+ with st.container(border=True):
296
+ st.markdown(title)
297
+ if rationale:
298
+ st.caption(rationale)
299
+ if metrics:
300
+ sub = st.columns(4)
301
+ sub[0].metric("tokens/sec", f"{metrics.get('tokens_per_sec', 0):,.0f}")
302
+ sub[1].metric("MFU %", f"{metrics.get('mfu_pct', 0):.2f}")
303
+ sub[2].metric("HBM peak GB", f"{metrics.get('hbm_peak_gb', 0):.1f}")
304
+ sub[3].metric("Δ vs best", f"{delta:+.2f}%" if delta is not None else "—")
305
+ wb = metrics.get("waste_budget") or {}
306
+ non_zero = {k: v for k, v in wb.items() if k != "useful_gpu" and v > 0}
307
+ if non_zero:
308
+ wb_text = ", ".join(
309
+ f"`{k}={v:.3f}`"
310
+ for k, v in sorted(non_zero.items(), key=lambda kv: kv[1], reverse=True)
311
+ )
312
+ st.caption(f"Waste: useful_gpu=`{wb.get('useful_gpu', 0):.3f}`, {wb_text}")
313
+ if reason and outcome != "evaluated":
314
+ st.caption(f"Reason: {reason}")
315
+
316
+
317
+ # ---------------------------------------------------------------------------
318
+ # Run
319
+ # ---------------------------------------------------------------------------
320
+
321
+
322
+ def _build_command(events_file: Path) -> list[str]:
323
+ cmd: list[str] = [sys.executable, "-u", str(AUTO_TUNE_SCRIPT)]
324
+ if workload_source == "Model id":
325
+ if not model_id.strip():
326
+ raise ValueError("Model id is required.")
327
+ cmd.extend(["--model", model_id.strip()])
328
+ else:
329
+ wp = workload_path.strip()
330
+ if not wp:
331
+ raise ValueError("Workload path is required.")
332
+ full = (REPO_ROOT / wp).resolve() if not Path(wp).is_absolute() else Path(wp)
333
+ if not full.exists():
334
+ raise ValueError(f"Workload not found: {full}")
335
+ cmd.append(str(full))
336
+ cmd.extend([
337
+ "--mode", mode,
338
+ "--steps", str(steps),
339
+ "--max-iterations", str(max_iterations),
340
+ "--early-stop-after", str(early_stop_after),
341
+ "--max-crashes", str(max_crashes),
342
+ "--improvement-threshold", str(improvement_threshold),
343
+ "--events", str(events_file),
344
+ ])
345
+ if mode == "llm-explore":
346
+ cmd.extend(["--candidates-per-iteration", str(candidates_per_iteration)])
347
+ return cmd
348
+
349
+
350
+ def _build_request_body() -> dict:
351
+ """Same config as _build_command, but as a JSON body for POST /auto-tune."""
352
+ body: dict[str, Any] = {
353
+ "mode": mode,
354
+ "steps": steps,
355
+ "max_iterations": max_iterations,
356
+ "early_stop_after": early_stop_after,
357
+ "max_crashes": max_crashes,
358
+ "improvement_threshold": float(improvement_threshold),
359
+ }
360
+ if mode == "llm-explore":
361
+ body["candidates_per_iteration"] = candidates_per_iteration
362
+ if workload_source == "Model id":
363
+ if not model_id.strip():
364
+ raise ValueError("Model id is required.")
365
+ body["model"] = model_id.strip()
366
+ else:
367
+ if not workload_path.strip():
368
+ raise ValueError("Workload path is required.")
369
+ body["workload"] = workload_path.strip()
370
+ return body
371
+
372
+
373
+ def _read_events(path: Path, seen: int) -> tuple[list[dict], int]:
374
+ """Read events from byte position `seen` onward; return (new_events, new_seen)."""
375
+ if not path.exists():
376
+ return [], seen
377
+ try:
378
+ with path.open("r") as f:
379
+ f.seek(seen)
380
+ chunk = f.read()
381
+ new_seen = f.tell()
382
+ except OSError:
383
+ return [], seen
384
+ if not chunk:
385
+ return [], new_seen
386
+ out: list[dict] = []
387
+ # The last line may be partial if the writer is mid-flush; drop it
388
+ # and back the pointer up so we re-read it next tick.
389
+ pieces = chunk.splitlines(keepends=True)
390
+ if pieces and not pieces[-1].endswith("\n"):
391
+ partial = pieces.pop()
392
+ new_seen -= len(partial.encode("utf-8"))
393
+ for line in pieces:
394
+ line = line.strip()
395
+ if not line:
396
+ continue
397
+ try:
398
+ out.append(json.loads(line))
399
+ except json.JSONDecodeError:
400
+ continue
401
+ return out, new_seen
402
+
403
+
404
+ # ---------------------------------------------------------------------------
405
+ # Initial state — empty page
406
+ # ---------------------------------------------------------------------------
407
+
408
+ if not run_pressed:
409
+ st.info(
410
+ "👈 Configure inputs in the sidebar and press **Run auto-tune**. "
411
+ "The script will profile your workload, try MI300X-specific tuning "
412
+ "changes, and report what worked."
413
+ )
414
+ _render_metrics_row(None, None)
415
+ st.subheader("How it works")
416
+ st.markdown(
417
+ """
418
+ 1. **Baseline benchmark** — runs your workload as-is (or as generated
419
+ from `--model`) on the GPU and measures tokens/sec, MFU, HBM peak,
420
+ and a per-bucket waste budget.
421
+ 2. **Iterative tuning** — applies one MI300X-specific change at a
422
+ time (e.g. `bf16`, `batch_size=16`, `TORCH_BLAS_PREFER_HIPBLASLT=1`)
423
+ and re-benchmarks. Keeps changes that beat the current best.
424
+ 3. **Live progress** — every accepted/rejected/crashed candidate
425
+ shows up here as it happens.
426
+ 4. **Final report** — overall improvement, which experiments won,
427
+ how much wastage was recovered, and a diff against the baseline
428
+ workload.
429
+ """
430
+ )
431
+ st.stop()
432
+
433
+
434
+ # ---------------------------------------------------------------------------
435
+ # Run path: launch subprocess + tail events
436
+ # ---------------------------------------------------------------------------
437
+
438
+ st.subheader("Live run")
439
+ header = st.empty()
440
+
441
+ # Validate inputs early so the spinner doesn't show before a clear error
442
+ try:
443
+ if backend_mode == "Local subprocess":
444
+ events_file = Path(tempfile.NamedTemporaryFile(
445
+ prefix="auto_tune_events_", suffix=".ndjson", delete=False
446
+ ).name)
447
+ cmd = _build_command(events_file)
448
+ with header.container():
449
+ st.code(" ".join(cmd), language="bash")
450
+ st.caption(f"Events stream: `{events_file}`")
451
+ else:
452
+ if not backend_url.strip():
453
+ raise ValueError("Backend URL is required for Remote GPU server mode.")
454
+ body = _build_request_body()
455
+ events_file = None
456
+ cmd = None
457
+ with header.container():
458
+ st.code(
459
+ f"POST {backend_url.rstrip('/')}/auto-tune\n"
460
+ + json.dumps(body, indent=2),
461
+ language="bash",
462
+ )
463
+ st.caption(
464
+ "Events stream: SSE from remote server. "
465
+ "All GPU work happens on that host."
466
+ )
467
+ except ValueError as exc:
468
+ st.error(str(exc))
469
+ st.stop()
470
+
471
+ baseline_metrics: dict | None = None
472
+ best_metrics: dict | None = None # most recent accepted iteration's metrics
473
+ final_summary: dict | None = None
474
+ iter_idx_to_container: dict[int, Any] = {}
475
+
476
+ metrics_row = st.empty()
477
+ with metrics_row.container():
478
+ _render_metrics_row(None, None)
479
+
480
+ progress_bar = st.progress(0, text="Starting auto_tune…")
481
+
482
+ iters_section = st.container()
483
+ iters_section.subheader("Iterations")
484
+
485
+ stdout_buffer: list[str] = []
486
+ expected_iters = max(1, max_iterations)
487
+ return_code: int | None = None
488
+
489
+
490
+ def _handle_event(event: dict) -> None:
491
+ """Render one event into the live UI. Mutates module-level state for
492
+ baseline/best_metrics/final_summary so the summary block below has
493
+ the data it needs after the run."""
494
+ global baseline_metrics, best_metrics, final_summary # noqa: PLW0603
495
+ etype = event.get("type")
496
+ if etype == "started":
497
+ st.session_state["started_event"] = event
498
+ elif etype == "baseline":
499
+ baseline_metrics = event["metrics"]
500
+ best_metrics = baseline_metrics
501
+ with metrics_row.container():
502
+ _render_metrics_row(baseline_metrics, best_metrics)
503
+ with iters_section:
504
+ with st.container(border=True):
505
+ st.markdown("**Baseline**")
506
+ sub = st.columns(4)
507
+ sub[0].metric("tokens/sec", f"{baseline_metrics.get('tokens_per_sec', 0):,.0f}")
508
+ sub[1].metric("MFU %", f"{baseline_metrics.get('mfu_pct', 0):.2f}")
509
+ sub[2].metric("HBM peak GB", f"{baseline_metrics.get('hbm_peak_gb', 0):.1f}")
510
+ sub[3].metric("GPU util %", f"{baseline_metrics.get('gpu_util_pct', 0):.1f}")
511
+ wb = baseline_metrics.get("waste_budget") or {}
512
+ non_zero = {k: v for k, v in wb.items() if k != "useful_gpu" and v > 0}
513
+ if non_zero:
514
+ wb_str = ", ".join(
515
+ f"`{k}={v:.3f}`"
516
+ for k, v in sorted(non_zero.items(), key=lambda kv: kv[1], reverse=True)
517
+ )
518
+ st.caption(f"Recoverable waste: {wb_str}")
519
+ elif etype == "iter_start":
520
+ i = event["iteration"]
521
+ with iters_section:
522
+ container = st.container(border=True)
523
+ iter_idx_to_container[i] = container
524
+ n_cand = len(event.get("candidates") or [])
525
+ cand_summary = ", ".join(
526
+ c.get("name", "?") for c in event.get("candidates") or []
527
+ )
528
+ container.markdown(
529
+ f"**Iteration {i}** · {n_cand} candidate{'s' if n_cand != 1 else ''}: {cand_summary}"
530
+ )
531
+ progress_bar.progress(
532
+ min(0.99, (i - 1) / expected_iters),
533
+ text=f"Iteration {i} of up to {expected_iters} — proposed: {cand_summary}",
534
+ )
535
+ elif etype == "candidate":
536
+ i = event["iteration"]
537
+ container = iter_idx_to_container.get(i)
538
+ if container is not None:
539
+ with container:
540
+ _render_candidate_card(event)
541
+ progress_bar.progress(
542
+ min(0.99, (i - 1) / expected_iters + 0.05),
543
+ text=f"Iter {i} · candidate {event.get('candidate_index')}/"
544
+ f"{event.get('n_candidates')}: {event.get('name')} "
545
+ f"({event.get('outcome')})",
546
+ )
547
+ elif etype == "merge_attempt":
548
+ i = event["iteration"]
549
+ container = iter_idx_to_container.get(i)
550
+ if container is not None:
551
+ with container:
552
+ outcome = event.get("outcome", "?")
553
+ names = ", ".join(event.get("candidate_names") or [])
554
+ if outcome == "wins":
555
+ delta = event.get("delta_vs_best", 0)
556
+ container.success(
557
+ f"🔗 MERGE WINS — combined ({names}) hit Δ {delta:+.2f}% "
558
+ f"(beats individual best `{event.get('individual_best_name')}`)"
559
+ )
560
+ elif outcome == "lost":
561
+ delta = event.get("delta_vs_best", 0)
562
+ container.info(
563
+ f"🔗 Merge tested ({names}) — Δ {delta:+.2f}%, didn't beat "
564
+ f"individual `{event.get('individual_best_name')}`"
565
+ )
566
+ elif outcome == "crashed":
567
+ container.warning(f"💥 Merge crashed: {names}")
568
+ else:
569
+ container.info(f"🔗 Merge skipped: {event.get('reason', '?')}")
570
+ elif etype == "iter_done":
571
+ i = event["iteration"]
572
+ container = iter_idx_to_container.get(i)
573
+ outcome = event.get("outcome")
574
+ if container is not None:
575
+ with container:
576
+ if outcome == "accepted":
577
+ container.success(
578
+ f"✅ ACCEPTED — `{event.get('winner_name')}` "
579
+ f"(Δ {event.get('winner_delta', 0):+.2f}%)"
580
+ )
581
+ else:
582
+ container.warning(
583
+ f"⏭️ ALL REJECTED — best was `{event.get('winner_name')}` "
584
+ f"(Δ {event.get('winner_delta', 0):+.2f}%, below threshold)"
585
+ )
586
+ if outcome == "accepted" and event.get("best_metrics"):
587
+ best_metrics = event["best_metrics"]
588
+ with metrics_row.container():
589
+ _render_metrics_row(baseline_metrics, best_metrics)
590
+ progress_bar.progress(
591
+ min(0.99, i / expected_iters),
592
+ text=f"Iter {i} done — best so far: {event.get('best_tps', 0):,.0f} tok/s",
593
+ )
594
+ elif etype == "summary":
595
+ final_summary = event
596
+ progress_bar.progress(1.0, text="Auto-tune complete")
597
+ elif etype == "error":
598
+ st.error(event.get("message", "unknown error"))
599
+ elif etype == "process_exit":
600
+ rc = event.get("returncode", "?")
601
+ st.error(
602
+ f"Backend subprocess exited (code {rc}): "
603
+ + event.get("message", "")
604
+ )
605
+
606
+
607
+ # ---- Source the events from either local subprocess or remote SSE ----
608
+ if backend_mode == "Local subprocess":
609
+ proc = subprocess.Popen(
610
+ cmd,
611
+ stdout=subprocess.PIPE,
612
+ stderr=subprocess.STDOUT,
613
+ text=True,
614
+ bufsize=1,
615
+ cwd=str(REPO_ROOT),
616
+ env={**os.environ},
617
+ )
618
+
619
+ seen_bytes = 0
620
+ try:
621
+ while True:
622
+ # Pull stdout incrementally so the user sees the raw script log too
623
+ while True:
624
+ try:
625
+ line = proc.stdout.readline() if proc.stdout else ""
626
+ except Exception:
627
+ line = ""
628
+ if not line:
629
+ break
630
+ stdout_buffer.append(line.rstrip())
631
+
632
+ new_events, seen_bytes = _read_events(events_file, seen_bytes)
633
+ for event in new_events:
634
+ _handle_event(event)
635
+
636
+ if proc.poll() is not None and not new_events:
637
+ new_events, seen_bytes = _read_events(events_file, seen_bytes)
638
+ if not new_events:
639
+ break
640
+ time.sleep(0.4)
641
+ finally:
642
+ if proc.poll() is None:
643
+ proc.terminate()
644
+ try:
645
+ proc.wait(timeout=3)
646
+ except subprocess.TimeoutExpired:
647
+ proc.kill()
648
+ return_code = proc.returncode
649
+ else:
650
+ # Remote SSE: stream from the FastAPI server
651
+ url = backend_url.rstrip("/") + "/auto-tune"
652
+ try:
653
+ response = requests.post(
654
+ url,
655
+ json=body,
656
+ stream=True,
657
+ timeout=(10, None), # connect timeout 10s, read timeout indefinite
658
+ headers={"Accept": "text/event-stream"},
659
+ )
660
+ response.raise_for_status()
661
+ except requests.RequestException as exc:
662
+ st.error(f"Failed to reach backend at {url}: {exc}")
663
+ st.stop()
664
+
665
+ try:
666
+ for raw_line in response.iter_lines(decode_unicode=True):
667
+ if not raw_line:
668
+ continue
669
+ stdout_buffer.append(raw_line)
670
+ # SSE framing: lines starting with `data: <json>`
671
+ if raw_line.startswith("data:"):
672
+ payload = raw_line[len("data:"):].strip()
673
+ try:
674
+ event = json.loads(payload)
675
+ except json.JSONDecodeError:
676
+ continue
677
+ _handle_event(event)
678
+ if event.get("type") == "summary":
679
+ # Final event; the server may still send a process_exit
680
+ # but we can stop blocking the UI.
681
+ pass
682
+ elif raw_line.startswith("event:") or raw_line.startswith(":"):
683
+ # SSE event-name line or comment — ignored
684
+ continue
685
+ except requests.RequestException as exc:
686
+ st.error(f"Stream interrupted: {exc}")
687
+ return_code = 0 if final_summary is not None else 1
688
+
689
+ # ---------------------------------------------------------------------------
690
+ # Final summary
691
+ # ---------------------------------------------------------------------------
692
+
693
+ st.divider()
694
+ st.subheader("Summary")
695
+
696
+ if final_summary is None:
697
+ st.error(
698
+ f"Auto-tune subprocess exited with code {return_code} but emitted "
699
+ "no summary event. Check the raw stdout below for details."
700
+ )
701
+ else:
702
+ base_tps = float(final_summary.get("baseline_tps") or 0)
703
+ best_tps = float(final_summary.get("best_tps") or 0)
704
+ improvement_pct = float(final_summary.get("improvement_pct") or 0)
705
+ base = final_summary.get("baseline_metrics") or {}
706
+ best = final_summary.get("best_metrics") or {}
707
+
708
+ summary_cols = st.columns(4)
709
+ summary_cols[0].metric(
710
+ "Baseline tokens/sec",
711
+ f"{base_tps:,.0f}",
712
+ )
713
+ summary_cols[1].metric(
714
+ "Best tokens/sec",
715
+ f"{best_tps:,.0f}",
716
+ delta=f"{best_tps - base_tps:+,.0f}",
717
+ )
718
+ summary_cols[2].metric(
719
+ "Improvement",
720
+ f"{improvement_pct:+.2f}%",
721
+ )
722
+ summary_cols[3].metric(
723
+ "MFU baseline → best",
724
+ f"{base.get('mfu_pct', 0):.1f} → {best.get('mfu_pct', 0):.1f} %",
725
+ )
726
+
727
+ chart = _waste_chart(base, best)
728
+ if chart is not None:
729
+ st.markdown("**Waste reduction (seconds/step, by bucket)**")
730
+ st.altair_chart(chart, use_container_width=True)
731
+ # Per-bucket reduction table
732
+ diff_rows = []
733
+ bwb = base.get("waste_budget") or {}
734
+ cwb = best.get("waste_budget") or {}
735
+ for bucket in WASTE_BUCKETS:
736
+ bv = float(bwb.get(bucket, 0.0))
737
+ cv = float(cwb.get(bucket, 0.0))
738
+ if bv == 0 and cv == 0:
739
+ continue
740
+ diff_rows.append({
741
+ "bucket": bucket,
742
+ "baseline (s)": round(bv, 4),
743
+ "best (s)": round(cv, 4),
744
+ "Δ (s)": round(cv - bv, 4),
745
+ })
746
+ if diff_rows:
747
+ st.dataframe(pd.DataFrame(diff_rows), use_container_width=True)
748
+
749
+ accepted = final_summary.get("accepted") or []
750
+ rejected = final_summary.get("rejected") or []
751
+
752
+ col_a, col_r = st.columns(2)
753
+ with col_a:
754
+ st.markdown(f"**✅ Accepted ({len(accepted)})**")
755
+ if accepted:
756
+ st.dataframe(
757
+ pd.DataFrame(accepted)[["name", "tps", "delta_pct"]].rename(
758
+ columns={"tps": "tokens/sec", "delta_pct": "Δ %"}
759
+ ),
760
+ use_container_width=True,
761
+ hide_index=True,
762
+ )
763
+ else:
764
+ st.caption("(no experiments accepted)")
765
+ with col_r:
766
+ st.markdown(f"**❌ Rejected ({len(rejected)})**")
767
+ if rejected:
768
+ st.dataframe(
769
+ pd.DataFrame(rejected),
770
+ use_container_width=True,
771
+ hide_index=True,
772
+ )
773
+ else:
774
+ st.caption("(none)")
775
+
776
+ env_vars = final_summary.get("best_env_vars") or {}
777
+ if env_vars:
778
+ st.markdown("**Required env vars for best config**")
779
+ st.code("\n".join(f"export {k}={v}" for k, v in env_vars.items()), language="bash")
780
+
781
+ best_path = final_summary.get("best_workload_path")
782
+ base_path = final_summary.get("baseline_workload_path")
783
+ if best_path and base_path:
784
+ st.markdown("**Best workload script**")
785
+ st.code(f"diff {base_path} {best_path}", language="bash")
786
+ try:
787
+ best_text = Path(best_path).read_text()
788
+ st.download_button(
789
+ "⬇️ Download best.py",
790
+ data=best_text,
791
+ file_name="best.py",
792
+ mime="text/x-python",
793
+ )
794
+ except OSError:
795
+ pass
796
+
797
+ # Raw stdout always at the bottom for debugging
798
+ with st.expander("Raw subprocess output"):
799
+ st.code("\n".join(stdout_buffer), language="text")
workloads/_runtime.py CHANGED
@@ -86,6 +86,22 @@ def parse_runtime_args() -> RuntimeArgs:
86
  )
87
 
88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89
  def emit_torch_profile(
90
  path: str,
91
  *,
@@ -94,13 +110,16 @@ def emit_torch_profile(
94
  per_device_batch: int,
95
  grad_accum: int = 1,
96
  seq_len_cap: int = 512,
 
97
  ) -> None:
98
  """Write the smallest torch_profile-shape JSON profile_parser will read.
99
 
100
  profile_parser._read_torch_profile looks for these top-level fields under
101
  ``metadata``: tokens_per_sec, mfu_pct, step_time_seconds, pytorch_version.
102
- We supply the first three and pytorch_version (mfu_pct is optional and
103
- estimated downstream).
 
 
104
 
105
  No-ops when ``path`` is empty (script run outside goblin_runner.sh) or
106
  when ``n_steps`` is 0 (training crashed before finishing a step).
@@ -113,43 +132,21 @@ def emit_torch_profile(
113
  global_batch = max(1, per_device_batch) * max(1, grad_accum)
114
  approx_tokens = n_steps * global_batch * seq_len_cap
115
  tokens_per_sec = approx_tokens / elapsed if elapsed > 0 else 0.0
116
- payload = {
117
- "metadata": {
118
- "tokens_per_sec": round(tokens_per_sec, 2),
119
- "step_time_seconds": round(elapsed / n_steps, 4),
120
- "pytorch_version": torch.__version__,
121
- "n_steps": n_steps,
122
- }
123
  }
 
 
 
 
 
 
124
  with open(path, "w") as f:
125
  json.dump(payload, f)
126
  except Exception as exc: # pragma: no cover — diagnostic only
127
  # Don't tank the run on a profile-emit failure; the agent will
128
  # just see "fake" metrics for this step instead of "live".
129
  print(f"[workloads._runtime] failed to write {path}: {exc}")
130
-
131
-
132
- def trainer_tokenizer_kwargs(trainer_cls, tokenizer) -> dict:
133
- """Return the right kwarg for handing a tokenizer to ``Trainer``.
134
-
135
- transformers ≥ 4.46 renamed ``tokenizer=`` to ``processing_class=`` (the
136
- old name is still accepted with a DeprecationWarning, but a future
137
- release drops it). Older versions only know ``tokenizer=``. We
138
- introspect ``Trainer.__init__`` so the workloads run on either
139
- generation without pinning a transformers version.
140
-
141
- Use site:
142
-
143
- trainer = Trainer(
144
- model=model,
145
- args=training_args,
146
- train_dataset=dataset,
147
- **trainer_tokenizer_kwargs(Trainer, tokenizer),
148
- data_collator=_toy_collate,
149
- )
150
- """
151
- import inspect
152
-
153
- if "processing_class" in inspect.signature(trainer_cls.__init__).parameters:
154
- return {"processing_class": tokenizer}
155
- return {"tokenizer": tokenizer}
 
86
  )
87
 
88
 
89
+ # MI300X (CDNA3) peak throughput, dense, bf16/fp16 — both arrive at the
90
+ # same number on this arch since the matrix engine is the same. Source:
91
+ # AMD Instinct MI300X datasheet. With sparsity it's ~2.6 PFLOPS, but
92
+ # transformers training rarely hits the sparse path so we use dense as
93
+ # the realistic peak.
94
+ _MI300X_PEAK_FLOPS_DENSE_BF16 = 1.307e15
95
+
96
+ # FLOPs per token for forward + backward. The standard 6N approximation
97
+ # (forward 2N + backward 4N for full fine-tuning) slightly overestimates
98
+ # LoRA — pure LoRA backward only computes weight gradients for the small
99
+ # adapter matrices, not the frozen base — so true LoRA flops/token is
100
+ # closer to 4N. We use 6N as the conventional choice and accept a ~30%
101
+ # pessimistic MFU for LoRA. Still useful as a relative metric run-to-run.
102
+ _FLOPS_PER_TOKEN_FACTOR = 6
103
+
104
+
105
  def emit_torch_profile(
106
  path: str,
107
  *,
 
110
  per_device_batch: int,
111
  grad_accum: int = 1,
112
  seq_len_cap: int = 512,
113
+ model_params: int = 0,
114
  ) -> None:
115
  """Write the smallest torch_profile-shape JSON profile_parser will read.
116
 
117
  profile_parser._read_torch_profile looks for these top-level fields under
118
  ``metadata``: tokens_per_sec, mfu_pct, step_time_seconds, pytorch_version.
119
+
120
+ `model_params` is optional — pass `sum(p.numel() for p in
121
+ model.parameters())` from the workload to get a populated `mfu_pct`.
122
+ Without it, mfu_pct stays unset (profile_parser will default to 0).
123
 
124
  No-ops when ``path`` is empty (script run outside goblin_runner.sh) or
125
  when ``n_steps`` is 0 (training crashed before finishing a step).
 
132
  global_batch = max(1, per_device_batch) * max(1, grad_accum)
133
  approx_tokens = n_steps * global_batch * seq_len_cap
134
  tokens_per_sec = approx_tokens / elapsed if elapsed > 0 else 0.0
135
+ metadata = {
136
+ "tokens_per_sec": round(tokens_per_sec, 2),
137
+ "step_time_seconds": round(elapsed / n_steps, 4),
138
+ "pytorch_version": torch.__version__,
139
+ "n_steps": n_steps,
 
 
140
  }
141
+ if model_params > 0 and tokens_per_sec > 0:
142
+ flops_per_token = _FLOPS_PER_TOKEN_FACTOR * model_params
143
+ mfu_pct = (flops_per_token * tokens_per_sec) / _MI300X_PEAK_FLOPS_DENSE_BF16 * 100
144
+ metadata["mfu_pct"] = round(mfu_pct, 2)
145
+ metadata["model_params"] = model_params
146
+ payload = {"metadata": metadata}
147
  with open(path, "w") as f:
148
  json.dump(payload, f)
149
  except Exception as exc: # pragma: no cover — diagnostic only
150
  # Don't tank the run on a profile-emit failure; the agent will
151
  # just see "fake" metrics for this step instead of "live".
152
  print(f"[workloads._runtime] failed to write {path}: {exc}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
workloads/scenarios/README.md ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test workloads
2
+
3
+ Hand-built training scripts that demo different misconfiguration patterns.
4
+ The agent should produce **different audit reports** for each one — useful
5
+ for stress-testing the KB, for live demos with variety, and for sanity-
6
+ checking that recommendations track the actual configuration.
7
+
8
+ | File | Headline misconfiguration | KB rules expected to fire |
9
+ |---|---|---|
10
+ | `train_qwen_full_ft.py` | Full fine-tune (no LoRA), fp32, no gradient checkpointing | `precision.fp32_default_wastes_matrix_cores`, `memory.gradient_checkpointing_for_long_seq`, `attention.sdpa_over_eager` |
11
+ | `train_qwen_long_context.py` | seq_len=8192, no gradient checkpointing, fp16, eager attention | `memory.gradient_checkpointing_for_long_seq`, `attention.flash_rocm_over_eager`, `precision.bf16_over_fp16_on_mi300x` |
12
+ | `train_qwen_distributed_bad.py` | One process driving 8 GPUs (ROCm launch serialization antipattern) | `collectives.one_process_per_gpu`, `env.nccl_min_nchannels`, `env.numa_auto_balancing_disable` |
13
+ | `train_qwen_bnb.py` | Uses bitsandbytes 8-bit Adam (not officially ROCm-supported) | `optimizer.bitsandbytes_not_supported_warning`, `precision.bf16_over_fp16_on_mi300x` |
14
+ | `train_qwen_well_tuned.py` | Already optimized (bf16, flash_rocm, prefetch=4, persistent_workers) | None or only minor suggestions — sanity check that the agent says "nothing to do" instead of inventing problems |
15
+
16
+ Run any of them through the agent:
17
+
18
+ ```bash
19
+ export GOBLIN_AGENT_BACKEND=qwen-vllm
20
+ export GOBLIN_QWEN_VLLM_URL=http://localhost:8001/v1
21
+ export GOBLIN_QWEN_VLLM_MODEL=Qwen/Qwen3-32B
22
+ python -m agent workloads/scenarios/train_qwen_long_context.py
23
+ ```
24
+
25
+ These are **AST-parseable, optionally executable**. Each one redirects on
26
+ the same fix — `parse_config` only walks the AST, so even if the script
27
+ doesn't actually train cleanly, the agent's audit still works. If you
28
+ want rocprofv3 to capture real numbers (vs FakeRunner fallback), the
29
+ scripts marked "executable: yes" below are runnable; the others rely on
30
+ FakeRunner replay.
workloads/scenarios/train_qwen_bnb.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test workload: bitsandbytes 8-bit Adam — the classic "ported from NVIDIA"
2
+ # script that doesn't realize bitsandbytes isn't officially supported on ROCm.
3
+ #
4
+ # Headline misconfigurations the agent should catch:
5
+ # - optimizer.bitsandbytes_not_supported_warning (optim="paged_adamw_8bit")
6
+ # - precision.bf16_over_fp16_on_mi300x (fp16=True)
7
+ # - attention.flash_rocm_over_eager (attn_implementation="eager")
8
+ # - data.dataloader_workers_zero (num_workers=0)
9
+ #
10
+ # Executable: not on ROCm — bitsandbytes will fail at import or on the
11
+ # first 8-bit GEMM. The whole point of this fixture is that the agent
12
+ # warns the user BEFORE they hit that wall. AST parse always works.
13
+
14
+ import os
15
+ import sys
16
+ import time
17
+
18
+ sys.path.insert(
19
+ 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
20
+ )
21
+
22
+ import torch
23
+ from datasets import load_dataset
24
+ from peft import LoraConfig, get_peft_model
25
+ from torch.utils.data import DataLoader
26
+ from transformers import (
27
+ AutoModelForCausalLM,
28
+ AutoTokenizer,
29
+ Trainer,
30
+ TrainingArguments,
31
+ )
32
+
33
+ from workloads._runtime import emit_torch_profile, parse_runtime_args
34
+
35
+ _runtime = parse_runtime_args()
36
+
37
+ os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
38
+ HF_TOKEN = os.environ["HF_TOKEN"]
39
+
40
+ MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
41
+
42
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
43
+ model = AutoModelForCausalLM.from_pretrained(
44
+ MODEL_ID,
45
+ torch_dtype=torch.float16,
46
+ attn_implementation="eager",
47
+ token=HF_TOKEN,
48
+ load_in_8bit=True, # bitsandbytes load — ROCm support unofficial
49
+ )
50
+
51
+ lora_config = LoraConfig(
52
+ r=16,
53
+ lora_alpha=32,
54
+ target_modules=["q_proj", "v_proj"],
55
+ lora_dropout=0.05,
56
+ bias="none",
57
+ task_type="CAUSAL_LM",
58
+ )
59
+ model = get_peft_model(model, lora_config)
60
+
61
+ dataset = load_dataset("yahma/alpaca-cleaned", split="train")
62
+
63
+ train_loader = DataLoader(
64
+ dataset,
65
+ batch_size=4,
66
+ num_workers=0,
67
+ pin_memory=False,
68
+ persistent_workers=False,
69
+ )
70
+
71
+ _RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
72
+
73
+ training_args = TrainingArguments(
74
+ output_dir="./out",
75
+ per_device_train_batch_size=4,
76
+ gradient_accumulation_steps=8,
77
+ num_train_epochs=3,
78
+ learning_rate=2e-4,
79
+ warmup_steps=100,
80
+ fp16=True,
81
+ optim="paged_adamw_8bit", # bitsandbytes optimizer — ROCm unofficial
82
+ logging_steps=10,
83
+ save_steps=500,
84
+ dataloader_num_workers=0,
85
+ dataloader_pin_memory=False,
86
+ gradient_checkpointing=False,
87
+ torch_compile=False,
88
+ report_to="none",
89
+ push_to_hub=False,
90
+ max_steps=_RUNTIME_MAX_STEPS,
91
+ )
92
+
93
+ trainer = Trainer(
94
+ model=model,
95
+ args=training_args,
96
+ train_dataset=dataset,
97
+ tokenizer=tokenizer,
98
+ )
99
+
100
+ if __name__ == "__main__":
101
+ _t0 = time.time()
102
+ trainer.train()
103
+ emit_torch_profile(
104
+ _runtime.torch_profile_out,
105
+ elapsed=time.time() - _t0,
106
+ n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
107
+ per_device_batch=training_args.per_device_train_batch_size,
108
+ grad_accum=training_args.gradient_accumulation_steps,
109
+ seq_len_cap=512,
110
+ )
workloads/scenarios/train_qwen_distributed_bad.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test workload: multi-GPU LoRA with the "one-process-many-GPUs" antipattern.
2
+ #
3
+ # A single Python process (no torchrun, no accelerate launch) tries to drive
4
+ # multiple MI300X devices via DataParallel. ROCm serializes kernel launches
5
+ # across devices when one process owns multiple HIP streams, so every
6
+ # collective becomes a launch-queue bottleneck.
7
+ #
8
+ # Headline misconfigurations the agent should catch:
9
+ # - collectives.one_process_per_gpu (single launcher, multiple devices)
10
+ # - env.nccl_min_nchannels (NCCL_MIN_NCHANNELS not set)
11
+ # - env.numa_auto_balancing_disable (NUMA balancing left on by default)
12
+ # - precision.bf16_over_fp16_on_mi300x (fp16=True)
13
+ # - attention.flash_rocm_over_eager
14
+ #
15
+ # Executable: requires multiple MI300X visible to one process. Even when
16
+ # runnable, ROCm warns on the launch serialization. AST parse is fine.
17
+
18
+ import os
19
+ import sys
20
+ import time
21
+
22
+ sys.path.insert(
23
+ 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
24
+ )
25
+
26
+ import torch
27
+ import torch.nn as nn
28
+ from datasets import load_dataset
29
+ from peft import LoraConfig, get_peft_model
30
+ from torch.utils.data import DataLoader
31
+ from transformers import (
32
+ AutoModelForCausalLM,
33
+ AutoTokenizer,
34
+ Trainer,
35
+ TrainingArguments,
36
+ )
37
+
38
+ from workloads._runtime import emit_torch_profile, parse_runtime_args
39
+
40
+ _runtime = parse_runtime_args()
41
+
42
+ os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
43
+ HF_TOKEN = os.environ["HF_TOKEN"]
44
+ os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
45
+ # Notably absent: NCCL_MIN_NCHANNELS, GOBLIN_HINT_NUMA_AUTO_BALANCING.
46
+
47
+ MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
48
+
49
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
50
+ model = AutoModelForCausalLM.from_pretrained(
51
+ MODEL_ID,
52
+ torch_dtype=torch.float16,
53
+ attn_implementation="eager",
54
+ token=HF_TOKEN,
55
+ )
56
+
57
+ lora_config = LoraConfig(
58
+ r=16,
59
+ lora_alpha=32,
60
+ target_modules=["q_proj", "v_proj"],
61
+ lora_dropout=0.05,
62
+ bias="none",
63
+ task_type="CAUSAL_LM",
64
+ )
65
+ model = get_peft_model(model, lora_config)
66
+
67
+ # THE ANTIPATTERN: DataParallel from a single process across all visible GPUs.
68
+ # Production code should use torchrun --nproc_per_node=N or accelerate launch.
69
+ if torch.cuda.device_count() > 1:
70
+ model = nn.DataParallel(model)
71
+
72
+ dataset = load_dataset("yahma/alpaca-cleaned", split="train")
73
+
74
+ train_loader = DataLoader(
75
+ dataset,
76
+ batch_size=8,
77
+ num_workers=4,
78
+ pin_memory=True,
79
+ persistent_workers=True,
80
+ )
81
+
82
+ _RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
83
+
84
+ training_args = TrainingArguments(
85
+ output_dir="./out",
86
+ per_device_train_batch_size=8,
87
+ gradient_accumulation_steps=4,
88
+ num_train_epochs=1,
89
+ learning_rate=2e-4,
90
+ warmup_steps=100,
91
+ fp16=True,
92
+ optim="adamw_torch",
93
+ logging_steps=10,
94
+ save_steps=500,
95
+ dataloader_num_workers=4,
96
+ dataloader_pin_memory=True,
97
+ dataloader_persistent_workers=True,
98
+ gradient_checkpointing=False,
99
+ torch_compile=False,
100
+ report_to="none",
101
+ push_to_hub=False,
102
+ max_steps=_RUNTIME_MAX_STEPS,
103
+ )
104
+
105
+ trainer = Trainer(
106
+ model=model,
107
+ args=training_args,
108
+ train_dataset=dataset,
109
+ tokenizer=tokenizer,
110
+ )
111
+
112
+ if __name__ == "__main__":
113
+ _t0 = time.time()
114
+ trainer.train()
115
+ emit_torch_profile(
116
+ _runtime.torch_profile_out,
117
+ elapsed=time.time() - _t0,
118
+ n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
119
+ per_device_batch=training_args.per_device_train_batch_size,
120
+ grad_accum=training_args.gradient_accumulation_steps,
121
+ seq_len_cap=512,
122
+ )
workloads/scenarios/train_qwen_full_ft.py ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test workload: full fine-tune (no LoRA), fp32, no gradient checkpointing.
2
+ #
3
+ # Headline misconfigurations the agent should catch:
4
+ # - precision.fp32_default_wastes_matrix_cores (fp16=False, bf16=False)
5
+ # - memory.gradient_checkpointing_for_long_seq (gradient_checkpointing=False, seq_len=2048)
6
+ # - attention.sdpa_over_eager OR flash_rocm (attn_implementation="eager")
7
+ # - data.dataloader_workers_zero (num_workers=0)
8
+ # - data.pin_memory_false (pin_memory=False)
9
+ #
10
+ # Executable: not really — full FT of a 7B model OOMs on a single MI300X
11
+ # without LoRA. AST parse is fine; rocprofv3 will fail and FakeRunner kicks in.
12
+
13
+ import os
14
+ import sys
15
+ import time
16
+
17
+ # Bootstrap repo root so `from workloads._runtime import ...` resolves
18
+ # regardless of the cwd goblin_runner.sh launches us from.
19
+ sys.path.insert(
20
+ 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
21
+ )
22
+
23
+ import torch
24
+ from datasets import load_dataset
25
+ from torch.utils.data import DataLoader
26
+ from transformers import (
27
+ AutoModelForCausalLM,
28
+ AutoTokenizer,
29
+ Trainer,
30
+ TrainingArguments,
31
+ )
32
+
33
+ from workloads._runtime import emit_torch_profile, parse_runtime_args
34
+
35
+ _runtime = parse_runtime_args()
36
+
37
+ os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
38
+ HF_TOKEN = os.environ["HF_TOKEN"]
39
+ os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
40
+
41
+ MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
42
+
43
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
44
+ model = AutoModelForCausalLM.from_pretrained(
45
+ MODEL_ID,
46
+ torch_dtype=torch.float32, # fp32 — wastes CDNA3 matrix cores
47
+ attn_implementation="eager",
48
+ token=HF_TOKEN,
49
+ )
50
+
51
+ dataset = load_dataset("yahma/alpaca-cleaned", split="train")
52
+
53
+ train_loader = DataLoader(
54
+ dataset,
55
+ batch_size=2, # tiny because fp32 + full FT eats HBM
56
+ num_workers=0,
57
+ pin_memory=False,
58
+ persistent_workers=False,
59
+ )
60
+
61
+ # Literal kwargs so parse_config picks them all up (see canonical workload).
62
+ _RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
63
+
64
+ training_args = TrainingArguments(
65
+ output_dir="./out",
66
+ per_device_train_batch_size=2,
67
+ gradient_accumulation_steps=16,
68
+ num_train_epochs=1,
69
+ learning_rate=1e-5,
70
+ warmup_steps=200,
71
+ fp16=False, # no precision flag → fp32 path
72
+ bf16=False,
73
+ optim="adamw_torch",
74
+ max_seq_length=2048,
75
+ logging_steps=10,
76
+ save_steps=500,
77
+ dataloader_num_workers=0,
78
+ dataloader_pin_memory=False,
79
+ gradient_checkpointing=False, # at seq=2048 this leaves a lot of HBM tied up
80
+ torch_compile=False,
81
+ report_to="none",
82
+ push_to_hub=False,
83
+ max_steps=_RUNTIME_MAX_STEPS,
84
+ )
85
+
86
+ trainer = Trainer(
87
+ model=model,
88
+ args=training_args,
89
+ train_dataset=dataset,
90
+ tokenizer=tokenizer,
91
+ )
92
+
93
+ if __name__ == "__main__":
94
+ _t0 = time.time()
95
+ trainer.train()
96
+ emit_torch_profile(
97
+ _runtime.torch_profile_out,
98
+ elapsed=time.time() - _t0,
99
+ n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
100
+ per_device_batch=training_args.per_device_train_batch_size,
101
+ grad_accum=training_args.gradient_accumulation_steps,
102
+ seq_len_cap=2048,
103
+ )
workloads/scenarios/train_qwen_long_context.py ADDED
@@ -0,0 +1,111 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test workload: long-context LoRA fine-tune.
2
+ #
3
+ # Headline misconfigurations the agent should catch:
4
+ # - memory.gradient_checkpointing_for_long_seq (seq_len=8192, gradient_checkpointing=False)
5
+ # - attention.flash_rocm_over_eager (attn_implementation="eager" — KILLS at seq 8K)
6
+ # - precision.bf16_over_fp16_on_mi300x (fp16=True)
7
+ # - data.persistent_workers_false (persistent_workers=False)
8
+ # - data.pin_memory_false (pin_memory=False)
9
+ #
10
+ # Executable: only with flash attention installed. Without it, eager attention
11
+ # at seq=8192 blows up HBM. AST parse always works.
12
+
13
+ import os
14
+ import sys
15
+ import time
16
+
17
+ sys.path.insert(
18
+ 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
19
+ )
20
+
21
+ import torch
22
+ from datasets import load_dataset
23
+ from peft import LoraConfig, get_peft_model
24
+ from torch.utils.data import DataLoader
25
+ from transformers import (
26
+ AutoModelForCausalLM,
27
+ AutoTokenizer,
28
+ Trainer,
29
+ TrainingArguments,
30
+ )
31
+
32
+ from workloads._runtime import emit_torch_profile, parse_runtime_args
33
+
34
+ _runtime = parse_runtime_args()
35
+
36
+ os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
37
+ HF_TOKEN = os.environ["HF_TOKEN"]
38
+ os.environ["MIOPEN_FIND_MODE"] = "3"
39
+
40
+ MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
41
+
42
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
43
+ model = AutoModelForCausalLM.from_pretrained(
44
+ MODEL_ID,
45
+ torch_dtype=torch.float16,
46
+ attn_implementation="eager", # at seq=8192 this is catastrophic
47
+ token=HF_TOKEN,
48
+ )
49
+
50
+ lora_config = LoraConfig(
51
+ r=32,
52
+ lora_alpha=64,
53
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
54
+ lora_dropout=0.05,
55
+ bias="none",
56
+ task_type="CAUSAL_LM",
57
+ )
58
+ model = get_peft_model(model, lora_config)
59
+
60
+ dataset = load_dataset("yahma/alpaca-cleaned", split="train")
61
+
62
+ train_loader = DataLoader(
63
+ dataset,
64
+ batch_size=1, # forced down by long context
65
+ num_workers=4, # workers OK; problem is elsewhere
66
+ pin_memory=False,
67
+ persistent_workers=False,
68
+ )
69
+
70
+ _RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
71
+
72
+ training_args = TrainingArguments(
73
+ output_dir="./out",
74
+ per_device_train_batch_size=1,
75
+ gradient_accumulation_steps=32,
76
+ num_train_epochs=2,
77
+ learning_rate=2e-4,
78
+ warmup_steps=50,
79
+ fp16=True, # bf16 is the right call on CDNA3
80
+ optim="adamw_torch",
81
+ max_seq_length=8192, # the long-context part
82
+ logging_steps=10,
83
+ save_steps=500,
84
+ dataloader_num_workers=4,
85
+ dataloader_pin_memory=False,
86
+ dataloader_persistent_workers=False,
87
+ gradient_checkpointing=False, # missing — at seq 8K activations dominate HBM
88
+ torch_compile=False,
89
+ report_to="none",
90
+ push_to_hub=False,
91
+ max_steps=_RUNTIME_MAX_STEPS,
92
+ )
93
+
94
+ trainer = Trainer(
95
+ model=model,
96
+ args=training_args,
97
+ train_dataset=dataset,
98
+ tokenizer=tokenizer,
99
+ )
100
+
101
+ if __name__ == "__main__":
102
+ _t0 = time.time()
103
+ trainer.train()
104
+ emit_torch_profile(
105
+ _runtime.torch_profile_out,
106
+ elapsed=time.time() - _t0,
107
+ n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
108
+ per_device_batch=training_args.per_device_train_batch_size,
109
+ grad_accum=training_args.gradient_accumulation_steps,
110
+ seq_len_cap=8192,
111
+ )
workloads/scenarios/train_qwen_well_tuned.py ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test workload: already well-tuned LoRA fine-tune.
2
+ #
3
+ # This is the negative-control fixture. The agent should produce a SHORT
4
+ # audit with few or no recommendations — the goal is to catch hallucinated
5
+ # rule applications. If GPU Goblin invents problems on a clean config, the
6
+ # audit isn't trustworthy.
7
+ #
8
+ # What's already correct here:
9
+ # - precision = bf16
10
+ # - attention_impl = flash_rocm (Optimum-AMD path)
11
+ # - dataloader_num_workers = 8, pin_memory=True, prefetch_factor=4,
12
+ # persistent_workers=True
13
+ # - gradient_checkpointing = True (long context)
14
+ # - torch_compile = True
15
+ # - NCCL_MIN_NCHANNELS, HSA_FORCE_FINE_GRAIN_PCIE both set
16
+ #
17
+ # Executable: yes (fastest path; meant to actually run cleanly).
18
+
19
+ import os
20
+ import sys
21
+ import time
22
+
23
+ sys.path.insert(
24
+ 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
25
+ )
26
+
27
+ import torch
28
+ from datasets import load_dataset
29
+ from peft import LoraConfig, get_peft_model
30
+ from torch.utils.data import DataLoader
31
+ from transformers import (
32
+ AutoModelForCausalLM,
33
+ AutoTokenizer,
34
+ Trainer,
35
+ TrainingArguments,
36
+ )
37
+
38
+ from workloads._runtime import emit_torch_profile, parse_runtime_args
39
+
40
+ _runtime = parse_runtime_args()
41
+
42
+ os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
43
+ HF_TOKEN = os.environ["HF_TOKEN"]
44
+
45
+ # All the env knobs from the KB, set correctly.
46
+ os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
47
+ os.environ["MIOPEN_FIND_MODE"] = "3"
48
+ os.environ["NCCL_MIN_NCHANNELS"] = "112"
49
+
50
+ MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
51
+
52
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
53
+ model = AutoModelForCausalLM.from_pretrained(
54
+ MODEL_ID,
55
+ torch_dtype=torch.bfloat16, # bf16 ✓
56
+ attn_implementation="flash_attention_2", # parser maps to flash / flash_rocm
57
+ token=HF_TOKEN,
58
+ )
59
+
60
+ # Compile path on — Qwen is on the eager-friendly list.
61
+ model = torch.compile(model, mode="reduce-overhead")
62
+
63
+ lora_config = LoraConfig(
64
+ r=16,
65
+ lora_alpha=32,
66
+ target_modules=["q_proj", "v_proj"],
67
+ lora_dropout=0.05,
68
+ bias="none",
69
+ task_type="CAUSAL_LM",
70
+ )
71
+ model = get_peft_model(model, lora_config)
72
+ model.gradient_checkpointing_enable()
73
+
74
+ dataset = load_dataset("yahma/alpaca-cleaned", split="train")
75
+
76
+ train_loader = DataLoader(
77
+ dataset,
78
+ batch_size=12, # comfortable on 192 GB HBM
79
+ num_workers=8,
80
+ pin_memory=True,
81
+ prefetch_factor=4,
82
+ persistent_workers=True,
83
+ )
84
+
85
+ _RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
86
+
87
+ training_args = TrainingArguments(
88
+ output_dir="./out",
89
+ per_device_train_batch_size=12,
90
+ gradient_accumulation_steps=2,
91
+ num_train_epochs=1,
92
+ learning_rate=2e-4,
93
+ warmup_steps=100,
94
+ bf16=True, # bf16 ✓
95
+ optim="adamw_torch",
96
+ max_seq_length=4096,
97
+ logging_steps=10,
98
+ save_steps=500,
99
+ dataloader_num_workers=8,
100
+ dataloader_pin_memory=True,
101
+ dataloader_prefetch_factor=4,
102
+ dataloader_persistent_workers=True,
103
+ gradient_checkpointing=True,
104
+ torch_compile=True,
105
+ report_to="none",
106
+ push_to_hub=False,
107
+ remove_unused_columns=False,
108
+ max_steps=_RUNTIME_MAX_STEPS,
109
+ )
110
+
111
+
112
+ def _toy_collate(rows):
113
+ texts = [
114
+ (r.get("instruction") or "")
115
+ + ("\n" + r["input"] if r.get("input") else "")
116
+ + "\n"
117
+ + (r.get("output") or "")
118
+ for r in rows
119
+ ]
120
+ enc = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
121
+ enc["labels"] = enc["input_ids"].clone()
122
+ return enc
123
+
124
+
125
+ trainer = Trainer(
126
+ model=model,
127
+ args=training_args,
128
+ train_dataset=dataset,
129
+ tokenizer=tokenizer,
130
+ data_collator=_toy_collate,
131
+ )
132
+
133
+ if __name__ == "__main__":
134
+ _t0 = time.time()
135
+ trainer.train()
136
+ emit_torch_profile(
137
+ _runtime.torch_profile_out,
138
+ elapsed=time.time() - _t0,
139
+ n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
140
+ per_device_batch=training_args.per_device_train_batch_size,
141
+ grad_accum=training_args.gradient_accumulation_steps,
142
+ seq_len_cap=4096,
143
+ )
workloads/train_qwen_lora.py CHANGED
@@ -30,11 +30,7 @@ from transformers import (
30
  TrainingArguments,
31
  )
32
 
33
- from workloads._runtime import (
34
- emit_torch_profile,
35
- parse_runtime_args,
36
- trainer_tokenizer_kwargs,
37
- )
38
 
39
  # Parse the goblin_runner.sh injected flags (--max_steps, --torch_profile_out).
40
  _runtime = parse_runtime_args()
@@ -150,12 +146,17 @@ trainer = Trainer(
150
  model=model,
151
  args=training_args,
152
  train_dataset=dataset,
153
- **trainer_tokenizer_kwargs(Trainer, tokenizer),
154
  data_collator=_toy_collate,
155
  )
156
 
157
 
158
  if __name__ == "__main__":
 
 
 
 
 
159
  _t0 = time.time()
160
  trainer.train()
161
  _elapsed = time.time() - _t0
@@ -170,4 +171,5 @@ if __name__ == "__main__":
170
  per_device_batch=training_args.per_device_train_batch_size,
171
  grad_accum=training_args.gradient_accumulation_steps,
172
  seq_len_cap=512,
 
173
  )
 
30
  TrainingArguments,
31
  )
32
 
33
+ from workloads._runtime import emit_torch_profile, parse_runtime_args
 
 
 
 
34
 
35
  # Parse the goblin_runner.sh injected flags (--max_steps, --torch_profile_out).
36
  _runtime = parse_runtime_args()
 
146
  model=model,
147
  args=training_args,
148
  train_dataset=dataset,
149
+ tokenizer=tokenizer,
150
  data_collator=_toy_collate,
151
  )
152
 
153
 
154
  if __name__ == "__main__":
155
+ # Total parameter count drives MFU. For a peft-wrapped model this
156
+ # includes both the frozen base (which still does forward) and the
157
+ # tiny LoRA adapters. The forward+backward FLOPs estimate uses the
158
+ # full count (the standard 6N rule lives inside emit_torch_profile).
159
+ _model_params = sum(p.numel() for p in model.parameters())
160
  _t0 = time.time()
161
  trainer.train()
162
  _elapsed = time.time() - _t0
 
171
  per_device_batch=training_args.per_device_train_batch_size,
172
  grad_accum=training_args.gradient_accumulation_steps,
173
  seq_len_cap=512,
174
+ model_params=_model_params,
175
  )