SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.
This commit is contained in:
Bruno Domingues
227
scripts/run_training_test.py
Normal file
227
scripts/run_training_test.py
Normal file
@@ -0,0 +1,227 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Utility script to trigger the training workflow end-to-end for testing.
|
||||
|
||||
Steps performed:
|
||||
1. Upload the CSV test dataset to MinIO using the configured `mc` alias.
|
||||
2. Insert a new experiment_run record in Postgres and capture the generated ID.
|
||||
3. Trigger the Temporal `train_model` workflow with the correct payload.
|
||||
|
||||
Prerequisites:
|
||||
- `mc` CLI configured with alias defined in MINIO_ALIAS.
|
||||
- PostgreSQL accessible with credentials in environment variables or defaults.
|
||||
- Temporal server reachable without TLS on TEMPORAL_HOST / TEMPORAL_NAMESPACE.
|
||||
- Python dependencies installed (see requirements.txt / requirements-dev.txt).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import uuid
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import psycopg2
|
||||
from psycopg2.extras import Json
|
||||
from temporalio import client
|
||||
|
||||
DOCS_PATH = Path('docs/test-model-data.csv')
|
||||
MINIO_ALIAS = os.getenv('MINIO_ALIAS', 'suse')
|
||||
MINIO_BUCKET = os.getenv('MINIO_BUCKET', 'model-training')
|
||||
|
||||
POSTGRES_CONFIG = {
|
||||
'host': os.getenv('POSTGRES_HOST', 'localhost'),
|
||||
'port': os.getenv('POSTGRES_PORT', '55432'),
|
||||
'user': os.getenv('POSTGRES_USER', 'postgres'),
|
||||
'password': os.getenv(
|
||||
'POSTGRES_PASSWORD',
|
||||
'nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3',
|
||||
),
|
||||
'dbname': os.getenv('POSTGRES_DBNAME', 'sientia-core-mlops-bff'),
|
||||
}
|
||||
|
||||
TEMPORAL_HOST = os.getenv('TEMPORAL_HOST', 'localhost:37463')
|
||||
TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE', 'model-manager')
|
||||
TEMPORAL_TASK_QUEUE = os.getenv('TEMPORAL_TASK_QUEUE', 'train_model-queue')
|
||||
TEMPORAL_WORKFLOW = os.getenv('TEMPORAL_WORKFLOW', 'train_model')
|
||||
|
||||
BASE_REQUEST_DATA = {
|
||||
'experimentName': 'model-manager-test-01',
|
||||
'username': 'bruno.domingues@aignosi.com.br',
|
||||
'modelType': 'Linear Regression',
|
||||
'targetVariable': '03CV020/CORRENTE_N_M1_PV(Value)',
|
||||
'variableColumns': ['303-WIT-200(Value)'],
|
||||
'lagTrain': 0,
|
||||
'lagVal': 0,
|
||||
'remStaticWin': False,
|
||||
'lowLim': {},
|
||||
'uppLim': {},
|
||||
'window': 0,
|
||||
'useScaler': False,
|
||||
'includeAr': False,
|
||||
'trainSize': 80,
|
||||
'shuffle': True,
|
||||
'lineSeparator': ',',
|
||||
'decimalSeparator': '.',
|
||||
'removedIntervals': [],
|
||||
}
|
||||
|
||||
|
||||
def _ensure_source_file(path: Path) -> None:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f'Test dataset not found at {path.resolve()}')
|
||||
|
||||
|
||||
def upload_to_minio(source_path: Path) -> str:
|
||||
"""Upload the CSV to MinIO using the mc CLI and return the object name."""
|
||||
_ensure_source_file(source_path)
|
||||
timestamp = datetime.utcnow().strftime('%Y%m%d-%H%M%S')
|
||||
object_name = f'test-model-data-{timestamp}.csv'
|
||||
target_uri = f'{MINIO_ALIAS}/{MINIO_BUCKET}/{object_name}'
|
||||
|
||||
subprocess.run(
|
||||
['mc', 'cp', str(source_path), target_uri],
|
||||
check=True,
|
||||
)
|
||||
return object_name
|
||||
|
||||
|
||||
def insert_experiment_run(file_name: str, request_data: dict) -> int:
|
||||
"""Insert experiment_run record and return the generated ID."""
|
||||
now = datetime.utcnow()
|
||||
payload = {
|
||||
**request_data,
|
||||
'fileName': file_name,
|
||||
'bucketName': MINIO_BUCKET,
|
||||
}
|
||||
|
||||
insert_sql = """
|
||||
INSERT INTO experiment_run (
|
||||
experiment_name,
|
||||
username,
|
||||
status,
|
||||
created_at,
|
||||
updated_at,
|
||||
bucket_name,
|
||||
file_name,
|
||||
request_data
|
||||
)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
||||
RETURNING id;
|
||||
"""
|
||||
|
||||
with psycopg2.connect(**POSTGRES_CONFIG) as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
insert_sql,
|
||||
(
|
||||
request_data['experimentName'],
|
||||
request_data['username'],
|
||||
'MAGE_REQUEST_SENT',
|
||||
now,
|
||||
now,
|
||||
MINIO_BUCKET,
|
||||
file_name,
|
||||
Json(payload),
|
||||
),
|
||||
)
|
||||
experiment_run_id = cur.fetchone()[0]
|
||||
|
||||
return experiment_run_id
|
||||
|
||||
|
||||
def build_workflow_payload(
|
||||
experiment_run_id: int,
|
||||
file_name: str,
|
||||
request_data: dict,
|
||||
) -> dict:
|
||||
"""Convert camelCase request data to snake_case and enrich with runtime values."""
|
||||
return {
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'experiment_name': request_data['experimentName'],
|
||||
'username': request_data['username'],
|
||||
'model_type': request_data['modelType'],
|
||||
'target_variable': request_data['targetVariable'],
|
||||
'variable_columns': request_data['variableColumns'],
|
||||
'lag_train': request_data['lagTrain'],
|
||||
'lag_val': request_data['lagVal'],
|
||||
'rem_static_win': request_data['remStaticWin'],
|
||||
'low_lim': request_data['lowLim'],
|
||||
'upp_lim': request_data['uppLim'],
|
||||
'window': request_data['window'],
|
||||
'use_scaler': request_data['useScaler'],
|
||||
'include_ar': request_data['includeAr'],
|
||||
'train_size': request_data['trainSize'],
|
||||
'shuffle': request_data['shuffle'],
|
||||
'bucket_name': MINIO_BUCKET,
|
||||
'file_name': file_name,
|
||||
'line_separator': request_data['lineSeparator'],
|
||||
'decimal_separator': request_data['decimalSeparator'],
|
||||
'removed_intervals': request_data['removedIntervals'],
|
||||
}
|
||||
|
||||
|
||||
async def trigger_temporal_workflow(workflow_input: dict) -> str:
|
||||
"""Connect to Temporal and trigger the training workflow."""
|
||||
temporal_client = await client.Client.connect(
|
||||
target_host=TEMPORAL_HOST,
|
||||
namespace=TEMPORAL_NAMESPACE,
|
||||
)
|
||||
|
||||
workflow_id = f'train-model-test-{uuid.uuid4()}'
|
||||
await temporal_client.execute_workflow(
|
||||
TEMPORAL_WORKFLOW,
|
||||
workflow_input,
|
||||
id=workflow_id,
|
||||
task_queue=TEMPORAL_TASK_QUEUE,
|
||||
execution_timeout=timedelta(minutes=5),
|
||||
run_timeout=timedelta(minutes=5),
|
||||
task_timeout=timedelta(minutes=5),
|
||||
)
|
||||
return workflow_id
|
||||
|
||||
|
||||
def main() -> None:
|
||||
try:
|
||||
uploaded_file_name = upload_to_minio(DOCS_PATH)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
print(f'Failed to upload file to MinIO: {exc}', file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except FileNotFoundError as exc:
|
||||
print(str(exc), file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
experiment_request = BASE_REQUEST_DATA.copy()
|
||||
|
||||
try:
|
||||
experiment_run_id = insert_experiment_run(uploaded_file_name, experiment_request)
|
||||
except psycopg2.Error as exc:
|
||||
print(f'Database error while inserting experiment_run: {exc}', file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
workflow_payload = build_workflow_payload(
|
||||
experiment_run_id=experiment_run_id,
|
||||
file_name=uploaded_file_name,
|
||||
request_data=experiment_request,
|
||||
)
|
||||
|
||||
try:
|
||||
workflow_id = asyncio.run(trigger_temporal_workflow(workflow_payload))
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f'Failed to start Temporal workflow: {exc}', file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(json.dumps(
|
||||
{
|
||||
'experiment_run_id': experiment_run_id,
|
||||
's3_object_name': uploaded_file_name,
|
||||
'workflow_id': workflow_id,
|
||||
},
|
||||
indent=2,
|
||||
))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user