gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1231

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2025-10-06 13:32:36 -03:00
parent 9b71ad7556
commit 36af84f056
3 changed files with 123 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
laborious/activities/mlflow.py
View File

@@ -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(

85
laborious/utils/repository/model_repository.py
View File

@@ -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:

50
laborious/worker/worker.py
View File

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