dftracer-ml-annotate

name	dftracer-ml-annotate
description	Annotate ML/DL Python workloads end-to-end with dftracer AI/ML region decorators. Detects frameworks and ROCm/HIP requirements, installs dftracer with correct flags, annotates all Python files, builds/runs smoke test, collects traces, and auto-updates the lessons file with any new pitfalls discovered during the session.

Lessons file: workspaces/.agents/skills/dftracer-annotation-lessons/LESSONS_LOG.md (workspace-local staging copy — shared with the general dftracer-pipeline skill; sync it back to the source repo with session_lessons_sync_preview / session_lessons_sync_pr, see Step 10). Load the compact rules/instructions with skill_load(name="dftracer-annotation-lessons"); load the accumulated log entries with skill_load(name="dftracer-annotation-lessons", file="LESSONS_LOG.md").

Read the lessons file before doing anything else. Apply every standing rule (ML-R1 through ML-R16) and check every session log entry for context matching the current application or framework.

DLIO benchmark (dlio_benchmark) is the canonical reference pattern. When in doubt about how to annotate a function, match what DLIO does.

══════════════════════════════════════════════════════════════════════ STEP 0 — SYSTEM DETECTION ══════════════════════════════════════════════════════════════════════

Invoke the system-detect skill to load module/compiler environment.

Store environment as ENV (module paths, MPI wrappers, ROCm paths).

══════════════════════════════════════════════════════════════════════ STEP 1 — GATHER INPUTS (if not supplied via arguments) ══════════════════════════════════════════════════════════════════════

If the user invoked this with named arguments (run_id=…, url=…, etc.), use those directly. Otherwise ask one question at a time:

Q1: "What is the Git URL of the ML application?" → APP_URL Q2: "Branch or tag? (default: main)" → REF Q3: "Smoke test command? (leave blank to auto-detect)" → SMOKE_CMD Q4: "Extra build flags? (leave blank to skip)" → EXTRA_FLAGS

If run_id was supplied, skip Q1–Q4 and jump to Step 3.

Print: "Starting ML annotation pipeline for <APP_URL> @ "

Track pitfalls found this session in a list: PITFALLS = []

══════════════════════════════════════════════════════════════════════ STEP 2 — SESSION SETUP ══════════════════════════════════════════════════════════════════════

2a. Create session and clone:

session_create(url=APP_URL, ref=REF)
→ store RUN_ID, WS

2b. Detect ML workload:

session_detect_ml_workload(run_id=RUN_ID)

Store:
  FRAMEWORKS       = result.frameworks
  HIP_NEEDED       = result.hip_tracing_needed
  ROCM_INFO        = result.rocm_info
  INSTALL_FLAGS    = result.install_flags
  DISTRIBUTED      = result.distributed
  HAS_DATALOADER   = result.has_dataloader
  CAPABILITIES     = result.capabilities

Print:
  "Frameworks: <FRAMEWORKS>"
  "ROCm: <ROCM_INFO.found> (<ROCM_INFO.path>)"
  "HIP tracing needed: <HIP_NEEDED>"
  "Distributed: <DISTRIBUTED>"

If HIP_NEEDED but ROCM_PATH is not set in ENV:
  PITFALLS.append({phase: "install", error: "ROCM_PATH not set",
    root_cause: "ROCm detected but ROCM_PATH env var absent",
    fix: "export ROCM_PATH=<ROCM_INFO.path> before install"})
  Set ROCM_PATH manually before Step 3.

2c. Configure and build original source:

session_configure(run_id=RUN_ID, extra_cmake_flags=EXTRA_FLAGS)
session_build_install(run_id=RUN_ID)

2d. Install dftracer with correct flags:

session_install_dftracer(run_id=RUN_ID)

The install reads INSTALL_FLAGS (which includes DFTRACER_ENABLE_HIP_TRACING=ON
when HIP_NEEDED) automatically from the session state set by
session_detect_ml_workload. No manual env override needed.

On failure → record pitfall in PITFALLS, then stop.

Print: "Setup complete. RUN_ID=<RUN_ID> dftracer installed."

2e. Validate session structure before annotating anything (STRICT RULE — see dftracer-cheatsheet S0: always use the MCP tool's own directory structure, never a hand-built path):

session_validate_structure(run_id=RUN_ID)

If clean=false → session_reorganize_structure(run_id=RUN_ID, dry_run=False)
then re-validate before proceeding to Step 3.

══════════════════════════════════════════════════════════════════════ STEP 3 — COPY SOURCE AND DISCOVER FILES ══════════════════════════════════════════════════════════════════════

3a. Copy source to annotated workspace:

session_copy_annotated(run_id=RUN_ID)

3b. Discover Python source files:

find_source_files(run_id=RUN_ID, folder="annotated", language="python",
  exclude_patterns=["**/test*", "**/__pycache__/**", "**/setup.py",
                    "**/conftest.py", "**/docs/**"])

Categorize files into:
  ENTRY_FILES    = files with if __name__ == "__main__" or def main(
  TRAIN_FILES    = files with train / fit / backward / optimizer.step
  DATA_FILES     = files with __getitem__ / Dataset / DataLoader
  CKPT_FILES     = files with save_checkpoint / load_checkpoint / state_dict
  COMM_FILES     = files with all_reduce / dist.barrier / horovod / hvd.
  OTHER_FILES    = all remaining .py files

Print: "Files: entry= train= data= ckpt= comm= other="

══════════════════════════════════════════════════════════════════════ STEP 4 — AI/ML REGION ANNOTATION ══════════════════════════════════════════════════════════════════════

FAST PATH (default) — one MCP call for the whole project

The per-file recipe in 4a–4g below is now implemented as deterministic MCP tools. Use them first; they are faster, reproducible, and do not drift.

ml_annotate_plan(run_id=RUN_ID, threshold=20)      # review, writes nothing
ml_annotate_project(run_id=RUN_ID, threshold=20)   # execute everything

ml_annotate_project categorizes every file (Step 3b), then runs the AI/ML cost estimator on EVERY file and passes its allow-list down. Semantic buckets (entry / train / data / ckpt / comm) go through python_annotate_ai_file, the leftovers through python_annotate_file.

The cost gate always runs. The one exception is the AI API: functions that map to a dft_ai.* region (data item, checkpoint, training step, comm, forward, backward) are annotated for what they MEAN and are kept regardless of score. Everything else must clear the threshold, so a training script does not get a decorator on every getter. ml_categorize_files exposes the bucketing alone.

It also validates at the end (validate=True) and returns validation.passed.

Then ALWAYS:

annotate_add_app_metadata(run_id=RUN_ID, filepath=<entry file>,
  language="python",
  params_json='{"app":"...","ranks":"...","batch_size":"..."}')

validate_annotations(run_id=RUN_ID, language="python")

Do not proceed to the build until validate_annotations passes. Also dispatch the dftracer-validate-python agent for an independent check.

BACKUP PATH — the manual recipe (4a–4g)

Use the steps below ONLY when: a tool is missing or errors; a file needs judgement the classifier cannot make; or you are repairing a file the tools got wrong. They document exactly what the fast path does.

══════════════════════════════════════════════════════════════════════ STEP 4 (manual) — AI/ML REGION ANNOTATION (python_annotate_ai_file) ══════════════════════════════════════════════════════════════════════

MANDATORY: Use python_annotate_ai_file for all files. NEVER manually write decorators with Edit/Write. The tool handles idempotency, loop wrapping, and import injection automatically.

Before annotating any function you are unsure about, call:

dftracer_get_ai_annotation(
  function_name=<fn_name>,
  context=<what it does>,
  phase=<compute|data|dataloader|comm|device|checkpoint|pipeline>
)

This returns the exact decorator and a ready-to-use code example.

4a. Entry-point files (is_entry=True)

python_annotate_ai_file(
  run_id=RUN_ID, filepath=<file>,
  category=<module_stem>,
  is_entry=True, annotate_loops=True,
)

Verify the result contains:
  - initialize_log injected near the top
  - finalize() injected before program exit
  - @dft_ai on main() / run() / __call__()

4b. Training files

python_annotate_ai_file(run_id=RUN_ID, filepath=<file>,
  category=<module_stem>, annotate_loops=True)

Expected decorators:
  train / fit / run_epoch     → @dft_ai.pipeline.train
  evaluate / validate         → @dft_ai.pipeline.evaluate
  test                        → @dft_ai.pipeline.test
  forward                     → @dft_ai.compute.forward
  backward / loss.backward    → @dft_ai.compute.backward
  for epoch in ...:           → dft_ai.pipeline.epoch.iter(...)
  for batch in ...:           → dft_ai.dataloader.fetch.iter(...)

Optimizer step — MUST use start/stop style (ML-R4):
  dft_ai.compute.step.start()
  optimizer.step()
  dft_ai.compute.step.stop()

After annotating each file, add inside the batch loop:
  ai.update(step=step, epoch=epoch)   ← ML-R8

4c. Data/Dataset files

python_annotate_ai_file(run_id=RUN_ID, filepath=<file>,
  category=<module_stem>, annotate_loops=True)

Expected decorators:
  __getitem__ / read_index / load_sample  → @dft_ai.data.item
  preprocess / transform / augment        → @dft_ai.data.preprocess.derive(name="<op>")
  collate                                 → @dft_ai.data.preprocess.derive(name="collate")
  to_device / .cuda() / .to(device)       → @dft_ai.device.transfer

**Data I/O rules (ML-R25) — use `ai.data.io.*` for explicit open/read/write/close:**

Annotate each I/O phase separately instead of lumping everything into `data.item`.
ALWAYS compute and pass `image_size` (bytes) as metadata to every I/O region.

Phase mapping:
  open file / open dataset / h5py.File(...)  → @dft_ai.data.io.open
  np.load / f.read() / dataset[idx]          → @dft_ai.data.io.read  + image_size
  f.write() / np.save()                      → @dft_ai.data.io.write + image_size
  f.close() / file handle cleanup            → @dft_ai.data.io.close

image_size MUST be the byte size of the actual DATA ARRAY, computed from the
loaded/written object — NEVER from the path string or file metadata:

  ✅ numpy array:   image_size=array.nbytes          (in-memory array bytes)
  ✅ torch tensor:  image_size=tensor.element_size() * tensor.nelement()
  ✅ bytes object:  image_size=len(buf)
  ✅ checkpoint:    image_size=sum(t.nbytes for t in state_dict.values() if hasattr(t, "nbytes"))
  ❌ WRONG:         image_size=os.path.getsize(path) (file-system metadata, not data)
  ❌ WRONG:         image_size=len(path)             (path string length)

Pass image_size via update() AFTER the read/write so the value is known:

```python
from dftracer.python import ai

# numpy example
def load_sample(path: str):
    with ai.data.io.open:
        pass
    with ai.data.io.read:
        data = np.load(path)
        ai.data.io.read.update(image_size=data.nbytes)
    with ai.data.io.close:
        pass
    return data
```

Context-manager style is preferred when the function mixes phases.
Decorator style is preferred when the method maps 1:1 to a phase.

4d. Checkpoint files

python_annotate_ai_file(run_id=RUN_ID, filepath=<file>,
  category=<module_stem>)

Expected decorators:
  save / save_checkpoint / write_ckpt     → @dft_ai.checkpoint.capture
  load / load_checkpoint / restore_ckpt  → @dft_ai.checkpoint.restart

**Checkpoint I/O rules (ML-R26) — use `ai.checkpoint.io.*` for explicit phases:**

Wrap the four I/O phases inside the outer capture/restart context.
ALWAYS compute and pass `image_size` (bytes of the checkpoint) as metadata.

Phase mapping:
  open file for checkpoint               → @ai.checkpoint.io.open
  torch.load / pickle.load / f.read()    → @ai.checkpoint.io.read  + image_size
  torch.save / pickle.dump / f.write()   → @ai.checkpoint.io.write + image_size
  f.close() / os.remove()               → @ai.checkpoint.io.close

image_size for checkpoints = total bytes of all tensors in the state dict.
Compute from the in-memory data — NOT from os.path.getsize or len(path):

  read:  sum(t.nbytes for t in checkpoint.get("model_state_dict", {}).values() if hasattr(t, "nbytes"))
  write: sum(t.nbytes for t in state_dict.get("model_state_dict", {}).values() if hasattr(t, "nbytes"))

```python
from dftracer.python import ai
import torch

def save_checkpoint(model, path: str):
    state_dict = {"model_state_dict": model.state_dict()}
    ckpt_bytes = sum(t.nbytes for t in state_dict["model_state_dict"].values() if hasattr(t, "nbytes"))
    with ai.checkpoint.capture:
        with ai.checkpoint.io.open:
            f = open(path, "wb")
        with ai.checkpoint.io.write:
            torch.save(state_dict, f)
            ai.checkpoint.io.write.update(image_size=ckpt_bytes)
        with ai.checkpoint.io.close:
            f.close()

def load_checkpoint(path: str):
    with ai.checkpoint.restart:
        with ai.checkpoint.io.open:
            f = open(path, "rb")
        with ai.checkpoint.io.read:
            state = torch.load(f)
            ckpt_bytes = sum(t.nbytes for t in state.get("model_state_dict", {}).values() if hasattr(t, "nbytes"))
            ai.checkpoint.io.read.update(image_size=ckpt_bytes)
        with ai.checkpoint.io.close:
            f.close()
    return state
```

4e. Other I/O files (config reads, stats writes, utility I/O — ML-R27)

Any I/O that is NOT inside the DataLoader path or checkpoint save/restore goes under ai.other.io.*. This covers: config file reads, CSV/stats file writes, rendezvous/coordination files, datagen scripts, logging helpers.

Phase mapping:
  open file / open dataset           → dft_ai.other.io.open
  f.read() / yaml.load / np.load()   → dft_ai.other.io.read  + image_size
  f.write() / outfile.write(row)     → dft_ai.other.io.write + image_size
  f.close()                          → dft_ai.other.io.close

Use `ai.other.log` for logging/print sinks that should be traced but
carry no I/O bytes.

image_size rules are identical to ML-R25:
  bytes object:  image_size=len(buf)
  numpy array:   image_size=array.nbytes
  encoded str:   image_size=len(s.encode())
  NEVER os.path.getsize or len(path)

Category decision tree:
  Is the I/O inside __getitem__ / DataLoader path?  → data.io.*
  Is the I/O torch.save / torch.load of model weights?  → checkpoint.io.*
  Everything else  → other.io.*

```python
from dftracer.python import ai

# config read example
with ai.other.io.open:
    f = open(config_path, "rb")
with ai.other.io.read:
    raw = f.read()
    cfg = yaml.safe_load(raw)
    ai.other.io.read.update(image_size=len(raw))
with ai.other.io.close:
    f.close()

# stats CSV write example
with ai.other.io.write:
    outfile.write(row)
    ai.other.io.write.update(image_size=len(row.encode()))
```

4f. Distributed communication files (when DISTRIBUTED=True)