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

SIENTIAPDE-1253: Expose workflow activity timeouts as environment variables and document them.

Browse Source
This commit is contained in:
Bruno Domingues
2025-10-15 16:03:01 -03:00
parent 1ec40bfb2a
commit 780184769d
3 changed files with 45 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
model_manager/workflows/train_model.py
View File

@@ -12,6 +12,7 @@ This workflow orchestrates the complete ML model training process, including:
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import os
from datetime import timedelta
from typing import Any
@@ -23,6 +24,16 @@ with workflow.unsafe.imports_passed_through():
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
# Activity Timeouts (in seconds) - Configurable via environment variables
# Defaults are designed to handle large files (up to 200MB)
TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
TIMEOUT_DOWNLOAD_FILE = int(os.getenv('TIMEOUT_DOWNLOAD_FILE', '600'))
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '1800'))
TIMEOUT_SAVE_MODEL = int(os.getenv('TIMEOUT_SAVE_MODEL', '300'))
TIMEOUT_CLEANUP_DIRECTORY = int(os.getenv('TIMEOUT_CLEANUP_DIRECTORY', '60'))
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '60'))
TIMEOUT_UPDATE_DATABASE = int(os.getenv('TIMEOUT_UPDATE_DATABASE', '30'))
@workflow.defn(name='train_model')
class TrainModel:
@@ -173,7 +184,7 @@ class TrainModel:
Activities.validate_train_params,
validation_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
)
# Validation succeeded: Update status
@@ -243,7 +254,7 @@ class TrainModel:
Activities.fetch_file_from_minio,
download_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
start_to_close_timeout=timedelta(seconds=TIMEOUT_DOWNLOAD_FILE),
)
# Step 2: Train model with downloaded file
@@ -257,7 +268,7 @@ class TrainModel:
Activities.train_model,
train_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
)
# Training succeeded: Update status
@@ -330,7 +341,7 @@ class TrainModel:
Activities.save_model,
save_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=120),
start_to_close_timeout=timedelta(seconds=TIMEOUT_SAVE_MODEL),
)
# Model saved successfully: Update status to MLFLOW_SENT with run_name
@@ -396,7 +407,7 @@ class TrainModel:
Activities.cleanup_run_directory,
cleanup_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_DIRECTORY),
)
workflow.logger.info(
@@ -414,7 +425,7 @@ class TrainModel:
Activities.delete_file_from_minio,
delete_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
)
# Cleanup succeeded: Update status to FILE_DELETED
@@ -486,5 +497,5 @@ class TrainModel:
Activities.update_experiment_run,
update_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE),
)