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

2
tests/activities/test_activities.py
View File

@@ -31,6 +31,7 @@ def _minio(endpoint_url: str):
'secret_key': 's',
'region': 'r',
'use_ssl': True,
'default_bucket': 'test-bucket',
}
@@ -64,6 +65,7 @@ def test_activities_strips_minio_endpoint_scheme(endpoint, expected_endpoint):
)
m_minio.assert_called_once()
assert m_minio.call_args.kwargs['endpoint'] == expected_endpoint
assert m_minio.call_args.kwargs['bucket'] == 'test-bucket'
m_mlflow.assert_called_once()

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'

4
tests/utils/test_connectors_config.py
View File

@@ -127,6 +127,7 @@ def test_build_minio_config_with_env_vars():
environ['MINIO_RETRY_MODE'] = 'standard'
environ['MINIO_CONNECT_TIMEOUT'] = '20'
environ['MINIO_READ_TIMEOUT'] = '120'
environ['MINIO_DEFAULT_BUCKET'] = 'my-bucket'
# Act
config = build_minio_config()
@@ -141,6 +142,7 @@ def test_build_minio_config_with_env_vars():
assert config['retry_mode'] == 'standard'
assert config['connect_timeout'] == 20
assert config['read_timeout'] == 120
assert config['default_bucket'] == 'my-bucket'
def test_build_plugin_store_config_cache_ttl_seconds():
@@ -162,6 +164,7 @@ def test_build_minio_config_with_defaults():
environ.pop('MINIO_RETRY_MODE', None)
environ.pop('MINIO_CONNECT_TIMEOUT', None)
environ.pop('MINIO_READ_TIMEOUT', None)
environ.pop('MINIO_DEFAULT_BUCKET', None)
# Act
config = build_minio_config()
@@ -176,3 +179,4 @@ def test_build_minio_config_with_defaults():
assert config['retry_mode'] == 'adaptive'
assert config['connect_timeout'] == 10
assert config['read_timeout'] == 60
assert config['default_bucket'] == 'streamlit-connectors'

6
tests/worker/test_prepare_worker.py
View File

@@ -29,12 +29,13 @@ def test_prepare_worker_train_queue_uses_train_limits():
activities=[],
temporal_client=fake_client,
logger=fake_logger,
runtime='model-manager-worker',
)
assert worker is fake_worker
worker_class.assert_called_once()
kwargs = worker_class.call_args.kwargs
assert kwargs['task_queue'] == 'train_model-queue'
assert kwargs['task_queue'] == 'train_model-model-manager-worker-queue'
assert kwargs['max_concurrent_activities'] == 6
assert kwargs['max_concurrent_workflow_tasks'] == 10
assert kwargs['activity_executor']._max_workers == 3
@@ -66,11 +67,12 @@ def test_prepare_worker_cleanup_queue_uses_cleanup_limits():
activities=[],
temporal_client=fake_client,
logger=fake_logger,
runtime='model-manager-worker',
)
assert worker is fake_worker
kwargs = worker_class.call_args.kwargs
assert kwargs['task_queue'] == 'cleanup_files-queue'
assert kwargs['task_queue'] == 'cleanup_files-model-manager-worker-queue'
assert kwargs['max_concurrent_activities'] == 7
assert kwargs['activity_executor']._max_workers == 5
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)

43
tests/worker/test_worker.py
View File

@@ -754,7 +754,7 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
@patch('model_manager.worker.worker.get_logger')
@patch('model_manager.worker.worker.start_prometheus_server')
@patch('model_manager.worker.worker.metrics')
async def test_main_missing_runtime_fails_fast(
async def test_main_missing_runtime_uses_single_fallback(
mock_metrics,
mock_start_prometheus,
mock_get_logger,
@@ -771,7 +771,7 @@ async def test_main_missing_runtime_fails_fast(
mock_prepare_worker,
mock_logger,
):
"""Test that main() fails fast when RUNTIME is missing."""
"""Test that main() uses single runtime fallback when RUNTIME is missing."""
from model_manager.worker.worker import main
mock_get_logger.return_value = mock_logger
@@ -792,11 +792,44 @@ async def test_main_missing_runtime_fails_fast(
mock_activities.shutdown = Mock()
mock_activities_class.return_value = mock_activities
with pytest.raises(ValueError, match='RUNTIME environment variable is required'):
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(return_value={'plugins': [], 'interceptors': []})
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
mock_prepare_worker.return_value = mock_worker_instance
mock_plugin_store_instance = AsyncMock()
mock_plugin_store_instance.install_runtime = AsyncMock(
return_value={'runtime': 'single', 'installed': []},
)
mock_plugin_store_class.return_value = mock_plugin_store_instance
mock_build_plugin_store_config.return_value = {
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
'owner': 'sientia',
'repo': 'model-library-store',
'branch': 'main',
'username': 'gitea-user',
'password': 'gitea-password',
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
'cache_ttl_seconds': None,
}
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
with pytest.raises(SystemExit):
await main()
mock_prepare_worker.assert_not_called()
mock_start_prometheus.assert_not_called()
assert mock_prepare_worker.call_count == 2
assert mock_prepare_worker.call_args_list[0].kwargs['runtime'] == 'single'
assert mock_prepare_worker.call_args_list[1].kwargs['runtime'] == 'single'
@patch('model_manager.worker.worker.asyncio.run')