Converting MLFlow Logs to AIM

[vishalbakshi]

I was able to track an experiment in MLFlow using composer (and its MLFlowLogger) and convert its logs to AIM.

Install MLFlow and Start Server

pip install mlflow

Then:

mlflow ui

This will give you a localhost URL where you can see their UI. It should by default to http://localhost:5000.

Training Script Changes

In composer_track.py add the following:

from composer.loggers import MLFlowLogger

...

aim_logger = MLFlowLogger(experiment_name="<experiment name>", tracking_uri="http://localhost:5000")

...

trainer = Trainer(
            ...
            loggers=aim_logger,
            ....
        )

...

MLFlowLogger doesn't have a ._run method so comment out the finally block:

finally:
        # Ensure proper cleanup
        print("Done.")
        # if aim_logger and aim_logger._run:
        #     print("Cleaning up Aim resources...")
        #     logger.info("Cleaning up Aim resources...")
        #     aim_logger.close()

Run the script with python compose_track.py.

Note: I had to install pytorch in my environment with:

mamba install pytorch::pytorch torchvision torchaudio -c pytorch

I also had to manually downgrade numpy to "numpy<2" otherwise it was throwing "Numpy not found" type error.

Training should happen quickly (took ~8 minutes on CPU) and logging should also happen quickly (no lag). The logs should be in a folder titled mlruns and have a default experiment 0 folder as well as a folder for the experiment you just ran (folder name will be a string of digits like 592053435467629112. Inside this folder you'll have a folder for each run, labeled with another string like 561facb27d6d46d9a4d24189e2f85187. Inside that folder is all the logged information like metrics. There's also YAML files with metadata.

Convert MLFlow Logs to AIM

I had to modify the CLI convert script so do not install AIM with pip (see this GitHub issue). Instead, clone the repo down and pip install -e . after making the following change in aim/convert/processors/mlflow.py:

Change:

experiments = client.list_experiments()

to

experiments = client.search_experiments()

You also have to set the tracking URI environment variable so you can do that in the terminal or add the following line after import mlflow:

mlflow.set_tracking_uri('http://localhost:5000')

Make sure the current directory is where you want your AIM stuff to be stored. Initialize AIM repository with:

aim init

Note that this creates a .aim folder which may be hidden in your file explorer.

Then run the following line, with the local path to your mlruns folder after --tracking_uri:

aim convert mlflow --tracking_uri 'file:///Users/username/path/to/mlruns'

This will create the necessary files in .aim which contains the new logs. Run aim up to see the AIM explorer and you should be able to see your logs!

[vishalbakshi]

MLFlow charts:

Corresponding AIM charts after conversion: