feat(operators): add mlflow integration

Author: Timon Viola

pyproject.toml

@@ -243,9 +243,7 @@ ignore = [
     "D214",
     "D215",
     "E731", # Do not assign a lambda expression, use a def
-    "TCH003", # Do not move imports from stdlib to TYPE_CHECKING block
-    "PT004", # Fixture does not return anything, add leading underscore
-    "PT005", # Fixture returns a value, remove leading underscore
+    "TC003", # Do not move imports from stdlib to TYPE_CHECKING block
     "PT006", # Wrong type of names in @pytest.mark.parametrize
     "PT007", # Wrong type of values in @pytest.mark.parametrize
     "PT011", # pytest.raises() is too broad, set the match parameter

src/dagcellent/operators/mlflow/_utils.py

@@ -1,5 +1,6 @@
-"""Mlflow operator."""
+"""Mlflow helpers."""
 # ruff: noqa: G004
+# typing: ignore
 from __future__ import annotations
 
 from enum import Enum
@@ -11,15 +12,15 @@
 
 
 class SlimModelVersion(TypedDict):
-    """Slim, json serializable type of mlflow.entities.model_registry.ModelVersion."""
+    """Slim, JSON serializable type of mlflow.entities.model_registry.ModelVersion."""
 
     name: str
     version: str
     run_id: str
     tags: dict[str, str]
 
 
-def _serialize_model_version(
+def serialize_model_version(
     model_version: mlflow.entities.model_registry.ModelVersion,
 ) -> SlimModelVersion:
     return {

src/dagcellent/operators/mlflow/get_latest_model_version.py

@@ -14,11 +14,11 @@
 from typing_extensions import ParamSpec
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
-
     # NOTE ruff fails for this check
     import mlflow.entities.model_registry  # noqa: TCH004
 
+    from dagcellent.operators.mlflow._utils import MlflowModelStage
+
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -120,7 +120,7 @@ def transition_model_version_stage(
     def get_latest_versions(
         self: MlflowHook,
         name: str,
-        stages: Sequence[str],
+        stages: list[str],
     ) -> list[mlflow.entities.model_registry.ModelVersion]:
         """Get latest model version. See MLFlow docs.