From c5cd382350dbd7749642f3820b7b2e93c5c4ae58 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 7 Apr 2026 12:57:29 -0300 Subject: [PATCH] feat: enhance configuration and scheduling for cleanup processes - 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. --- .env.example | 15 +- model_manager/activities/activities.py | 3 +- model_manager/schedules/cleanup_schedule.py | 106 +++++++++++--- model_manager/utils/connectors_config.py | 2 + model_manager/worker/prepare_worker.py | 23 ++- model_manager/worker/worker.py | 36 +++-- requirements.txt | 2 +- tests/activities/test_activities.py | 2 + tests/schedules/test_cleanup_schedule.py | 146 +++++++++++++++----- tests/utils/test_connectors_config.py | 4 + tests/worker/test_prepare_worker.py | 6 +- tests/worker/test_worker.py | 43 +++++- values.yaml | 14 ++ 13 files changed, 326 insertions(+), 76 deletions(-) diff --git a/.env.example b/.env.example index 64f08f9..c8e0e7e 100644 --- a/.env.example +++ b/.env.example @@ -17,9 +17,19 @@ PROJECT_NAME=sientia-model-manager TEMPORAL_HOST=temporal-frontend.temporal.svc.cluster.local:7233 TEMPORAL_NAMESPACE=model-manager -TRAIN_TASK_QUEUE=train_model-queue -CLEANUP_TASK_QUEUE=cleanup-queue TEMPORAL_USE_TLS=false +RUNTIME=basic + +STORE_BASE_URL=http://gitea-http.gitea.svc.cluster.local +STORE_OWNER=aignosi +STORE_REPO=suse-model-store +STORE_BRANCH=main +STORE_USERNAME= +STORE_PASSWORD= +STORE_CACHE_TTL_SECONDS=3600 +PYPI_SERVER=http://library-distribution-server.library.svc.cluster.local:5000 +PYPI_USERNAME= +PYPI_PASSWORD= MONGODB_USERNAME=mongo_user MONGODB_PASSWORD=mongo_db_password @@ -36,6 +46,7 @@ MINIO_MAX_RETRY_ATTEMPTS=3 MINIO_RETRY_MODE=adaptive MINIO_CONNECT_TIMEOUT=10 MINIO_READ_TIMEOUT=60 +MINIO_DEFAULT_BUCKET=model-training TIMEOUT_VALIDATE_PARAMS=30 TIMEOUT_TRAIN_MODEL=2700 diff --git a/model_manager/activities/activities.py b/model_manager/activities/activities.py index 18f6cd5..5cdb61f 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -59,7 +59,7 @@ class Activities(ExperimentTracking, Training, Cleanup): mlflow_config: MLFlow server configuration dictionary Required keys: host, port, username, password minio_config: MinIO storage configuration dictionary - Required keys: endpoint_url, access_key, secret_key, region, use_ssl + Required keys: endpoint_url, access_key, secret_key, region, use_ssl, default_bucket logger: Logger instance for observability and debugging notification_handler: Notification handler for alerts and monitoring @@ -108,6 +108,7 @@ class Activities(ExperimentTracking, Training, Cleanup): notification_handler=notification_handler, metrics_controller=metrics_controller, secure=minio_config['use_ssl'], + bucket=minio_config['default_bucket'], ) Training.__init__( diff --git a/model_manager/schedules/cleanup_schedule.py b/model_manager/schedules/cleanup_schedule.py index 4a54d1f..c41b4f3 100644 --- a/model_manager/schedules/cleanup_schedule.py +++ b/model_manager/schedules/cleanup_schedule.py @@ -2,6 +2,7 @@ import os from datetime import timedelta +from typing import Any from sientia_do.observability.logger import Logger as SientiaLogger from temporalio.client import ( @@ -11,14 +12,68 @@ from temporalio.client import ( ScheduleSpec, ) +from model_manager.worker.prepare_worker import build_queue_name + # 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')) +def build_cleanup_schedule_id(runtime: str | None) -> str: + """ + Build cleanup schedule ID using runtime-derived naming. + + Args: + runtime: Runtime suffix used by workers + + Return: + str: Cleanup schedule ID + """ + normalized_runtime = runtime.strip() if runtime else '' + return f'cleanup-files-{normalized_runtime or "single"}-daily' + + +async def _needs_schedule_reconcile( + schedule_handle: Any, + cleanup_task_queue: str, + logger: SientiaLogger, + metadata: dict[str, str | None], +) -> bool: + """ + Compare configured cleanup schedule against current expected values. + + Args: + schedule_handle: Temporal schedule handle for current schedule ID + cleanup_task_queue: Expected cleanup task queue + logger: Logger instance for errors + metadata: Metadata dictionary for logging context + + Return: + bool: True when schedule should be recreated to apply current config + """ + try: + schedule_description = await schedule_handle.describe() + schedule = getattr(schedule_description, 'schedule', None) + action = getattr(schedule, 'action', None) + spec = getattr(schedule, 'spec', None) + + current_task_queue = getattr(action, 'task_queue', None) + current_execution_timeout = getattr(action, 'execution_timeout', None) + current_cron = getattr(spec, 'cron_expressions', None) + current_timezone = getattr(spec, 'time_zone_name', None) + + return ( + current_task_queue != cleanup_task_queue + or current_execution_timeout != timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS) + or current_cron != [CLEANUP_CRON] + or current_timezone != CLEANUP_TIMEZONE + ) + except Exception as e: # noqa: BLE001 + logger.custom_error(f'Error describing cleanup schedule for reconcile: {e}', metadata) + return True + + async def schedule_exists( client: Client, schedule_id: str, logger: SientiaLogger, metadata: dict[str, str | None] ) -> bool: @@ -58,21 +113,32 @@ async def create_cleanup_schedule( logger: Logger instance for logging schedule operations metadata: Metadata dictionary for logging context """ - # 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 + runtime = (os.getenv('RUNTIME') or 'single').strip() + cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single') + schedule_id = build_cleanup_schedule_id(runtime) + created = False + updated = False + + if await schedule_exists(client, schedule_id, logger, metadata): + handle = client.get_schedule_handle(schedule_id) + if await _needs_schedule_reconcile(handle, cleanup_task_queue, logger, metadata): + await handle.delete() + updated = True + else: + logger.custom_info( + f"Schedule '{schedule_id}' is already up to date, no-op reconcile", + metadata, + ) + return await client.create_schedule( - SCHEDULE_ID, + 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, + id=f'cleanup-files-scheduled-{schedule_id}', + task_queue=cleanup_task_queue, execution_timeout=timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS), ), spec=ScheduleSpec( @@ -81,9 +147,17 @@ async def create_cleanup_schedule( ), ), ) + created = not updated - logger.custom_info( - f"Schedule '{SCHEDULE_ID}' created successfully. " - f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', - metadata, - ) + if updated: + logger.custom_info( + f"Schedule '{schedule_id}' reconciled successfully. " + f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', + metadata, + ) + elif created: + logger.custom_info( + f"Schedule '{schedule_id}' created successfully. " + f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', + metadata, + ) diff --git a/model_manager/utils/connectors_config.py b/model_manager/utils/connectors_config.py index 61853f0..f78b417 100644 --- a/model_manager/utils/connectors_config.py +++ b/model_manager/utils/connectors_config.py @@ -106,6 +106,7 @@ def build_minio_config() -> dict[str, Any]: MINIO_RETRY_MODE: Retry mode - standard, legacy, or adaptive (default: adaptive) MINIO_CONNECT_TIMEOUT: Connection timeout in seconds (default: 10) MINIO_READ_TIMEOUT: Read timeout in seconds (default: 60) + MINIO_DEFAULT_BUCKET: Default S3 bucket for MinioRepository (default: model-training) Returns: dict: MinIO configuration dictionary with all required parameters @@ -120,6 +121,7 @@ def build_minio_config() -> dict[str, Any]: 'retry_mode': getenv('MINIO_RETRY_MODE', 'adaptive'), 'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')), 'read_timeout': int(getenv('MINIO_READ_TIMEOUT', '60')), + 'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'model-training'), } diff --git a/model_manager/worker/prepare_worker.py b/model_manager/worker/prepare_worker.py index f522249..2855ce3 100644 --- a/model_manager/worker/prepare_worker.py +++ b/model_manager/worker/prepare_worker.py @@ -39,6 +39,23 @@ def camel_to_snake(text: str) -> str: return text.lower() +def build_queue_name(workflow_name: str, runtime: str | None = None) -> str: + """ + Build Temporal queue name from workflow name and runtime. + + Args: + - workflow_name: str, workflow class name in CamelCase format + - runtime: str | None, runtime suffix for environment-specific queues + + Return: + str: queue name in the format --queue or -queue + """ + snake_workflow_name = camel_to_snake(workflow_name) + if runtime: + return f'{snake_workflow_name}-{runtime}-queue' + return f'{snake_workflow_name}-queue' + + def prepare_worker( main_workflow: type, other_workflows: Sequence[type], @@ -62,11 +79,7 @@ def prepare_worker( Worker: fully configured Temporal worker instance ready to run """ main_workflow_name = main_workflow.__name__.upper() - queue_name = ( - f'{camel_to_snake(main_workflow.__name__)}-{runtime}-queue' - if runtime - else f'{camel_to_snake(main_workflow.__name__)}-queue' - ) + queue_name = build_queue_name(main_workflow.__name__, runtime) local_workflow_parameters: dict[str, int] = {} for parameter_name, default_value in parameters: diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index 74f1e22..e381ea8 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -5,8 +5,8 @@ It orchestrates Temporal workers, manages task queues, and handles the lifecycle model training and cleanup workflows. The worker supports two task queues: -- train_model-queue: For ML model training workflows -- cleanup-queue: For file cleanup workflows +- train_model--queue: For ML model training workflows +- cleanup_files--queue: For file cleanup workflows Key Features: - Automatic scaling with PollerBehaviorAutoscaling @@ -20,8 +20,7 @@ Environment Variables: - TEMPORAL_HOST: Temporal server address (default: localhost:7233) - TEMPORAL_NAMESPACE: Temporal namespace (default: model-manager) - TEMPORAL_USE_TLS: Enable TLS for Temporal connection (default: false) -- TRAIN_TASK_QUEUE: Task queue for training workflows (default: train_model-queue) -- CLEANUP_TASK_QUEUE: Task queue for cleanup workflows (default: cleanup-queue) +- RUNTIME: Runtime identifier used in queue naming (default: single) - POD_ID: Kubernetes pod identifier for metrics - HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090) - HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091) @@ -60,8 +59,20 @@ with workflow.unsafe.imports_passed_through(): POD_ID = os.getenv('POD_ID') RUNTIME = os.getenv('RUNTIME') SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) -TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue') -CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue') + + +def _get_runtime(runtime: str | None) -> str: + """ + Resolve runtime using fallback when missing. + + Args: + - runtime: str | None, runtime value from environment + + Return: + str: normalized runtime value + """ + normalized_runtime = runtime.strip() if runtime else '' + return normalized_runtime or 'single' async def main(): @@ -84,8 +95,7 @@ async def main(): SystemExit: On graceful shutdown or error conditions """ - if not RUNTIME: - raise ValueError('RUNTIME environment variable is required') + runtime = _get_runtime(RUNTIME) host = os.getenv('TEMPORAL_HOST', 'localhost:7233') use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true' @@ -93,7 +103,7 @@ async def main(): metadata = { 'pod_id': POD_ID, - 'runtime': RUNTIME, + 'runtime': runtime, } start_prometheus_server(logger, metadata) @@ -112,7 +122,7 @@ async def main(): metrics_controller = MetricsController(logger=logger) - logger.custom_info(f'Installing runtime {RUNTIME}', metadata) + logger.custom_info(f'Installing runtime {runtime}', metadata) plugin_store_parameters = build_plugin_store_config() plugin_store = PluginStore( @@ -131,7 +141,7 @@ async def main(): metrics_controller=metrics_controller, ) - await plugin_store.install_runtime(runtime_name=RUNTIME) + await plugin_store.install_runtime(runtime_name=runtime) activities = Activities( postgres_config=build_postgres_config(), @@ -181,7 +191,7 @@ async def main(): ], temporal_client=temporal_client, logger=logger, - runtime=RUNTIME, + runtime=runtime, ), prepare_worker( main_workflow=CleanupFiles, @@ -191,7 +201,7 @@ async def main(): ], temporal_client=temporal_client, logger=logger, - runtime=RUNTIME, + runtime=runtime, ), ] diff --git a/requirements.txt b/requirements.txt index 86c0310..cf6c0f7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,4 @@ botocore==1.40.55 /home/grezewave/Documents/projects/sientia/sientia-dataops-library prometheus-client==0.23.1 beautifulsoup4==4.12.3 -evidently \ No newline at end of file +evidently==0.4.39 \ No newline at end of file diff --git a/tests/activities/test_activities.py b/tests/activities/test_activities.py index 2e3b4b5..ace21b3 100644 --- a/tests/activities/test_activities.py +++ b/tests/activities/test_activities.py @@ -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() diff --git a/tests/schedules/test_cleanup_schedule.py b/tests/schedules/test_cleanup_schedule.py index 3f4277c..63e9ef1 100644 --- a/tests/schedules/test_cleanup_schedule.py +++ b/tests/schedules/test_cleanup_schedule.py @@ -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' diff --git a/tests/utils/test_connectors_config.py b/tests/utils/test_connectors_config.py index 680abd2..84a2b6a 100644 --- a/tests/utils/test_connectors_config.py +++ b/tests/utils/test_connectors_config.py @@ -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' diff --git a/tests/worker/test_prepare_worker.py b/tests/worker/test_prepare_worker.py index bc1b0a7..0e0616d 100644 --- a/tests/worker/test_prepare_worker.py +++ b/tests/worker/test_prepare_worker.py @@ -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) diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 028a351..3781a84 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -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') diff --git a/values.yaml b/values.yaml index 7c999f8..cdb2bb3 100644 --- a/values.yaml +++ b/values.yaml @@ -160,6 +160,8 @@ global: value: "10" - name: MINIO_READ_TIMEOUT value: "60" + - name: MINIO_DEFAULT_BUCKET + value: "model-training" - name: TIMEOUT_VALIDATE_PARAMS value: "30" @@ -195,6 +197,18 @@ global: - name: PYPI_SERVER value: "http://library-distribution-server.library.svc.cluster.local:5000" + - name: PYPI_USERNAME + valueFrom: + secretKeyRef: + name: sientia-plugin-store-credentials + key: pypi_username + optional: true + - name: PYPI_PASSWORD + valueFrom: + secretKeyRef: + name: sientia-plugin-store-credentials + key: pypi_password + optional: true # ----------------------------------------------------------------------------- # Runtimes configuration