Add Mlflow to docker compose setup

Also add a test command to show communication between the django web
server and the new Mlflow service.

Author: naglepuff

bats_ai/core/management/commands/makeexperiment.py

@@ -0,0 +1,38 @@
+from django.conf import settings
+from django.contrib.auth.models import User
+import djclick as click
+from mlflow import MlflowClient
+
+
+@click.option(
+    '--username',
+    type=click.STRING,
+    required=True,
+    help='superuser username for mlflow experiment creation',
+)
+@click.option('--name', type=click.STRING, required=True, help='a name for the mlflow experiment')
+@click.option(
+    '--description',
+    type=click.STRING,
+    required=False,
+    help='a description for the mlflow experiment',
+)
+@click.command()
+def command(username, name, description: str | None = None):
+    if username:
+        user = User.objects.get(username=username)
+    else:
+        first_user = User.objects.first()
+        if first_user:
+            user = user
+
+    if user:
+        mlflow_client = MlflowClient(tracking_uri=settings.MLFLOW_ENDPOINT)
+
+        experiment_tags = {
+            'project-name': 'batsai',
+        }
+        if description:
+            experiment_tags['mlflow.note.content'] = description
+
+        mlflow_client.create_experiment(name=name, tags=experiment_tags)

bats_ai/settings.py

@@ -81,6 +81,8 @@ class DevelopmentConfiguration(BatsAiMixin, DevelopmentBaseConfiguration):
     MINIO_STORAGE_MEDIA_USE_PRESIGNED = True
     MINIO_STORAGE_MEDIA_URL = 'http://127.0.0.1:9000/django-storage'
 
+    MLFLOW_ENDPOINT = values.Value('http://localhost:5000')
+
 
 class TestingConfiguration(BatsAiMixin, TestingBaseConfiguration):
     pass

dev/.env.docker-compose

@@ -13,3 +13,4 @@ DOCKER_MINIO_PORT=9000
 DOCKER_MINIO_CONSOLE_PORT=9001
 APPLICATION_CLIENT_ID=HSJWFZ2cIpWQOvNyCXyStV9hiOd7DfWeBOCzo4pP # Application Identification
 NABAT_API_URL=https://api.sciencebase.gov/nabat-graphql/graphql
+DJANGO_MLFLOW_ENDPOINT=mlflow:5000

dev/.env.docker-compose-native

@@ -7,3 +7,4 @@ DJANGO_MINIO_STORAGE_SECRET_KEY=minioSecretKey
 DJANGO_STORAGE_BUCKET_NAME=django-storage
 APPLICATION_CLIENT_ID=HSJWFZ2cIpWQOvNyCXyStV9hiOd7DfWeBOCzo4pP # Application Identification
 NABAT_API_URL=https://api.sciencebase.gov/nabat-graphql/graphql
+DJANGO_MLFLOW_ENDPOINT=http://localhost:5000

docker-compose.yml

@@ -31,6 +31,16 @@ services:
     volumes:
       - minio:/data
 
+  mlflow:
+    image: ghcr.io/mlflow/mlflow:latest
+    ports:
+      - "5000:5000"
+    environment:
+      - MLFLOW_S3_ENDPOINT_URL=http://minio:9000
+      - AWS_ACCESS_KEY=${DJANGO_MINIO_STORAGE_ACCESS_KEY:-minioAccessKey}
+      - AWS_SECRET_ACCESS_KEY=${DJANGO_DATABASE_PASSWORD:-minioSecretKey}
+    command: mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri postgresql://postgres:postgres@postgres:5432/mlflow --artifacts-destination s3://bucket
+
 volumes:
   postgres:
   sourcedb:

setup.py

@@ -71,6 +71,8 @@
         # guano metadata
         'guano',
         'django_celery_results',
+        # ML Ops
+        'mlflow',
     ],
     extras_require={
         'dev': [