vitor-aignosi
192834bde3
QTZPOC-20: Refactor model import workflow, introduce atomic model ID generation, and clarify MinIO expiry. Fuses bundle download, decryption, MLflow experiment creation, and artifact upload into a single activity to prevent cross-worker issues. Implements an atomic counter for MongoDB model IDs, replacing the max(id) + 1 scan to avoid collisions during concurrent imports. Adds comprehensive documentation for MinIO object expiry mechanisms and refines lifecycle rule application. Improves error handling with clearer responsibility classification and centralizes MLflow repository injection.
395 lines
13 KiB
Python
395 lines
13 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
from pytest import mark
|
|
|
|
from laborious.activities.activities import Activities
|
|
from laborious.activities.api import API
|
|
from laborious.activities.gates import Gates
|
|
from laborious.activities.mlflow import MLFlow
|
|
from laborious.activities.model_metrics import ModelMetrics
|
|
from laborious.activities.opc import OPC
|
|
from laborious.activities.storage import Storage
|
|
|
|
|
|
@patch('laborious.activities.activities.Storage.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
|
@patch('laborious.activities.activities.API.__init__')
|
|
@patch('laborious.activities.activities.MLFlowRepository')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test___init__(
|
|
mock_metrics_controller,
|
|
mock_minio_repository,
|
|
mock_mlflow_repository,
|
|
mock_api_init,
|
|
mock_model_metrics_init,
|
|
mock_gates_init,
|
|
mock_opc_init,
|
|
mock_mlflow_init,
|
|
mock_storage_init,
|
|
):
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10,
|
|
}
|
|
|
|
minio_config = {
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'default_bucket': 'test',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
}
|
|
|
|
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group',
|
|
}
|
|
|
|
pi_web_api_config = {
|
|
'base_url': 'https://test-pi-server.com',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
mlflow_config=mlflow_config,
|
|
minio_config=minio_config,
|
|
opc_config=opc_config,
|
|
pi_web_api_config=pi_web_api_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
assert isinstance(activities, Activities)
|
|
assert isinstance(activities, Storage)
|
|
assert isinstance(activities, MLFlow)
|
|
assert isinstance(activities, OPC)
|
|
assert isinstance(activities, Gates)
|
|
assert isinstance(activities, ModelMetrics)
|
|
assert isinstance(activities, API)
|
|
|
|
mock_storage_init.assert_called_once_with(
|
|
ANY,
|
|
host=postgres_config['host'],
|
|
port=postgres_config['port'],
|
|
user=postgres_config['user'],
|
|
password=postgres_config['password'],
|
|
dbname=postgres_config['dbname'],
|
|
min_connections=postgres_config['min_connections'],
|
|
max_connections=postgres_config['max_connections'],
|
|
retention_hours=minio_config['retention_hours'],
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_mlflow_init.assert_called_once_with(
|
|
ANY,
|
|
mlflow_repository=mock_mlflow_repository.return_value,
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_mlflow_repository.assert_called_once_with(
|
|
f'{mlflow_config["host"]}:{mlflow_config["port"]}',
|
|
mlflow_config['username'],
|
|
mlflow_config['password'],
|
|
logger,
|
|
notification_handler,
|
|
mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_opc_init.assert_called_once_with(
|
|
ANY,
|
|
opc_servers=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_gates_init.assert_called_once_with(
|
|
ANY,
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_model_metrics_init.assert_called_once_with(
|
|
ANY,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_api_init.assert_called_once_with(
|
|
ANY,
|
|
base_url=pi_web_api_config['base_url'],
|
|
auth_type=pi_web_api_config['auth_type'],
|
|
auth_token=pi_web_api_config['auth_token'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint=minio_config['endpoint_url'],
|
|
access_key=minio_config['access_key'],
|
|
secret_key=minio_config['secret_key'],
|
|
bucket=minio_config['default_bucket'],
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=mock_metrics_controller.return_value,
|
|
secure=minio_config['secure'],
|
|
)
|
|
|
|
|
|
BASE_CONFIGS = {
|
|
'postgres_config': {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10,
|
|
},
|
|
'minio_config': {
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'default_bucket': 'test',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
},
|
|
'mlflow_config': {'host': 'localhost', 'port': 5000, 'username': 'm', 'password': 'm'},
|
|
'opc_config': {},
|
|
'pi_web_api_config': {'base_url': 'https://pi', 'auth_type': 'bearer', 'auth_token': 't'},
|
|
}
|
|
|
|
IMPORT_CONFIG = {'bucket': 'sientia', 'prefix': 'imported_models/'}
|
|
IMPORT_STATUS_CONFIG = {
|
|
'host': 'bff-db',
|
|
'port': 5432,
|
|
'dbname': 'sientia-core-mlops-bff',
|
|
'user': 'u',
|
|
'password': 'p',
|
|
'schema': 'public',
|
|
}
|
|
|
|
|
|
@patch('laborious.activities.activities.ModelImport.__init__')
|
|
@patch('laborious.activities.activities.Storage.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
|
@patch('laborious.activities.activities.API.__init__')
|
|
@patch('laborious.activities.activities.MongoDBRepository')
|
|
@patch('laborious.activities.activities.MLFlowRepository')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test_mlflow_and_import_activities_receive_the_same_mlflow_repository_instance( # noqa: PLR0913
|
|
_mock_metrics_controller,
|
|
mock_minio_repository,
|
|
mock_mlflow_repository,
|
|
mock_mongo_repository,
|
|
_mock_api_init,
|
|
_mock_model_metrics_init,
|
|
_mock_gates_init,
|
|
_mock_opc_init,
|
|
mock_mlflow_init,
|
|
_mock_storage_init,
|
|
mock_model_import_init,
|
|
):
|
|
"""Setting the MLflow tracking URI is a process-wide side effect, so two repository instances
|
|
would be two chances to disagree about where MLflow is."""
|
|
Activities(
|
|
**BASE_CONFIGS,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
mongo_config={'connection_string': 'mongodb://mongo', 'database_name': 'sientia'},
|
|
import_config=IMPORT_CONFIG,
|
|
import_status_config=IMPORT_STATUS_CONFIG,
|
|
)
|
|
|
|
repository = mock_mlflow_repository.return_value
|
|
assert mock_mlflow_init.call_args.kwargs['mlflow_repository'] is repository
|
|
|
|
passed = mock_model_import_init.call_args.kwargs
|
|
assert passed['mlflow_repository'] is repository
|
|
assert passed['minio_repository'] is mock_minio_repository.return_value
|
|
assert passed['mongo_repository'] is mock_mongo_repository.return_value
|
|
assert passed['import_config'] is IMPORT_CONFIG
|
|
assert passed['status_config'] is IMPORT_STATUS_CONFIG
|
|
|
|
|
|
@patch('laborious.activities.activities.ModelImport.__init__')
|
|
@patch('laborious.activities.activities.Storage.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
|
@patch('laborious.activities.activities.API.__init__')
|
|
@patch('laborious.activities.activities.MongoDBRepository')
|
|
@patch('laborious.activities.activities.MLFlowRepository')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test_without_a_mongo_config_the_import_gets_no_repository( # noqa: PLR0913
|
|
_mock_metrics_controller,
|
|
_mock_minio_repository,
|
|
_mock_mlflow_repository,
|
|
mock_mongo_repository,
|
|
_mock_api_init,
|
|
_mock_model_metrics_init,
|
|
_mock_gates_init,
|
|
_mock_opc_init,
|
|
_mock_mlflow_init,
|
|
_mock_storage_init,
|
|
mock_model_import_init,
|
|
):
|
|
"""It refuses rather than guessing: the notification handler keeps its own client either way."""
|
|
Activities(
|
|
**BASE_CONFIGS,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
import_config=IMPORT_CONFIG,
|
|
import_status_config=IMPORT_STATUS_CONFIG,
|
|
)
|
|
|
|
assert mock_model_import_init.call_args.kwargs['mongo_repository'] is None
|
|
mock_mongo_repository.assert_not_called()
|
|
|
|
|
|
@patch('laborious.activities.activities.ModelImport.__init__')
|
|
@patch('laborious.activities.activities.Storage.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
@patch('laborious.activities.activities.ModelMetrics.__init__')
|
|
@patch('laborious.activities.activities.API.__init__')
|
|
@patch('laborious.activities.activities.MongoDBRepository')
|
|
@patch('laborious.activities.activities.MLFlowRepository')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test_exactly_one_mlflow_repository_is_built_even_with_import_active( # noqa: PLR0913
|
|
_mock_metrics_controller,
|
|
_mock_minio_repository,
|
|
mock_mlflow_repository,
|
|
_mock_mongo_repository,
|
|
_mock_api_init,
|
|
_mock_model_metrics_init,
|
|
_mock_gates_init,
|
|
_mock_opc_init,
|
|
_mock_mlflow_init,
|
|
_mock_storage_init,
|
|
_mock_model_import_init,
|
|
):
|
|
"""No second `MLFlowRepository` is ever created on any path: `MLFlow` and `ModelImport` only
|
|
receive the one instance built here, they never build their own."""
|
|
Activities(
|
|
**BASE_CONFIGS,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
mongo_config={'connection_string': 'mongodb://mongo', 'database_name': 'sientia'},
|
|
import_config=IMPORT_CONFIG,
|
|
import_status_config=IMPORT_STATUS_CONFIG,
|
|
)
|
|
|
|
mock_mlflow_repository.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.activities.Storage')
|
|
@patch('laborious.activities.activities.MLFlow')
|
|
@patch('laborious.activities.activities.OPC')
|
|
@patch('laborious.activities.activities.Gates')
|
|
@patch('laborious.activities.activities.ModelMetrics')
|
|
@patch('laborious.activities.activities.API')
|
|
@patch('laborious.activities.activities.MLFlowRepository')
|
|
@patch('laborious.activities.activities.MinioRepository')
|
|
async def test_shutdown(
|
|
_mock_minio_repository,
|
|
_mock_mlflow_repository,
|
|
mock_api_init,
|
|
mock_model_metrics_init,
|
|
mock_gates_init,
|
|
mock_opc_init,
|
|
mock_mlflow_init,
|
|
mock_storage_init,
|
|
):
|
|
mock_opc_init.aclose = AsyncMock()
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10,
|
|
}
|
|
|
|
minio_config = {
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'default_bucket': 'test',
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
}
|
|
|
|
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group',
|
|
}
|
|
|
|
pi_web_api_config = {
|
|
'base_url': 'https://test-pi-server.com',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'test_token',
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
mlflow_config=mlflow_config,
|
|
minio_config=minio_config,
|
|
opc_config=opc_config,
|
|
pi_web_api_config=pi_web_api_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
await activities.shutdown()
|
|
mock_opc_init.aclose.assert_called_once()
|
|
mock_storage_init.close.assert_called_once()
|
|
mock_mlflow_init.close.assert_called_once()
|
|
mock_gates_init.close.assert_called_once()
|
|
mock_model_metrics_init.close.assert_called_once()
|
|
mock_api_init.close.assert_called_once()
|