feat: Added JOB_NAME to MLFlow tag, lambda-rised to set_tag function …

…for future tagging in MLFlow
aisingapore · Sep 10, 2024 · 6e6e507 · 6e6e507
1 parent 787217b
commit 6e6e507
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/{{cookiecutter.repo_name}}/src/{{cookiecutter.src_package_name}}/general_utils.py b/{{cookiecutter.repo_name}}/src/{{cookiecutter.src_package_name}}/general_utils.py
@@ -80,10 +80,13 @@ def mlflow_init(args, run_name='train-model', setup_mlflow=False, autolog=False)
 
             mlflow.start_run(run_name=run_name)
 
-            if "MLFLOW_HPTUNING_TAG" in os.environ:
-                mlflow.set_tag("hptuning_tag", os.environ.get("MLFLOW_HPTUNING_TAG"))
-            if "JOB_UUID" in os.environ:
-                mlflow.set_tag("job_uuid", os.environ.get("JOB_UUID"))
+            set_tag = lambda env_var, tag_name='': mlflow.set_tag(
+                tag_name if tag_name != '' else env_var.lower(), 
+                os.environ.get(env_var)
+            ) if env_var in os.environ else None
+            set_tag("MLFLOW_HP_TUNING_TAG", "hptuning_tag")
+            set_tag("JOB_UUID")
+            set_tag("JOB_NAME")
 
             mlflow_run = mlflow.active_run()
             init_success = True