Update Mlflow

Author: naglepuff

bats_ai/core/management/commands/examplelog.py

@@ -14,7 +14,6 @@ def command(experiment_name):
     if experiment:
         click.echo(f'Creating a log for experiment {experiment_name}')
         example_train.delay(experiment_name)
-        # train_body(experiment_name)
     else:
         click.echo(
             f'Could not find experiment {experiment_name}.'

bats_ai/core/management/commands/makeexperiment.py

@@ -19,6 +19,7 @@
 )
 @click.command()
 def command(username, name, description: str | None = None):
+    user = None
     if username:
         user = User.objects.get(username=username)
     else:

bats_ai/core/management/commands/registeronnxmodel.py

@@ -3,7 +3,7 @@
 from django.conf import settings
 import djclick as click
 import mlflow
-import mlflow.onnx
+import mlflow.onnx as mlflow_onnx
 import onnx
 
 
@@ -19,7 +19,7 @@ def command():
     with mlflow.start_run() as run:
         run_id = run.info.run_id
         click.echo(f'Run ID: {run_id}')
-        mlflow.onnx.log_model(
+        mlflow_onnx.log_model(
             onnx_model=onnx_model,
             artifact_path='onnx_model',
             # save_as_external_data=True,

bats_ai/core/management/commands/setupmlflow.py

@@ -7,7 +7,7 @@
 
 @click.command()
 def setupmlflow():
-    db_name = settings.MLFLOW_PG_DB if settings.MLFLOW_PG_DB else 'mlflow'
+    db_name = settings.MLFLOW_DB if settings.MLFLOW_DB else 'mlflow'
     bucket_name = settings.MLFLOW_BUCKET if settings.MLFLOW_BUCKET else 'mlflow'
 
     click.echo(f'Creating database {db_name} for mlflow')

bats_ai/settings/development.py

@@ -15,6 +15,18 @@
 staticfiles_index = INSTALLED_APPS.index('django.contrib.staticfiles')
 INSTALLED_APPS.insert(staticfiles_index, 'whitenoise.runserver_nostatic')
 
+DATABASES = {
+    **DATABASES,
+    'mlflow': {
+        **env.db_url('DJANGO_MLFLOW_DB_URL', engine='django.contrib.gis.db.backends.postgis'),
+        'CONN_MAX_AGE': timedelta(minutes=10).total_seconds(),
+    }
+}
+
+MLFLOW_BUCKET: str = env.str('DJANGO_MLFLOW_BUCKET')
+MLFLOW_DB: str = env.str('DJANGO_MLFLOW_DB')
+MLFLOW_ENDPOINT: str = env.str('DJANGO_MLFLOW_ENDPOINT')
+
 # Include Debug Toolbar middleware as early as possible in the list.
 # However, it must come after any other middleware that encodes the response's content,
 # such as GZipMiddleware.

bats_ai/tasks/tasks.py

@@ -264,6 +264,10 @@ def train_body(experiment_name: str):
         mlflow.log_metric('accuracy', accuracy)
         mlflow.set_tag('Training Info', 'Basic LR model for iris data')
 
+        print("ENV AWS_ACCESS_KEY_ID =", os.getenv("AWS_ACCESS_KEY_ID"))
+        print("ENV AWS_SECRET_ACCESS_KEY =", os.getenv("AWS_SECRET_ACCESS_KEY"))
+        print("ENV MLFLOW_S3_ENDPOINT_URL =", os.getenv("MLFLOW_S3_ENDPOINT_URL"))
+
         signature = infer_signature(X_train, lr.predict(X_train))
         _ = mlflow.sklearn.log_model(
             sk_model=lr,

dev/.env.docker-compose

@@ -3,6 +3,13 @@ DJANGO_DATABASE_URL=postgres://postgres:postgres@postgres:5432/django
 DJANGO_CELERY_BROKER_URL=amqp://rabbitmq:5672/
 DJANGO_MINIO_STORAGE_URL=http://minioAccessKey:minioSecretKey@minio:9000/django-storage
 DJANGO_MINIO_STORAGE_MEDIA_URL=http://localhost:9000/django-storage
+
+# Mlflow
+DJANGO_MLFLOW_ENDPOINT=http://localhost:5000
+DJANGO_MLFLOW_DB_URL=postgres://postgres:postgres@localhost:5432/mlflow
+DJANGO_MLFLOW_DB=mlflow
+DJANGO_MLFLOW_BUCKET=mlflow
+
 # When in Docker, the bridge network sends requests from the host machine exclusively via a
 # dedicated IP address. Since there's no way to determine the real origin address,
 # consider any IP address (though actually this will only be the single dedicated address) to

dev/.env.docker-compose-native

@@ -3,6 +3,15 @@ DJANGO_DATABASE_URL=postgres://postgres:postgres@localhost:5432/django
 DJANGO_CELERY_BROKER_URL=amqp://localhost:5672/
 DJANGO_MINIO_STORAGE_URL=http://minioAccessKey:minioSecretKey@localhost:9000/django-storage
 
+# MLFLOW
+DJANGO_MLFLOW_ENDPOINT=http://localhost:5000
+DJANGO_MLFLOW_DB_URL=postgres://postgres:postgres@localhost:5432/mlflow
+DJANGO_MLFLOW_DB=mlflow
+DJANGO_MLFLOW_BUCKET=mlflow
+AWS_ACCESS_KEY_ID=minioAccessKey  # Needed for celery, not for docker
+AWS_SECRET_ACCESS_KEY=minioSecretKey
+MLFLOW_S3_ENDPOINT_URL=http://localhost:900
+
 SERVERHOSTNAME=localhost
 
 APPLICATION_CLIENT_ID=HSJWFZ2cIpWQOvNyCXyStV9hiOd7DfWeBOCzo4pP # Application Identification
@@ -11,7 +20,6 @@ VITE_APP_API_ROOT=http://localhost:8000/api/v1
 VITE_APP_OAUTH_API_ROOT=http://localhost:8000/oauth/
 VITE_APP_OAUTH_CLIENT_ID=HSJWFZ2cIpWQOvNyCXyStV9hiOd7DfWeBOCzo4pP
 VITE_APP_LOGIN_REDIRECT=http://localhost:3000
-DJANGO_MLFLOW_ENDPOINT=http://localhost:5000
 # Needed for AWS/minio access from mlflow
 AWS_ACCESS_KEY_ID=minioAccessKey
 AWS_SECRET_ACCESS_KEY=minioSecretKey

dev/mlflow.Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/mlflow/mlflow:v2.20.3
+FROM ghcr.io/mlflow/mlflow:v3.3.0
 
 
 # hadolint ignore=DL3013

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "gunicorn",
     "geopandas>=1.1.1",
     "mlflow",
+    "psycopg2>=2.9.10",
 ]
 
 [project.optional-dependencies]