Lightning-AI / Lightning-AI/pytorch-lightning

MLFlowLogger does not save config.yaml for each run

Open
#20,184 0 comments 0 reactions 0 assignees View on GitHub
bug logger: mlflow ver: 2.4.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

The `MLFlowLogger` seems to save the `config.yaml` in the top-level `save_dir` (e.g. `./mlruns`) directory (not even inside the experiment directory), instead of the specific run directory as for the other loggers. See below for minimal example. When running the same experiment twice, this results in an error because the `config.yaml` already exists.

Here is an example folder structure where you can see the `config.yaml` being at the top-level.
```shell
mlruns/
├── 557060468949431600 (experiment ID)
│ ├── 14625fca5e654f7faff19061b1ed44fa (run ID)
│ ├── 8b0a025336d6492391929adb37c18d2b (run ID)
│ └── meta.yaml
└── config.yaml
```

**Expected behavior:** just like with the default logger, we expect the `config.yaml` to be saved for inside the directory of each run of the given experiment.

```shell
mlruns/
└── 519079607625374876 (experiment ID)
├── 71d8f4b93eac490c8046d07bf7b49d31 (run ID)
│ ├── ...
│ └── config.yaml
├── 81a4e345f552487ea0d591e6bc14c881 (run ID)
│ ├── ...
│ └── config.yaml
└── meta.yaml
```

**Solution idea:** two lines of interest seem to be:
- [lightning/pytorch/loggers/mlflow.py#L302](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/loggers/mlflow.py#L302)
- [lightning/pytorch/trainer/trainer.py#L1227](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/trainer/trainer.py#L1227)

**Workaround 1:** we can just avoid the error with `LightningCLI(save_config_kwargs={"overwrite": True})` as suggested in the error message. However this does not save the config per-run.

**Workaround 2:** We can override [cli.SaveConfigCallback.save_config](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.cli.SaveConfigCallback.html#lightning.pytorch.cli.SaveConfigCallback) to set `save_to_log_dir=False`, and implement logic to save in the correct folder by using the experiment ID and run ID.

```python
from pathlib import Path

from lightning.fabric.utilities.cloud_io import get_filesystem
from lightning.pytorch.cli import LightningCLI, SaveConfigCallback
from lightning.pytorch.demos.boring_classes import DemoModel, BoringDataModule

class MLFlowSaveConfigCallback(SaveConfigCallback):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.save_to_log_dir = False

def save_config(self, trainer, pl_module, stage):
dir_runs = Path(trainer.logger.save_dir)
dir_run = dir_runs / trainer.logger.experiment_id / trainer.logger.run_id
path_config = dir_run / self.config_filename

fs = get_filesystem(dir_run)
fs.makedirs(dir_run, exist_ok=True)

self.parser.save(
self.config, path_config, skip_none=False, overwrite=self.overwrite, multifile=self.multifile
)

def cli_main():
LightningCLI(DemoModel, BoringDataModule,
save_config_callback=MLFlowSaveConfigCallback)

if __name__ == "__main__":
cli_main()
```

### What version are you seeing the problem on?

v2.4

### How to reproduce the bug
With the files below, run `python main.py fit --config config.yaml` twice. The first run will succeed, and the second one will fail with the error message below.

**main.py**
```python
from lightning.pytorch.cli import LightningCLI
from lightning.pytorch.demos.boring_classes import DemoModel, BoringDataModule

def cli_main():
LightningCLI(DemoModel, BoringDataModule)

if __name__ == "__main__":
cli_main()
```

**config.yaml**
```yaml
# lightning.pytorch==2.4.0
trainer:
logger:
class_path: lightning.pytorch.loggers.MLFlowLogger
```

### Error messages and logs

```shell
RuntimeError: SaveConfigCallback expected ./mlruns/config.yaml to NOT exist. Aborting to avoid overwriting results of a previous run. You can delete the previous config file, set `LightningCLI(save_config_callback=None)` to disable config saving, or set `LightningCLI(save_config_kwargs={"overwrite": True})` to overwrite the config file.
```

### Environment

Current environment

* CUDA:
- GPU:
- NVIDIA RTX 2000 Ada Generation Laptop GPU
- available: True
- version: 12.1
* Lightning:
- efficientnet-pytorch: 0.7.1
- lightning: 2.4.0
- lightning-utilities: 0.11.3.post0
- pytorch-lightning: 2.3.1
- segmentation-models-pytorch: 0.3.3
- torch: 2.3.1
- torchgeo: 0.5.2
- torchmetrics: 1.4.0.post0
- torchvision: 0.18.1
* Packages:
- aenum: 3.1.15
- affine: 2.4.0
- aiohttp: 3.9.5
- aiosignal: 1.3.1
- albucore: 0.0.12
- albumentations: 1.4.10
- alembic: 1.13.2
- aniso8601: 9.0.1
- annotated-types: 0.7.0
- antlr4-python3-runtime: 4.9.3
- asttokens: 2.4.1
- async-timeout: 4.0.3
- attrs: 23.2.0
- basemap: 1.4.1
- basemap-data: 1.3.2
- bitsandbytes: 0.43.1
- blinker: 1.8.2
- cachetools: 5.3.3
- certifi: 2024.6.2
- charset-normalizer: 3.3.2
- click: 8.1.7
- click-plugins: 1.1.1
- cligj: 0.7.2
- cloudpickle: 3.0.0
- comm: 0.2.2
- contourpy: 1.2.1
- cycler: 0.12.1
- databricks-sdk: 0.29.0
- debugpy: 1.8.2
- decorator: 5.1.1
- deprecated: 1.2.14
- docker: 7.1.0
- docstring-parser: 0.16
- efficientnet-pytorch: 0.7.1
- einops: 0.8.0
- entrypoints: 0.4
- exceptiongroup: 1.2.1
- executing: 2.0.1
- filelock: 3.15.4
- fiona: 1.9.6
- flask: 3.0.3
- fonttools: 4.53.0
- frozenlist: 1.4.1
- fsspec: 2024.6.1
- gitdb: 4.0.11
- gitpython: 3.1.43
- google-auth: 2.33.0
- graphene: 3.3
- graphql-core: 3.2.3
- graphql-relay: 3.2.0
- greenlet: 3.0.3
- gunicorn: 22.0.0
- huggingface-hub: 0.23.4
- hydra-core: 1.3.2
- idna: 3.7
- imageio: 2.34.2
- importlib-metadata: 7.2.1
- importlib-resources: 6.4.0
- ipykernel: 6.29.5
- ipython: 8.26.0
- itsdangerous: 2.2.0
- jedi: 0.19.1
- jinja2: 3.1.4
- joblib: 1.4.2
- jsonargparse: 4.31.0
- jupyter-client: 8.6.2
- jupyter-core: 5.7.2
- kiwisolver: 1.4.5
- kornia: 0.7.3
- kornia-rs: 0.1.4
- lazy-loader: 0.4
- lightly: 1.5.8
- lightly-utils: 0.0.2
- lightning: 2.4.0
- lightning-utilities: 0.11.3.post0
- mako: 1.3.5
- markdown: 3.6
- markdown-it-py: 3.0.0
- markupsafe: 2.1.5
- matplotlib: 3.8.4
- matplotlib-inline: 0.1.7
- mdurl: 0.1.2
- mlflow: 2.15.1
- mlflow-skinny: 2.15.1
- mpmath: 1.3.0
- multidict: 6.0.5
- munch: 4.0.0
- nest-asyncio: 1.6.0
- networkx: 3.3
- numpy: 1.26.4
- nvidia-cublas-cu12: 12.1.3.1
- nvidia-cuda-cupti-cu12: 12.1.105
- nvidia-cuda-nvrtc-cu12: 12.1.105
- nvidia-cuda-runtime-cu12: 12.1.105
- nvidia-cudnn-cu12: 8.9.2.26
- nvidia-cufft-cu12: 11.0.2.54
- nvidia-curand-cu12: 10.3.2.106
- nvidia-cusolver-cu12: 11.4.5.107
- nvidia-cusparse-cu12: 12.1.0.106
- nvidia-ml-py: 12.535.161
- nvidia-nccl-cu12: 2.20.5
- nvidia-nvjitlink-cu12: 12.5.82
- nvidia-nvtx-cu12: 12.1.105
- nvitop: 1.3.2
- omegaconf: 2.3.0
- opencv-python-headless: 4.10.0.84
- opentelemetry-api: 1.26.0
- opentelemetry-sdk: 1.26.0
- opentelemetry-semantic-conventions: 0.47b0
- packaging: 23.2
- pandas: 2.2.2
- parso: 0.8.4
- pexpect: 4.9.0
- pillow: 10.4.0
- pip: 24.1.1
- platformdirs: 4.2.2
- pretrainedmodels: 0.7.4
- prompt-toolkit: 3.0.47
- protobuf: 5.27.2
- psutil: 6.0.0
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pyarrow: 15.0.2
- pyasn1: 0.6.0
- pyasn1-modules: 0.4.0
- pydantic: 2.8.0
- pydantic-core: 2.20.0
- pygments: 2.18.0
- pyparsing: 3.1.2
- pyproj: 3.6.1
- pyshp: 2.3.1
- python-dateutil: 2.9.0.post0
- pytorch-lightning: 2.3.1
- pytz: 2024.1
- pyyaml: 6.0.1
- pyzmq: 26.0.3
- querystring-parser: 1.2.4
- rasterio: 1.3.10
- requests: 2.32.3
- rich: 13.7.1
- rsa: 4.9
- rtree: 1.2.0
- safetensors: 0.4.3
- scikit-image: 0.24.0
- scikit-learn: 1.5.0
- scipy: 1.14.0
- segmentation-models-pytorch: 0.3.3
- setuptools: 65.5.0
- shapely: 2.0.4
- six: 1.16.0
- smmap: 5.0.1
- snuggs: 1.4.7
- sqlalchemy: 2.0.32
- sqlparse: 0.5.1
- stack-data: 0.6.3
- sympy: 1.12.1
- tensorboardx: 2.6.2.2
- termcolor: 2.4.0
- threadpoolctl: 3.5.0
- tifffile: 2024.6.18
- timm: 0.9.2
- tomli: 2.0.1
- torch: 2.3.1
- torchgeo: 0.5.2
- torchmetrics: 1.4.0.post0
- torchvision: 0.18.1
- tornado: 6.4.1
- tqdm: 4.66.4
- traitlets: 5.14.3
- triton: 2.3.1
- typeshed-client: 2.5.1
- typing-extensions: 4.12.2
- tzdata: 2024.1
- urllib3: 2.2.2
- wcwidth: 0.2.13
- werkzeug: 3.0.3
- wrapt: 1.16.0
- yarl: 1.9.4
- zipp: 3.19.2
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.14
- release: 6.5.0-1025-oem
- version: #26-Ubuntu SMP PREEMPT_DYNAMIC Tue Jun 18 12:35:22 UTC 2024

### More info

_No response_

Contributor guide

Open the contributing guide

Research direction

Start with lightning/pytorch/loggers/mlflow.py around line 302 and lightning/pytorch/trainer/trainer.py around line 1227, then reproduce the issue with main.py and config.yaml using `python main.py fit --config config.yaml` twice. Done means each MLflow run stores config.yaml inside its own run directory, and repeating the experiment no longer fails because of a top-level config file.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
observability-sre
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.