diff --git a/.env.example b/.env.example index 6efb9cb..22e8021 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/model_manager/schedules/__init__.py b/model_manager/schedules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/model_manager/schedules/cleanup_schedule.py b/model_manager/schedules/cleanup_schedule.py new file mode 100644 index 0000000..084051b --- /dev/null +++ b/model_manager/schedules/cleanup_schedule.py @@ -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, + ) diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index 29200aa..4cda4b1 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -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, diff --git a/run_local.sh b/run_local.sh index 4c88728..4fa3fba 100755 --- a/run_local.sh +++ b/run_local.sh @@ -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." diff --git a/tests/activities/test_cleanup.py b/tests/activities/test_cleanup.py index b86f0d8..2d8f7b3 100644 --- a/tests/activities/test_cleanup.py +++ b/tests/activities/test_cleanup.py @@ -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, diff --git a/tests/schedules/__init__.py b/tests/schedules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 1a3d5d0..1b81da1 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -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 diff --git a/todo-list.txt b/todo-list.txt index 2ff06bd..d8c18c2 100644 --- a/todo-list.txt +++ b/todo-list.txt @@ -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. diff --git a/values.yaml b/values.yaml index eeee7ac..e488d00 100644 --- a/values.yaml +++ b/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"