Lightning-AI / Lightning-AI/pytorch-lightning

MLFlowLogger used with server crashes training

Open
#12,833 8 comments 5 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature logger: mlflow
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🐛 Bug

Exceptions in MLFlowLogger are not caught, these produce the break of the training and loss of the model. This happens most often when using a server to log data. Since the MLFlowLogger uses MLFlow, and MLFlow has an option to log data over a server, exceptions can happen for various reasons concerning the third party server components.

### To Reproduce

In order to reproduce the bug most effectively, you must run a mlflow server with some database. I used MySQL database and the server uses the local HTTP protocol.

#### Prepare the MySQL database and run the server

1. Install MySQL in your system.
```
sudo apt install mysql-server
```

2. Create a database to use as an MLflow backend tracking server.

```
CREATE DATABASE mlflow_tracking;
```

Under Linux mlflow server can be run with the following command:

```bash
user@pc:~/ModelFolder$ mlflow server --backend-store-uri mysql+pymysql://root@localhost/mlflow_tracking --default-artifact-root file:/home/user/ModelFolder/mlruns -h 0.0.0.0 -p 5000
```
Info:

- `ModelFolder` is the place of the training script
- `mlflow_tracking` is the MySQL database (it can be empty, mlflow will create the needed structure)
- `ModelFolder/mlruns` is the default dir where mlflow stores data

In order to run locally the server, you need to disable passwords for the MySQL user `root@localhost`.

You can use the following gist link to the training script (it is based on the default `bug_report_model.ipynb`):

[bug_report_model_mlflowlogger_server.ipynb ](https://gist.github.com/GinkoBalboa/ba16f2f08496a32d0136b7cc13ee1ce9)

### Expected behavior

Whenever the server is not working the training crashes.

### Proposed solution

I've proposed a solution to the given problem by catching exceptions. I've tested three scenarios with the server:

1. Server is turned off, and the script is started.
2. Server is on, the script is started, and the server is killed during the training.
3. Server is on, the script is started, and the server is killed after the first training finishes but before it writes to the database. For this, I used a sleep timer of 60 sec. in the script.

I've investigated these three scenarios since the exception is caught in different functions so a solution to only one scenario is not sufficient.

### Additional context

Similar issues were raised previously:

- https://github.com/PyTorchLightning/pytorch-lightning/issues/6641
- https://github.com/PyTorchLightning/pytorch-lightning/issues/9607

**NOTE**: PR with the proposed solution will follow quickly.

cc @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with MLFlowLogger and the linked bug_report_model_mlflowlogger_server.ipynb, then reproduce the three listed server-failure scenarios using an MLflow server backed by MySQL. Trace where each failure is raised and verify that training completes without losing the model when logging fails.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.