Lightning-AI / Lightning-AI/pytorch-lightning
MLFlowLogger does not save config.yaml for each run
- 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
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