Modularized the post-promotion code; [skip ci]

Modularized both inference and monitoring logic; added some docstrings;
removed unused imports.

Author: Sebastijan-Dominis

ml/metadata/schemas/post_promotion/__init__.py

@@ -0,0 +1 @@
+"""A package for post-promotion metadata schemas."""

ml/metadata/schemas/post_promotion/infer.py

@@ -0,0 +1,21 @@
+"""Models for inference metadata."""
+from typing import Literal
+
+from pydantic import BaseModel
+
+from ml.modeling.models.feature_lineage import FeatureLineage
+
+
+class InferenceMetadata(BaseModel):
+    problem_type: str
+    segment: str
+    model_version: str
+    model_stage: str
+    run_id: str
+    timestamp: str
+    columns: list[str]
+    snapshot_bindings_id: str
+    feature_lineage: list[FeatureLineage]
+    artifact_type: Literal["pipeline", "model"]
+    artifact_hash: str
+    inference_latency_seconds: float

ml/metadata/validation/post_promotion/__init__.py

@@ -0,0 +1 @@
+"""A package for post-promotion metadata validation."""

ml/metadata/validation/post_promotion/infer.py

@@ -0,0 +1,26 @@
+"""A module for validating inference metadata in the inference pipeline."""
+import logging
+
+from ml.exceptions import RuntimeMLError
+from ml.metadata.schemas.post_promotion.infer import InferenceMetadata
+
+logger = logging.getLogger(__name__)
+
+def validate_inference_metadata(metadata: dict) -> InferenceMetadata:
+    """
+    Validate the inference metadata against the InferenceMetadata schema.
+
+    Args:
+        metadata (dict): The metadata dictionary to validate.
+
+    Returns:
+        InferenceMetadata: The validated metadata object.
+    """
+    try:
+        validated_metadata = InferenceMetadata.model_validate(metadata)
+        logger.debug("Successfully validated inference metadata.")
+        return validated_metadata
+    except Exception as e:
+        msg = "Error validating inference metadata."
+        logger.exception(msg)
+        raise RuntimeMLError(msg) from e

ml/metadata/validation/promotion/promote.py

@@ -1,3 +1,4 @@
+"""A module for validating promotion metadata against Pydantic models for staging and production."""
 import logging
 from typing import Literal
 

ml/post_promotion/inference/__init__.py

@@ -0,0 +1 @@
+"""A package for inference-related code."""

ml/post_promotion/inference/classes/__init__.py

@@ -0,0 +1 @@
+"""A package for classes in the inference pipeline."""

ml/post_promotion/inference/classes/function_returns.py

@@ -0,0 +1,18 @@
+"""A module for defining function return types in the inference pipeline."""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+
+@dataclass
+class ArtifactLoadingReturn:
+    """Return type for the artifact loading function."""
+    artifact: Any
+    artifact_hash: str
+    artifact_type: Literal["pipeline", "model"]
+
+@dataclass
+class PredictionStoringReturn:
+    """Return type for the prediction storing function."""
+    file_path: Path
+    cols: list[str]

ml/post_promotion/inference/classes/predictions_schema.py

@@ -0,0 +1,21 @@
+"""A module for defining the schema of predictions in the inference pipeline."""
+
+# Hardcoded, since these are not expected to change very often.
+# If changes become common, consider a refactor to make these configurable.
+# For now, simply update this file when changes are needed, and ensure that
+# all components of the pipeline are updated accordingly.
+SCHEMA_VERSION = "V1"
+
+BASE_EXPECTED_COLUMNS = [
+    "run_id",
+    "prediction_id",
+    "timestamp",
+    "model_stage",
+    "model_version",
+    "entity_id",
+    "input_hash",
+    "prediction",
+    "schema_version"
+]
+
+PROBA_PREFIX = "proba_"

ml/post_promotion/inference/execution/__init__.py

@@ -0,0 +1 @@
+"""A package for the execution code of the inference pipeline."""