339 lines
11 KiB
Python
339 lines
11 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 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
|
|
""",
|
|
)
|
|
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.',
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
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)
|
|
|
|
# Require scenario argument if not listing
|
|
if not args.scenario:
|
|
print('Error: --scenario is required. Use --list to see available scenarios.', file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Load scenario
|
|
try:
|
|
experiment_request = load_scenario(args.scenario)
|
|
print(f"Loaded scenario: {args.scenario}")
|
|
except FileNotFoundError as exc:
|
|
print(str(exc), file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Upload CSV to MinIO
|
|
try:
|
|
uploaded_file_name = upload_to_minio(args.csv)
|
|
print(f"Uploaded CSV to MinIO: {uploaded_file_name}")
|
|
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)
|
|
|
|
# Insert experiment run
|
|
try:
|
|
experiment_run_id = insert_experiment_run(uploaded_file_name, experiment_request)
|
|
print(f"Created experiment_run with ID: {experiment_run_id}")
|
|
except psycopg2.Error as exc:
|
|
print(f'Database error while inserting experiment_run: {exc}', file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# 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))
|
|
except Exception as exc: # noqa: BLE001
|
|
print(f'Failed to start Temporal workflow: {exc}', file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
'scenario': args.scenario,
|
|
'experiment_run_id': experiment_run_id,
|
|
's3_object_name': uploaded_file_name,
|
|
'workflow_id': workflow_id,
|
|
},
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|