gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: update environment configuration and remove deprecated model serving

Browse Source 
- Modified `.env.example` to set local defaults for PostgreSQL, MLflow, and MinIO configurations.
- Added MongoDB configuration parameters to the environment setup.
- Updated `README.md` to reflect changes in workflow input parameters and task queue naming conventions.
- Removed the `ModelServing` class to streamline the codebase, as it was deemed unnecessary.
- Adjusted `connectors_config.py` to align with new environment variable names and improve clarity.
- Updated tests to reflect changes in configuration handling and removed tests related to the deleted `ModelServing` class.
This commit is contained in:
vitor-aignosi
2026-04-07 16:58:50 -03:00
parent c5cd382350
commit 0ae03b246f
11 changed files with 301 additions and 1303 deletions
Download Patch File Download Diff File Expand all files Collapse all files

750
scripts/run_training_test.py
View File

@@ -1,37 +1,31 @@
#!/usr/bin/env python3
"""Utility script to trigger the training workflow end-to-end for testing.
# ---
# jupyter:
# jupytext:
# formats: py:percent
# text_representation:
# extension: .py
# format_name: percent
# kernelspec:
# display_name: Python 3
# language: python
# name: python3
# ---
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.
Alternatives for local diagnosis (--local / --validate-only):
- example: python scripts/run_training_test.py --scenario 01-linear-regression-basic --local --csv docs/test-model-data.csv
- --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).
"""
# %% [markdown]
# # Training smoke test (`input-sample.md`)
#
# Run cells top to bottom in VS Code / Cursor (**Run Cell** on each `# %%` block).
#
# Steps mirror `input-sample.md`: optional DB delete + insert, `mc cp` to MinIO, Temporal `train_model`.
# Set `POSTGRES_*`, `TEMPORAL_*`, `TRAIN_TASK_QUEUE`, and configure the `mc` alias (default `suse`).
# %%
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
@@ -39,27 +33,20 @@ 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
# %%
# --- configuration (edit here or use `.env` at repo root) ---
PROJECT_ROOT = Path(__file__).resolve().parent.parent
ENV_PATH = PROJECT_ROOT / '.env'
if ENV_PATH.exists():
load_dotenv(dotenv_path=ENV_PATH)
load_dotenv(PROJECT_ROOT / '.env')
EXPERIMENT_RUN_ID = 1001
MINIO_MC_ALIAS = os.getenv('MINIO_MC_ALIAS', 'suse')
MINIO_BUCKET = os.getenv('MINIO_DEFAULT_BUCKET', 'model-training')
OBJECT_NAME = f'training-sample-dataset-{EXPERIMENT_RUN_ID}.csv'
LOCAL_CSV = PROJECT_ROOT / 'input_dataset.csv'
DEFAULT_CSV_PATH = Path('docs/test-model-data.csv')
TEST_SCENARIOS_DIR = PROJECT_ROOT / 'docs' / 'test-scenarios'
MINIO_ALIAS = 'suse'
MINIO_BUCKET = 'model-training'
# Mapeamento de CSV específico por cenário
SCENARIO_CSV_MAPPING = {
'12-angular-test-date-format': Path('docs/DB_CV022_WIT230.csv'),
'13-angular-test-double-date-column': Path('docs/DB_CV022_WIT230 _double_date_column.csv'),
}
POSTGRES_CONFIG = {
PG = {
'host': os.getenv('POSTGRES_HOST'),
'port': os.getenv('POSTGRES_PORT'),
'port': int(os.getenv('POSTGRES_PORT', '5432')),
'user': os.getenv('POSTGRES_USER'),
'password': os.getenv('POSTGRES_PASSWORD'),
'dbname': os.getenv('POSTGRES_DBNAME'),
@@ -68,597 +55,108 @@ POSTGRES_CONFIG = {
TEMPORAL_HOST = os.getenv('TEMPORAL_HOST')
TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE')
TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE')
TEMPORAL_WORKFLOW = 'train_model'
TEMPORAL_TLS = os.getenv('TEMPORAL_USE_TLS', 'false').lower() in ('1', 'true', 'yes')
# %%
# --- 1) database: delete previous row (same id), then insert `experiment_run` ---
# Primary key column is `id` (see `experiment_tracking` updates). `request_data` matches the SQL sample in `input-sample.md`.
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')])
now = datetime.utcnow()
request_data = {
'experiment_run_id': EXPERIMENT_RUN_ID,
'variable_columns': ['feature_a', 'feature_b'],
'target_variable': 'target',
'bucket_name': MINIO_BUCKET,
'file_name': OBJECT_NAME,
'line_separator': ',',
'decimal_separator': '.',
'train_size': 80,
'shuffle': True,
'random_state': 42,
'model_name': 'test-runtime-linear-regression-model',
'model_type': 'linear_regression',
'model_id': EXPERIMENT_RUN_ID,
'data_model_kwargs': {},
'model_kwargs': {},
'opt_params': {},
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
}
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),
),
with psycopg2.connect(**PG) as conn:
with conn.cursor() as cur:
cur.execute('DELETE FROM public.experiment_run WHERE id = %s', (EXPERIMENT_RUN_ID,))
cur.execute(
"""
INSERT INTO public.experiment_run (
id, experiment_name, run_name, username, status, error_message,
created_at, updated_at, bucket_name, file_name, request_data, orchestrator_response_data
)
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 + metrics).
Reads CSV from disk, runs DataManagerRepository.prepare_training_data and
compute_regression_metrics. 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.data_manager_repository import DataManagerRepository
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
# Use scenario-specific CSV if mapped, otherwise use provided csv_path
if scenario_name in SCENARIO_CSV_MAPPING:
csv_path = SCENARIO_CSV_MAPPING[scenario_name]
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__)
data_manager_repository = DataManagerRepository(logger)
with open(csv_path, 'rb') as f:
file_content = BytesIO(f.read())
try:
train_result = data_manager_repository.prepare_training_data(
train_file_bytes=file_content.getvalue(),
validation_file_bytes=None,
params=train_params,
metadata={'source': 'run_local_pipeline', 'scenario': scenario_name},
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
(
EXPERIMENT_RUN_ID,
'test-experiment-name',
'test-run-name',
'test-username',
'ORCHESTRATOR_WAITING_PROC',
None,
now,
now,
MINIO_BUCKET,
OBJECT_NAME,
Json(request_data),
None,
),
)
train_result = data_manager_repository.compute_regression_metrics(
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
# %%
# --- 2) MinIO: upload local CSV (requires `mc` CLI and alias configured) ---
subprocess.run(
['mc', 'cp', str(LOCAL_CSV), f'{MINIO_MC_ALIAS}/{MINIO_BUCKET}/{OBJECT_NAME}'],
check=True,
)
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)
# %%
# --- 3) Temporal: start `train_model` (flat payload; worker fills `model_metadata` in `load_model_metadata`) ---
result['success'] = True
result['run_name'] = getattr(train_result, 'run_name', None)
result['run_dir'] = getattr(train_result, 'run_dir', None)
return result
if not TEMPORAL_HOST or not TEMPORAL_NAMESPACE or not TRAIN_TASK_QUEUE:
raise RuntimeError('Set TEMPORAL_HOST, TEMPORAL_NAMESPACE, and TRAIN_TASK_QUEUE')
TH, TN, TQ = TEMPORAL_HOST, TEMPORAL_NAMESPACE, TRAIN_TASK_QUEUE
_workflow_input = {
'experiment_run_id': EXPERIMENT_RUN_ID,
'variable_columns': ['feature_a', 'feature_b'],
'target_variable': 'target',
'bucket_name': MINIO_BUCKET,
'file_name': OBJECT_NAME,
'line_separator': ',',
'decimal_separator': '.',
'train_size': 80,
'shuffle': True,
'random_state': 42,
'model_name': 'test-runtime-linear-regression-model',
'model_type': 'linear_regression',
'model_id': EXPERIMENT_RUN_ID,
'data_model_kwargs': {},
'model_kwargs': {},
'opt_params': {},
}
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
# Use scenario-specific CSV if mapped, otherwise use provided csv_path
if scenario_name in SCENARIO_CSV_MAPPING:
csv_path = SCENARIO_CSV_MAPPING[scenario_name]
print(f' Using scenario-specific CSV: {csv_path}')
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 _handle_list_scenarios() -> None:
"""Print available scenarios and exit."""
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)
def _handle_run_all(args: argparse.Namespace) -> None:
"""Run all scenarios and exit with appropriate code."""
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)
failed_count = sum(1 for r in results if not r['success'])
sys.exit(1 if failed_count > 0 else 0)
def _handle_validate_only(args: argparse.Namespace) -> None:
"""Validate scenario parameters only and exit."""
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)
sys.exit(0)
def _handle_local(args: argparse.Namespace) -> None:
"""Run local pipeline and exit."""
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))
def _handle_single_scenario(args: argparse.Namespace) -> None:
"""Run one scenario (MinIO + Postgres + Temporal) and print result."""
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,
)
)
def main() -> None:
args = parse_args()
if args.list:
_handle_list_scenarios()
if args.all:
_handle_run_all(args)
if args.validate_only:
_handle_validate_only(args)
if args.local:
_handle_local(args)
return
if not args.scenario:
print(
'Error: --scenario or --all is required. Use --list to see available scenarios.',
file=sys.stderr,
)
sys.exit(1)
_handle_single_scenario(args)
if __name__ == '__main__':
main()
c = await client.Client.connect(
target_host=TH,
namespace=TN,
tls=TEMPORAL_TLS,
)
wid = f'train-model-test-{uuid.uuid4()}'
await c.execute_workflow( # type: ignore[call-overload]
'train_model',
_workflow_input,
id=wid,
task_queue=TQ,
execution_timeout=timedelta(minutes=5),
run_timeout=timedelta(minutes=5),
task_timeout=timedelta(minutes=5),
)
print(wid)