238 lines
7.2 KiB
Python
238 lines
7.2 KiB
Python
#!/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
|
|
|
|
from dotenv import load_dotenv
|
|
import psycopg2
|
|
from psycopg2.extras import Json
|
|
from temporalio import client
|
|
|
|
|
|
# Carrega variáveis de ambiente do arquivo .env na raiz do projeto
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
ENV_PATH = PROJECT_ROOT / '.env'
|
|
if ENV_PATH.exists():
|
|
load_dotenv(dotenv_path=ENV_PATH)
|
|
|
|
|
|
DOCS_PATH = Path('docs/test-model-data.csv')
|
|
MINIO_ALIAS = 'suse'
|
|
MINIO_BUCKET = 'model-training'
|
|
|
|
POSTGRES_CONFIG = {
|
|
'host': os.getenv('POSTGRES_HOST'),
|
|
'port': os.getenv('POSTGRES_PORT'),
|
|
'user': os.getenv('POSTGRES_USER'),
|
|
'password': os.getenv('POSTGRES_PASSWORD'),
|
|
'dbname': os.getenv('POSTGRES_DBNAME'),
|
|
}
|
|
|
|
TEMPORAL_HOST = os.getenv('TEMPORAL_HOST')
|
|
TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE')
|
|
TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE')
|
|
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( # noqa: S603
|
|
['mc', 'cp', str(source_path), target_uri], # noqa: S607
|
|
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'],
|
|
'ORCHESTRATOR_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,
|
|
tls=os.getenv('TEMPORAL_USE_TLS', False),
|
|
)
|
|
|
|
workflow_id = f'train-model-test-{uuid.uuid4()}'
|
|
await temporal_client.execute_workflow(
|
|
TEMPORAL_WORKFLOW,
|
|
workflow_input,
|
|
id=workflow_id,
|
|
task_queue=TRAIN_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()
|