ddi-fw 0.0.217__py3-none-any.whl → 0.0.218__py3-none-any.whl

ddi_fw/ml/tensorflow_wrapper.py

@@ -7,18 +7,16 @@ from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, Callback
 from sklearn.model_selection import train_test_split, KFold, StratifiedKFold
 import numpy as np
 from tensorflow.keras import Model
-import mlflow
-from mlflow.utils.autologging_utils import batch_metrics_logger
-
-from mlflow.models import infer_signature
 from ddi_fw.ml.evaluation_helper import Metrics, evaluate
 
 # import tf2onnx
 # import onnx
 
+from ddi_fw.ml.tracking_service import TrackingService
 import ddi_fw.utils as utils
 import os
 
+
 def convert_to_categorical(arr, num_classes):
     """
     This function takes an array of labels and converts them to one-hot encoding 
@@ -33,7 +31,7 @@ def convert_to_categorical(arr, num_classes):
     - The one-hot encoded array if the original array was binary or label encoded
     - The original array if it doesn't require any conversion
     """
-    
+
     try:
         # First, check if the array is binary-encoded
         if not utils.is_binary_encoded(arr):
@@ -45,7 +43,7 @@ def convert_to_categorical(arr, num_classes):
     except Exception as e:
         # If binary encoding check raises an error, print it and continue to label encoding check
         print(f"Error while checking binary encoding: {e}")
-    
+
     try:
         # Check if the array is label-encoded
         if utils.is_label_encoded(arr):
@@ -56,21 +54,21 @@ def convert_to_categorical(arr, num_classes):
         print(f"Error while checking label encoding: {e}")
         # If the arr labels don't match any of the known encodings, raise an error
         raise ValueError("Unknown label encoding format.")
-    
+
     # If no conversion was needed, return the original array
-    
+
     return arr
 
 
 class TFModelWrapper(ModelWrapper):
 
-    def __init__(self, date, descriptor, model_func, use_mlflow=False, **kwargs):
+    def __init__(self, date, descriptor, model_func, tracking_service: TrackingService | None = None, **kwargs):
         super().__init__(date, descriptor, model_func, **kwargs)
         self.batch_size = kwargs.get('batch_size', 128)
         self.epochs = kwargs.get('epochs', 100)
-        self.use_mlflow = use_mlflow
+        self.tracking_service = tracking_service
 
-    # TODO think different settings for num_classes 
+    # TODO think different settings for num_classes
     def fit_model(self, X_train, y_train, X_valid, y_valid):
         self.kwargs['input_shape'] = self.train_data.shape
         self.num_classes = len(np.unique(y_train, axis=0))
@@ -85,7 +83,11 @@ class TFModelWrapper(ModelWrapper):
         )
         early_stopping = EarlyStopping(
             monitor='val_loss', patience=10, mode='auto')
-        custom_callback = CustomCallback(self.use_mlflow)
+        if self.tracking_service:
+            custom_callback = CustomCallback(self.tracking_service)
+            callbacks=[early_stopping, checkpoint, custom_callback]
+        else:
+            callbacks=[early_stopping, checkpoint]
         train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
         train_dataset = train_dataset.batch(batch_size=self.batch_size)
 
@@ -100,7 +102,7 @@ class TFModelWrapper(ModelWrapper):
             train_dataset,
             epochs=self.epochs,
             # validation_data=val_dataset,
-            callbacks=[early_stopping, checkpoint, custom_callback]
+            callbacks=callbacks
         )
         # ex
         # history = model.fit(
@@ -123,38 +125,51 @@ class TFModelWrapper(ModelWrapper):
         if self.train_idx_arr and self.val_idx_arr:
             for i, (train_idx, val_idx) in enumerate(zip(self.train_idx_arr, self.val_idx_arr)):
                 print(f"Validation {i}")
+                X_train_cv = self.train_data[train_idx]
+                y_train_cv = self.train_label[train_idx]
+                X_valid_cv = self.train_data[val_idx]
+                y_valid_cv = self.train_label[val_idx]
 
-                if self.use_mlflow:
-                    with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
-                        X_train_cv = self.train_data[train_idx]
-                        y_train_cv = self.train_label[train_idx]
-                        X_valid_cv = self.train_data[val_idx]
-                        y_valid_cv = self.train_label[val_idx]
-                        model, checkpoint = self.fit_model(
-                            X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
-                        models[f'{self.descriptor}_validation_{i}'] = model
-                        models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
-                else:
-                    X_train_cv = self.train_data[train_idx]
-                    y_train_cv = self.train_label[train_idx]
-                    X_valid_cv = self.train_data[val_idx]
-                    y_valid_cv = self.train_label[val_idx]
+                def fit_model_cv_func():
                     model, checkpoint = self.fit_model(
                         X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
-                    models[f'{self.descriptor}_validation_{i}'] = model
-                    models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
+                    return model, checkpoint
+
+                if self.tracking_service:
+                    model, checkpoint = self.tracking_service.run(
+                        run_name=f'Validation {i}', description='CV models', nested_run=True, func=fit_model_cv_func)
+                    # with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
+
+                    #     model, checkpoint = self.fit_model(
+                    #         X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
+                    #     models[f'{self.descriptor}_validation_{i}'] = model
+                    #     models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
+                else:
+                    model, checkpoint = fit_model_cv_func()
+                    # model, checkpoint = self.fit_model(
+                    #     X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
+                models[f'{self.descriptor}_validation_{i}'] = model
+                models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
         else:
-            if self.use_mlflow:
-                with mlflow.start_run(run_name=f'Training', description='Training', nested=True) as cv_fit:
-                    model, checkpoint = self.fit_model(
-                        self.train_data, self.train_label, None, None)
-                    models[self.descriptor] = model
-                    models_val_acc[self.descriptor] = checkpoint.best
-            else:
+            def fit_model_func():
                 model, checkpoint = self.fit_model(
                     self.train_data, self.train_label, None, None)
-                models[self.descriptor] = model
-                models_val_acc[self.descriptor] = checkpoint.best
+                return model, checkpoint
+
+            if self.tracking_service:
+                model, checkpoint = self.tracking_service.run(
+                    run_name=f'Training', description='Training', nested_run=True, func=fit_model_func)
+                # with mlflow.start_run(run_name=f'Training', description='Training', nested=True) as cv_fit:
+                #     model, checkpoint = self.fit_model(
+                #         self.train_data, self.train_label, None, None)
+                #     models[self.descriptor] = model
+                #     models_val_acc[self.descriptor] = checkpoint.best
+            else:
+                model, checkpoint = fit_model_func()
+                # models[self.descriptor] = model
+                # models_val_acc[self.descriptor] = checkpoint.best
+            models[self.descriptor] = model
+            models_val_acc[self.descriptor] = checkpoint.best
         if models_val_acc == {}:
             return model, None
         best_model_key = max(models_val_acc, key=lambda k: models_val_acc[k])
@@ -172,62 +187,113 @@ class TFModelWrapper(ModelWrapper):
         pred = self.best_model.predict(test_dataset)
         return pred
 
+    # def fit_and_evaluate(self, print_detail=False) -> tuple[dict[str, Any], Metrics, Any]:
+    #     if self.use_mlflow:
+    #         with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
+    #             best_model, best_model_key = self.fit()
+    #             self.best_model: Model = best_model
+    #             pred = self.predict()
+    #             actual = self.test_label
+    #             # if not utils.is_binary_encoded(pred):
+    #             #     pred = tf.keras.utils.to_categorical(np.argmax(pred,axis=1), num_classes=self.num_classes)
+    #             pred_as_cat = convert_to_categorical(pred, self.num_classes)
+    #             actual_as_cat = convert_to_categorical(
+    #                 actual, self.num_classes)
+
+    #             logs, metrics = evaluate(
+    #                 actual=actual_as_cat, pred=pred_as_cat, info=self.descriptor, print_detail=print_detail)
+    #             metrics.format_float()
+    #             mlflow.log_metrics(logs)
+    #             mlflow.log_param('best_cv', best_model_key)
+    #             utils.compress_and_save_data(
+    #                 metrics.__dict__, run.info.artifact_uri, f'{self.date}_metrics.gzip')
+    #             mlflow.log_artifact(
+    #                 f'{run.info.artifact_uri}/{self.date}_metrics.gzip')
+
+    #             return logs, metrics, pred
+    #     else:
+    #         best_model, best_model_key = self.fit()
+    #         self.best_model = best_model
+    #         pred = self.predict()
+    #         actual = self.test_label
+
+    #         pred_as_cat = convert_to_categorical(pred, self.num_classes)
+    #         actual_as_cat = convert_to_categorical(actual, self.num_classes)
+    #         logs, metrics = evaluate(
+    #             actual=actual_as_cat, pred=pred_as_cat, info=self.descriptor)
+    #         metrics.format_float()
+    #         return logs, metrics, pred
+
     def fit_and_evaluate(self, print_detail=False) -> tuple[dict[str, Any], Metrics, Any]:
-        if self.use_mlflow:
-            with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
-                print(run.info.artifact_uri)
-                best_model, best_model_key = self.fit()
-                print(best_model_key)
-                self.best_model: Model = best_model
-                pred = self.predict()
-                actual = self.test_label
-                # if not utils.is_binary_encoded(pred):
-                #     pred = tf.keras.utils.to_categorical(np.argmax(pred,axis=1), num_classes=self.num_classes)
-                pred_as_cat= convert_to_categorical(pred, self.num_classes)
-                actual_as_cat= convert_to_categorical(actual, self.num_classes)
-
-                logs, metrics = evaluate(
-                    actual=actual_as_cat, pred=pred_as_cat, info=self.descriptor, print_detail=print_detail)
-                metrics.format_float()
-                mlflow.log_metrics(logs)
-                mlflow.log_param('best_cv', best_model_key)
-                utils.compress_and_save_data(
-                    metrics.__dict__, run.info.artifact_uri, f'{self.date}_metrics.gzip')
-                mlflow.log_artifact(
-                    f'{run.info.artifact_uri}/{self.date}_metrics.gzip')
-
-                return logs, metrics, pred
-        else:
+        """
+        Fit the model, evaluate it, and log results using the tracking service.
+
+        Args:
+            print_detail (bool): Whether to print detailed evaluation logs.
+
+        Returns:
+            tuple: A tuple containing logs, metrics, and predictions.
+        """
+        self.best_model: Model = None
+        def evaluate_and_log(artifact_uri=None):
+            # Fit the model
             best_model, best_model_key = self.fit()
-            print(best_model_key)
             self.best_model = best_model
+
+            # Make predictions
             pred = self.predict()
             actual = self.test_label
-            # if not utils.is_binary_encoded(pred):
-            #     pred = tf.keras.utils.to_categorical(np.argmax(pred,axis=1), num_classes=self.num_classes)
-            # if not utils.is_binary_encoded(actual):
-            #     actual = tf.keras.utils.to_categorical(actual, num_classes=self.num_classes)
-            pred= convert_to_categorical(pred, self.num_classes)
-            actual= convert_to_categorical(actual, self.num_classes)
+
+            # Convert predictions and actual labels to categorical format
+            pred_as_cat = convert_to_categorical(pred, self.num_classes)
+            actual_as_cat = convert_to_categorical(actual, self.num_classes)
+
+            # Evaluate the model
             logs, metrics = evaluate(
-                actual=actual, pred=pred, info=self.descriptor)
+                actual=actual_as_cat, pred=pred_as_cat, info=self.descriptor, print_detail=print_detail
+            )
             metrics.format_float()
-            return logs, metrics, pred
 
+            if self.tracking_service:
+                # Log metrics and parameters
+                self.tracking_service.log_metrics(logs)
+                self.tracking_service.log_param('best_cv', best_model_key)
 
-"""
-    Custom Keras callback for logging training metrics and model summary to MLflow.
-"""
+                # Save metrics to the artifact URI if provided
+                if artifact_uri:
+                    utils.compress_and_save_data(
+                        metrics.__dict__, artifact_uri, f'{self.date}_metrics.gzip'
+                    )
+                    self.tracking_service.log_artifact(
+                        f'{artifact_uri}/{self.date}_metrics.gzip'
+                    )
+
+            return logs, metrics, pred
+
+        # Use the tracking service to run the evaluation
+        if self.tracking_service:
+            return self.tracking_service.run(
+                run_name=self.descriptor,
+                description="Fit and evaluate the model",
+                nested_run=True,
+                func=evaluate_and_log
+            )
+        else:
+            # If no tracking service is provided, run the evaluation directly
+            return evaluate_and_log()
 
 
 class CustomCallback(Callback):
-    def __init__(self, use_mlflow: bool = True):
+    """
+    Custom Keras callback for logging training metrics and model summary to MLflow.
+    """
+    def __init__(self, tracking_service: TrackingService):
         super().__init__()
-        self.use_mlflow = use_mlflow
+        self.tracking_service = tracking_service
 
-    def _mlflow_log(self, func: Callable):
-        if self.use_mlflow:
-            func()
+    # def _mlflow_log(self, func: Callable):
+    #     if self.use_mlflow:
+    #         func()
 
     def on_train_begin(self, logs=None):
         if logs is None:
@@ -236,24 +302,32 @@ class CustomCallback(Callback):
             raise TypeError("self.model must be an instance of Model")
 
         keys = list(logs.keys())
-        self._mlflow_log(lambda: mlflow.log_param("train_begin_keys", keys))
+
+        self.tracking_service.log_param("train_begin_keys", keys)
+        # self._mlflow_log(lambda: mlflow.log_param("train_begin_keys", keys))
+
         # config = self.model.optimizer.get_config()
         config = self.model.get_config()
         for attribute in config:
-            self._mlflow_log(lambda: mlflow.log_param(
-                "opt_" + attribute, config[attribute]))
+            self.tracking_service.log_param(
+                "opt_" + attribute, config[attribute])
+            # self._mlflow_log(lambda: mlflow.log_param(
+            #     "opt_" + attribute, config[attribute]))
 
         sum_list = []
         self.model.summary(print_fn=sum_list.append)
         summary = "\n".join(sum_list)
-        self._mlflow_log(lambda: mlflow.log_text(
-            summary, artifact_file="model_summary.txt"))
+        self.tracking_service.log_text(
+            summary, file_name="model_summary.txt")
+        # self._mlflow_log(lambda: mlflow.log_text(
+        #     summary, artifact_file="model_summary.txt"))
 
     def on_train_end(self, logs=None):
         if logs is None:
             logs = {}
         print(logs)
-        self._mlflow_log(lambda: mlflow.log_metrics(logs))
+        self.tracking_service.log_metrics(logs)
+        # self._mlflow_log(lambda: mlflow.log_metrics(logs))
 
     def on_epoch_begin(self, epoch, logs=None):
         if logs is None:
@@ -273,7 +347,8 @@ class CustomCallback(Callback):
     def on_test_end(self, logs=None):
         if logs is None:
             logs = {}
-        self._mlflow_log(lambda: mlflow.log_metrics(logs))
+        self.tracking_service.log_metrics(logs)
+        # self._mlflow_log(lambda: mlflow.log_metrics(logs))
         print(logs)
 
     def on_predict_begin(self, logs=None):
@@ -285,7 +360,8 @@ class CustomCallback(Callback):
         if logs is None:
             logs = {}
         keys = list(logs.keys())
-        self._mlflow_log(lambda: mlflow.log_metrics(logs))
+        self.tracking_service.log_metrics(logs)
+        # self._mlflow_log(lambda: mlflow.log_metrics(logs))
 
     def on_train_batch_begin(self, batch, logs=None):
         if logs is None:

ddi_fw/ml/tracking_service.py

@@ -0,0 +1,194 @@
+import inspect
+import os
+from typing import Optional, Dict, Any
+import logging
+from urllib.parse import urlparse
+import mlflow
+from abc import ABC, abstractmethod
+from typing import Callable, Optional, Dict, Any
+
+
+def normalize_artifact_uri(artifact_uri: str) -> str:
+    """
+    Normalize the artifact URI to a standard file path.
+
+    Args:
+        artifact_uri (str): The artifact URI to normalize.
+
+    Returns:
+        str: The normalized file path.
+    """
+    if artifact_uri.startswith("file:///"):
+        parsed_uri = urlparse(artifact_uri)
+        return os.path.abspath(os.path.join(parsed_uri.path.lstrip('/')))
+    return artifact_uri
+
+class Tracking(ABC):
+    def __init__(self,  experiment_name: str, tracking_params: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the tracking backend.
+
+        Args:
+            experiment_name (str): The name of the experiment.
+            experiment_tags (dict, optional): Tags for the experiment.
+        """
+        self.experiment_name = experiment_name
+        self.tracking_params = tracking_params or {}
+
+    @abstractmethod
+    def setup_experiment(self):
+        """Set up the experiment in the tracking backend."""
+        pass
+
+    @abstractmethod
+    def run(self, run_name: str, description:str, func: Callable, nested_run: bool = False):
+        """Run the experiment with the given function."""
+        pass
+
+    @abstractmethod
+    def log_text(self, content:str, file_name: str):
+        """Log parameters to the tracking backend."""
+        pass
+    @abstractmethod
+    def log_param(self, key:str, value: Any):
+        """Log parameters to the tracking backend."""
+        pass
+    @abstractmethod
+    def log_params(self, params: Dict[str, Any]):
+        """Log parameters to the tracking backend."""
+        pass
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+        """Log metrics to the tracking backend."""
+        pass
+
+    @abstractmethod
+    def log_artifact(self, artifact_path: str):
+        """Log an artifact to the tracking backend."""
+        pass
+
+
+logger = logging.getLogger(__name__)
+
+
+class MLFlowTracking(Tracking):
+    def __init__(self, experiment_name: str, tracking_params: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the MLFlowTracking backend.
+
+        Args:
+            experiment_name (str): The name of the experiment.
+            tracking_params (dict, optional): Parameters for MLflow tracking.
+        """
+        super().__init__(experiment_name, tracking_params)
+        if tracking_params:
+            self.experiment_tags = tracking_params.get("experiment_tags", {})
+
+    def setup_experiment(self):
+        """Set up an experiment in MLflow."""
+        tracking_uri = self.tracking_params.get("tracking_uri")
+        if not tracking_uri:
+            raise ValueError("Tracking URI must be specified for MLflow.")
+
+        mlflow.set_tracking_uri(tracking_uri)
+
+        if mlflow.get_experiment_by_name(self.experiment_name) is None:
+            artifact_location = self.tracking_params.get("artifact_location")
+            mlflow.create_experiment(self.experiment_name, artifact_location)
+            logger.info(
+                f"Created new MLflow experiment: {self.experiment_name}")
+
+        mlflow.set_experiment(self.experiment_name)
+
+        if self.experiment_tags:
+            mlflow.set_experiment_tags(self.experiment_tags)
+            logger.info(
+                f"Set tags for MLflow experiment '{self.experiment_name}': {self.experiment_tags}")
+
+    def run(self, run_name: str, description:str, func: Callable, nested_run: bool = False):
+        """Run the experiment with the given function."""
+        func_signature = inspect.signature(func)
+       
+        if nested_run:
+            with mlflow.start_run(run_name=run_name, description= description, nested=True) as run:
+                if "artifact_uri" in func_signature.parameters:
+                    artifact_uri = normalize_artifact_uri(run.info.artifact_uri) if run.info.artifact_uri else ""
+                    return func(artifact_uri=artifact_uri)
+                else:
+                    return func()
+        else:
+            with mlflow.start_run(run_name=run_name, description= description) as run:
+                 if "artifact_uri" in func_signature.parameters:
+                    artifact_uri = normalize_artifact_uri(run.info.artifact_uri) if run.info.artifact_uri else ""
+                    return func(artifact_uri=artifact_uri)
+                 else:
+                    return func()
+    
+    def log_text(self, content: str, file_name: str):
+        mlflow.log_text(
+            content, artifact_file=file_name)
+    
+    def log_param(self, key: str, value: Any):
+         mlflow.log_param(key, value)
+
+    def log_params(self, params: Dict[str, Any]):
+        """Log parameters to MLflow."""
+        mlflow.log_params(params)
+
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+        """Log metrics to MLflow."""
+        mlflow.log_metrics(metrics, step=step)
+
+    def log_artifact(self, artifact_path: str):
+        """Log an artifact to MLflow."""
+        mlflow.log_artifact(artifact_path)
+
+
+class TrackingService:
+    def __init__(self, experiment_name: str, backend: str = "mlflow", tracking_params: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the TrackingService.
+
+        Args:
+            backend (str): The tracking backend to use (e.g., "mlflow").
+            tracking_params (dict, optional): Parameters for the tracking backend.
+        """
+        self.experiment_name = experiment_name
+        self.backend = backend.lower()
+        self.tracking_params = tracking_params or {}
+        self.tracking_instance = self._initialize_backend()
+
+    def _initialize_backend(self) -> Tracking:
+        """Initialize the appropriate tracking backend."""
+        if self.backend == "mlflow":
+            return MLFlowTracking(self.experiment_name, self.tracking_params)
+        else:
+            raise ValueError(f"Unsupported tracking backend: {self.backend}")
+
+    def setup(self):
+        """Set up the experiment in the tracking backend."""
+        self.tracking_instance.setup_experiment()
+
+    def run(self, run_name: str, description:str ,func: Callable, nested_run: bool = False) -> Any:
+        """Run the experiment with the given function."""
+        return self.tracking_instance.run(run_name, description , func, nested_run=nested_run)
+        
+    def log_text(self, content: str, file_name: str):
+        self.tracking_instance.log_text(content, file_name)
+        
+    def log_param(self, key: str, value: Any):
+        """Log a parameter to the tracking backend."""
+        self.tracking_instance.log_param(key, value)
+
+    def log_params(self, params: Dict[str, Any]):
+        """Log parameters to the tracking backend."""
+        self.tracking_instance.log_params(params)
+
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+        """Log metrics to the tracking backend."""
+        self.tracking_instance.log_metrics(metrics, step=step)
+
+    def log_artifact(self, artifact_path: str):
+        """Log an artifact to the tracking backend."""
+        self.tracking_instance.log_artifact(artifact_path)