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

feat: enhance configuration and scheduling for cleanup processes

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

146
tests/schedules/test_cleanup_schedule.py
View File

@@ -14,6 +14,15 @@ def mock_temporal_client():
client = AsyncMock()
client.list_schedules = AsyncMock()
client.create_schedule = AsyncMock()
handle = AsyncMock()
handle.delete = AsyncMock()
schedule = MagicMock()
schedule.action.task_queue = 'cleanup_files-model-manager-worker-queue'
schedule.action.execution_timeout = timedelta(hours=1)
schedule.spec.cron_expressions = ['0 0 * * *']
schedule.spec.time_zone_name = 'UTC'
handle.describe = AsyncMock(return_value=MagicMock(schedule=schedule))
client.get_schedule_handle = MagicMock(return_value=handle)
return client
@@ -120,10 +129,16 @@ async def test_schedule_exists_handles_exception(mock_temporal_client, mock_logg
@pytest.mark.asyncio
async def test_create_cleanup_schedule_skips_when_exists(
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'RUNTIME': 'model-manager-worker',
},
)
async def test_create_cleanup_schedule_reconciles_when_exists(
mock_temporal_client, mock_logger, metadata
):
"""Test that create_cleanup_schedule skips creation when schedule already exists."""
"""Test that create_cleanup_schedule recreates schedule when it already exists."""
import model_manager.schedules.cleanup_schedule
reload(model_manager.schedules.cleanup_schedule)
@@ -131,7 +146,47 @@ async def test_create_cleanup_schedule_skips_when_exists(
# Mock schedule already exists
mock_schedule = MagicMock()
mock_schedule.id = 'cleanup-files-daily'
mock_schedule.id = 'cleanup-files-model-manager-worker-daily'
async def mock_list_schedules():
yield mock_schedule
mock_temporal_client.list_schedules.return_value = mock_list_schedules()
# Force reconcile by diverging task queue
mock_temporal_client.get_schedule_handle.return_value.describe.return_value.schedule.action.task_queue = 'different-queue'
await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata)
# Verify schedule was reconciled via delete + create
mock_temporal_client.get_schedule_handle.assert_called_once_with(
'cleanup-files-model-manager-worker-daily'
)
mock_temporal_client.get_schedule_handle.return_value.delete.assert_called_once()
mock_temporal_client.create_schedule.assert_called_once()
mock_logger.custom_info.assert_called_once()
assert 'reconciled successfully' in mock_logger.custom_info.call_args[0][0]
@pytest.mark.asyncio
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'RUNTIME': 'model-manager-worker',
},
)
async def test_create_cleanup_schedule_noop_when_schedule_is_up_to_date(
mock_temporal_client, mock_logger, metadata
):
"""Test no-op reconcile when existing schedule already matches current config."""
import model_manager.schedules.cleanup_schedule
reload(model_manager.schedules.cleanup_schedule)
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
mock_schedule = MagicMock()
mock_schedule.id = 'cleanup-files-model-manager-worker-daily'
async def mock_list_schedules():
yield mock_schedule
@@ -140,22 +195,22 @@ async def test_create_cleanup_schedule_skips_when_exists(
await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata)
# Verify schedule creation was NOT called
mock_temporal_client.get_schedule_handle.assert_called_once_with(
'cleanup-files-model-manager-worker-daily'
)
mock_temporal_client.get_schedule_handle.return_value.delete.assert_not_called()
mock_temporal_client.create_schedule.assert_not_called()
# Verify info log was called
mock_logger.custom_info.assert_called_once()
assert 'already configured' in mock_logger.custom_info.call_args[0][0]
assert 'no-op reconcile' in mock_logger.custom_info.call_args[0][0]
@pytest.mark.asyncio
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'CLEANUP_SCHEDULE_ID': 'test-cleanup-schedule',
'RUNTIME': 'model-manager-worker',
'CLEANUP_CRON': '0 2 * * *',
'CLEANUP_TIMEZONE': 'America/Sao_Paulo',
'CLEANUP_TASK_QUEUE': 'test-cleanup-queue',
'CLEANUP_EXECUTION_TIMEOUT_HOURS': '2',
},
)
@@ -185,9 +240,9 @@ async def test_create_cleanup_schedule_creates_with_custom_config(
schedule_id = call_args[0][0]
schedule_obj = call_args[0][1]
assert schedule_id == 'test-cleanup-schedule'
assert schedule_id == 'cleanup-files-model-manager-worker-daily'
assert schedule_obj.action.workflow == 'cleanup_files'
assert schedule_obj.action.task_queue == 'test-cleanup-queue'
assert schedule_obj.action.task_queue == 'cleanup_files-model-manager-worker-queue'
assert schedule_obj.action.execution_timeout == timedelta(hours=2)
assert schedule_obj.spec.cron_expressions == ['0 2 * * *']
assert schedule_obj.spec.time_zone_name == 'America/Sao_Paulo'
@@ -201,7 +256,7 @@ async def test_create_cleanup_schedule_creates_with_custom_config(
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'CLEANUP_SCHEDULE_ID': 'default-schedule',
'RUNTIME': 'model-manager-worker',
},
)
async def test_create_cleanup_schedule_uses_defaults(mock_temporal_client, mock_logger, metadata):
@@ -212,7 +267,6 @@ async def test_create_cleanup_schedule_uses_defaults(mock_temporal_client, mock_
for key in [
'CLEANUP_CRON',
'CLEANUP_TIMEZONE',
'CLEANUP_TASK_QUEUE',
'CLEANUP_EXECUTION_TIMEOUT_HOURS',
]:
os.environ.pop(key, None)
@@ -238,11 +292,17 @@ async def test_create_cleanup_schedule_uses_defaults(mock_temporal_client, mock_
assert schedule_obj.spec.cron_expressions == ['0 0 * * *'] # Default midnight
assert schedule_obj.spec.time_zone_name == 'UTC' # Default UTC
assert schedule_obj.action.task_queue == 'cleanup-queue' # Default queue
assert schedule_obj.action.task_queue == 'cleanup_files-model-manager-worker-queue'
assert schedule_obj.action.execution_timeout == timedelta(hours=1) # Default 1 hour
@pytest.mark.asyncio
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'RUNTIME': 'model-manager-worker',
},
)
async def test_create_cleanup_schedule_workflow_id_format(
mock_temporal_client, mock_logger, metadata
):
@@ -250,10 +310,7 @@ async def test_create_cleanup_schedule_workflow_id_format(
import model_manager.schedules.cleanup_schedule
reload(model_manager.schedules.cleanup_schedule)
from model_manager.schedules.cleanup_schedule import (
SCHEDULE_ID,
create_cleanup_schedule,
)
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
# Mock schedule does not exist (empty list)
async def mock_list_schedules():
@@ -268,11 +325,17 @@ async def test_create_cleanup_schedule_workflow_id_format(
call_args = mock_temporal_client.create_schedule.call_args
schedule_obj = call_args[0][1]
expected_workflow_id = f'cleanup-files-scheduled-{SCHEDULE_ID}'
expected_workflow_id = 'cleanup-files-scheduled-cleanup-files-model-manager-worker-daily'
assert schedule_obj.action.id == expected_workflow_id
@pytest.mark.asyncio
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'RUNTIME': 'model-manager-worker',
},
)
async def test_create_cleanup_schedule_empty_workflow_args(
mock_temporal_client, mock_logger, metadata
):
@@ -305,10 +368,9 @@ async def test_create_cleanup_schedule_empty_workflow_args(
@patch.dict(
'model_manager.schedules.cleanup_schedule.os.environ',
{
'CLEANUP_SCHEDULE_ID': 'custom-id',
'RUNTIME': 'model-manager-worker',
'CLEANUP_CRON': '30 3 * * 1',
'CLEANUP_TIMEZONE': 'Europe/London',
'CLEANUP_TASK_QUEUE': 'custom-queue',
'CLEANUP_EXECUTION_TIMEOUT_HOURS': '3',
},
)
@@ -320,15 +382,16 @@ def test_environment_variables_loaded_correctly():
from model_manager.schedules.cleanup_schedule import (
CLEANUP_CRON,
CLEANUP_EXECUTION_TIMEOUT_HOURS,
CLEANUP_TASK_QUEUE,
CLEANUP_TIMEZONE,
SCHEDULE_ID,
build_cleanup_schedule_id,
)
assert SCHEDULE_ID == 'custom-id'
assert (
build_cleanup_schedule_id('model-manager-worker')
== 'cleanup-files-model-manager-worker-daily'
)
assert CLEANUP_CRON == '30 3 * * 1'
assert CLEANUP_TIMEZONE == 'Europe/London'
assert CLEANUP_TASK_QUEUE == 'custom-queue'
assert CLEANUP_EXECUTION_TIMEOUT_HOURS == 3
@@ -338,10 +401,9 @@ def test_environment_variables_use_defaults_when_not_set():
# Remove all env vars
for key in [
'CLEANUP_SCHEDULE_ID',
'RUNTIME',
'CLEANUP_CRON',
'CLEANUP_TIMEZONE',
'CLEANUP_TASK_QUEUE',
'CLEANUP_EXECUTION_TIMEOUT_HOURS',
]:
os.environ.pop(key, None)
@@ -350,13 +412,35 @@ def test_environment_variables_use_defaults_when_not_set():
from model_manager.schedules.cleanup_schedule import (
CLEANUP_CRON,
CLEANUP_EXECUTION_TIMEOUT_HOURS,
CLEANUP_TASK_QUEUE,
CLEANUP_TIMEZONE,
SCHEDULE_ID,
build_cleanup_schedule_id,
)
assert SCHEDULE_ID == 'cleanup-files-daily'
assert build_cleanup_schedule_id(None) == 'cleanup-files-single-daily'
assert CLEANUP_CRON == '0 0 * * *'
assert CLEANUP_TIMEZONE == 'UTC'
assert CLEANUP_TASK_QUEUE == 'cleanup-queue'
assert CLEANUP_EXECUTION_TIMEOUT_HOURS == 1
@pytest.mark.asyncio
async def test_create_cleanup_schedule_uses_single_runtime_when_runtime_missing(
mock_temporal_client, mock_logger, metadata
):
"""Test create_cleanup_schedule uses single runtime fallback."""
import model_manager.schedules.cleanup_schedule
os.environ.pop('RUNTIME', None)
reload(model_manager.schedules.cleanup_schedule)
from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
async def mock_list_schedules():
return
yield
mock_temporal_client.list_schedules.return_value = mock_list_schedules()
await create_cleanup_schedule(mock_temporal_client, mock_logger, metadata)
call_args = mock_temporal_client.create_schedule.call_args
schedule_obj = call_args[0][1]
assert schedule_obj.action.task_queue == 'cleanup_files-single-queue'