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,