SIENTIAPDE-1350: Implement scheduled cleanup workflow using Temporal schedules. Adds schedule creation to worker startup and configures environment variables.
This commit is contained in:
@@ -49,4 +49,9 @@ TIMEOUT_CLEANUP_LOCAL=120
|
||||
MAX_KEYS_CLEANUP=1000
|
||||
DEFAULT_CLEANUP_BUCKET=model-training
|
||||
|
||||
CLEANUP_SCHEDULE_ID=cleanup-files-daily
|
||||
CLEANUP_CRON="0 0 * * *"
|
||||
CLEANUP_TIMEZONE=UTC
|
||||
CLEANUP_EXECUTION_TIMEOUT_HOURS=1
|
||||
|
||||
EXTRA_PIP_REQUIREMENTS=git+https://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@github.com/Aignosi/sientia-mlops-library.git
|
||||
|
||||
0
model_manager/schedules/__init__.py
Normal file
0
model_manager/schedules/__init__.py
Normal file
85
model_manager/schedules/cleanup_schedule.py
Normal file
85
model_manager/schedules/cleanup_schedule.py
Normal file
@@ -0,0 +1,85 @@
|
||||
"""Schedule configuration for cleanup workflow."""
|
||||
|
||||
import os
|
||||
from datetime import timedelta
|
||||
|
||||
from sientia_do.observability.logger import Logger as SientiaLogger
|
||||
from temporalio.client import (
|
||||
Client,
|
||||
Schedule,
|
||||
ScheduleActionStartWorkflow,
|
||||
ScheduleSpec,
|
||||
)
|
||||
|
||||
# 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'))
|
||||
|
||||
|
||||
async def schedule_exists(
|
||||
client: Client, schedule_id: str, logger: SientiaLogger, metadata: dict[str, str | None]
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a schedule already exists.
|
||||
|
||||
Args:
|
||||
client: Temporal client instance
|
||||
schedule_id: ID of the schedule to check
|
||||
|
||||
Returns:
|
||||
True if schedule exists, False otherwise
|
||||
"""
|
||||
try:
|
||||
async for schedule in await client.list_schedules():
|
||||
if schedule.id == schedule_id:
|
||||
return True
|
||||
return False
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.custom_error(f'Error checking if schedule exists: {e}', metadata)
|
||||
return False
|
||||
|
||||
|
||||
async def create_cleanup_schedule(
|
||||
client: Client, logger: SientiaLogger, metadata: dict[str, str | None]
|
||||
) -> None:
|
||||
"""
|
||||
Create or update the cleanup files schedule.
|
||||
|
||||
This function is idempotent and can be called multiple times safely.
|
||||
It will only create the schedule if it doesn't already exist.
|
||||
|
||||
Args:
|
||||
client: Temporal client instance
|
||||
"""
|
||||
# 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
|
||||
|
||||
await client.create_schedule(
|
||||
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,
|
||||
execution_timeout=timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS),
|
||||
),
|
||||
spec=ScheduleSpec(
|
||||
cron_expressions=[CLEANUP_CRON],
|
||||
time_zone_name=CLEANUP_TIMEZONE,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
logger.custom_info(
|
||||
f"Schedule '{SCHEDULE_ID}' created successfully. "
|
||||
f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
|
||||
metadata,
|
||||
)
|
||||
@@ -37,6 +37,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
|
||||
from model_manager import metrics
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
|
||||
from model_manager.utils.connectors_config import (
|
||||
build_minio_config,
|
||||
build_mlflow_config,
|
||||
@@ -117,6 +118,14 @@ async def main():
|
||||
|
||||
logger.custom_info(f'Temporal client initialized at {host}', metadata)
|
||||
|
||||
# Create cleanup schedule (idempotent - only creates if doesn't exist)
|
||||
try:
|
||||
await create_cleanup_schedule(temporal_client, logger, metadata)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.custom_error(f'Failed to configure cleanup schedule: {e}', metadata)
|
||||
# Don't fail the worker startup if schedule creation fails
|
||||
# The schedule can be created manually if needed
|
||||
|
||||
workers = [
|
||||
Worker(
|
||||
temporal_client,
|
||||
|
||||
@@ -4,7 +4,9 @@
|
||||
set -e
|
||||
|
||||
if [ -f .env ]; then
|
||||
export $(cat .env | grep -v '^#' | xargs)
|
||||
set -a
|
||||
source <(cat .env | grep -v '^#' | grep -v '^$')
|
||||
set +a
|
||||
echo "Environment variables loaded from .env"
|
||||
else
|
||||
echo "Warning: .env file not found. Continuing without environment variables."
|
||||
|
||||
@@ -54,6 +54,14 @@ def temp_dir():
|
||||
# --- Initialization Tests ---
|
||||
|
||||
|
||||
@patch.dict(
|
||||
'model_manager.activities.cleanup.os.environ',
|
||||
{
|
||||
'CLEANUP_RETENTION_HOURS': '24',
|
||||
'CLEANUP_DRY_RUN': 'false',
|
||||
'MAX_KEYS_CLEANUP': '1000',
|
||||
},
|
||||
)
|
||||
def test_cleanup_init_default_values(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
@@ -61,6 +69,9 @@ def test_cleanup_init_default_values(
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test Cleanup initialization uses default environment values."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
@@ -140,6 +151,7 @@ def test_cleanup_minio_files_missing_bucket_name(
|
||||
asyncio.run(cleanup.cleanup_minio_files({'metadata': {}}))
|
||||
|
||||
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
def test_cleanup_minio_files_success_with_deletions(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
@@ -207,6 +219,7 @@ def test_cleanup_minio_files_dry_run(
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
def test_cleanup_minio_files_delete_error(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
@@ -294,6 +307,7 @@ def test_cleanup_temp_directories_nonexistent_path(
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
def test_cleanup_temp_directories_success_with_deletions(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
@@ -362,6 +376,7 @@ def test_cleanup_temp_directories_dry_run(
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
def test_cleanup_temp_directories_delete_error(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
|
||||
0
tests/schedules/__init__.py
Normal file
0
tests/schedules/__init__.py
Normal file
@@ -335,6 +335,7 @@ async def test_main_handles_exception(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
||||
@patch('model_manager.worker.worker.Worker')
|
||||
@patch('model_manager.worker.worker.client.Client')
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@@ -360,11 +361,14 @@ async def test_main_temporal_client_configuration(
|
||||
mock_runtime_class,
|
||||
mock_client_class,
|
||||
mock_worker_class,
|
||||
mock_create_cleanup_schedule,
|
||||
mock_logger,
|
||||
):
|
||||
"""Test that Temporal client is configured correctly."""
|
||||
from model_manager.worker.worker import main
|
||||
|
||||
mock_create_cleanup_schedule.return_value = AsyncMock()
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production'},
|
||||
@@ -410,11 +414,12 @@ async def test_main_temporal_client_configuration(
|
||||
target_host='temporal.example.com:7233',
|
||||
namespace='production',
|
||||
runtime=mock_runtime,
|
||||
tls=False,
|
||||
tls=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
||||
@patch('model_manager.worker.worker.Worker')
|
||||
@patch('model_manager.worker.worker.client.Client')
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@@ -440,12 +445,15 @@ async def test_main_worker_configuration(
|
||||
mock_runtime_class,
|
||||
mock_client_class,
|
||||
mock_worker_class,
|
||||
mock_create_cleanup_schedule,
|
||||
mock_env_vars,
|
||||
mock_logger,
|
||||
):
|
||||
"""Test that Temporal worker is configured with correct parameters."""
|
||||
from model_manager.worker.worker import main
|
||||
|
||||
mock_create_cleanup_schedule.return_value = AsyncMock()
|
||||
|
||||
# Setup mocks
|
||||
mock_get_logger.return_value = mock_logger
|
||||
mock_build_mongodb.return_value = {
|
||||
@@ -488,10 +496,10 @@ async def test_main_worker_configuration(
|
||||
# Verify Worker was created with correct configuration
|
||||
assert mock_worker_class.call_count == 2
|
||||
|
||||
# Primeira chamada: worker de treinamento (train_model-queue)
|
||||
# Primeira chamada: worker de treinamento (train_model-local_queue)
|
||||
train_call_args = mock_worker_class.call_args_list[0]
|
||||
assert train_call_args[0][0] == mock_client_instance # temporal_client
|
||||
assert train_call_args[1]['task_queue'] == 'train_model-queue'
|
||||
assert train_call_args[1]['task_queue'] == 'train_model-local_queue'
|
||||
assert train_call_args[1]['max_concurrent_workflow_tasks'] == 10
|
||||
assert train_call_args[1]['max_concurrent_activities'] == 10
|
||||
assert train_call_args[1]['max_concurrent_local_activities'] == 10
|
||||
@@ -503,10 +511,10 @@ async def test_main_worker_configuration(
|
||||
assert mock_activities.train_model in train_activities_list
|
||||
assert mock_activities.cleanup_resources in train_activities_list
|
||||
|
||||
# Segunda chamada: worker de cleanup (cleanup-queue)
|
||||
# Segunda chamada: worker de cleanup (cleanup-local_queue)
|
||||
cleanup_call_args = mock_worker_class.call_args_list[1]
|
||||
assert cleanup_call_args[0][0] == mock_client_instance # temporal_client
|
||||
assert cleanup_call_args[1]['task_queue'] == 'cleanup-queue'
|
||||
assert cleanup_call_args[1]['task_queue'] == 'cleanup-local_queue'
|
||||
assert cleanup_call_args[1]['max_concurrent_workflow_tasks'] == 20
|
||||
assert cleanup_call_args[1]['max_concurrent_activities'] == 20
|
||||
assert cleanup_call_args[1]['max_concurrent_local_activities'] == 20
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
- Adcionar configuração do cron job ao helm chart
|
||||
- Refatorar o arquivo .dockerignore para só deixar copiar os arquivos que forem necessários para a execução do container, pois ele está copiando muitos arquivos desnecessários.
|
||||
- Criar um gráfico no grafana para cada nova atividade.
|
||||
- Atualizar a documentação dos métodos alterados.
|
||||
|
||||
10
values.yaml
10
values.yaml
@@ -246,6 +246,16 @@ env:
|
||||
- name: DEFAULT_CLEANUP_BUCKET
|
||||
value: "model-training"
|
||||
|
||||
# Cleanup Schedule Configuration
|
||||
- name: CLEANUP_SCHEDULE_ID
|
||||
value: "cleanup-files-daily"
|
||||
- name: CLEANUP_CRON
|
||||
value: "0 0 * * *" # Midnight UTC
|
||||
- name: CLEANUP_TIMEZONE
|
||||
value: "UTC"
|
||||
- name: CLEANUP_EXECUTION_TIMEOUT_HOURS
|
||||
value: "1"
|
||||
|
||||
- name: EXTRA_PIP_REQUIREMENTS
|
||||
value: "git+https://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@github.com/Aignosi/sientia-mlops-library.git"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user