SIENTIAPDE-1646

Enhance MLFlow run ID resolution with error handling for missing and invalid source URIs

- Added checks in `get_model_run_id` method to raise exceptions for models with missing or invalid source URIs.
- Introduced new test cases to validate error handling for these scenarios.
- Updated `requirements-light.txt` to include `mlflow` as a dependency.

laborious/utils/repository/model_repository.py

@@ -170,8 +170,19 @@ class MLFlowRepository(SientiaMonitoring):
 
         # Sort by version number to get the latest
         latest_version = max(stage_versions, key=lambda v: int(v.version))
-        run_id = latest_version.source.split('/')
-        return run_id[2]
+        source = latest_version.source
+        if source is None:
+            raise mlflow.exceptions.MlflowException(
+                f"Model '{model_name}' version '{latest_version.version}' in stage '{stage}' "
+                'has no source URI to resolve run ID.'
+            )
+        parts = source.split('/')
+        if len(parts) <= 2 or not parts[2]:
+            raise mlflow.exceptions.MlflowException(
+                f"Model '{model_name}' version '{latest_version.version}' in stage '{stage}' "
+                f"has invalid source URI '{source}' for run ID resolution."
+            )
+        return parts[2]
 
     def get_next_run_name(self, model_name: str) -> str:
         """

requirements-light.txt

@@ -4,6 +4,7 @@ sqlalchemy
 asyncua==1.0.6
 redis
 sientia_do>=1.12.1
+mlflow
 prometheus-client
 botocore
 boto3

tests/laborious/utils/repository/test_model_repository.py

@@ -142,6 +142,41 @@ def test_get_model_run_id_success(mlflow_repository):
     assert output == '1'
 
 
+def test_get_model_run_id_missing_source(mlflow_repository):
+    mlflow_repository.client.search_registered_models.return_value = [MagicMock(name='test')]
+
+    mlflow_repository.client.search_model_versions.return_value = [
+        MagicMock(current_stage='Production', version='1', source='runs/test/0'),
+        MagicMock(current_stage='Production', version='2', source=None),
+    ]
+
+    with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
+        mlflow_repository.get_model_run_id('test')
+
+    assert (
+        str(exc_info.value)
+        == "Model 'test' version '2' in stage 'Production' has no source URI to resolve run ID."
+    )
+
+
+def test_get_model_run_id_invalid_source(mlflow_repository):
+    mlflow_repository.client.search_registered_models.return_value = [MagicMock(name='test')]
+
+    mlflow_repository.client.search_model_versions.return_value = [
+        MagicMock(current_stage='Production', version='1', source='runs/test/0'),
+        MagicMock(current_stage='Production', version='2', source='runs/test'),
+    ]
+
+    with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
+        mlflow_repository.get_model_run_id('test')
+
+    assert (
+        str(exc_info.value)
+        == "Model 'test' version '2' in stage 'Production' has invalid source URI "
+        "'runs/test' for run ID resolution."
+    )
+
+
 def test_get_next_run_name(mlflow, mlflow_repository):
     mlflow.search_runs.return_value = [1, 2, 3]
     output = mlflow_repository.get_next_run_name('run')