feat: enhance configuration and error handling in project setup
- 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.
This commit is contained in:
76
tests/worker/test_prepare_worker.py
Normal file
76
tests/worker/test_prepare_worker.py
Normal file
@@ -0,0 +1,76 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user