Integrate with Ray¶

Ray is an open-source framework for scaling Python and AI workloads — from distributed model training to hyperparameter tuning and production serving.

Instrument your runs with Comet to track hyperparameters and metrics, monitor the system and GPU usage of every worker, and version model checkpoints — for faster, easier reproducibility and collaboration.

Comet integrates with both Ray Train (distributed training) and Ray Tune (hyperparameter tuning).

Ray Train¶

Ray Train scales model training across many workers and GPUs for frameworks such as PyTorch, TensorFlow, and XGBoost.

A distributed run spans many processes, but the Comet integration collects everything into a single experiment:

hyperparameters and the metrics you report during training;
per-worker system and GPU metrics, prefixed with each worker's rank — so you can confirm your expensive GPUs are fully used and your CPUs aren't the bottleneck;
optionally, model checkpoints, versioned as Comet Artifacts.

The integration has two public pieces, used together:

CometTrainLoggerCallback — added to the trainer's RunConfig on the driver (the script that launches the run).
comet_worker_logger — a context manager wrapped around your training function, which runs on each worker.

Ray Train V1 vs V2

Ray Train V2 (the ray.train/UserCallback API) is the default since Ray 2.51. Earlier Ray versions — or setting RAY_TRAIN_V2_ENABLED=0 — use V1 (the older ray.air/ray.tune-based API).

The Comet API is identical on both (CometTrainLoggerCallback + comet_worker_logger). The integration detects which Ray Train version is active and selects the correct callback automatically — you don't configure anything. Only the surrounding Ray APIs differ:

	Ray Train V2 (Ray ≥ 2.51, default)	Ray Train V1 (Ray ≤ 2.50, or `RAY_TRAIN_V2_ENABLED=0`)
Imports	`from ray.train import RunConfig, ScalingConfig`	`from ray.air.config import RunConfig, ScalingConfig`
Worker rank	`ray.train.get_context().get_world_rank()`	`from ray.air import session` → `session.get_world_rank()`
Report metrics	`ray.train.report({...})`	`session.report({...})`

The examples below are written for V2; a full V1 version is in Complete runnable example.

Connect Comet¶

Two changes connect Comet to a Ray Train script — wrap your training function with comet_worker_logger, and add CometTrainLoggerCallback to the RunConfig:

import comet_ml
from comet_ml.integration.ray import CometTrainLoggerCallback, comet_worker_logger

import ray.train
from ray.train import RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer

comet_ml.login()


def train_func(config):
    with comet_worker_logger(config) as experiment:
        # Your distributed training code here.
        ...
        # Report metrics to Ray Train; Comet logs them to the experiment.
        ray.train.report({"accuracy": accuracy})


config = {"lr": 1e-3, "epochs": 10}
callback = CometTrainLoggerCallback(config, project_name="comet-ray-example")

trainer = TorchTrainer(
    train_func,
    train_loop_config=config,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
    run_config=RunConfig(callbacks=[callback]),
)
result = trainer.fit()

As a shorthand, you can decorate the training function with @comet_worker instead of using the context manager.

Pass the same config object to both

The callback writes Comet's connection details (the experiment key, and the API key if shared) into the config dict you give it, and each worker reads them back from the train_loop_config it receives. So you must pass the same dict instance to CometTrainLoggerCallback(config, ...) and to train_loop_config=config — don't copy it or build a separate one, or the workers won't join the same experiment.

Note

comet_ml.login() picks up your API key, workspace, and project from the environment or config file. See Configure Comet for the other ways to set them.

Configuration options¶

Pass extra arguments to the callback on the driver to control logging:

callback = CometTrainLoggerCallback(
    config,
    project_name="comet-ray-example",
    parse_args=False,        # don't log CLI args as hyperparameters
    save_checkpoints=True,   # log reported Ray checkpoints as Comet Artifacts (online only)
    online=False,            # create an offline experiment instead of online
)

On each worker, pass any Experiment parameter to the context manager — for example, to change the metric reporting frequency:

with comet_worker_logger(config, auto_metric_step_rate=1) as experiment:
    ...

Sharing your API key with workers

Workers need a Comet API key to log. CometTrainLoggerCallback(share_api_key_to_workers=True) will distribute your key to the workers through the Ray config, which is convenient but insecure. For shared or multi-node clusters, prefer setting the API key in each worker's environment instead — see Distributed Training.

What gets logged¶

Everything from a run lands in a single Comet experiment, however many workers you use:

Hyperparameters — your train_loop_config (plus CLI arguments, unless parse_args=False).
Reported metrics — whatever you pass to ray.train.report({...}) (session.report({...}) on V1). These are logged once, from rank 0.
Per-worker system & GPU metrics — captured on every worker automatically and prefixed with the worker's rank (0.sys.*, 1.sys.*, …), so you can see each GPU's utilization separately.
Model checkpoints — each checkpoint you report is versioned as a Comet Artifact when save_checkpoints=True. This requires an online experiment.
Run environment — installed packages, source code, and git metadata, among other details Comet logs automatically.

Metrics normally reach Comet through report(...). comet_worker_logger also returns the run's Experiment, which you can use to log other data — images, audio, artifacts — directly from a worker. If you do log a metric this way, note that values logged from a worker are not automatically namespaced (only the system and GPU metrics above are rank-prefixed), so workers logging the same name overwrite each other — include the rank in the name to keep a value per worker (f"grad_norm_rank_{rank}").

You can turn each automatically-logged item on or off — see Configuration options.

Migrating from Ray Tune's CometLoggerCallback

Ray's built-in CometLoggerCallback is for Ray Tune, not Ray Train. For Ray Train, use Comet's CometTrainLoggerCallback — it additionally logs per-worker system and GPU metrics.

Note

Don't see what you need to log here? You can manually log any kind of data using the Experiment object — for example experiment.log_image for images or experiment.log_audio for audio.

Complete runnable example¶

A minimal, self-contained run (synthetic data, so there's nothing dataset-specific to set up). It reports metrics and checkpoints the model from rank 0.

Ray Train V2 (recommended)¶

python -m pip install -U comet_ml "ray[train]>=2.51" torch pillow

import os
import tempfile

import comet_ml
from comet_ml.integration.ray import CometTrainLoggerCallback, comet_worker_logger

import numpy as np
import ray.train
import ray.train.torch
import torch
from PIL import Image
from ray.train import Checkpoint, RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer
from torch import nn

comet_ml.login()


def train_func(config):
    with comet_worker_logger(config) as experiment:
        rank = ray.train.get_context().get_world_rank()

        # --- your training code; kept trivial here on purpose ---
        model = ray.train.torch.prepare_model(nn.Linear(20, 2))
        optimizer = torch.optim.SGD(model.parameters(), lr=config["lr"])
        loss_fn = nn.CrossEntropyLoss()

        for epoch in range(config["epochs"]):
            features, labels = torch.randn(64, 20), torch.randint(0, 2, (64,))
            loss = loss_fn(model(features), labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            # ---------------------------------------------------------

            # Report metrics, and from rank 0 also report a checkpoint. With
            # save_checkpoints=True (below) the callback logs each reported
            # checkpoint to Comet as an Artifact.
            with tempfile.TemporaryDirectory() as ckpt_dir:
                checkpoint = None
                if rank == 0:
                    torch.save(model.state_dict(), os.path.join(ckpt_dir, "model.pt"))
                    checkpoint = Checkpoint.from_directory(ckpt_dir)
                ray.train.report({"loss": loss.item(), "epoch": epoch}, checkpoint=checkpoint)

        # The worker's Experiment is for richer per-run data that doesn't go
        # through report() — images, audio, tables, ... Here, log a sample
        # image (e.g. a prediction) from rank 0.
        if rank == 0:
            sample = Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8))
            experiment.log_image(sample, name="sample_prediction")


config = {"lr": 1e-3, "epochs": 10}
callback = CometTrainLoggerCallback(
    config,
    project_name="comet-ray-example",
    save_checkpoints=True,  # log reported Ray checkpoints to Comet as Artifacts
)

trainer = TorchTrainer(
    train_func,
    train_loop_config=config,
    scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
    run_config=RunConfig(callbacks=[callback]),
)
trainer.fit()

Ray Train V1 (legacy)¶

The same run on the V1 API. Only the Ray pieces change (ray.air imports, session.report, session.get_world_rank) — the Comet code is unchanged.

python -m pip install -U comet_ml "ray[train]<=2.50" torch pillow

import os
import tempfile

import comet_ml
from comet_ml.integration.ray import CometTrainLoggerCallback, comet_worker_logger

import numpy as np
import ray.train.torch
import torch
from PIL import Image
from ray.air import session                       # V1: AIR session + config
from ray.air.config import RunConfig, ScalingConfig
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer
from torch import nn

comet_ml.login()


def train_func(config):
    with comet_worker_logger(config) as experiment:
        rank = session.get_world_rank()            # V1: rank via AIR session

        model = ray.train.torch.prepare_model(nn.Linear(20, 2))
        optimizer = torch.optim.SGD(model.parameters(), lr=config["lr"])
        loss_fn = nn.CrossEntropyLoss()

        for epoch in range(config["epochs"]):
            features, labels = torch.randn(64, 20), torch.randint(0, 2, (64,))
            loss = loss_fn(model(features), labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            with tempfile.TemporaryDirectory() as ckpt_dir:
                checkpoint = None
                if rank == 0:
                    torch.save(model.state_dict(), os.path.join(ckpt_dir, "model.pt"))
                    checkpoint = Checkpoint.from_directory(ckpt_dir)
                session.report({"loss": loss.item(), "epoch": epoch}, checkpoint=checkpoint)  # V1: AIR session

        if rank == 0:
            sample = Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8))
            experiment.log_image(sample, name="sample_prediction")


config = {"lr": 1e-3, "epochs": 10}
callback = CometTrainLoggerCallback(
    config,
    project_name="comet-ray-example",
    save_checkpoints=True,
)

trainer = TorchTrainer(
    train_loop_per_worker=train_func,
    train_loop_config=config,
    scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
    run_config=RunConfig(callbacks=[callback]),
)
trainer.fit()

Multi-GPU and multi-worker¶

However many GPUs/workers you scale to, the integration works unchanged. A few things to know about a distributed run:

Per-GPU resource monitoring. Every worker's system and GPU metrics are logged under its rank (0.sys.*, 1.sys.*, …) in the one experiment, so you can see utilization for each GPU at a glance.
Checkpoints reported from rank 0 are logged as Comet Artifacts when you set save_checkpoints=True (online experiments only).
Multiple nodes: give every worker the Comet API key via its environment rather than share_api_key_to_workers=True (see Distributed Training), and point offline runs at shared storage with offline_directory.

Try it out!¶

More examples: our examples GitHub repository.
Run the Keras example in Colab:
Preview the result of an example Ray Train experiment.

Ray Tune¶

Ray Tune is a Python library for experiment execution and hyperparameter tuning at any scale.

Connect Comet¶

Use Ray's built-in CometLoggerCallback to log Tune trials to Comet — add it to your RunConfig:

import comet_ml

from ray import tune
from ray.air.integrations.comet import CometLoggerCallback

comet_ml.login()

tuner = tune.Tuner(
    train_function,
    run_config=tune.RunConfig(callbacks=[CometLoggerCallback()]),
    # ... tune_config / param_space ...
)
tuner.fit()

You can configure the callback by passing it the same arguments you would pass to the Experiment object (project_name, tags, …). Learn more in the Ray documentation.

What gets logged¶

For every trial, Comet automatically logs:

Hyperparameters — from the Ray config / param_space.
Metrics — reported through tune.report(...) / ray.train.report(...).

When using rllib, the callback also logs episode-level metrics to Comet as curves.

End-to-end example¶

python -m pip install -U "ray[tune]" comet_ml

import comet_ml

import numpy as np
from ray import tune
from ray.air.integrations.comet import CometLoggerCallback

comet_ml.login()


def train_function(config):
    for _ in range(30):
        loss = config["mean"] + config["sd"] * np.random.randn()
        tune.report({"loss": loss})


tuner = tune.Tuner(
    train_function,
    tune_config=tune.TuneConfig(metric="loss", mode="min"),
    run_config=tune.RunConfig(callbacks=[CometLoggerCallback(tags=["my-trial"])]),
    param_space={
        "mean": tune.grid_search([1, 2, 3]),
        "sd": tune.uniform(0.2, 0.8),
    },
)
results = tuner.fit()

# Hyperparameters and metrics for every trial are now in Comet.
print("Best loss:", results.get_best_result().metrics["loss"])

Try it out!¶

Here's an example Colab Notebook for using Comet with Ray Tune.

Jul. 3, 2026