Skip to content

feat: Add --disable-traceml flag for true zero-overhead no-op mode#51

Merged
abhinavsriva merged 5 commits intotraceopt-ai:mainfrom
ppraneth:dis
Mar 15, 2026
Merged

feat: Add --disable-traceml flag for true zero-overhead no-op mode#51
abhinavsriva merged 5 commits intotraceopt-ai:mainfrom
ppraneth:dis

Conversation

@ppraneth
Copy link
Copy Markdown
Contributor

@ppraneth ppraneth commented Mar 12, 2026

Resolves #50

This PR introduces a true no-op mode via the --disable-traceml CLI flag (and TRACEML_DISABLED environment variable).

Previously, users had to manually remove import traceml and instrumentation wrappers (like @trace_step or Lightning/HF callbacks) from their code to run clean benchmarks or execute natively. With this PR, users can leave all TraceML instrumentation exactly as-is. Passing the flag safely short-circuits all TraceML logic, providing exactly native PyTorch performance.

Changes Implemented

  • CLI Launcher (cli.py): Bypasses the background aggregator subprocess and directly launches the user script via torchrun.
  • Runtime (executor.py): Intercepts the disable flag to skip TraceMLRuntime initialization, falling back to a safe NoOpRuntime to prevent background threads/TCP clients from spawning.
  • Decorators (decorators.py, timing.py): Public context managers (@trace_step, timed_region) and monkey-patches (patch_dataloader) safely yield or return immediately.
  • Integrations (hf_decorators.py, lightning.py): Hugging Face and PyTorch Lightning callbacks return instantly, completely skipping PyTorch hook attachments and gradient accumulation logic.
  • Memory Tracking (step_memory.py, flush_buffers.py): CUDA max-memory trackers and global buffer flush methods safely short-circuit to ensure zero network/disk I/O.

Native Execution (TraceML Disabled):
time traceml run --disable-traceml src/examples/simple_cnn.py

Traced Execution (TraceML Enabled):
time traceml run src/examples/simple_cnn.py

Benchmark Results

  • Traced Time: ~21.29s
  • Disabled Time: ~17.28s (Reclaimed ~4s)

Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
Signed-off-by: ppraneth <pranethparuchuri@gmail.com>
@abhinavsriva abhinavsriva merged commit b5f13cb into traceopt-ai:main Mar 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add --disable-traceml flag for no-op runtime disable

2 participants