SIENTIAPDE-1231

Enhance MLFlow and MLFlowRepository with improved data handling and logging

- Refactored MLFlow class to sort data by 'created_at' and drop duplicates for better input preparation.
- Updated MLFlowRepository methods to include detailed logging for artifact downloads and model predictions.
- Introduced LzmaPayloadCodec for efficient payload compression in the worker, optimizing data handling for large payloads.
- Enhanced timestamp handling in treated data to ensure compatibility with model expectations.

laborious/activities/mlflow.py

@@ -232,14 +232,23 @@ class MLFlow(BaseActivity):
         timestamp = data['timestamp'].max()
         self.debug(f'Timestamp: {timestamp}', metadata)
 
+        # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
+        data = data.sort_values('created_at', ascending=False).drop_duplicates(
+            subset=['variable', 'timestamp'], keep='first'
+        )
+
         data.drop(columns=['model_id'], inplace=True, errors='ignore')
         data.drop(columns=['created_at'], inplace=True, errors='ignore')
 
-        data = data.pivot(index='timestamp', columns='variable',
+        # Pivot data for model input format
-                          values='value')
+        data = data.pivot(
-        data.sort_index(inplace=True)
+            index='timestamp', columns='variable',
-        data.reset_index(inplace=True)
+            values='value')
+        data.fillna(np.nan, inplace=True)
+        # data.reset_index(inplace=True)
+        data.columns.name = None
 
+        data['timestamp'] = data.index
         data['timestamp'] = to_datetime(
             data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
         data['timestamp'] = to_datetime(

laborious/utils/repository/model_repository.py

@@ -235,6 +235,9 @@ class MLFlowRepository():
         if not path.exists(output_dir):
             makedirs(output_dir)
 
+        self.logger.info(
+            f"Downloading artifacts from {run_id} to {output_dir}")
+
         return self.client.download_artifacts(
             run_id,
             artifact_path,
@@ -347,18 +350,13 @@ class MLFlowRepository():
         Raises:
             ValueError: If model cannot be loaded with any compression method
         """
+        code_path = path.join(
+            artifact_path, "code")
 
-        if type == "transformer":
+        pickle_file = "training_transformer.pkl" if type == "transformer" else "stacking_model.pkl"
-            code_path = path.join(
-                artifact_path, "transformer_pyfunc", "code")
-            pickle_path = path.join(
-                artifact_path, "transformer_pyfunc", "artifacts", "training_transformer.pkl")
 
-        elif type == "prediction":
+        pickle_path = path.join(
-            code_path = path.join(
+            artifact_path, "artifacts", pickle_file)
-                artifact_path, "stacking_model", "code")
-            pickle_path = path.join(
-                artifact_path, "stacking_model", "artifacts", "stacking_model.pkl")
 
         if code_path not in sys_path:
             sys_path.insert(0, code_path)
@@ -468,6 +466,9 @@ class MLFlowRepository():
                     f"{message}")
 
         elif index_type == datetime or index_type == pd.Timestamp:
+            if data.index.tz is None:
+                data.index = data.index.tz_localize('UTC')
+
             data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
         else:
             raise ValueError(
@@ -690,6 +691,9 @@ class MLFlowRepository():
         self.logger.custom_debug(
             f"Fitting transformation model with training data (shape: {data.shape})", metadata)
 
+        data.to_csv(
+            f"tmp/retrain_data_{model_name}.csv", index=True)
+
         data_model = data_model.fit(data)
 
         self.logger.custom_debug(
@@ -697,6 +701,20 @@ class MLFlowRepository():
 
         treated_data = data_model.predict(data)
 
+        # Stores current index as timestamp, Courier model expects a timestamp column
+        # with specific format
+        treated_data['timestamp'] = treated_data.index
+
+        # Parses timestamp column to datetime format to align with data
+        treated_data = self.detect_and_parse_datetime_index(
+            treated_data, metadata)
+
+        self.logger.custom_debug(
+            f"Treated data index: {treated_data.index}", metadata)
+
+        treated_data.to_csv(
+            f"tmp/retrain_treated_data_{model_name}.csv", index=True)
+
         self.logger.custom_debug(
             f"Transformed data shape: {treated_data.shape}", metadata)
 
@@ -708,12 +726,24 @@ class MLFlowRepository():
             self.logger.custom_debug(
                 f"Using provided target variable: {target_name}", metadata)
 
+        # Check if treated_data contains target variable
+        if target_name not in treated_data.columns:
+            self.logger.custom_debug(
+                f"Target variable {target_name} not found in treated data, aligning data with treated data indexes", metadata)
+            # Aligns data with treated data indexes to get target variable
+            aligned_data = data.loc[treated_data.index]
+        else:
+            # Uses target variable from treated data
+            self.logger.custom_debug(
+                f"Target variable {target_name} found in treated data, using it", metadata)
+            aligned_data = treated_data
+
         if fit_config.get('y_type', 'series').lower() == 'series':
-            y = data[target_name]
+            y = aligned_data[target_name]
             self.logger.custom_debug(
                 f"Extracting target as series, shape: {y.shape}", metadata)
         else:
-            y = data[[target_name]]
+            y = aligned_data[[target_name]]
             self.logger.custom_debug(
                 f"Extracting target as dataframe, shape: {y.shape}", metadata)
 
@@ -726,22 +756,25 @@ class MLFlowRepository():
             self.logger.custom_debug(
                 f"Combined data shape for prediction model fit: {treated_data.shape}", metadata)
 
+            treated_data.to_csv(
+                f"tmp/retrain_treated_data_combined_{model_name}.csv", index=True)
+
             prediction_model = prediction_model.fit(treated_data)
             self.logger.custom_debug(
                 "Prediction model fitted with combined data", metadata)
         else:
             # Keep data separated
             fit_order = fit_config.get('split_fit_first', 'x').lower()
+
+            y.to_csv(
+                f"tmp/retrain_y_{model_name}.csv", index=True)
+
             self.logger.custom_debug(
-                f"Fitting prediction model with separated data, order: {fit_order}", metadata)
+                f"Fitting prediction model with separated data, first arg: {fit_order}", metadata)
             if fit_order == 'x':
                 prediction_model = prediction_model.fit(treated_data, y)
-                self.logger.custom_debug(
-                    "Prediction model fitted with X, y order", metadata)
             else:
                 prediction_model = prediction_model.fit(y, treated_data)
-                self.logger.custom_debug(
-                    "Prediction model fitted with y, X order", metadata)
 
         experiment = self.get_experiment_by_run_id(latest_production_id)
         mlflow.set_experiment(experiment)
@@ -1012,30 +1045,32 @@ class MLFlowRepository():
             data.to_csv(
                 f"tmp/treated_data_{model_name}.csv", index=True)
 
-            data = self.get_cached_predict(
+            predict_data = self.get_cached_predict(
                 model_name, data, model_retention, flavor)
 
             end_time = datetime.now()
 
-            if isinstance(data, pd.DataFrame):
+            if isinstance(predict_data, pd.DataFrame):
                 self.logger.custom_debug(
                     f"Data received from model prediction: {data.head(5).to_csv()}", metadata)
 
-                data.to_csv(
+                predict_data.to_csv(
                     f"tmp/predicted_data_{model_name}.csv", index=True)
                 data.columns = ['prediction']
 
             else:
-                data = pd.DataFrame(data, columns=['prediction'])
+                predict_data = pd.DataFrame(
-                data.to_csv(
+                    predict_data, columns=['prediction'])
+                predict_data.to_csv(
                     f"tmp/predicted_data_{model_name}.csv", index=True)
 
-            data.index = input_index
+            predict_data.index = input_index
-            data['response_time'] = (end_time - start_time).total_seconds()
+            predict_data['response_time'] = (
+                end_time - start_time).total_seconds()
 
             return {
                 'success': True,
-                'content': data.to_dict()
+                'content': predict_data.to_dict()
             }
 
         except Exception as e:

laborious/worker/worker.py

@@ -28,6 +28,8 @@ Environment Variables:
 from temporalio import workflow, client
 from temporalio.worker import Worker, PollerBehaviorAutoscaling
 from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
+from temporalio.converter import PayloadCodec, DataConverter
+from temporalio.api.common.v1 import Payload
 
 with workflow.unsafe.imports_passed_through():
     import os
@@ -49,9 +51,54 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.observability.logger import get_logger
     from laborious import metrics
     from prometheus_client import start_http_server
+    import lzma
+    import dataclasses
+
 
 POD_ID = os.getenv('POD_ID')
 SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
+LZMA_MIN_MB = float(os.getenv('LZMA_MIN_MB', "1.5"))
+
+
+class LzmaPayloadCodec(PayloadCodec):
+    async def encode(self, payloads):
+        out = []
+        for p in payloads:
+            if p.data:
+                old_len = len(p.data) / 1000000
+                # Only compress payloads larger than 1.5 MB
+                if old_len > LZMA_MIN_MB:
+                    compressed_data = lzma.compress(p.data)
+                    new_len = len(compressed_data) / 1000000
+                    ratio = new_len / old_len if old_len else 0
+                    print(
+                        f"[codec] encode lzma: {old_len} MB -> {new_len} MB ({ratio:.2f}x)")
+                    meta = dict(p.metadata or {})
+                    meta[b"codec"] = b"lzma"
+                    out.append(Payload(metadata=meta, data=compressed_data))
+                else:
+                    out.append(p)
+
+            else:
+                out.append(p)
+        return out
+
+    async def decode(self, payloads):
+        out = []
+        for p in payloads:
+            if p.data and (p.metadata or {}).get(b"codec") == b"lzma":
+                # comp_len = len(p.data)
+                decomp = lzma.decompress(p.data)
+                # decomp_len = len(decomp)
+                # ratio = (decomp_len / comp_len) if comp_len else 0
+                # print(
+                #     f"[codec] decode lzma: {comp_len} B -> {decomp_len} B ({ratio:.2f}x)")
+                meta = dict(p.metadata or {})
+                meta.pop(b"codec", None)
+                out.append(Payload(metadata=meta, data=decomp))
+            else:
+                out.append(p)
+        return out
 
 
 async def main():
@@ -125,9 +172,12 @@ async def main():
 
     logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
 
+    codec_dc = DataConverter(payload_codec=LzmaPayloadCodec())
+
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
+        data_converter=codec_dc,
         runtime=new_runtime
     )