TraceML: logging metrics, inputs and artifacts for Polyaxon runs

What TraceML adds to the Polyaxon client

TraceML is described as an engine for ML and data tracking, visualization, explainability, drift detection and dashboards for Polyaxon. In practice the part of it that the documentation shows in detail is a tracking client. A run is opened with an init call that names a project, a run name, a description, a list of tags and an artifacts path, and then the run accumulates inputs, metrics, outputs and artifacts through a set of log functions. The intended users are machine learning engineers who want each training run to leave behind a record of its hyper-parameters, its loss curve and the model file, and who would rather not hand-roll a metrics writer for every project. The engine is a separate Python package from the polyaxon client, and the README states that you need polyaxon installed as well if you want to use the tracking features. The dashboard, drift and explainability claims live in the project description, not in the documented API.

Offline mode removes the API from the loop

The most useful thing in the README is also the least documented, because it arrives in a section of its own. TraceML can track runs without an API by enabling an offline mode, and there are two ways to do it. The first is an environment variable, which applies to a whole process or shell session:

export POLYAXON_OFFLINE="true"

The second is a flag passed to the init call, which is what every example in the README uses:

from traceml import tracking

tracking.init(..., is_offline=True, ...)

The difference in practice is scope. The environment variable turns off the API for everything in the environment, while the per-call flag makes a single run offline, which is useful in a script that mixes online and offline experiments. What offline means for the stored data is not spelled out beyond the artifacts path and the tracking calls, so the reasonable assumption is that the run record is written locally rather than pushed to a server.

tracking.init sets the shape of the run

Every example opens with the same block, and the arguments are worth reading individually because they decide what ends up in the record. The simplest usage example imports the package under the name tracking, which is an alias of the top-level package rather than the submodule:

import traceml as tracking

tracking.init(
    is_offline=True,
    project='quick-start',
    name="my-new-run",
    description="trying TraceML",
    tags=["examples"],
    artifacts_path="path/to/artifacts/repo"
)

The project groups runs, the name identifies one run inside it, the description and the tags give you something to search on later, and artifacts_path says where files belonging to the run are written. Other examples use the same structure with different values, such as project='tracking-project' and names like keras-run, pytorch-run and tf-run, which suggests the run is addressable by those fields. The tracking module also exposes a stop function that the example calls explicitly, described as optional, so a run can be finalised early rather than waiting for process exit.

Four kinds of payload: refs, inputs, metrics and outputs

The tracking surface divides into four kinds of call, and the difference matters for what you store. Data references log the content itself, as in log_data_ref(content=X_train, name='x_train'), so the training set is attached to the run under a name. Inputs log the hyper-parameters, for example a batch_size, dropout, learning_rate and optimizer call, so that a run can be reproduced from the record. Metrics log the numbers that change during training:

# Track metrics
for step in range(100):
    loss = get_loss(step)
    tracking.log_metrics(
    loss=loss,
    accuracy=(100 - loss) / 100.0,
)

Outputs log one-time results such as a validation score. The distinction between a data reference and an artifact reference is the one to keep straight: a data reference records the dataset itself, while an artifact reference points at a file that already exists on disk. Metric calls are plain keyword arguments, so a new metric needs no registration step, and the return value is not used in the examples, which means a failed log call would not be noticed unless you check the return value yourself.

Framework callbacks cover Keras, TensorFlow, Fastai, Lightning and HuggingFace

Rather than ask every project to log metrics by hand, TraceML ships integration callbacks that hook into the training loop of each supported framework. The Keras integration is imported from traceml.integrations.keras and documented as automatically saving all metrics and collecting outputs and models, while still allowing extra information through the logging methods. The same Callback class is passed into fit:

model.fit(
    x_train,
    y_train,
    validation_data=(X_test, y_test),
    epochs=epochs,
    batch_size=100,
    callbacks=[Callback()],
)

The TensorFlow integration is imported from traceml.integrations.tensorflow and its callback takes flags for what to collect:

estimator.train(hooks=[Callback(log_image=True, log_histo=True, log_tensor=True)])

Those flags are the interesting part, because they are the only place in the README where the content of the tracked payload is specified in any detail. Fastai uses traceml.integrations.fastai with the callback passed as learn.fit(..., cbs=[Callback()]), PyTorch Lightning uses traceml.integrations.pytorch_lightning with the callback supplied as logger=Callback() inside pl.Trainer, and HuggingFace uses traceml.integrations.hugging_face with callbacks=[Callback].

PyTorch is the manual path, and it is the clearest one

PyTorch is handled differently from the other frameworks, and the example explains why. There is no PyTorch callback in the README, only the statement that you can log metrics, inputs and outputs of PyTorch experiments using the tracking module. That means the training loop stays yours:

for batch_idx, (data, target) in enumerate(train_loader):
    output = model(data)
    loss = F.nll_loss(output, target)
    loss.backward()
    optimizer.step()
    tracking.log_metrics(loss=loss)

Saving a checkpoint then follows a two-step pattern. You ask the tracking module for a path inside the run's artifacts area, write the file there, and register it as an artifact:

asset_path = tracking.get_outputs_path('model.ckpt')
torch.save(model.state_dict(), asset_path)

# log model
tracking.log_artifact_ref(asset_path, framework="pytorch", ...)

This is arguably the cleanest illustration of the design in the whole README. The engine does not own the file, it only names a location and records a reference, and the framework argument tells the record what kind of artifact it is. The trade-off is that a metric logged after the optimizer step is only captured if you remember the call, and nothing in the documentation suggests the module can attach itself to an arbitrary PyTorch loop.

Artifact plotting, packaging and the unfinished local sandbox

The README's artifact section starts to show how plots are produced and then cuts off mid-function, with a matplotlib figure routine and imports for altair, plotly, bokeh and vega_datasets. Because the example is truncated, the API for attaching a plot to a run is not visible here, and the safe conclusion is that the plotting helpers are partially documented. On packaging, installation is a plain pip install of traceml, with polyaxon added when you want the tracking features, and the repository itself is small: a README, a LICENSE, a CODE_OF_CONDUCT.md, a CONTRIBUTING.md and the single traceml package directory. The licence is Apache 2.0. A local sandbox section is explicitly marked WIP and says coming soon, so the standalone local server experience is not there yet. There are no GitHub releases, and the last push was on 2026-09-17.

Editorial conclusion

TraceML is a reasonable choice for teams already running Polyaxon who want run metadata, metrics and artifacts collected with one pip install and no server, and for offline experiments where the artifacts directory is enough. It is the wrong tool when you need the drift detection, explainability and dashboard features named in the project description, because the README documents no API for them, and it is the wrong tool when you need a versioned release, since the repository publishes no GitHub releases and the last push was on 2026-09-17. Before adopting it, run the offline example with POLYAXON_OFFLINE set, check that your framework has a matching traceml.integrations callback, and confirm that the artifacts path is writable from the training process.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. polyaxon/traceml on GitHub
4. README