09ee92f100
- Added new ignore rule for Ruff to allow temporary paths in tests. - Introduced MyPy overrides for specific modules to ignore errors. - Refactored `Cleanup` and `ExperimentTracking` classes to remove async keywords from methods, improving consistency in method signatures. - Updated `Training` class methods to handle synchronous operations, enhancing performance and clarity. - Adjusted `requirements.txt` to remove unnecessary Git dependency, streamlining project setup.
77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
"""Unit tests for local worker factory."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
def test_prepare_worker_train_queue_uses_train_limits():
|
|
from model_manager.worker.prepare_worker import prepare_worker
|
|
from model_manager.workflows.train_model import TrainModel
|
|
|
|
fake_worker = MagicMock()
|
|
fake_client = MagicMock()
|
|
fake_logger = MagicMock()
|
|
|
|
with patch(
|
|
'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
|
|
) as worker_class:
|
|
with patch.dict(
|
|
'os.environ',
|
|
{
|
|
'TRAINMODEL_ACTIVITY_EXECUTOR_MAX_WORKERS': '3',
|
|
'TRAINMODEL_MAX_CONCURRENT_ACTIVITIES': '6',
|
|
'TRAINMODEL_MAX_CONCURRENT_WORKFLOW_TASKS': '10',
|
|
},
|
|
clear=False,
|
|
):
|
|
worker = prepare_worker(
|
|
main_workflow=TrainModel,
|
|
other_workflows=[],
|
|
activities=[],
|
|
temporal_client=fake_client,
|
|
logger=fake_logger,
|
|
)
|
|
|
|
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['max_concurrent_activities'] == 6
|
|
assert kwargs['max_concurrent_workflow_tasks'] == 10
|
|
assert kwargs['activity_executor']._max_workers == 3
|
|
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)
|
|
|
|
|
|
def test_prepare_worker_cleanup_queue_uses_cleanup_limits():
|
|
from model_manager.worker.prepare_worker import prepare_worker
|
|
from model_manager.workflows.cleanup_files import CleanupFiles
|
|
|
|
fake_worker = MagicMock()
|
|
fake_client = MagicMock()
|
|
fake_logger = MagicMock()
|
|
|
|
with patch(
|
|
'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
|
|
) as worker_class:
|
|
with patch.dict(
|
|
'os.environ',
|
|
{
|
|
'CLEANUPFILES_ACTIVITY_EXECUTOR_MAX_WORKERS': '5',
|
|
'CLEANUPFILES_MAX_CONCURRENT_ACTIVITIES': '7',
|
|
},
|
|
clear=False,
|
|
):
|
|
worker = prepare_worker(
|
|
main_workflow=CleanupFiles,
|
|
other_workflows=[],
|
|
activities=[],
|
|
temporal_client=fake_client,
|
|
logger=fake_logger,
|
|
)
|
|
|
|
assert worker is fake_worker
|
|
kwargs = worker_class.call_args.kwargs
|
|
assert kwargs['task_queue'] == 'cleanup_files-queue'
|
|
assert kwargs['max_concurrent_activities'] == 7
|
|
assert kwargs['activity_executor']._max_workers == 5
|
|
kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)
|