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
|
MAX_KEYS_CLEANUP=1000
|
||||||
DEFAULT_CLEANUP_BUCKET=model-training
|
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
|
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 import metrics
|
||||||
from model_manager.activities.activities import Activities
|
from model_manager.activities.activities import Activities
|
||||||
|
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
|
||||||
from model_manager.utils.connectors_config import (
|
from model_manager.utils.connectors_config import (
|
||||||
build_minio_config,
|
build_minio_config,
|
||||||
build_mlflow_config,
|
build_mlflow_config,
|
||||||
@@ -117,6 +118,14 @@ async def main():
|
|||||||
|
|
||||||
logger.custom_info(f'Temporal client initialized at {host}', metadata)
|
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 = [
|
workers = [
|
||||||
Worker(
|
Worker(
|
||||||
temporal_client,
|
temporal_client,
|
||||||
|
|||||||
@@ -4,7 +4,9 @@
|
|||||||
set -e
|
set -e
|
||||||
|
|
||||||
if [ -f .env ]; then
|
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"
|
echo "Environment variables loaded from .env"
|
||||||
else
|
else
|
||||||
echo "Warning: .env file not found. Continuing without environment variables."
|
echo "Warning: .env file not found. Continuing without environment variables."
|
||||||
|
|||||||
@@ -54,6 +54,14 @@ def temp_dir():
|
|||||||
# --- Initialization Tests ---
|
# --- 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(
|
def test_cleanup_init_default_values(
|
||||||
mock_storage_repository,
|
mock_storage_repository,
|
||||||
mock_logger,
|
mock_logger,
|
||||||
@@ -61,6 +69,9 @@ def test_cleanup_init_default_values(
|
|||||||
mock_metrics_controller,
|
mock_metrics_controller,
|
||||||
):
|
):
|
||||||
"""Test Cleanup initialization uses default environment values."""
|
"""Test Cleanup initialization uses default environment values."""
|
||||||
|
import model_manager.activities.cleanup
|
||||||
|
|
||||||
|
reload(model_manager.activities.cleanup)
|
||||||
from model_manager.activities.cleanup import Cleanup
|
from model_manager.activities.cleanup import Cleanup
|
||||||
|
|
||||||
cleanup = Cleanup(
|
cleanup = Cleanup(
|
||||||
@@ -140,6 +151,7 @@ def test_cleanup_minio_files_missing_bucket_name(
|
|||||||
asyncio.run(cleanup.cleanup_minio_files({'metadata': {}}))
|
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(
|
def test_cleanup_minio_files_success_with_deletions(
|
||||||
mock_storage_repository,
|
mock_storage_repository,
|
||||||
mock_logger,
|
mock_logger,
|
||||||
@@ -207,6 +219,7 @@ def test_cleanup_minio_files_dry_run(
|
|||||||
cleanup._emit_metrics.assert_called_once()
|
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(
|
def test_cleanup_minio_files_delete_error(
|
||||||
mock_storage_repository,
|
mock_storage_repository,
|
||||||
mock_logger,
|
mock_logger,
|
||||||
@@ -294,6 +307,7 @@ def test_cleanup_temp_directories_nonexistent_path(
|
|||||||
cleanup._emit_metrics.assert_called_once()
|
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(
|
def test_cleanup_temp_directories_success_with_deletions(
|
||||||
temp_dir,
|
temp_dir,
|
||||||
mock_storage_repository,
|
mock_storage_repository,
|
||||||
@@ -362,6 +376,7 @@ def test_cleanup_temp_directories_dry_run(
|
|||||||
cleanup._emit_metrics.assert_called_once()
|
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(
|
def test_cleanup_temp_directories_delete_error(
|
||||||
temp_dir,
|
temp_dir,
|
||||||
mock_storage_repository,
|
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
|
@pytest.mark.asyncio
|
||||||
|
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
||||||
@patch('model_manager.worker.worker.Worker')
|
@patch('model_manager.worker.worker.Worker')
|
||||||
@patch('model_manager.worker.worker.client.Client')
|
@patch('model_manager.worker.worker.client.Client')
|
||||||
@patch('model_manager.worker.worker.Runtime')
|
@patch('model_manager.worker.worker.Runtime')
|
||||||
@@ -360,11 +361,14 @@ async def test_main_temporal_client_configuration(
|
|||||||
mock_runtime_class,
|
mock_runtime_class,
|
||||||
mock_client_class,
|
mock_client_class,
|
||||||
mock_worker_class,
|
mock_worker_class,
|
||||||
|
mock_create_cleanup_schedule,
|
||||||
mock_logger,
|
mock_logger,
|
||||||
):
|
):
|
||||||
"""Test that Temporal client is configured correctly."""
|
"""Test that Temporal client is configured correctly."""
|
||||||
from model_manager.worker.worker import main
|
from model_manager.worker.worker import main
|
||||||
|
|
||||||
|
mock_create_cleanup_schedule.return_value = AsyncMock()
|
||||||
|
|
||||||
with patch.dict(
|
with patch.dict(
|
||||||
os.environ,
|
os.environ,
|
||||||
{'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production'},
|
{'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',
|
target_host='temporal.example.com:7233',
|
||||||
namespace='production',
|
namespace='production',
|
||||||
runtime=mock_runtime,
|
runtime=mock_runtime,
|
||||||
tls=False,
|
tls=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
@patch('model_manager.worker.worker.create_cleanup_schedule')
|
||||||
@patch('model_manager.worker.worker.Worker')
|
@patch('model_manager.worker.worker.Worker')
|
||||||
@patch('model_manager.worker.worker.client.Client')
|
@patch('model_manager.worker.worker.client.Client')
|
||||||
@patch('model_manager.worker.worker.Runtime')
|
@patch('model_manager.worker.worker.Runtime')
|
||||||
@@ -440,12 +445,15 @@ async def test_main_worker_configuration(
|
|||||||
mock_runtime_class,
|
mock_runtime_class,
|
||||||
mock_client_class,
|
mock_client_class,
|
||||||
mock_worker_class,
|
mock_worker_class,
|
||||||
|
mock_create_cleanup_schedule,
|
||||||
mock_env_vars,
|
mock_env_vars,
|
||||||
mock_logger,
|
mock_logger,
|
||||||
):
|
):
|
||||||
"""Test that Temporal worker is configured with correct parameters."""
|
"""Test that Temporal worker is configured with correct parameters."""
|
||||||
from model_manager.worker.worker import main
|
from model_manager.worker.worker import main
|
||||||
|
|
||||||
|
mock_create_cleanup_schedule.return_value = AsyncMock()
|
||||||
|
|
||||||
# Setup mocks
|
# Setup mocks
|
||||||
mock_get_logger.return_value = mock_logger
|
mock_get_logger.return_value = mock_logger
|
||||||
mock_build_mongodb.return_value = {
|
mock_build_mongodb.return_value = {
|
||||||
@@ -488,10 +496,10 @@ async def test_main_worker_configuration(
|
|||||||
# Verify Worker was created with correct configuration
|
# Verify Worker was created with correct configuration
|
||||||
assert mock_worker_class.call_count == 2
|
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]
|
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[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_workflow_tasks'] == 10
|
||||||
assert train_call_args[1]['max_concurrent_activities'] == 10
|
assert train_call_args[1]['max_concurrent_activities'] == 10
|
||||||
assert train_call_args[1]['max_concurrent_local_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.train_model in train_activities_list
|
||||||
assert mock_activities.cleanup_resources 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]
|
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[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_workflow_tasks'] == 20
|
||||||
assert cleanup_call_args[1]['max_concurrent_activities'] == 20
|
assert cleanup_call_args[1]['max_concurrent_activities'] == 20
|
||||||
assert cleanup_call_args[1]['max_concurrent_local_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.
|
- 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.
|
- Criar um gráfico no grafana para cada nova atividade.
|
||||||
- Atualizar a documentação dos métodos alterados.
|
- Atualizar a documentação dos métodos alterados.
|
||||||
|
|||||||
10
values.yaml
10
values.yaml
@@ -246,6 +246,16 @@ env:
|
|||||||
- name: DEFAULT_CLEANUP_BUCKET
|
- name: DEFAULT_CLEANUP_BUCKET
|
||||||
value: "model-training"
|
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
|
- name: EXTRA_PIP_REQUIREMENTS
|
||||||
value: "git+https://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@github.com/Aignosi/sientia-mlops-library.git"
|
value: "git+https://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@github.com/Aignosi/sientia-mlops-library.git"
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user