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 argparse
import asyncio
import json
import os
import subprocess
import sys
import uuid
from datetime import datetime, timedelta
from pathlib import Path

import psycopg2
from dotenv import load_dotenv
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)


DEFAULT_CSV_PATH = Path('docs/test-model-data.csv')
TEST_SCENARIOS_DIR = PROJECT_ROOT / 'docs' / 'test-scenarios'
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'


def list_available_scenarios() -> list[str]:
    """List all available test scenario files."""
    if not TEST_SCENARIOS_DIR.exists():
        return []
    return sorted([f.stem for f in TEST_SCENARIOS_DIR.glob('*.json')])


def load_scenario(scenario_name: str) -> dict:
    """Load a test scenario from JSON file.

    Args:
        scenario_name: Name of the scenario (without .json extension)
                       or full path to a JSON file.

    Returns:
        Dictionary with scenario data.

    Raises:
        FileNotFoundError: If scenario file doesn't exist.
    """
    # Check if it's a full path
    scenario_path = Path(scenario_name)
    if scenario_path.suffix == '.json' and scenario_path.exists():
        with open(scenario_path) as f:
            return json.load(f)

    # Otherwise, look in the test-scenarios directory
    scenario_file = TEST_SCENARIOS_DIR / f'{scenario_name}.json'
    if not scenario_file.exists():
        available = list_available_scenarios()
        available_str = ', '.join(available) if available else 'none'
        raise FileNotFoundError(
            f"Scenario '{scenario_name}' not found at {scenario_file}.\n"
            f'Available scenarios: {available_str}'
        )

    with open(scenario_file) as f:
        return json.load(f)


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_WAITING_PROC',
                    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'],
        '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'],
        # New parameters
        'model_name': request_data.get('modelName', 'Linear Regression'),
        'degree': request_data.get('degree', 1),
        'interaction_only': request_data.get('interactionOnly', False),
        'nan_treatment': request_data.get('nanTreatment', 'drop'),
        'start_date': request_data.get('startDate'),
        'end_date': request_data.get('endDate'),
        'scaler_name': request_data.get('scalerName', 'None'),
        'support_filters': request_data.get('supportFilters', {}),
    }


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 parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Run training workflow tests with different scenarios.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List available scenarios
  python scripts/run_training_test.py --list

  # Run a specific scenario
  python scripts/run_training_test.py --scenario linear-regression-basic

  # Run with a custom JSON file
  python scripts/run_training_test.py --scenario /path/to/custom-scenario.json

  # Run with a custom CSV data file
  python scripts/run_training_test.py --scenario linear-regression-basic --csv docs/other-data.csv

  # Run all scenarios
  python scripts/run_training_test.py --all
""",
    )
    parser.add_argument(
        '--scenario',
        '-s',
        type=str,
        help='Name of the test scenario (without .json) or path to a JSON file.',
    )
    parser.add_argument(
        '--csv',
        '-c',
        type=Path,
        default=DEFAULT_CSV_PATH,
        help=f'Path to the CSV data file (default: {DEFAULT_CSV_PATH}).',
    )
    parser.add_argument(
        '--list',
        '-l',
        action='store_true',
        help='List all available test scenarios and exit.',
    )
    parser.add_argument(
        '--all',
        '-a',
        action='store_true',
        help='Run all available test scenarios sequentially.',
    )
    return parser.parse_args()


def run_single_scenario(scenario_name: str, csv_path: Path) -> dict:
    """Run a single test scenario and return the result.

    Args:
        scenario_name: Name of the scenario to run.
        csv_path: Path to the CSV data file.

    Returns:
        Dictionary with scenario result including success status and details.
    """
    result = {
        'scenario': scenario_name,
        'success': False,
        'error': None,
        'experiment_run_id': None,
        's3_object_name': None,
        'workflow_id': None,
    }

    # Load scenario
    try:
        experiment_request = load_scenario(scenario_name)
        print(f"  Loaded scenario: {scenario_name}")
    except FileNotFoundError as exc:
        result['error'] = str(exc)
        return result

    # Upload CSV to MinIO
    try:
        uploaded_file_name = upload_to_minio(csv_path)
        result['s3_object_name'] = uploaded_file_name
        print(f"  Uploaded CSV to MinIO: {uploaded_file_name}")
    except subprocess.CalledProcessError as exc:
        result['error'] = f'Failed to upload file to MinIO: {exc}'
        return result
    except FileNotFoundError as exc:
        result['error'] = str(exc)
        return result

    # Insert experiment run
    try:
        experiment_run_id = insert_experiment_run(uploaded_file_name, experiment_request)
        result['experiment_run_id'] = experiment_run_id
        print(f"  Created experiment_run with ID: {experiment_run_id}")
    except psycopg2.Error as exc:
        result['error'] = f'Database error while inserting experiment_run: {exc}'
        return result

    # Build and trigger workflow
    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))
        result['workflow_id'] = workflow_id
        result['success'] = True
        print(f"  Workflow started: {workflow_id}")
    except Exception as exc:  # noqa: BLE001
        result['error'] = f'Failed to start Temporal workflow: {exc}'
        return result

    return result


def print_summary(results: list[dict]) -> None:
    """Print a summary of all scenario results.

    Args:
        results: List of result dictionaries from run_single_scenario.
    """
    passed = [r for r in results if r['success']]
    failed = [r for r in results if not r['success']]

    print('\n' + '=' * 60)
    print('SUMMARY')
    print('=' * 60)
    print(f"Total: {len(results)} | Passed: {len(passed)} | Failed: {len(failed)}")
    print('=' * 60)

    if passed:
        print('\n✓ PASSED:')
        for r in passed:
            print(f"  - {r['scenario']}")

    if failed:
        print('\n✗ FAILED:')
        for r in failed:
            print(f"  - {r['scenario']}")
            if r['error']:
                print(f"    Error: {r['error']}")

    print()


def main() -> None:
    args = parse_args()

    # List scenarios and exit if requested
    if args.list:
        scenarios = list_available_scenarios()
        if scenarios:
            print('Available test scenarios:')
            for scenario in scenarios:
                print(f'  - {scenario}')
        else:
            print(f'No scenarios found in {TEST_SCENARIOS_DIR}')
        sys.exit(0)

    # Run all scenarios if requested
    if args.all:
        scenarios = list_available_scenarios()
        if not scenarios:
            print(f'No scenarios found in {TEST_SCENARIOS_DIR}', file=sys.stderr)
            sys.exit(1)

        print(f'Running {len(scenarios)} scenarios...\n')
        results = []

        for i, scenario in enumerate(scenarios, 1):
            print(f'[{i}/{len(scenarios)}] Running scenario: {scenario}')
            result = run_single_scenario(scenario, args.csv)
            results.append(result)
            status = '✓' if result['success'] else '✗'
            print(f'[{i}/{len(scenarios)}] {status} {scenario}\n')

        print_summary(results)

        # Exit with error code if any scenario failed
        failed_count = sum(1 for r in results if not r['success'])
        sys.exit(1 if failed_count > 0 else 0)

    # Require scenario argument if not listing or running all
    if not args.scenario:
        print('Error: --scenario or --all is required. Use --list to see available scenarios.', file=sys.stderr)
        sys.exit(1)

    # Run single scenario
    print(f'Running scenario: {args.scenario}')
    result = run_single_scenario(args.scenario, args.csv)

    if not result['success']:
        print(f"Error: {result['error']}", file=sys.stderr)
        sys.exit(1)

    print(
        json.dumps(
            {
                'scenario': result['scenario'],
                'experiment_run_id': result['experiment_run_id'],
                's3_object_name': result['s3_object_name'],
                'workflow_id': result['workflow_id'],
            },
            indent=2,
        )
    )


if __name__ == '__main__':
    main()