feat: enhance configuration and scheduling for cleanup processes

- Updated `.env.example` to include new environment variables for MinIO and PyPI configuration.
- Refactored `create_cleanup_schedule` to utilize runtime-specific task queues and improve schedule reconciliation logic.
- Enhanced `Activities` class to require a default bucket in MinIO configuration.
- Adjusted `requirements.txt` to specify version for `evidently`.
- Updated tests to reflect changes in schedule creation and configuration handling.
This commit is contained in:
vitor-aignosi
2026-04-07 12:57:29 -03:00
parent 09ee92f100
commit c5cd382350
13 changed files with 326 additions and 76 deletions

View File

@@ -59,7 +59,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
mlflow_config: MLFlow server configuration dictionary
Required keys: host, port, username, password
minio_config: MinIO storage configuration dictionary
Required keys: endpoint_url, access_key, secret_key, region, use_ssl
Required keys: endpoint_url, access_key, secret_key, region, use_ssl, default_bucket
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
@@ -108,6 +108,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
notification_handler=notification_handler,
metrics_controller=metrics_controller,
secure=minio_config['use_ssl'],
bucket=minio_config['default_bucket'],
)
Training.__init__(

View File

@@ -2,6 +2,7 @@
import os
from datetime import timedelta
from typing import Any
from sientia_do.observability.logger import Logger as SientiaLogger
from temporalio.client import (
@@ -11,14 +12,68 @@ from temporalio.client import (
ScheduleSpec,
)
from model_manager.worker.prepare_worker import build_queue_name
# Schedule configuration from environment variables
SCHEDULE_ID = os.getenv('CLEANUP_SCHEDULE_ID', 'cleanup-files-daily')
CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 0 * * *') # Default: midnight UTC
CLEANUP_TIMEZONE = os.getenv('CLEANUP_TIMEZONE', 'UTC')
CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue')
CLEANUP_EXECUTION_TIMEOUT_HOURS = int(os.getenv('CLEANUP_EXECUTION_TIMEOUT_HOURS', '1'))
def build_cleanup_schedule_id(runtime: str | None) -> str:
"""
Build cleanup schedule ID using runtime-derived naming.
Args:
runtime: Runtime suffix used by workers
Return:
str: Cleanup schedule ID
"""
normalized_runtime = runtime.strip() if runtime else ''
return f'cleanup-files-{normalized_runtime or "single"}-daily'
async def _needs_schedule_reconcile(
schedule_handle: Any,
cleanup_task_queue: str,
logger: SientiaLogger,
metadata: dict[str, str | None],
) -> bool:
"""
Compare configured cleanup schedule against current expected values.
Args:
schedule_handle: Temporal schedule handle for current schedule ID
cleanup_task_queue: Expected cleanup task queue
logger: Logger instance for errors
metadata: Metadata dictionary for logging context
Return:
bool: True when schedule should be recreated to apply current config
"""
try:
schedule_description = await schedule_handle.describe()
schedule = getattr(schedule_description, 'schedule', None)
action = getattr(schedule, 'action', None)
spec = getattr(schedule, 'spec', None)
current_task_queue = getattr(action, 'task_queue', None)
current_execution_timeout = getattr(action, 'execution_timeout', None)
current_cron = getattr(spec, 'cron_expressions', None)
current_timezone = getattr(spec, 'time_zone_name', None)
return (
current_task_queue != cleanup_task_queue
or current_execution_timeout != timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS)
or current_cron != [CLEANUP_CRON]
or current_timezone != CLEANUP_TIMEZONE
)
except Exception as e: # noqa: BLE001
logger.custom_error(f'Error describing cleanup schedule for reconcile: {e}', metadata)
return True
async def schedule_exists(
client: Client, schedule_id: str, logger: SientiaLogger, metadata: dict[str, str | None]
) -> bool:
@@ -58,21 +113,32 @@ async def create_cleanup_schedule(
logger: Logger instance for logging schedule operations
metadata: Metadata dictionary for logging context
"""
# Check if schedule already exists
if await schedule_exists(client, SCHEDULE_ID, logger, metadata):
logger.custom_info(
f"Schedule '{SCHEDULE_ID}' already configured, skipping creation", metadata
)
return
runtime = (os.getenv('RUNTIME') or 'single').strip()
cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single')
schedule_id = build_cleanup_schedule_id(runtime)
created = False
updated = False
if await schedule_exists(client, schedule_id, logger, metadata):
handle = client.get_schedule_handle(schedule_id)
if await _needs_schedule_reconcile(handle, cleanup_task_queue, logger, metadata):
await handle.delete()
updated = True
else:
logger.custom_info(
f"Schedule '{schedule_id}' is already up to date, no-op reconcile",
metadata,
)
return
await client.create_schedule(
SCHEDULE_ID,
schedule_id,
Schedule(
action=ScheduleActionStartWorkflow(
'cleanup_files',
{}, # Empty input, will use default bucket from environment
id=f'cleanup-files-scheduled-{SCHEDULE_ID}',
task_queue=CLEANUP_TASK_QUEUE,
id=f'cleanup-files-scheduled-{schedule_id}',
task_queue=cleanup_task_queue,
execution_timeout=timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS),
),
spec=ScheduleSpec(
@@ -81,9 +147,17 @@ async def create_cleanup_schedule(
),
),
)
created = not updated
logger.custom_info(
f"Schedule '{SCHEDULE_ID}' created successfully. "
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
metadata,
)
if updated:
logger.custom_info(
f"Schedule '{schedule_id}' reconciled successfully. "
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
metadata,
)
elif created:
logger.custom_info(
f"Schedule '{schedule_id}' created successfully. "
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
metadata,
)

View File

@@ -106,6 +106,7 @@ def build_minio_config() -> dict[str, Any]:
MINIO_RETRY_MODE: Retry mode - standard, legacy, or adaptive (default: adaptive)
MINIO_CONNECT_TIMEOUT: Connection timeout in seconds (default: 10)
MINIO_READ_TIMEOUT: Read timeout in seconds (default: 60)
MINIO_DEFAULT_BUCKET: Default S3 bucket for MinioRepository (default: model-training)
Returns:
dict: MinIO configuration dictionary with all required parameters
@@ -120,6 +121,7 @@ def build_minio_config() -> dict[str, Any]:
'retry_mode': getenv('MINIO_RETRY_MODE', 'adaptive'),
'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')),
'read_timeout': int(getenv('MINIO_READ_TIMEOUT', '60')),
'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'model-training'),
}

View File

@@ -39,6 +39,23 @@ def camel_to_snake(text: str) -> str:
return text.lower()
def build_queue_name(workflow_name: str, runtime: str | None = None) -> str:
"""
Build Temporal queue name from workflow name and runtime.
Args:
- workflow_name: str, workflow class name in CamelCase format
- runtime: str | None, runtime suffix for environment-specific queues
Return:
str: queue name in the format <workflow>-<runtime>-queue or <workflow>-queue
"""
snake_workflow_name = camel_to_snake(workflow_name)
if runtime:
return f'{snake_workflow_name}-{runtime}-queue'
return f'{snake_workflow_name}-queue'
def prepare_worker(
main_workflow: type,
other_workflows: Sequence[type],
@@ -62,11 +79,7 @@ def prepare_worker(
Worker: fully configured Temporal worker instance ready to run
"""
main_workflow_name = main_workflow.__name__.upper()
queue_name = (
f'{camel_to_snake(main_workflow.__name__)}-{runtime}-queue'
if runtime
else f'{camel_to_snake(main_workflow.__name__)}-queue'
)
queue_name = build_queue_name(main_workflow.__name__, runtime)
local_workflow_parameters: dict[str, int] = {}
for parameter_name, default_value in parameters:

View File

@@ -5,8 +5,8 @@ It orchestrates Temporal workers, manages task queues, and handles the lifecycle
model training and cleanup workflows.
The worker supports two task queues:
- train_model-queue: For ML model training workflows
- cleanup-queue: For file cleanup workflows
- train_model-<runtime>-queue: For ML model training workflows
- cleanup_files-<runtime>-queue: For file cleanup workflows
Key Features:
- Automatic scaling with PollerBehaviorAutoscaling
@@ -20,8 +20,7 @@ Environment Variables:
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: model-manager)
- TEMPORAL_USE_TLS: Enable TLS for Temporal connection (default: false)
- TRAIN_TASK_QUEUE: Task queue for training workflows (default: train_model-queue)
- CLEANUP_TASK_QUEUE: Task queue for cleanup workflows (default: cleanup-queue)
- RUNTIME: Runtime identifier used in queue naming (default: single)
- POD_ID: Kubernetes pod identifier for metrics
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
@@ -60,8 +59,20 @@ with workflow.unsafe.imports_passed_through():
POD_ID = os.getenv('POD_ID')
RUNTIME = os.getenv('RUNTIME')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue')
CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue')
def _get_runtime(runtime: str | None) -> str:
"""
Resolve runtime using fallback when missing.
Args:
- runtime: str | None, runtime value from environment
Return:
str: normalized runtime value
"""
normalized_runtime = runtime.strip() if runtime else ''
return normalized_runtime or 'single'
async def main():
@@ -84,8 +95,7 @@ async def main():
SystemExit: On graceful shutdown or error conditions
"""
if not RUNTIME:
raise ValueError('RUNTIME environment variable is required')
runtime = _get_runtime(RUNTIME)
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
@@ -93,7 +103,7 @@ async def main():
metadata = {
'pod_id': POD_ID,
'runtime': RUNTIME,
'runtime': runtime,
}
start_prometheus_server(logger, metadata)
@@ -112,7 +122,7 @@ async def main():
metrics_controller = MetricsController(logger=logger)
logger.custom_info(f'Installing runtime {RUNTIME}', metadata)
logger.custom_info(f'Installing runtime {runtime}', metadata)
plugin_store_parameters = build_plugin_store_config()
plugin_store = PluginStore(
@@ -131,7 +141,7 @@ async def main():
metrics_controller=metrics_controller,
)
await plugin_store.install_runtime(runtime_name=RUNTIME)
await plugin_store.install_runtime(runtime_name=runtime)
activities = Activities(
postgres_config=build_postgres_config(),
@@ -181,7 +191,7 @@ async def main():
],
temporal_client=temporal_client,
logger=logger,
runtime=RUNTIME,
runtime=runtime,
),
prepare_worker(
main_workflow=CleanupFiles,
@@ -191,7 +201,7 @@ async def main():
],
temporal_client=temporal_client,
logger=logger,
runtime=RUNTIME,
runtime=runtime,
),
]