Public API

Import the stable core API from traceml_ai:

import traceml_ai as traceml

The reference below documents every symbol in traceml_ai.__all__. The old import traceml path remains a compatibility import and emits a FutureWarning; new code should use traceml_ai.

Stable Core API

traceml.__version__

A string identifying the installed TraceML version. It is useful when recording the environment for a run or bug report.

Lifecycle

init

init(*, mode: str = 'auto', patch_dataloader: Optional[bool] = None, patch_forward: Optional[bool] = None, patch_backward: Optional[bool] = None, patch_h2d: Optional[bool] = None, disabled: Optional[bool] = None, ui_mode: Optional[str] = None, interval: Optional[float] = None, logs_dir: Optional[str] = None, enable_logging: Optional[bool] = None, session_id: Optional[str] = None, aggregator_host: Optional[str] = None, aggregator_port: Optional[int] = None, connect_timeout_sec: float = 10.0, connect_retry_interval_sec: float = 0.25, on_missing_aggregator: Optional[str] = None) -> TraceMLInitConfig

Initialize TraceML and start its runtime for this process.

Parameters:

Name	Type	Description	Default
mode	str	Instrumentation policy: "auto" (default), "manual", or "selective". "custom" is accepted as an alias for selective.	'auto'
patch_dataloader	Optional[bool]	Per-feature automatic-instrumentation settings. They may be supplied only with mode="selective".	None
patch_forward	Optional[bool]	Per-feature automatic-instrumentation settings. They may be supplied only with mode="selective".	None
patch_backward	Optional[bool]	Per-feature automatic-instrumentation settings. They may be supplied only with mode="selective".	None
patch_h2d	Optional[bool]	Per-feature automatic-instrumentation settings. They may be supplied only with mode="selective".	None
disabled	Optional[bool]	Disable TraceML entirely for this process.	None
ui_mode	Optional[str]	Display mode: "summary", "cli", or "dashboard".	None
interval	Optional[float]	Sampling interval in seconds for runtime telemetry.	None
logs_dir	Optional[str]	Directory for session artifacts.	None
enable_logging	Optional[bool]	Enable TraceML file logging.	None
session_id	Optional[str]	Identifier shared by workers that write one run's artifacts.	None
aggregator_host	Optional[str]	Aggregator endpoint for direct launches. traceml run configures it automatically.	None
aggregator_port	Optional[str]	Aggregator endpoint for direct launches. traceml run configures it automatically.	None
connect_timeout_sec	float	Bounded connection wait and retry interval for the aggregator.	10.0
connect_retry_interval_sec	float	Bounded connection wait and retry interval for the aggregator.	10.0
on_missing_aggregator	Optional[str]	Missing-aggregator policy: "warn" prints one stderr warning and continues with tracing disabled; "raise" fails the call. Resolution is explicit argument, then TRACEML_ON_MISSING_AGGREGATOR, then "warn".	None

Returns:

Type	Description
TraceMLInitConfig	The effective instrumentation configuration. Fail-open initialization returns a disabled configuration.

Raises:

Type	Description
ValueError	If mode or selective patch settings are invalid.
RuntimeError	If initialization conflicts with an existing configuration, or when on_missing_aggregator="raise" is selected and setup fails.

start

start(*, mode: str = 'auto', patch_dataloader: Optional[bool] = None, patch_forward: Optional[bool] = None, patch_backward: Optional[bool] = None, patch_h2d: Optional[bool] = None, disabled: Optional[bool] = None, ui_mode: Optional[str] = None, interval: Optional[float] = None, logs_dir: Optional[str] = None, enable_logging: Optional[bool] = None, session_id: Optional[str] = None, aggregator_host: Optional[str] = None, aggregator_port: Optional[int] = None, connect_timeout_sec: float = 10.0, connect_retry_interval_sec: float = 0.25, on_missing_aggregator: Optional[str] = None) -> TraceMLInitConfig

Backward-compatible alias for :func:init.

This function accepts the same parameters and has the same missing- aggregator behavior as :func:init. New code may use either name; prefer init for consistency with framework integration entry points.

Step boundary

trace_step

trace_step(model: Module) -> ContextManager[None]

Return the context manager that defines one training-step boundary.

Parameters:

Name	Type	Description	Default
model	Module	The PyTorch model being trained. TraceML uses it to collect step-memory measurements and associate automatic phase timing with the step.	required

Returns:

Type	Description
ContextManager[None]	Use around the work from optimizer.zero_grad(...) through optimizer.step().

Notes

When tracing is disabled, the context manager is a no-op. Exceptions from the training body still propagate normally.

End-of-run summaries

summary

summary(*, timeout_sec: float = 30.0, poll_interval_sec: float = 0.1, print_text: bool = False, rank0_only: bool = True) -> Optional[Dict[str, Any]]

Return a compact tracker-friendly summary for the active session.

Parameters:

Name	Type	Description	Default
timeout_sec	float	Maximum time to wait for the aggregator to publish the summary.	30.0
poll_interval_sec	float	Delay between checks for the aggregator response.	0.1
print_text	bool	Print the matching final_summary.txt artifact when available.	False
rank0_only	bool	Return None on non-primary ranks when true.	True

Returns:

Type	Description
dict or None	A flat projection of the final summary for experiment trackers, or None for a non-primary rank when rank0_only=True.

Raises:

Type	Description
RuntimeError	If session history is unavailable, the aggregator returns an error, or the request times out.

final_summary

final_summary(*, timeout_sec: float = 30.0, poll_interval_sec: float = 0.1, print_text: bool = False, rank0_only: bool = True) -> Optional[Dict[str, Any]]

Return the full final-summary payload for the active session.

Parameters:

Name	Type	Description	Default
timeout_sec	float	Maximum time to wait for the aggregator to publish the summary.	30.0
poll_interval_sec	float	Delay between checks for the aggregator response.	0.1
print_text	bool	Print the matching final_summary.txt artifact when available.	False
rank0_only	bool	Return None on non-primary ranks when true.	True

Returns:

Type	Description
dict or None	The parsed final_summary.json payload, or None for a non-primary rank when rank0_only=True.

Raises:

Type	Description
RuntimeError	If session history is unavailable, the aggregator returns an error, or the request times out.

Manual instrumentation helpers

Use these only for manual or selective instrumentation. Automatic mode already times the matching PyTorch paths, and manual wrappers reject duplicate automatic instrumentation where double-counting would be possible.

wrap_dataloader_fetch

wrap_dataloader_fetch(obj: Any) -> Any

Wrap a loader or iterator to time step-scoped batch fetches.

Parameters:

Name	Type	Description	Default
obj	Any	A loader implementing __iter__ or an iterator implementing __next__.	required

Returns:

Type	Description
Any	A loader or iterator proxy that records fetch time.

Raises:

Type	Description
TypeError	If obj is neither a loader nor an iterator.
RuntimeError	If a torch DataLoader is already automatically instrumented.

wrap_forward

wrap_forward(model: Module) -> nn.Module

Wrap one model instance's forward(...) for step-scoped timing.

Parameters:

Name	Type	Description	Default
model	Module	The PyTorch module to instrument.	required

Returns:

Type	Description
Module	The same model instance after its forward method is wrapped in place.

Raises:

Type	Description
TypeError	If model is not a PyTorch module with a callable forward.
RuntimeError	If automatic forward instrumentation is active or the instance cannot be wrapped safely.

wrap_backward

wrap_backward(loss: Any) -> Any

Wrap a loss-like object so .backward(...) records timing.

Parameters:

Name	Type	Description	Default
loss	Any	An object with a callable backward method, normally a loss tensor.	required

Returns:

Type	Description
Any	A proxy that forwards attributes and times its backward call.

Raises:

Type	Description
TypeError	If loss has no callable backward method.
RuntimeError	If automatic backward instrumentation is active.

wrap_optimizer

wrap_optimizer(optimizer: Any) -> Any

Wrap one optimizer instance's .step(...) for timing.

Parameters:

Name	Type	Description	Default
optimizer	Any	An optimizer with a callable step method.	required

Returns:

Type	Description
Any	The same optimizer instance after its step method is wrapped in place.

Raises:

Type	Description
TypeError	If optimizer has no callable step method.
RuntimeError	If automatic optimizer instrumentation is active or the instance cannot be wrapped safely.

wrap_h2d

wrap_h2d(obj: Any) -> Any

Wrap an object so qualifying CPU-to-CUDA .to(...) calls are timed.

Parameters:

Name	Type	Description	Default
obj	Any	A tensor or batch object with a callable .to(...) method.	required

Returns:

Type	Description
Any	A proxy that forwards attributes and returns the transferred object on its first .to(...) call.

Raises:

Type	Description
TypeError	If obj has no callable .to(...) method.

Notes

Use this in manual or selective mode. If automatic H2D timing becomes active after wrapping, the proxy passes through to avoid double counting.

CLI

TraceML installs the traceml command-line entry point:

traceml run <script>                  # final summary JSON/text
traceml run <script> --mode=summary   # explicit summary mode
traceml run <script> --mode=cli       # live terminal view
traceml run <script> --mode=dashboard # live browser view
traceml watch <script>                # zero-code system/process summary
traceml serve                         # standalone aggregator

Summary mode is the default for every topology. Live cli and dashboard modes are intended for single-node runs. Use PyTorch Profiler or Nsight for operator- or kernel-level profiling.

Direct Launch with traceml serve

traceml run starts the aggregator and your script together. For direct python or torchrun launches, start the aggregator yourself and call traceml.init(...) inside the training script:

# terminal 1
traceml serve --aggregator-host 127.0.0.1 --aggregator-port 29765

# terminal 2
python train.py

traceml serve owns only the aggregator. It binds the endpoint, waits for a shutdown signal, and writes the final summary; it never launches or wraps the training script.

For multi-node workers, bind the aggregator on a reachable address and set the endpoint on every training node:

traceml serve --aggregator-bind-host 0.0.0.0 --aggregator-host <node0-ip> \
  --aggregator-port 29765 --nnodes <N> --nproc-per-node <M>

TRACEML_AGGREGATOR_HOST=<node0-ip> TRACEML_AGGREGATOR_PORT=29765 \
  torchrun ... train.py

Workers resolve TRACEML_AGGREGATOR_HOST as 127.0.0.1 by default, so every non-aggregator node needs the reachable node-0 address above.

traceml serve flags:

Flag	Meaning
--aggregator-host	Address workers connect to; default 127.0.0.1.
--aggregator-bind-host	Bind address; use 0.0.0.0 for multi-node.
--aggregator-port	Aggregator TCP port; default 29765.
--nnodes / --nproc-per-node	Expected world size; the aggregator waits for all ranks before finalizing.
--mode	summary (default), cli, or dashboard.
--logs-dir	Directory for session logs.
--run-name / --session-id	Shared run identity for worker artifacts.

Missing-aggregator behavior

If the aggregator cannot be reached, traceml.init(...) retries for its bounded timeout, writes one stderr warning, and continues with tracing disabled as a no-op. This is the default warn policy; it does not stop training.

Use strict behavior when telemetry is required, for example in CI:

traceml.init(on_missing_aggregator="raise")

The policy resolves in this order: the explicit on_missing_aggregator argument, TRACEML_ON_MISSING_AGGREGATOR, then warn. It is not read from traceml.yaml.

aggregator_host and aggregator_port are direct-launch settings, not traceml.yaml settings. Other runtime settings resolve as explicit traceml.init(...) arguments, then TRACEML_* environment variables, then traceml.yaml, then built-in defaults.

Matching display modes across processes

In direct-launch mode, set the aggregator display with traceml serve --mode and the worker display with traceml.init(ui_mode=...) or TRACEML_UI_MODE. Use cli for both when the live terminal panel should include worker output:

traceml serve --mode cli --run-name demo --aggregator-port 29765
TRACEML_UI_MODE=cli TRACEML_SESSION_ID=demo python train.py

If the modes differ, telemetry, diagnosis, and final artifacts are unaffected; only worker stdout mirroring into the live panel is skipped.

Framework Integrations

Framework integrations are separate from the stable core API above. Use the matching integration guide for installation and runtime requirements.

Hugging Face

Preferred path: call the integration init() once and register TraceMLTrainerCallback with your existing transformers.Trainer.

traceml_ai.integrations.huggingface.init

init()

Initialize TraceML for Hugging Face Trainer runs.

Call once before constructing the Trainer, then register TraceMLTrainerCallback. init() makes TraceML's process-wide instrumentation explicit: PyTorch DataLoader fetch timing, the H2D Tensor.to patch, and the forward/backward/optimizer auto-timers that trace_step arms inside each bracketed step.

The callback is a per-step bracket and cannot install these process-wide patches on its own; the auto-timers it arms are no-ops unless the matching patch is installed. init() is the recommended entry point so the DataLoader fetch patch in particular is installed deterministically rather than relying on import order. This mirrors the PyTorch Lightning integration's init(); HF uses mode="auto" because trace_step drives forward/backward timing through the patch-gated auto-timers, whereas Lightning's callback owns that timing directly.

traceml_ai.integrations.huggingface.TraceMLTrainerCallback

TraceMLTrainerCallback()

Bases: TrainerCallback if HAS_TRANSFORMERS else object

Preferred Hugging Face integration for TraceML.

Register with Trainer(..., callbacks=[TraceMLTrainerCallback()]).

The callback is a pure bracket around TraceML's trace_step context manager: it opens trace_step in on_step_begin and closes it in on_step_end. trace_step owns the step memory tracker, the step counter advance, the auto-timers for forward/backward/h2d, and the per-step flush. Nothing is duplicated here.

One TraceML step equals one optimizer step. With gradient_accumulation_steps > 1, forward and backward events from all accumulated micro-batches fold into a single TraceML step. See the HF integration docs for the full list of limitations vs. TraceMLTrainer.

Legacy compatibility: TraceMLTrainer

TraceMLTrainer remains supported for existing users. New code should prefer TraceMLTrainerCallback; see the Hugging Face guide for its trade-offs.

traceml_ai.integrations.huggingface.TraceMLTrainer

TraceMLTrainer(*args, traceml_enabled: bool = True, **kwargs)

Bases: Trainer if HAS_TRANSFORMERS else object

Thin wrapper around transformers.Trainer that auto-installs TraceMLTrainerCallback.

Kept for backward compatibility with users on the original TraceML HF integration API. New code should prefer Trainer(..., callbacks=[TraceMLTrainerCallback()]) directly.

PyTorch Lightning

traceml_ai.integrations.lightning.init

init()

Initialize TraceML for PyTorch Lightning runs.

Lightning owns the training loop, so TraceMLCallback owns step boundaries, flushing, and framework hook integration. The integration init enables DataLoader fetch timing plus the H2D Tensor.to patch. The callback turns H2D timing on only around Lightning's batch transfer hooks and wraps LightningModule.forward directly for model-forward timing.

traceml_ai.integrations.lightning.TraceMLCallback

TraceMLCallback()

Bases: Callback

Official TraceML Callback for PyTorch Lightning.

Captures full step time (forward + backward + optimizer) as well as individual phase timings. Safely handles gradient accumulation by treating each micro-batch as a step, providing 0-duration optimizer events on accumulating steps to preserve dashboard step alignment.

Ray Train

traceml_ai.integrations.ray.TraceMLTorchTrainer

TraceMLTorchTrainer(train_loop_per_worker: TrainLoop, *, train_loop_config: Optional[Dict[str, Any]] = None, traceml_config: Optional[TraceMLRayConfig] = None, **torch_trainer_kwargs: Any)

TraceML wrapper for ray.train.torch.TorchTrainer.

The wrapper deliberately uses composition instead of subclassing. Ray's TorchTrainer still owns training orchestration; TraceML only adds:

• one aggregator actor before fit()
• one runtime wrapper inside each worker
• best-effort aggregator shutdown after fit() completes or fails

last_endpoint property

last_endpoint: Optional[AggregatorEndpoint]

Aggregator endpoint used by the most recent fit() call.

fit

fit() -> Any

Run Ray Train with TraceML telemetry enabled.

traceml_ai.integrations.ray.TraceMLRayConfig dataclass

TraceMLRayConfig(mode: str = 'summary', profile: str = 'run', init_mode: str = 'auto', patch_dataloader: Optional[bool] = None, patch_forward: Optional[bool] = None, patch_backward: Optional[bool] = None, patch_h2d: Optional[bool] = None, logs_dir: str = './logs', session_id: str = '', sampler_interval_sec: float = DEFAULT_INTERVAL_SEC, summary_window_rows: int = DEFAULT_SUMMARY_WINDOW_ROWS, bind_host: str = '0.0.0.0', port: int = 0, stop_timeout_sec: float = 5.0)

TraceML settings used by the Ray Train integration.

Parameters:

Name	Type	Description	Default
mode	str	TraceML display/reporting mode. "summary" is the default for Ray because distributed worker logs are often noisy.	'summary'
profile	str	TraceML sampler profile. Public Ray integration uses the normal "run" profile.	'run'
init_mode	str	Instrumentation mode passed to traceml.init() inside each worker.	'auto'
patch_dataloader	Optional[bool]	Selective-mode-only override for DataLoader fetch patching.	None
patch_forward	Optional[bool]	Selective-mode-only override for forward-pass patching.	None
patch_backward	Optional[bool]	Selective-mode-only override for backward-pass patching.	None
patch_h2d	Optional[bool]	Selective-mode-only override for host-to-device transfer patching.	None
logs_dir	str	Directory where TraceML writes session logs and summary artifacts.	'./logs'
session_id	str	Optional explicit TraceML session id. If omitted, a unique Ray session id is generated for each fit() call.	''
sampler_interval_sec	float	Background sampler cadence in seconds. The aggregator actor uses the same value for its live UI refresh cadence; TCP ingestion is immediate.	DEFAULT_INTERVAL_SEC
summary_window_rows	int	Number of recent history rows used by final summary generation.	DEFAULT_SUMMARY_WINDOW_ROWS
bind_host	str	Host interface used by the aggregator actor. Use "0.0.0.0" for multi-node Ray clusters so workers on other nodes can connect.	'0.0.0.0'
port	int	Aggregator TCP port. 0 asks the OS to pick a free port.	0
stop_timeout_sec	float	Best-effort timeout for aggregator shutdown.	5.0