final

config.env

@@ -5,5 +5,5 @@ AZURE_STORAGE_CONNECTION_STRING=DefaultEndpointsProtocol=https;AccountName=azrml
 AZURE_STORAGE_ACCESS_KEY=tDGOKfiZ2svfoMvVmS0Fbpf0FTHfTq4wKYuDX7cAxlhve/3991QuzdvJHm9vWc+lo6mtC+x9yPSghWNR4+gacg==
 EXPERIMENT_NAME=Golden Batch Models Test
 RUN_NAME=r5-golden-batch-test
-MODEL_NAME=r5_model
-TOTAL_MODELS_NEEDED=2
\ No newline at end of file
+MODEL_NAME=r5_model/model.pkl
+TOTAL_MODELS_NEEDED=1
\ No newline at end of file

scripts/core/mlflow_util.py

@@ -20,14 +20,23 @@ mlflow.set_registry_uri(mlflow_tracking_uri)
 client = mlflow.tracking.MlflowClient()
 
 
-class MlFlowUtil:
+class MlflowCleanUp:
+    def __init__(self, experiment_name, parent_run_name, model_save_name):
+        self.experiment_name = experiment_name
+        self.parent_run_name = parent_run_name
+        self.model_save_name = model_save_name
+        self.total_models_needed = int(MlflowMetaData.TOTAL_MODELS_NEEDED)
+        self.model_name = MODEL_NAME
+        self.model_history_key = 'tags.mlflow.log-model.history'
+        self.model_parent_run_id_key = 'tags.mlflow.parentRunId'
+
     @staticmethod
-    def delete_artifact(run_id, parent_run_name, artifact_uri, file_path, model_name):
+    def delete_artifact(run_id, parent_run_name, artifact_uri, file_path):
         logger.info(f"Deleting artifact for {run_id} under {parent_run_name}")
         container_name = artifact_uri.split("//")[-1].split('@')[0]
         mlflow_name = artifact_uri.split("//")[-1].split('@')[-1].split('/')[1]
         mlflow_id = artifact_uri.split("//")[-1].split('@')[-1].split('/')[2]
-        path = f'{mlflow_name}/{mlflow_id}/{run_id}/artifacts/{file_path}/{model_name}'
+        path = f'{mlflow_name}/{mlflow_id}/{run_id}/artifacts/{file_path}'
         logger.info(f'Deleting artifact from path: {path}')
         blob_service_client = BlobServiceClient.from_connection_string(MlflowMetaData.AZURE_STORAGE_CONNECTION_STRING)
         container_client = blob_service_client.get_container_client(container_name)
@@ -39,18 +48,6 @@ class MlFlowUtil:
         else:
             logger.info(f"The blob {path} does not exist, which means its already deleted")
 
-
-class MlflowCleanUp:
-    def __init__(self, experiment_name, parent_run_name, model_save_name):
-        self.experiment_name = experiment_name
-        self.parent_run_name = parent_run_name
-        self.model_save_name = model_save_name
-        self._mfu_ = MlFlowUtil()
-        self.total_models_needed = int(MlflowMetaData.TOTAL_MODELS_NEEDED)
-        self.model_name = MODEL_NAME
-        self.model_history_key = 'tags.mlflow.log-model.history'
-        self.model_parent_run_id_key = 'tags.mlflow.parentRunId'
-
     def check_experiment(self):
         experiment_info = mlflow.get_experiment_by_name(self.experiment_name)
         if experiment_info is None:
@@ -92,8 +89,8 @@ class MlflowCleanUp:
                     for record in all_records:
                         run_id = record['run_id']
                         artifact_uri = record['artifact_uri']
-                        self._mfu_.delete_artifact(run_id, run_name_mapping[parent_run_id], artifact_uri,
-                                                   self.model_save_name, self.model_name)
+                        self.delete_artifact(run_id, run_name_mapping[parent_run_id], artifact_uri,
+                                             self.model_save_name)
                 else:
                     logger.info(f'No records to cleanup for Run {run_name_mapping[parent_run_id]}')
             else: