sientia-dataops-model-manager/scripts/run_training_test.py

#!/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(  # 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'],
                    '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()