Disable traceml

[abhinavsriva]

What changed

This hardens the TraceML disabled/no-op paths so training can run natively when TraceML is explicitly disabled or when instrumentation cannot be installed safely.

• Adds --disable_traceml as an alias for the existing --disable-traceml launcher flag.
• Makes TRACEML_DISABLED=1 behave as a launcher-level kill switch.
• When disabled, launches the target script directly through torchrun instead of starting the TraceML executor, aggregator, manifests, logs, stderr capture, or report generation.
• Strips stale TRACEML_* settings from the native child environment and keeps only TRACEML_DISABLED=1.
• Adds an SDK tracing gate, is_tracing_armed(), so auto-patched wrappers pass through natively until patch installation finishes cleanly.
• Converts patch-install failures to fail-open no-op behavior by default: TraceML stops any started runtime, sets TRACEML_DISABLED=1, warns once, and lets training continue.
• Preserves strict behavior with on_missing_aggregator='raise', which still raises on patch-install failure after cleanup.
• Makes low-level timing, step-memory, buffer flushing, and Lightning callback paths read TRACEML_DISABLED dynamically instead of caching it at import time.
• Updates launcher output wording from "TraceML executor" to the more accurate "training process".

Why

The disabled path should be a true native execution mode. If a user passes--disable-traceml or exports TRACEML_DISABLED=1, TraceML should not create telemetry side effects, should not parse TraceML config, and should not leave partial instrumentation active.

The same fail-open behavior is important when patch installation fails. Partial patching can produce inconsistent telemetry, so this change keeps the wrappers installed-but-unarmed and degrades to native pass-through behavior instead of half-tracing a training process.

Files changed vs version_0.3.5

• src/traceml_ai/launcher/cli.py
• src/traceml_ai/launcher/commands.py
• src/traceml_ai/launcher/process.py
• src/traceml_ai/sdk/initial.py
• src/traceml_ai/instrumentation/patches/dataloader_patch.py
• src/traceml_ai/instrumentation/patches/forward_auto_timer_patch.py
• src/traceml_ai/instrumentation/patches/backward_auto_timer_patch.py
• src/traceml_ai/instrumentation/patches/h2d_auto_timer_patch.py
• src/traceml_ai/integrations/lightning.py
• src/traceml_ai/utils/timing.py
• src/traceml_ai/utils/step_memory.py
• src/traceml_ai/utils/flush_buffers.py
• tests/runtime/test_launcher.py
• tests/sdk/test_init_runtime.py
• tests/sdk/test_init_and_wrappers.py
• tests/integrations/test_lightning.py
• tests/integrations/test_ray.py

How I tested

Not run in this pass.

Recommended focused tests:

PYTHONDONTWRITEBYTECODE=1 pytest -p no:cacheprovider \
  tests/runtime/test_launcher.py \
  tests/sdk/test_init_runtime.py \
  tests/sdk/test_init_and_wrappers.py \
  tests/integrations/test_lightning.py \
  tests/integrations/test_ray.py \
  -q

Runtime impact

• Disabled launches now bypass TraceML setup entirely and run the script natively.
• TRACEML_DISABLED=1 can disable tracing even if it is set after modules were imported.
• Auto-patch wrappers remain safe after failed initialization because they check is_tracing_armed() before recording.
• Existing enabled TraceML runs should keep the normal launcher, runtime, aggregator, and telemetry behavior.

Extra context

Base comparison used for this summary:version_0.3.5...disable_traceml.

Current branch commits included in this comparison:

80c4388 Merge branch 'version_0.3.5' into disable_traceml
6a3d917 Merge branch 'version_0.3.5' into disable_traceml
3f1561d diable hardned
b395b66 test corrected
ce390e6 added gating to 0 or less
2e2a1b0 straggler simplified

[Pendu]

Pushed a fix for the blocker here (114ae80): import traceml_ai.integrations.lightning raised an ImportError for every user.

The cause: a circular import. Four patch modules (h2d_auto_timer_patch.py, forward_auto_timer_patch.py, backward_auto_timer_patch.py, dataloader_patch.py) each import is_tracing_armed from traceml_ai.sdk.initial at module scope, and importing that submodule first runs traceml_ai/sdk/__init__.py. That in turn imports sdk.instrumentation, which imports back from the same patch module while it's still mid-initialization. The name isn't defined yet at that point, so the import fails.

Fix: moved is_tracing_armed / _set_tracing_armed into a new leaf module, traceml_ai/runtime/arming.py, with no imports beyond __future__. The four patch modules now import from there instead of sdk.initial. sdk/initial.py re-exports is_tracing_armed (added to its __all__) so every existing internal call site and the is_tracing_armed()/_set_tracing_armed() behavior stay unchanged, just relocated.

The branch was also a merge behind version_0.3.5; merged it up to date (0905191) the same way you did on the earlier branches in this stack, resolving one trivial line-wrap conflict in commands.py.

Verified on the pushed head:

• import traceml_ai.integrations.lightning now succeeds.
• pytest tests/sdk tests/runtime tests/integrations -> 172 passed, 3 failed, 1 skipped. The 3 failures (2 in test_deepspeed.py, 1 in test_hf_trainer.py) are unchanged from the branch state before this fix, so they are pre-existing and out of scope here.
• Lint clean at the pinned pre-commit versions (isort 5.13.2 / black 24.10.0 / ruff 0.9.4 / codespell), and mkdocs build --strict passes.