diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 3e4b492..87f6ea3 100644 --- a/laborious/activities/mlflow.py +++ b/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', - values='value') - data.sort_index(inplace=True) - data.reset_index(inplace=True) + # Pivot data for model input format + data = data.pivot( + index='timestamp', columns='variable', + 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( diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 89922a8..a58857d 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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": - code_path = path.join( - artifact_path, "transformer_pyfunc", "code") - pickle_path = path.join( - artifact_path, "transformer_pyfunc", "artifacts", "training_transformer.pkl") + pickle_file = "training_transformer.pkl" if type == "transformer" else "stacking_model.pkl" - elif type == "prediction": - code_path = path.join( - artifact_path, "stacking_model", "code") - pickle_path = path.join( - artifact_path, "stacking_model", "artifacts", "stacking_model.pkl") + pickle_path = path.join( + artifact_path, "artifacts", pickle_file) 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']) - data.to_csv( + predict_data = pd.DataFrame( + predict_data, columns=['prediction']) + predict_data.to_csv( f"tmp/predicted_data_{model_name}.csv", index=True) - data.index = input_index - data['response_time'] = (end_time - start_time).total_seconds() + predict_data.index = input_index + 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: diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 81c716b..ad2e6c0 100644 --- a/laborious/worker/worker.py +++ b/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 )