Skip to main content

Guangzhou, China

MLFlow Docker

MLFlow in Docker

This is just an experiment to see if I can use MLFlow inside my pytorch-jupyter Docker container with the latest version of YOLOv8.1. To do this I added MLFlow to the Dockerfile:

FROM pytorch/pytorch:latest

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive

# Install system dependencies
RUN apt-get update && \
    apt-get install -y \
        git \
        tini \
        python3-pip \
        python3-dev \
        python3-opencv \
        libglib2.0-0

# intall optional python deps
RUN python -m pip install --upgrade pip
# jupyter notebooks
RUN pip install jupyter
# fastdup https://github.com/visual-layer/fastdup
RUN pip install fastdup
RUN pip install opencv-python
RUN pip install matplotlib matplotlib-inline pandas
RUN pip install pillow
RUN pip install pyyaml
# YOLO 8.1
RUN pip install ultralytics Cython>=0.29.32 lapx>=0.5.5
# MLFlow 2.10
RUN pip install mlflow pytorch_lightning



# Set the working directory
WORKDIR /opt/app

# Start the notebook
RUN chmod +x /usr/bin/tini
ENTRYPOINT ["/usr/bin/tini", "--"]
CMD ["jupyter", "notebook", "--port=8888", "--no-browser", "--ip=0.0.0.0", "--allow-root"]

Let's build this custom image with:

docker build -t pytorch-jupyter . -f Dockerfile

I can now create the container and mount my working directory into the container WORKDIR to get started:

docker run --ipc=host --gpus all -ti --rm \
    -v $(pwd):/opt/app -p 8888:8888 -p 5000:5000 \
    --name pytorch-jupyter \
    pytorch-jupyter:latest

This will start Jupyter on Port 8888 and I can start MLFlow manually:

docker exec -ti pytorch-jupyter mlflow ui --host 0.0.0.0

The MLFlow is now available on localhost:5000 on my host system:

MLFlow in Docker

Just to be sure I stop MLFlow and try to run it directly from a Jupyter Notebook:

get_ipython().system_raw("mlflow ui --port 5000 --host 0.0.0.0 --backend-store-uri runs/mlflow &")

And the UI is still available on localhost:5000 - nice:

MLFlow in Docker

Create an MLFlow Experiment​

Enable MLflow Logging​

from ultralytics import settings

settings.update({
    'mlflow': True,
    'clearml': False,
    'comet': False,
    'dvc': False,
    'hub': False,
    'neptune': False,
    'raytune': False,
    'tensorboard': False,
    'wandb': False
})

Model Training​

from ultralytics import YOLO

# Load a pretrained YOLO model (recommended for training)
model = YOLO('yolov8n.pt')

# Train the model using the 'coco128.yaml' dataset for 3 epochs
results = model.train(data='coco128.yaml', epochs=3)

MLFlow in Docker