Lightning-AI / Lightning-AI/pytorch-lightning
Log key prefix controls for Trainer, WandbLogger, and CometLogger
- 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
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