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

15
.env.example
View File

@@ -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

3
model_manager/activities/activities.py
View File

@@ -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__(

106
model_manager/schedules/cleanup_schedule.py
View File

@@ -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,
)

2
model_manager/utils/connectors_config.py
View File

@@ -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'),
}

23
model_manager/worker/prepare_worker.py
View File

@@ -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 <workflow>-<runtime>-queue or <workflow>-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:

36
model_manager/worker/worker.py
View File

@@ -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-<runtime>-queue: For ML model training workflows
- cleanup_files-<runtime>-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,
),
]

2
requirements.txt
View File

@@ -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
evidently==0.4.39

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')

14
values.yaml
View File

@@ -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