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 (default):
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.

Local diagnosis (--local / --validate-only):
- --validate-only: Validates scenario parameters only (no MinIO, Postgres, Temporal).
- --local: Runs the same training pipeline locally (validate + load CSV + train +
  after_train_calculation). Use to get full Python tracebacks for debugging.
  Does not upload to MinIO, insert DB, or start Temporal. By default skips
  MLflow save; use --local-save-mlflow to also test saving to MLflow.

Prerequisites (default flow):
- `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 io import BytesIO
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 _resolve_csv_path(csv_path: Path) -> Path:
    """Resolve CSV path; if not found in project root, try docs/."""
    if csv_path.is_absolute():
        return csv_path
    resolved = PROJECT_ROOT / csv_path
    if resolved.exists():
        return resolved
    docs_path = PROJECT_ROOT / 'docs' / csv_path.name
    if docs_path.exists():
        return docs_path
    return resolved


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'],
        'date_column': request_data.get('dateColumn'),
        'date_format': request_data.get('dateFormat'),
        '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', {}),
        'static_threshold': request_data.get('staticThreshold'),
    }


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 _build_local_payload(request_data: dict, csv_path: Path) -> dict:
    """Build workflow payload for local run (no real experiment_run_id)."""
    return build_workflow_payload(
        experiment_run_id=0,
        file_name=csv_path.name,
        request_data=request_data,
    )


def run_validate_only(scenario_name: str) -> dict:
    """Validate scenario parameters only. No MinIO, Postgres, or Temporal.

    Returns:
        dict: {'success': bool, 'error': str | None, 'scenario': str}
    """
    from model_manager.utils.models.train_model_params import TrainModelParams

    result = {'scenario': scenario_name, 'success': False, 'error': None}
    try:
        request_data = load_scenario(scenario_name)
    except FileNotFoundError as exc:
        result['error'] = str(exc)
        return result

    payload = _build_local_payload(request_data, Path('local.csv'))
    try:
        train_params = TrainModelParams.from_dict(payload)
        train_params.validate_business_rules()
        result['success'] = True
    except (ValueError, TypeError, KeyError) as e:
        result['error'] = str(e)
    return result


def run_local_pipeline(
    scenario_name: str,
    csv_path: Path,
    save_mlflow: bool = False,
) -> dict:
    """Run the same training pipeline locally (validate + train + after_train).

    Reads CSV from disk, runs TrainingRepository.train and after_train_calculation.
    Optionally saves to MLflow if save_mlflow is True (requires MLflow env).

    Returns:
        dict: {'success': bool, 'error': str | None, 'scenario': str, ...}
    """
    from model_manager.utils.logger_helper import get_logger
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.repository.training_repository import TrainingRepository

    result = {
        'scenario': scenario_name,
        'success': False,
        'error': None,
    }
    try:
        request_data = load_scenario(scenario_name)
    except FileNotFoundError as exc:
        result['error'] = str(exc)
        return result

    csv_path = _resolve_csv_path(csv_path)
    _ensure_source_file(csv_path)
    payload = _build_local_payload(request_data, csv_path)
    try:
        train_params = TrainModelParams.from_dict(payload)
        train_params.validate_business_rules()
    except (ValueError, TypeError, KeyError) as e:
        result['error'] = f'Validation failed: {e}'
        return result

    logger = get_logger(__name__)
    training_repository = TrainingRepository(logger)

    with open(csv_path, 'rb') as f:
        file_content = BytesIO(f.read())

    try:
        train_result = training_repository.train(file_content, train_params)
        train_result = training_repository.after_train_calculation(
            train_params, train_result
        )
    except Exception as e:
        result['error'] = str(e)
        raise  # re-raise so caller gets full traceback for diagnosis

    if save_mlflow:
        from model_manager.utils.connectors_config import build_mlflow_config
        from model_manager.utils.repository.model_repository import ModelRepository

        mlflow_config = build_mlflow_config()
        model_repository = ModelRepository(
            url=mlflow_config['url'],
            username=mlflow_config['username'],
            password=mlflow_config['password'],
            logger=logger,
        )
        train_result = model_repository.save_model(train_result)

    result['success'] = True
    result['run_name'] = getattr(train_result, 'run_name', None)
    result['run_dir'] = getattr(train_result, 'run_dir', None)
    return result


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

  # Validate scenario parameters only (no external services)
  python scripts/run_training_test.py --scenario linear-regression-basic --validate-only

  # Run training pipeline locally to diagnose errors (full traceback)
  python scripts/run_training_test.py --scenario linear-regression-basic --local --csv docs/test-model-data.csv

  # Local run and save to MLflow (requires MLflow env)
  python scripts/run_training_test.py --scenario linear-regression-basic --local --local-save-mlflow
""",
    )
    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.',
    )
    parser.add_argument(
        '--validate-only',
        action='store_true',
        help='Only validate scenario parameters (no MinIO, Postgres, Temporal).',
    )
    parser.add_argument(
        '--local',
        action='store_true',
        help='Run training pipeline locally (validate + train from CSV) to get full tracebacks.',
    )
    parser.add_argument(
        '--local-save-mlflow',
        action='store_true',
        help='With --local, also save the model to MLflow (requires MLflow env).',
    )
    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

    csv_path = _resolve_csv_path(csv_path)

    # 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)

    # Validate-only mode: no external services
    if args.validate_only:
        if not args.scenario:
            print('Error: --scenario is required with --validate-only.', file=sys.stderr)
            sys.exit(1)
        result = run_validate_only(args.scenario)
        if result['success']:
            print(f"Validation OK: {result['scenario']}")
        else:
            print(f"Validation failed: {result['error']}", file=sys.stderr)
            sys.exit(1)
        return

    # Local pipeline: same code path as worker, full traceback on error
    if args.local:
        if not args.scenario:
            print('Error: --scenario is required with --local.', file=sys.stderr)
            sys.exit(1)
        print(f'Running local pipeline: {args.scenario} (CSV: {args.csv})')
        result = run_local_pipeline(
            args.scenario,
            args.csv,
            save_mlflow=args.local_save_mlflow,
        )
        if not result['success']:
            print(f"Error: {result['error']}", file=sys.stderr)
            sys.exit(1)
        out = {'scenario': result['scenario'], 'success': True}
        if result.get('run_name') is not None:
            out['run_name'] = result['run_name']
        if result.get('run_dir') is not None:
            out['run_dir'] = result['run_dir']
        print(json.dumps(out, indent=2))
        return

    # 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 (MinIO + Postgres + Temporal)
    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()