vitor-aignosi
dac4f1718f
QTZPOC-20: Implement model import workflow for .sientia bundles Introduces the import_model workflow and its associated activities to enable importing .sientia bundles into MLflow and MongoDB. This includes: * New core workflow and activity logic for handling bundle download, decryption, validation, MLflow provisioning (experiments, runs, registered versions), and MongoDB model document creation. * Comprehensive E2E and unit test suites, including a dedicated test harness and a manual run script for development. * Updated .env.example and README.md with new configurations and usage instructions. * Ensures correct MLflow model source resolution for legacy loader compatibility and robust MongoDB document management. * Refactors existing utility scripts into a scripts/ directory.
345 lines
11 KiB
Python
345 lines
11 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.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test___init__(
|
|
mock_metrics_controller,
|
|
mock_minio_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_host=mlflow_config['host'],
|
|
mlflow_port=mlflow_config['port'],
|
|
mlflow_username=mlflow_config['username'],
|
|
mlflow_password=mlflow_config['password'],
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=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.MinioRepository')
|
|
@patch('laborious.activities.activities.MetricsController')
|
|
def test_the_import_activities_share_the_worker_s_mlflow_repository( # noqa: PLR0913
|
|
_mock_metrics_controller,
|
|
mock_minio_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."""
|
|
repository = MagicMock(name='model_monitoring_repository')
|
|
mock_mlflow_init.side_effect = lambda instance, **_: setattr(
|
|
instance, 'model_monitoring_repository', repository
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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.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_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()
|
|
|
|
|
|
@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.MinioRepository')
|
|
async def test_shutdown(
|
|
_mock_minio_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()
|