triggerflow 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

triggerflow/mlflow_wrapper.py

@@ -1,5 +1,6 @@
 # trigger_mlflow.py
 import mlflow
+import os
 import mlflow.pyfunc
 import tempfile
 from pathlib import Path
@@ -8,6 +9,120 @@ from mlflow.tracking import MlflowClient
 from .core import TriggerModel
 
 
+def setup_mlflow(mlflow_uri: str = None, 
+                 model_name: str = None, 
+                 experiment_name: str = None, 
+                 run_name: str = None, 
+                 experiment_id: str = None,
+                 run_id: str = None,
+                 save_env_file: bool = True,
+                 auto_configure: bool = False
+                 ):
+
+    # Set the MLflow tracking URI
+    if mlflow_uri is None:
+        mlflow_uri = os.getenv('MLFLOW_URI', 'https://ngt.cern.ch/models')
+    mlflow.set_tracking_uri(mlflow_uri)
+    os.environ["MLFLOW_URI"] = mlflow_uri
+    print(f"Using MLflow tracking URI: {mlflow_uri}")
+
+
+    # Set the model name
+    if model_name is None:
+        if os.getenv('MLFLOW_MODEL_NAME'):
+            model_name = os.getenv('MLFLOW_MODEL_NAME')
+        else:
+            model_name = os.getenv('CI_COMMIT_BRANCH', 'Test-Model')
+    os.environ["MLFLOW_MODEL_NAME"] = model_name
+    print(f"Using model name: {model_name}")
+
+
+    # Set the experiment name
+    if experiment_name is None:
+        if os.getenv('MLFLOW_EXPERIMENT_NAME'):
+            experiment_name = os.getenv('MLFLOW_EXPERIMENT_NAME')
+        else:
+            experiment_name = os.getenv('CI_COMMIT_BRANCH', 'Test-Training-Torso')
+    os.environ["MLFLOW_EXPERIMENT_NAME"] = experiment_name
+    print(f"Using experiment name: {experiment_name}")
+
+
+    # Set the run name
+    if run_name is None:
+        if os.getenv('CI') == 'true':
+            if os.getenv('CI_PARENT_PIPELINE_ID'):
+                run_name = f"{os.getenv('CI_PARENT_PIPELINE_ID')}-{os.getenv('CI_PIPELINE_ID')}"
+            else:
+                run_name = f"{os.getenv('CI_PIPELINE_ID')}"
+        else:
+            import datetime
+            run_name = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+    os.environ["MLFLOW_RUN_NAME"] = run_name
+    print(f"Using run name: {run_name}")
+
+
+    # Create a new experiment or get the existing one
+    if experiment_id is None:
+        if os.getenv("MLFLOW_EXPERIMENT_ID"):
+            experiment_id = os.getenv("MLFLOW_EXPERIMENT_ID")
+        else:
+            try:
+                experiment_id = mlflow.create_experiment(experiment_name)
+            except mlflow.exceptions.MlflowException:
+                experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id
+    
+    check_experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id
+    if str(check_experiment_id) != str(experiment_id):
+        raise ValueError(f"Provided experiment_id {experiment_id} does not match the ID of experiment_name {experiment_name} ({check_experiment_id})")
+
+    # if mlflow.get_experiment_by_name(experiment_name).experiment_id is None:
+    #     experiment_id = mlflow.create_experiment(experiment_name)
+    # else:
+    #     experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id
+
+    mlflow.set_experiment(experiment_id=experiment_id)
+    os.environ["MLFLOW_EXPERIMENT_ID"] = experiment_id
+    print(f"Using experiment ID: {experiment_id}")
+
+
+    # Start a new MLflow run
+    if run_id is None:
+        if os.getenv("MLFLOW_RUN_ID"):
+            run_id = os.getenv("MLFLOW_RUN_ID")
+        else:
+            with mlflow.start_run(experiment_id=experiment_id, run_name=run_name) as run:
+                run_id = run.info.run_id
+
+    check_run_info = mlflow.get_run(run_id)
+    if str(check_run_info.info.experiment_id) != str(experiment_id):
+        raise ValueError(f"Provided run_id {run_id} does not belong to experiment_id {experiment_id} (found {check_run_info.info.experiment_id})")
+    
+    os.environ["MLFLOW_RUN_ID"] = run_id
+    print(f"Started run with ID: {run_id}")
+
+    # Save environment variables to a file for later steps in CI/CD pipelines
+    if save_env_file and os.getenv("CI") == "true":
+        print(f"Saving MLflow environment variables to {os.getenv('CI_ENV_FILE', 'mlflow.env')}")
+        with open(os.getenv('CI_ENV_FILE', 'mlflow.env'), 'a') as f:
+            f.write(f"MLFLOW_URI={mlflow_uri}\n")
+            f.write(f"MLFLOW_MODEL_NAME={model_name}\n")
+            f.write(f"MLFLOW_EXPERIMENT_NAME={experiment_name}\n")
+            f.write(f"MLFLOW_RUN_NAME={run_name}\n")
+            f.write(f"MLFLOW_EXPERIMENT_ID={experiment_id}\n")
+            f.write(f"MLFLOW_RUN_ID={run_id}\n")
+
+            if auto_configure:
+                print("Auto_configure is set to true. Exporting AUTO_CONFIGURE=true")
+                f.write(f"AUTO_CONFIGURE=true\n")
+
+    return experiment_name, run_name, experiment_id, run_id
+
+if os.getenv("AUTO_CONFIGURE") == "true":
+    print("AUTO_CONFIGURE is true and running in CI environment. Setting up mlflow...")
+    setup_mlflow()
+else:
+    print("AUTO_CONFIGURE is not set. Skipping mlflow run setup")
+
 class MLflowWrapper(mlflow.pyfunc.PythonModel):
     """PyFunc wrapper for TriggerModel; backend can be set at runtime."""
     def load_context(self, context):

{triggerflow-0.1.6.dist-info → triggerflow-0.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: triggerflow
-Version: 0.1.6
+Version: 0.1.8
 Summary: Utilities for ML models targeting hardware triggers
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License

{triggerflow-0.1.6.dist-info → triggerflow-0.1.8.dist-info}/RECORD

@@ -1,11 +1,11 @@
 triggerflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 triggerflow/core.py,sha256=MRr2M2_9kvANx8ZW0wBnS0f6L3uAQBDPQle9NXsPbvs,20639
-triggerflow/mlflow_wrapper.py,sha256=2z-278Z1gvoEdVsNw6M1S7UoDw14Vlr2FCJWBVMLwSg,3889
+triggerflow/mlflow_wrapper.py,sha256=BJhkVMn7-H4AnjgnpHiprLUqvdXUXWvN4NPdL7anZhQ,8652
 triggerflow/templates/makefile,sha256=VL39isTUBewrs8zTSDzdP6LLln7zpGoCZnLadpMu7CA,808
 triggerflow/templates/makefile_version,sha256=Tmu0tyAopJbiBQVMMOa6l2Cz5GkEn20mwgzIi0CfhyM,338
 triggerflow/templates/model_template.cpp,sha256=eGwY5ca_HgjoIvqorOBPSJspP0wngpjJheq3meb48r4,1616
 triggerflow/templates/scales.h,sha256=5bq6lVF36SRQKE2zg9RpBG6K5orpPlnJ8g125nbtFow,365
-triggerflow-0.1.6.dist-info/METADATA,sha256=_EsPMT8F3BTu1hA9Ff0oxisqB2qeqpgZdIqf_k_PLIQ,1942
-triggerflow-0.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-triggerflow-0.1.6.dist-info/top_level.txt,sha256=g4M0nqpVPFZcmVmsoLExDtJFLDBK4fzobCIBqo13BEw,12
-triggerflow-0.1.6.dist-info/RECORD,,
+triggerflow-0.1.8.dist-info/METADATA,sha256=_ER4GTO1DCk4dgS_eALiHOeFflnNleuir8Q7_3xwAxM,1942
+triggerflow-0.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+triggerflow-0.1.8.dist-info/top_level.txt,sha256=g4M0nqpVPFZcmVmsoLExDtJFLDBK4fzobCIBqo13BEw,12
+triggerflow-0.1.8.dist-info/RECORD,,