From e9030c867d9d16dfe1152b5a0729a4297f05cd46 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20D=C5=82uski?= <tomasz.dluski@medicover.pl>
Date: Sun, 5 Nov 2023 21:15:28 +0100
Subject: [PATCH] feat: add built-in run_test_experiment job to test the mlflow
 server internally so its super easy to test if it actually works e2e

---
 docker-compose.yml                     | 19 +++++++++++++++++++
 quickstart/mlflow_tracking.py          | 12 ++++++------
 {mlflow => test_experiment}/Dockerfile |  5 ++---
 test_experiment/mlflow_tracking.py     | 22 ++++++++++++++++++++++
 4 files changed, 49 insertions(+), 9 deletions(-)
 rename {mlflow => test_experiment}/Dockerfile (50%)
 create mode 100644 test_experiment/mlflow_tracking.py

diff --git a/docker-compose.yml b/docker-compose.yml
index 18bc7bd..8cf4f2b 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -67,6 +67,25 @@ services:
     command: tcp db:3306 -t 90s -i 250ms
     networks:
       - internal
+  run_test_experiment:
+    build:
+      context: ./test_experiment
+      dockerfile: Dockerfile
+    depends_on:
+      - "mlflow"
+    environment:
+      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
+      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
+      - AWS_DEFAULT_REGION=${AWS_REGION}
+      - MLFLOW_S3_ENDPOINT_URL=http://s3:9000
+      - MLFLOW_TRACKING_URI=http://mlflow:5000
+    entrypoint: >
+      /bin/sh -c "
+      python3 mlflow_tracking.py;
+      exit 0;
+      "
+    networks:
+      - internal
 networks:
   internal:
   public:
diff --git a/quickstart/mlflow_tracking.py b/quickstart/mlflow_tracking.py
index 386a2d8..9cd9fa0 100644
--- a/quickstart/mlflow_tracking.py
+++ b/quickstart/mlflow_tracking.py
@@ -1,22 +1,22 @@
 import os
 from random import random, randint
 
-from mlflow import mlflow,log_metric, log_param, log_artifacts
+import mlflow
 
 if __name__ == "__main__":
     with mlflow.start_run() as run:
         mlflow.set_tracking_uri('http://localhost:5000')
         print("Running mlflow_tracking.py")
 
-        log_param("param1", randint(0, 100))
+        mlflow.log_param("param1", randint(0, 100))
         
-        log_metric("foo", random())
-        log_metric("foo", random() + 1)
-        log_metric("foo", random() + 2)
+        mlflow.log_metric("foo", random())
+        mlflow.log_metric("foo", random() + 1)
+        mlflow.log_metric("foo", random() + 2)
 
         if not os.path.exists("outputs"):
             os.makedirs("outputs")
         with open("outputs/test.txt", "w") as f:
             f.write("hello world!")
 
-        log_artifacts("outputs")
+        mlflow.log_artifacts("outputs")
diff --git a/mlflow/Dockerfile b/test_experiment/Dockerfile
similarity index 50%
rename from mlflow/Dockerfile
rename to test_experiment/Dockerfile
index f75ec2a..15a7467 100644
--- a/mlflow/Dockerfile
+++ b/test_experiment/Dockerfile
@@ -1,7 +1,6 @@
 FROM continuumio/miniconda3:latest
 
-RUN pip install mlflow boto3 pymysql
+RUN pip install mlflow boto3
 
-ADD . /app
 WORKDIR /app
-
+COPY . .
diff --git a/test_experiment/mlflow_tracking.py b/test_experiment/mlflow_tracking.py
new file mode 100644
index 0000000..be5236b
--- /dev/null
+++ b/test_experiment/mlflow_tracking.py
@@ -0,0 +1,22 @@
+import os
+from random import random, randint
+
+import mlflow
+
+if __name__ == "__main__":
+    with mlflow.start_run() as run:
+        mlflow.set_tracking_uri('http://mlflow:5000')
+        print("Running mlflow_tracking.py")
+
+        mlflow.log_param("param1", randint(0, 100))
+        
+        mlflow.log_metric("foo", random())
+        mlflow.log_metric("foo", random() + 1)
+        mlflow.log_metric("foo", random() + 2)
+
+        if not os.path.exists("outputs"):
+            os.makedirs("outputs")
+        with open("outputs/test.txt", "w") as f:
+            f.write("hello world!")
+
+        mlflow.log_artifacts("outputs")