Lightning-AI / Lightning-AI/pytorch-lightning

Log key prefix controls for Trainer, WandbLogger, and CometLogger

Open
#21,782 0 comments 0 reactions 0 assignees View on GitHub
feature
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

Lightning automatically injects metric keys like `epoch` (via `Trainer`) and `global_step` (via
`WandbLogger`) into the logged metrics stream. These keys are currently hardcoded and cannot be
prefixed or renamed without monkey-patching internals.

This creates a friction point when users want a clean, consistent metric namespace. For example,
when using W&B or Comet, it is common to organise metrics under prefixes like `trainer/`,
`train/`, `val/`, etc. User-logged metrics can already be namespaced freely via `self.log(...)`,
but Lightning's own injected keys (`epoch`, `trainer/global_step`) sit outside that control.
They are either hardcoded into a different namespace or into no namespace at all.

A concrete pain point: `WandbLogger` hardcodes `trainer/global_step` as the step axis key, but
`Trainer` logs the epoch metric as bare `epoch`. There is no way to align these without modifying
Lightning internals. This means manual work on each new experiment to align the dashboards in
W&B.

### Pitch

Add controls so that each component that *generates* metric keys exposes a parameter to configure
those keys without touching any metric keys the user provides.

The design principle:

> `log_key_prefix` prefixes the keys *generated by that component*.

### Proposed API

**`Trainer(log_key_prefix=...)`**

Prefixes Trainer-generated metric keys. Currently the only such key is the automatic `epoch`
metric injected at each logging step.

```python
Trainer(log_key_prefix="trainer/")
# logs: "trainer/epoch"

Trainer(log_key_prefix="") # produces bare "epoch", the previous behaviour, backwards compatible.
Trainer(log_key_prefix=None) # produces bare "epoch", the previous behaviour, backwards compatible.
Trainer() # produces bare "epoch", the previous behaviour, backwards compatible.
```

**`WandbLogger(log_key_prefix=...)`**

Controls the W&B-generated step key used in `define_metric` and `experiment.log`. The existing
`prefix=` argument is unchanged and continues to apply to the pass-through metric keys.

```python
WandbLogger(log_key_prefix="trainer/") # Produces "trainer/global_step", the previous behaviour.

WandbLogger(log_key_prefix="") # Produces bare "global_step"
WandbLogger(log_key_prefix=None) # Produces bare "global_step"
WandbLogger() # Produces bare "global_step", NOT backwards compatible.
```

**`CometLogger(epoch_key=...)`**

Controls which incoming metric key is extracted and forwarded to Comet's dedicated `epoch=`
argument. `None` disables epoch extraction entirely.

```python
CometLogger() # Matches Trainer(), extracting "epoch", so it's backward compatible.
CometLogger(epoch_key="epoch") # Extracts "epoch".
CometLogger(epoch_key="trainer/epoch") # Extracts "trainer/epoch", matching Trainer(log_key_prefix="trainer/")
CometLogger(epoch_key=None) # Disables special epoch extraction.
```

### Aligning all Lightning-generated keys under one namespace

```python
Trainer(log_key_prefix="trainer/")
WandbLogger(log_key_prefix="trainer/")
CometLogger(epoch_key="trainer/epoch")
```

This places `trainer/epoch` and `trainer/global_step` in the same namespace while leaving all
user metrics untouched.

### Alternatives

I explored having a dictionary of namespaces in Trainer that other modules could pick from so
they would be automatically aligned, but it was creating a lot of coupling. WandbLogger needed
a backreference to the trainer to pick up the log_key_prefix for example.

If this is an approach that interests you, I can explore it.

### Additional context

_No response_

cc @lantiga

Contributor guide

Open the contributing guide

Research direction

Start by locating the Trainer, WandbLogger, and CometLogger entry points and the existing metric-key handling around define_metric and experiment.log. Trace how generated keys are separated from user-provided metrics, then verify that the proposed prefix and epoch_key options preserve the stated defaults and produce the aligned example keys without changing user metrics.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.