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/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
     )