collie-mlops 0.1.0b0__py3-none-any.whl

collie/core/evaluator/evaluator.py

@@ -0,0 +1,147 @@
+from typing import (
+    Dict, 
+    Any,
+    List,
+    Optional
+)
+
+from collie.contracts.event import (
+    Event, 
+    EventHandler, 
+    EventType
+)
+from collie.contracts.mlflow import MLFlowComponentABC
+from collie._common.decorator import type_checker
+from collie.core.models import (
+    EvaluatorArtifact,
+    EvaluatorPayload
+)
+from collie._common.exceptions import EvaluatorError
+
+class Evaluator(EventHandler, MLFlowComponentABC):
+    
+    def __init__(
+        self, 
+        description: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None
+    ) -> None:
+        
+        """
+        Initializes the Evaluator.
+
+        Args:
+            description (Optional[str], optional): Description for the MLflow run. Defaults to None.
+            tags (Optional[Dict[str, str]], optional): Tags to associate with the MLflow run. Defaults to None.
+
+        """
+        super().__init__()
+        self._registered_model_name = None
+        self.description = description
+        self.tags = tags or {"component": "Evaluator"}
+
+        self.model_uri = None
+        self.metrics = None
+    
+    @property
+    def registered_model_name(self) -> str:
+        if not self._registered_model_name:
+            raise EvaluatorError("Registered model name is not set.")
+        return self._registered_model_name
+    
+    @registered_model_name.setter
+    def registered_model_name(self, name: str) -> None:
+        self._registered_model_name = name
+        
+    def run(self, event: Event) -> Event:
+       
+        with self.start_run(
+            tags=self.tags,
+            run_name="Evaluator",
+            log_system_metrics=True,
+            description=self.description,
+            nested=True,
+        ) as run:
+            
+            try:
+                evaluator_event = self._handle(event)
+                payload = self._get_evaluator_payload(evaluator_event)
+
+                self.metrics: List[Dict[str, Any]] = payload.metrics
+                self.model_uri = event.context.get('model_uri')
+        
+                self._log_metrics()
+                self._log_summary(payload)
+
+                event.context.set("evaluator_report_uri", self.artifact_uri(run))
+
+                if self.experiment_is_better(payload):
+                    event.context.set("pass_evaluation", True)
+                    self.mlflow.log_param("evaluation_result", "passed")
+                else:
+                    event.context.set("pass_evaluation", False)
+                    self.mlflow.log_param("evaluation_result", "failed")
+               
+            except Exception as e:
+                raise EvaluatorError(f"Evaluator failed: {e}") from e
+
+            return Event(
+                type=EventType.EVALUATION_DONE,
+                payload=payload,
+                context=event.context,
+            )
+
+    @staticmethod
+    def experiment_is_better(
+        payload: EvaluatorPayload
+    ) -> bool:
+        """
+        Decide if experiment model is better based on score direction.
+        """
+        return payload.is_better_than_production
+    
+    def _flatten_metrics(
+        self,
+        metrics: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        """
+        Flatten list of metric dictionaries into a single dictionary.
+        Useful for logging and analysis.
+        """
+        flattened = {}
+        for metric_dict in metrics:
+            flattened.update(metric_dict)
+        return flattened
+
+    def _log_summary(
+        self, 
+        payload: EvaluatorPayload
+    ) -> None:
+        """
+        Generate summary statistics from metrics.
+        """
+        flattened = self._flatten_metrics(payload.metrics)
+        
+        summary = {
+            "total_metrics": len(flattened),
+            "is_better": payload.is_better_than_production,
+            "metrics": flattened
+        }
+        self.mlflow.log_dict(
+            dictionary=summary, 
+            artifact_file=f"{EvaluatorArtifact().report}"
+        )
+    
+    def _log_metrics(self) -> None:
+        """
+        Log individual metrics to MLflow.
+        """
+        for metric_dict in self.metrics:
+            for metric_name, metric_value in metric_dict.items():
+                self.mlflow.log_metric(metric_name, metric_value)
+    
+    @type_checker((EvaluatorPayload,), "EvaluatorPayload must be of type EvaluatorPayload.")
+    def _get_evaluator_payload(self, event: Event) -> EvaluatorPayload: 
+        return event.payload
+    
+    def artifact_uri(self, run) -> str:
+        return f"{run.info.artifact_uri}/{EvaluatorArtifact().report}"

collie/core/models.py

@@ -0,0 +1,125 @@
+from typing import (
+    Dict, 
+    Any,
+    Optional,
+    List,
+)
+from pydantic import BaseModel, ConfigDict, Field
+
+import pandas as pd
+
+
+class BasePayload(BaseModel):
+    """Base class for all payload types with optional extra_data functionality."""
+    
+    extra_data: Optional[Dict[str, Any]] = None
+    
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    
+    def set_extra(self, key: str, value: Any):
+        """
+        Set a value in extra_data and return self for chaining.
+        
+        Args:
+            key: The key to set
+            value: The value to store
+            
+        Returns:
+            Self for method chaining
+            
+        Example:
+            >>> payload.set_extra("feature_names", ["age", "income"])
+            >>> payload.set_extra("n_classes", 3).set_extra("version", "1.0")
+        """
+        if self.extra_data is None:
+            self.extra_data = {}
+        self.extra_data[key] = value
+        return self
+    
+    def get_extra(self, key: str, default: Any = None) -> Any:
+        """
+        Get a value from extra_data with optional default.
+        
+        Args:
+            key: The key to retrieve
+            default: Default value if key doesn't exist
+            
+        Returns:
+            The value or default
+            
+        Example:
+            >>> features = payload.get_extra("feature_names", [])
+        """
+        if self.extra_data is None:
+            return default
+        return self.extra_data.get(key, default)
+    
+    def has_extra(self, key: str) -> bool:
+        """
+        Check if a key exists in extra_data.
+        
+        Args:
+            key: The key to check
+            
+        Returns:
+            True if key exists, False otherwise
+            
+        Example:
+            >>> if payload.has_extra("feature_names"):
+            ...     features = payload.get_extra("feature_names")
+        """
+        if self.extra_data is None:
+            return False
+        return key in self.extra_data
+
+
+class TransformerArtifact(BaseModel):
+    train_data: str = "train.csv"
+    validation_data: str = "validation.csv"
+    test_data: str = "test.csv"
+
+
+class TunerArtifact(BaseModel):
+    hyperparameters: str = "hyperparameters.json"
+
+
+class TrainerArtifact(BaseModel):
+    model: str = "model"
+
+
+class EvaluatorArtifact(BaseModel):
+    report: str = "report.json"
+
+
+class TransformerPayload(BasePayload):
+    train_data: Optional[pd.DataFrame] = None
+    validation_data: Optional[pd.DataFrame]  = None
+    test_data: Optional[pd.DataFrame]  = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class TrainerPayload(BasePayload):
+    model: Any = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class TunerPayload(BasePayload):
+    hyperparameters: Dict[str, Any]
+    train_data: Optional[pd.DataFrame] = None
+    validation_data: Optional[pd.DataFrame]  = None
+    test_data: Optional[pd.DataFrame]  = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class EvaluatorPayload(BasePayload):
+    metrics: List[Dict[str, Any]]
+    is_better_than_production: bool
+
+
+class PusherPayload(BasePayload):
+    model_uri: str
+    status: Optional[str] = None
+    model_version: Optional[str] = None

collie/core/orchestrator/orchestrator.py

@@ -0,0 +1,47 @@
+from typing import (
+    Optional,
+    Dict,
+)
+
+from collie.contracts.orchestrator import OrchestratorBase
+from collie.core.enums.components import CollieComponentType
+from collie._common.utils import get_logger
+
+logger = get_logger()
+
+
+class Orchestrator(OrchestratorBase):
+
+    def __init__(
+        self,
+        components: CollieComponentType,
+        tracking_uri: str,
+        registered_model_name: str,
+        mlflow_tags: Optional[Dict[str, str]] = None,
+        experiment_name: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            components=components,
+            tracking_uri=tracking_uri,
+            mlflow_tags=mlflow_tags,
+            experiment_name=experiment_name,
+            description=description
+        )
+        self.registered_model_name = registered_model_name
+
+    def orchestrate_pipeline(self) -> None:
+        """Run the pipeline sequentially without Airflow."""
+
+        logger.info("Pipeline started.")
+        incoming_event = self.initialize_event()
+
+        for idx, component in enumerate(self.components):
+            logger.info(f"Running component {idx}: {type(component).__name__}")
+            component.mlflow_config = self.mlflow_config
+            if hasattr(component, "_registered_model_name"):
+                component.registered_model_name = self.registered_model_name
+        
+            incoming_event = component.run(incoming_event)
+
+            logger.info(f"Component {idx} finished: {type(component).__name__}")

collie/core/pusher/pusher.py

@@ -0,0 +1,98 @@
+from typing import Optional
+from collie.contracts.event import (
+    Event, 
+    EventHandler, 
+    EventType
+)
+from collie.contracts.mlflow import MLFlowComponentABC
+from collie.core.models import PusherPayload
+from collie._common.decorator import type_checker
+from collie._common.exceptions import PusherError
+from collie.core.enums.ml_models import MLflowModelStage
+
+
+class Pusher(EventHandler, MLFlowComponentABC):
+    def __init__(
+        self,
+        target_stage: MLflowModelStage = MLflowModelStage.PRODUCTION,
+        archive_existing_versions: bool = True,
+        description: Optional[str] = None,
+        tags: Optional[dict] = None
+    ) -> None:
+        """
+        Initializes the Pusher.
+
+        Args:
+            target_stage (Optional[MLflowModelStage], optional): The stage to transition the model to after evaluation. Defaults to None.
+            archive_existing_versions (bool, optional): Whether to archive existing versions at the target stage. Defaults to True.
+            description (Optional[str], optional): Description for the MLflow run. Defaults to None.
+            tags (Optional[dict], optional): Tags to associate with the MLflow run. Defaults to None.
+
+        """
+        super().__init__()
+        self._registered_model_name = None
+        self.target_stage = target_stage
+        self.archive_existing_versions = archive_existing_versions
+        self.description = description
+        self.tags = tags or {"component": "Pusher"}
+    
+    @property
+    def registered_model_name(self) -> str:
+        if not self._registered_model_name:
+            raise PusherError("Registered model name is not set.")
+        return self._registered_model_name
+    
+    @registered_model_name.setter
+    def registered_model_name(self, name: str) -> None:
+        self._registered_model_name = name
+
+    def run(self, event: Event) -> Event:
+        with self.start_run(
+            tags=self.tags,
+            run_name="Pusher",
+            log_system_metrics=False,
+            nested=True,
+            description=self.description
+        ):
+            try:
+                pusher_event = self._handle(event)
+                payload = self._get_pusher_payload(pusher_event)
+                pass_evaluation = event.context.get("pass_evaluation")
+                
+                if pass_evaluation:
+                    model_uri = event.context.get('model_uri')
+                    if not model_uri:
+                        raise PusherError("model_uri not found in event context")
+                    
+                    version = self.register_model(
+                        model_name=self.registered_model_name,
+                        model_uri=model_uri
+                    )
+                    self.mlflow.log_param("model_registered", "true")
+                    event.context.set("registered_version", str(version))
+                    self.mlflow.log_param("model_version", version)
+                    
+                    self.mlflow.log_param("target_stage", self.target_stage.value)
+                    self.transition_model_version(
+                        registered_model_name=self.registered_model_name,
+                        version=str(version),
+                        desired_stage=self.target_stage,
+                        archive_existing_versions_at_stage=self.archive_existing_versions,
+                    )
+
+                else:
+                    self.mlflow.log_param("skip_reason", "evaluation_failed")
+                    self.mlflow.log_param("model_registered", "false")
+
+                return Event(
+                    type=EventType.PUSHER_DONE,
+                    payload=payload,
+                    context=event.context,
+                )
+
+            except Exception as e:
+                raise PusherError(f"Pusher failed with error: {e}")
+
+    @type_checker((PusherPayload,), "PusherPayload must be of type PusherPayload.")
+    def _get_pusher_payload(self, event: Event) -> PusherPayload:
+        return event.payload

collie/core/trainer/trainer.py

@@ -0,0 +1,78 @@
+from typing import Optional
+from collie.contracts.event import (
+    Event, 
+    EventHandler, 
+    EventType
+)
+from collie.contracts.mlflow import MLFlowComponentABC
+from collie.core.models import (
+    TrainerPayload,
+    TrainerArtifact
+)
+from collie._common.decorator import type_checker
+from collie._common.exceptions import TrainerError
+
+
+class Trainer(EventHandler, MLFlowComponentABC):
+
+    def __init__(
+        self, 
+        description: Optional[str] = None,
+        tags: Optional[dict] = None
+    ) -> None:
+        """
+        Initializes the Trainer component.
+
+        Args:
+            description (Optional[str], optional): Description for the MLflow run. Defaults to None.
+            tags (Optional[dict], optional): Tags to associate with the MLflow run. Defaults to None.
+        """
+        super().__init__()
+        self.description = description
+        self.tags = tags or {"component": "Trainer"}
+    
+    def run(self, event: Event) -> Event:
+        """
+        Run the trainer component.
+
+        This method starts a new MLflow run, trains the model,
+        logs metrics, and sets the outputs.
+        """
+        
+        with self.start_run(
+            run_name="Trainer", 
+            tags=self.tags,
+            log_system_metrics=True, 
+            nested=True,
+            description=self.description
+        ) as run:
+            try:
+                trainer_event = self._handle(event)
+
+                trainer_payload = self._trainer_payload(trainer_event)
+                event_type = EventType.TRAINING_DONE
+                
+                model = trainer_payload.model
+                self.log_model(
+                    model=model, 
+                    name=TrainerArtifact().model
+                )
+                event.context.set("model_uri", self.artifact_uri(run))
+
+                return Event(
+                    type=event_type,
+                    payload=trainer_payload,
+                    context=event.context
+                )
+            except Exception as e:
+                raise TrainerError(f"Trainer failed with error: {e}") from e
+
+    @type_checker((TrainerPayload,) , 
+        "TrainerPayload must be of type TrainerPayload."
+    )    
+    def _trainer_payload(self, event: Event) -> TrainerPayload: 
+
+        return event.payload
+    
+    def artifact_uri(self, run) -> str:
+        return f"runs:/{run.info.run_id}/{TrainerArtifact().model}"

collie/core/transform/transform.py

@@ -0,0 +1,87 @@
+from typing import Optional
+from collie.contracts.event import (
+    Event, 
+    EventHandler, 
+    EventType
+)
+from collie.contracts.mlflow import MLFlowComponentABC
+from collie.core.models import (
+    TransformerPayload,
+    TransformerArtifact
+)
+from collie._common.decorator import type_checker
+from collie._common.exceptions import TransformerError
+
+
+class Transformer(EventHandler, MLFlowComponentABC):
+
+    def __init__(
+        self,
+        description: Optional[str] = None,
+        tags: Optional[dict] = None
+    ) -> None:
+        """
+        Initializes the Transformer component.
+
+        Args:
+            description (Optional[str], optional): Description for the MLflow run. Defaults to None.
+            tags (Optional[dict], optional): Tags to associate with the MLflow run. Defaults to None.
+        """
+        super().__init__()
+        self.description = description
+        self.tags = tags or {"component": "Transformer"}
+    
+    def run(self, event: Event) -> Event:
+        """
+        Run the transformer component.
+
+        This method starts a new MLflow run, transforms the input data,
+        logs metrics, and sets the outputs.
+        """
+        with self.start_run(
+            tags=self.tags,
+            run_name="Transformer",
+            log_system_metrics=True,
+            nested=True,
+            description=self.description
+        ) as run:
+            try:
+                transformer_event = self._handle(event)
+
+                transformer_payload = self._transformer_payload(transformer_event)
+                event_type = EventType.DATA_READY
+
+                data = transformer_payload.model_dump()
+                for data_type, data in data.items():
+
+                    if data is not None:
+                        source = self.artifact_uri(run, data_type)
+                        event.context.set(
+                            f"{data_type}_uri",
+                            source
+                        )
+                        
+                        self.log_pd_data(
+                            data=data, 
+                            context=data_type,
+                            source=source
+                        )
+
+                return Event(
+                    type=event_type,
+                    payload=transformer_payload,
+                    context=event.context
+                )
+
+            except Exception as e:
+                raise TransformerError(f"Transformer failed with error: {e}")
+        
+    @type_checker((TransformerPayload,) , 
+        "TransformerPayload must be of type TransformerPayload."
+    )    
+    def _transformer_payload(self, event: Event) -> TransformerPayload: 
+
+        return event.payload
+    
+    def artifact_uri(self, run, data_type) -> str:
+         return f"{run.info.artifact_uri}/{TransformerArtifact().model_dump()[data_type]}" 

collie/core/tuner/tuner.py

@@ -0,0 +1,84 @@
+from typing import Optional
+from collie.contracts.event import (
+    Event, 
+    EventHandler, 
+    EventType
+)
+from collie.contracts.mlflow import MLFlowComponentABC
+from collie._common.decorator import type_checker
+from collie.core.models import (
+    TunerArtifact,
+    TunerPayload
+)
+from collie._common.exceptions import TunerError
+
+
+class Tuner(EventHandler, MLFlowComponentABC):
+
+    def __init__(
+        self,
+        description: Optional[str] = None,
+        tags: Optional[dict] = None
+    ) -> None:
+        """
+        Initializes the Tuner component.
+
+        Args:
+            description (Optional[str], optional): Description for the MLflow run. Defaults to None.
+            tags (Optional[dict], optional): Tags to associate with the MLflow run. Defaults to None.
+        """
+        super().__init__()
+        self.description = description
+        self.tags = tags or {"component": "Tuner"}
+    
+    def run(self, event: Event) -> None:
+        """
+        Run the hyperparameter tuner component.
+
+        This method starts a new MLflow run, tunes the hyperparameters,
+        logs metrics, and sets the outputs.
+        """
+
+        with self.start_run(
+            tags=self.tags,
+            run_name="Tuner",
+            log_system_metrics=True,
+            nested=True,
+            description=self.description
+        ) as run:
+            try:
+                tuner_event = self._handle(event)
+
+                tuner_payload = self._tuner_payload(tuner_event)
+                hyperparameters = tuner_payload.model_dump()
+
+                self.mlflow.log_dict(
+                    dictionary=hyperparameters, 
+                    artifact_file=TunerArtifact().hyperparameters
+                )
+                event.context.set(
+                    "hyperparameters_uri",
+                    self.artifact_path(run)
+                )
+
+                event_type = EventType.TUNING_DONE
+
+                return Event(
+                    type=event_type,
+                    payload=tuner_payload,
+                    context=event.context
+                )
+            except Exception as e:
+                raise TunerError(f"Tuner failed with error: {e}") from e
+
+    @type_checker((TunerPayload,) , 
+        "TunerPayload must be of type TunerPayload."
+    )    
+    def _tuner_payload(self, event: Event) -> TunerPayload: 
+
+        tuner_payload: TunerPayload = event.payload
+
+        return tuner_payload
+    
+    def artifact_path(self, run) -> str:
+        return f"{run.info.artifact_uri}/{TunerArtifact().hyperparameters}"