SIENTIAPDE-1273

Enhance MLFlowRepository and Activities classes with new methods and metrics

- Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry.
- Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models.
- Updated Activities class to integrate ModelMetrics for improved metrics handling.
- Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities.
- Updated various workflows to include `transform_table_name` in input data for better data handling.
This commit is contained in:
vitor-aignosi
2025-11-13 16:40:23 -03:00
parent 6ac0f38d59
commit b68674fe64
15 changed files with 1639 additions and 68 deletions

View File

@@ -29,6 +29,7 @@ from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import PollerBehaviorAutoscaling, Worker
with workflow.unsafe.imports_passed_through():
import asyncio
import os
@@ -49,6 +50,7 @@ with workflow.unsafe.imports_passed_through():
)
from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.drift import Drift
from laborious.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)
@@ -156,6 +158,23 @@ async def main():
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
Worker(
temporal_client,
task_queue='drift-queue',
workflows=[Drift],
activities=[
activities.load_custom_query,
activities.get_reference_data,
activities.calculate_drift,
activities.export_data_to_postgres,
],
max_concurrent_workflow_tasks=50,
max_concurrent_activities=50,
max_concurrent_local_activities=50,
max_cached_workflows=2,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
Worker(
temporal_client,
task_queue='predictions_batch-queue',
@@ -168,6 +187,7 @@ async def main():
activities.input_gate,
activities.mlflow_response_gate,
activities.mlflow_content_gate,
activities.format_transformed_data,
activities.format_prediction,
activities.format_default_prediction,
activities.get_last_timestamp,