diff --git a/laborious/activities/model_import.py b/laborious/activities/model_import.py index 1a6f547..166962c 100644 --- a/laborious/activities/model_import.py +++ b/laborious/activities/model_import.py @@ -40,7 +40,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.repository.mongodb_repository import MongoDBRepository from sientia_do.temporal.constants import now from sqlalchemy import create_engine, text - from sqlalchemy.engine import URL + from sqlalchemy.engine import URL, Engine from sqlalchemy.exc import DataError, IntegrityError, ProgrammingError from sqlalchemy.sql import Executable from temporalio.common import RetryPolicy @@ -321,25 +321,37 @@ class ModelImport(SientiaMonitoring): self.import_minio_repository = minio_repository self.import_mlflow_repository = mlflow_repository self.mongo_repository = mongo_repository - self.status_engine = create_engine( - URL.create( - 'postgresql', - username=status_config['user'], - password=status_config['password'], - host=status_config['host'], - port=status_config['port'], - database=status_config['dbname'], - ), - pool_size=2, - max_overflow=3, - pool_pre_ping=True, - ) + self._status_engine: Engine | None = None + + @property + def status_engine(self) -> Engine: + """The BFF status database's engine, created on first access and cached from then on. + + A worker pod where the import feature is never exercised never opens a connection pool + against that Postgres: nothing in this module reaches `status_engine` unless an import + activity actually runs. + """ + if self._status_engine is None: + self._status_engine = create_engine( + URL.create( + 'postgresql', + username=self.status_config['user'], + password=self.status_config['password'], + host=self.status_config['host'], + port=self.status_config['port'], + database=self.status_config['dbname'], + ), + pool_size=2, + max_overflow=3, + pool_pre_ping=True, + ) + return self._status_engine def close(self) -> None: - """Dispose the status engine. The import holds no other connection.""" - engine = getattr(self, 'status_engine', None) - if engine is not None: - engine.dispose() + """Dispose the status engine, when one was ever created. The import holds no other + connection.""" + if self._status_engine is not None: + self._status_engine.dispose() # ------------------------------------------------------------------ helpers @@ -1128,7 +1140,7 @@ class ModelImport(SientiaMonitoring): upsert=True, return_document=ReturnDocument.AFTER, ) - return str(doc['seq']), 0 + return str(doc['seq']), 0 # type: ignore[index] # ----------------------------------------------------------- cleanup and retention @@ -1217,7 +1229,12 @@ class ModelImport(SientiaMonitoring): # The read still runs every time (it is how drift — someone editing the rule outside # this code — gets detected); only the write is skipped when nothing has drifted. self.info(f'import retention rule already correct on {bucket} ({days} days)', metadata) - return {'bucket': bucket, 'prefix': prefix, 'days': days, 'rules': len(existing.rules)} + return { + 'bucket': bucket, + 'prefix': prefix, + 'days': days, + 'rules': len(existing.rules), # type: ignore[union-attr] + } rules = [rule for rule in (existing.rules if existing else []) if rule.rule_id != rule_id] rules.append( diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 04ed18e..2c14cf5 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -106,13 +106,21 @@ class OPC(SientiaMonitoring): 4. Logs connection status for operational visibility Raises: - Exception: If OPC repository initialization fails or connection - establishment encounters critical errors + ValueError: A server whose `cert_path` is set but whose `server_uri` + (or private key) is missing fails `OpcRepository.set_security()` + with no try/except between here and `main()`. This is a + configuration error, not a connectivity one, and is intentionally + fatal to the whole worker process — fix the deployment config + rather than catch it here. Note: - Connection failures are logged and reported but do not prevent - the initialization of other OPC servers. Each server is handled - independently to ensure maximum availability. + Connection failures (server unreachable, timeout, etc.) are logged + and reported but do not prevent the initialization of other OPC + servers — each server's `connect()` call catches those on its own. + This isolation does NOT extend to configuration errors raised before + a connection is attempted (see `Raises` above): a single + misconfigured server intentionally stops every other server, and the + whole worker, from starting. """ self.logger.info('Initializing OPC servers...') for opc_id, server in self.opc_servers.items(): diff --git a/laborious/utils/bundle/reader.py b/laborious/utils/bundle/reader.py index aea9646..5d223f4 100644 --- a/laborious/utils/bundle/reader.py +++ b/laborious/utils/bundle/reader.py @@ -35,6 +35,9 @@ import nacl.exceptions from laborious.utils.bundle import format as fmt from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep +# Gate 5 — chunk size for the byte-counted extraction copy loop. +_EXTRACT_CHUNK_BYTES = 1024 * 1024 + # Gate 6 — the layout `bundle.build_zip` writes and `model_repository.py` reads back. METADATA_FILE_NAME = 'metadata.json' ARTIFACTS_DIR_NAME = 'artifacts' @@ -389,8 +392,23 @@ class BundleReader: resolved.parent.mkdir(parents=True, exist_ok=True) # Written as a plain file: the archive's stored permissions are never restored, so no # member can arrive executable. + # + # Copied in chunks, counting bytes actually written, instead of `shutil.copyfileobj`: + # `entry.file_size` is central-directory metadata gate 4 already checked in aggregate, but a + # DEFLATE stream can be crafted to decompress to far more bytes than its own header + # declares. Aborting as soon as the running count exceeds `file_size` catches that member + # before the excess ever reaches disk, rather than trusting the declared size alone. + written = 0 with archive.open(entry) as source, resolved.open('wb') as sink: - shutil.copyfileobj(source, sink) + while chunk := source.read(_EXTRACT_CHUNK_BYTES): + written += len(chunk) + if written > entry.file_size: + raise BundleGateError( + f'archive member {entry.filename!r} decompresses to more than its ' + f'declared {entry.file_size} bytes', + gate=5, + ) + sink.write(chunk) # ------------------------------------------------------------------ gate 6 diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 1c32baf..118b69e 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -139,9 +139,9 @@ class MLFlowRepository(SientiaMonitoring): """ run_info = mlflow.get_run(run_id) if prediction: - model_uri = run_info.info.artifact_uri + '/prediction_model' + model_uri = run_info.info.artifact_uri + '/prediction_model' # type: ignore[operator] else: - model_uri = run_info.info.artifact_uri + '/data_model' + model_uri = run_info.info.artifact_uri + '/data_model' # type: ignore[operator] return model_uri def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str: @@ -290,7 +290,11 @@ class MLFlowRepository(SientiaMonitoring): """ async def dowload_artifacts( - self, model_name: str, metadata: dict[str, Any], artifact_path: str = 'data_model' + self, + model_name: str, + metadata: dict[str, Any], + artifact_path: str = 'data_model', + run_id: str | None = None, ) -> str: """ Download artifacts from the latest production run of a model. @@ -299,11 +303,14 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): Registered model name. metadata (dict[str, Any]): Metadata used for structured logging. artifact_path (str): Relative path to artifacts within the run. + run_id (str | None): Run already resolved by the caller. When `None`, + the Production run is resolved here. Returns: str: Local filesystem path where artifacts are saved. """ - run_id = self.get_model_run_id(model_name=model_name, stage='Production') + if run_id is None: + run_id = self.get_model_run_id(model_name=model_name, stage='Production') output_dir = f'{ARTIFACTS_PATH}/{model_name}' full_path = path.join(output_dir, artifact_path) @@ -472,6 +479,7 @@ class MLFlowRepository(SientiaMonitoring): artifact_name: str, flavor: str, metadata: dict[str, Any] | None = None, + run_id: str | None = None, ) -> str | None: """ Download only an artifact's `code` directory and put it on `sys.path`. @@ -489,6 +497,8 @@ class MLFlowRepository(SientiaMonitoring): flavor (str): Flavor doing the load, for the log lines only. metadata (dict[str, Any] | None): Metadata used for structured logging. + run_id (str | None): Run already resolved by the caller. When `None`, + the Production run is resolved here. Returns: str | None: Local directory the `code` directory was downloaded @@ -500,7 +510,8 @@ class MLFlowRepository(SientiaMonitoring): Exception: Any error raised while resolving the Production run, listing the artifact or downloading `code`. """ - run_id = self.get_model_run_id(model_name=model_name, stage='Production') + if run_id is None: + run_id = self.get_model_run_id(model_name=model_name, stage='Production') remote_code_path = f'{artifact_name}/code' if not self.check_artifact_exists(run_id, remote_code_path, metadata): @@ -623,7 +634,11 @@ class MLFlowRepository(SientiaMonitoring): return model, local_dir async def load_predict_model( - self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' + self, + model_name: str, + metadata: dict[str, Any], + flavor: str = 'sklearn', + run_id: str | None = None, ) -> tuple[Any, str | None]: """ Load a predictive model from the MLflow Model Registry. @@ -632,6 +647,10 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): The name of the model to download from the registry. metadata (dict[str, Any]): Metadata used for structured logging. flavor (str): Model flavor ('pyfunc', 'sklearn', 'joblib') + run_id (str | None): Run already resolved by the caller. When `None`, + the Production run is resolved where needed (only the `joblib` + branch downloads by run; the `pyfunc`/`sklearn` branches load + from the symbolic `models://production` URI regardless). Returns: tuple[Any, str | None]: The loaded predictive model and the local @@ -667,6 +686,7 @@ class MLFlowRepository(SientiaMonitoring): model_name=model_name, artifact_name='prediction_model', metadata=metadata, + run_id=run_id, ) elif flavor == 'pyfunc': model, local_dir = mlflow.pyfunc.load_model(model_uri), None @@ -681,7 +701,11 @@ class MLFlowRepository(SientiaMonitoring): return model, local_dir async def load_transform_model( - self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' + self, + model_name: str, + metadata: dict[str, Any], + flavor: str = 'sklearn', + run_id: str | None = None, ) -> tuple[Any, str | None]: """ Load the latest Production version of a transformation model. @@ -693,6 +717,8 @@ class MLFlowRepository(SientiaMonitoring): model_name (str): The name of the model to download. metadata (dict[str, Any]): Metadata used for structured logging. flavor (str): Model flavor ('sklearn', 'pyfunc', 'joblib') + run_id (str | None): Run already resolved by the caller. When `None`, + the Production run is resolved here. Returns: tuple[Any, str | None]: The loaded model object, depending on the @@ -719,7 +745,9 @@ class MLFlowRepository(SientiaMonitoring): if flavor not in ('sklearn', 'pyfunc', 'joblib'): raise ValueError(INVALID_FLAVOR_MESSAGE) - latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production') + latest_production_id = run_id or self.get_model_run_id( + model_name=model_name, stage='Production' + ) model_uri = ( None @@ -802,8 +830,10 @@ class MLFlowRepository(SientiaMonitoring): target = 'prediction_model' if model_type == 'predict' else 'data_model' + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + code_artifact_dir = await self.download_artifact_code_dir( - model_name, target, flavor, metadata + model_name, target, flavor, metadata, run_id=run_id ) if load_wrapper: @@ -811,7 +841,9 @@ class MLFlowRepository(SientiaMonitoring): self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}') try: - artifact_path = await self.dowload_artifacts(model_name, metadata, target) + artifact_path = await self.dowload_artifacts( + model_name, metadata, target, run_id=run_id + ) self.info( f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' @@ -832,9 +864,13 @@ class MLFlowRepository(SientiaMonitoring): return model, code_artifact_dir or artifact_path if model_type == 'predict': - model, local_dir = await self.load_predict_model(model_name, metadata, flavor) + model, local_dir = await self.load_predict_model( + model_name, metadata, flavor, run_id=run_id + ) else: - model, local_dir = await self.load_transform_model(model_name, metadata, flavor) + model, local_dir = await self.load_transform_model( + model_name, metadata, flavor, run_id=run_id + ) return model, code_artifact_dir or local_dir diff --git a/laborious/workflows/import_model.py b/laborious/workflows/import_model.py index 3a81df7..0d8ea0d 100644 --- a/laborious/workflows/import_model.py +++ b/laborious/workflows/import_model.py @@ -26,6 +26,7 @@ with workflow.unsafe.imports_passed_through(): from typing import Any from sientia_do.temporal.policies import retry_policy + from temporalio.exceptions import is_cancelled_exception from laborious.activities.activities import Activities from laborious.activities.model_import import ( @@ -309,6 +310,8 @@ class ImportModel: start_to_close_timeout=STATUS_WRITE_TIMEOUT, ) except BaseException as error: # NOSONAR - a lost hint must not fail a working import + if is_cancelled_exception(error): + raise await self._report_status_write_failure(base, error) async def _record_terminal( @@ -344,6 +347,8 @@ class ImportModel: except ( BaseException ) as error: # NOSONAR - the outcome stands even when it cannot be written + if is_cancelled_exception(error): + raise await self._report_status_write_failure( base, error, @@ -377,8 +382,9 @@ class ImportModel: retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY, start_to_close_timeout=STATUS_WRITE_TIMEOUT, ) - except BaseException: # noqa: S110 # NOSONAR - the report of a report is where this stops - pass + except BaseException as error: # NOSONAR - the report of a report is where this stops + if is_cancelled_exception(error): + raise async def _report_status_write_failure( self, @@ -394,10 +400,12 @@ class ImportModel: retry_policy=retry_policy, start_to_close_timeout=STATUS_WRITE_TIMEOUT, ) - except BaseException: # noqa: S110 # NOSONAR - the report of a report is where this stops - # Deliberately silent: the channel this was trying to reach is the thing that logs, so - # there is nowhere left to say it, and an import that worked must not fail over it. - pass + except BaseException as error: # NOSONAR - the report of a report is where this stops + if is_cancelled_exception(error): + raise + # Deliberately silent otherwise: the channel this was trying to reach is the thing that + # logs, so there is nowhere left to say it, and an import that worked must not fail + # over it. # ----------------------------------------------------------- side-effect steps @@ -418,12 +426,13 @@ class ImportModel: retry_policy=retry_policy, start_to_close_timeout=RETENTION_TIMEOUT, ) - except BaseException: # noqa: S110 # NOSONAR - retention is not a precondition of an import - # Dropped, not reported: the activity's own failure is already in the event history, - # and a bucket policy the importer cannot set is an operational problem rather than a - # bad import. Failing a valid model import over a lifecycle API would be the wrong - # trade. - pass + except BaseException as error: # NOSONAR - retention is not a precondition of an import + if is_cancelled_exception(error): + raise + # Dropped, not reported otherwise: the activity's own failure is already in the event + # history, and a bucket policy the importer cannot set is an operational problem rather + # than a bad import. Failing a valid model import over a lifecycle API would be the + # wrong trade. @staticmethod async def _cleanup(base: dict[str, Any]) -> dict[str, Any]: @@ -435,5 +444,7 @@ class ImportModel: retry_policy=retry_policy, start_to_close_timeout=CLEANUP_TIMEOUT, ) - except BaseException: # NOSONAR - cleanup never replaces the failure that brought us here + except BaseException as error: # NOSONAR - cleanup never replaces the failure that brought us here + if is_cancelled_exception(error): + raise return {'cleaned': False} diff --git a/tests/laborious/activities/test_model_import_bundle.py b/tests/laborious/activities/test_model_import_bundle.py index 8a41ddb..e4940b1 100644 --- a/tests/laborious/activities/test_model_import_bundle.py +++ b/tests/laborious/activities/test_model_import_bundle.py @@ -72,24 +72,25 @@ def minio_client() -> MagicMock: def import_activity(tmp_path: Path, minio_client: MagicMock) -> ModelImport: repository = MagicMock() repository.minio_client = minio_client - with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()): - activity = ModelImport( - status_config={ - 'host': 'h', - 'port': 5432, - 'dbname': 'sientia-core-mlops-bff', - 'user': 'u', - 'password': 'p', - 'schema': 'public', - }, - import_config=import_config(tmp_path), - minio_repository=repository, - mlflow_repository=MagicMock(), - mongo_repository=AsyncMock(), - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), - ) + activity = ModelImport( + status_config={ + 'host': 'h', + 'port': 5432, + 'dbname': 'sientia-core-mlops-bff', + 'user': 'u', + 'password': 'p', + 'schema': 'public', + }, + import_config=import_config(tmp_path), + minio_repository=repository, + mlflow_repository=MagicMock(), + mongo_repository=AsyncMock(), + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + # `status_engine` is created lazily on first access; pre-seed it with the recording double. + activity._status_engine = FakeStatusEngine() activity.emit_metric = AsyncMock() activity.send_notification_async = AsyncMock() return activity diff --git a/tests/laborious/activities/test_model_import_provisioning.py b/tests/laborious/activities/test_model_import_provisioning.py index e025898..3054ada 100644 --- a/tests/laborious/activities/test_model_import_provisioning.py +++ b/tests/laborious/activities/test_model_import_provisioning.py @@ -100,36 +100,37 @@ def import_activity(tmp_path: Path, mlflow_client: MagicMock, mongo: AsyncMock) minio_repository = MagicMock() mlflow_repository = MagicMock() mlflow_repository.client = mlflow_client - with patch('laborious.activities.model_import.create_engine', return_value=FakeStatusEngine()): - activity = ModelImport( - status_config={ - 'host': 'h', - 'port': 5432, - 'dbname': 'sientia-core-mlops-bff', - 'user': 'u', - 'password': 'p', - 'schema': 'public', - }, - import_config={ - 'password_key': None, - 'bucket': 'sientia', - 'prefix': 'imported_models/', - 'work_dir': str(tmp_path), - 'max_object_bytes': 1024, - 'max_entries': 10, - 'max_uncompressed_bytes': 4096, - 'max_compression_ratio': 200.0, - 'models_collection': 'models', - 'counters_collection': 'model_id_counters', - 'retention_days': 7, - }, - minio_repository=minio_repository, - mlflow_repository=mlflow_repository, - mongo_repository=mongo, - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), - ) + activity = ModelImport( + status_config={ + 'host': 'h', + 'port': 5432, + 'dbname': 'sientia-core-mlops-bff', + 'user': 'u', + 'password': 'p', + 'schema': 'public', + }, + import_config={ + 'password_key': None, + 'bucket': 'sientia', + 'prefix': 'imported_models/', + 'work_dir': str(tmp_path), + 'max_object_bytes': 1024, + 'max_entries': 10, + 'max_uncompressed_bytes': 4096, + 'max_compression_ratio': 200.0, + 'models_collection': 'models', + 'counters_collection': 'model_id_counters', + 'retention_days': 7, + }, + minio_repository=minio_repository, + mlflow_repository=mlflow_repository, + mongo_repository=mongo, + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + # `status_engine` is created lazily on first access; pre-seed it with the recording double. + activity._status_engine = FakeStatusEngine() activity.emit_metric = AsyncMock() activity.send_notification_async = AsyncMock() return activity diff --git a/tests/laborious/activities/test_model_import_status.py b/tests/laborious/activities/test_model_import_status.py index 893d945..684c715 100644 --- a/tests/laborious/activities/test_model_import_status.py +++ b/tests/laborious/activities/test_model_import_status.py @@ -164,17 +164,19 @@ def test_build_update_with_no_casts_or_extra_where_renders_a_bare_update(): def build_activity(engine: FakeStatusEngine) -> ModelImport: """A `ModelImport` whose only connection is the recording engine double.""" - with patch('laborious.activities.model_import.create_engine', return_value=engine): - activity = ModelImport( - status_config=STATUS_CONFIG, - import_config=IMPORT_CONFIG, - minio_repository=MagicMock(), - mlflow_repository=MagicMock(), - mongo_repository=AsyncMock(), - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), - ) + activity = ModelImport( + status_config=STATUS_CONFIG, + import_config=IMPORT_CONFIG, + minio_repository=MagicMock(), + mlflow_repository=MagicMock(), + mongo_repository=AsyncMock(), + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + # `status_engine` is created lazily on first access; setting the private attribute directly + # pre-seeds it with the recording double instead of exercising the real `create_engine`. + activity._status_engine = engine activity.emit_metric = AsyncMock() activity.send_notification_async = AsyncMock() return activity @@ -983,8 +985,62 @@ async def test_the_engine_is_disposed_on_shutdown(import_activity: ModelImport, @mark.asyncio async def test_close_is_a_no_op_without_an_engine(import_activity: ModelImport, engine): """A partially-constructed instance (or a second close) must not blow up on a missing engine.""" - del import_activity.status_engine + import_activity._status_engine = None import_activity.close() assert engine.disposed is False + + +# ------------------------------------------------------------- the lazy status engine + + +def test_no_engine_is_created_when_status_engine_is_never_accessed(): + """A worker pod that never runs an import activity never opens a pool against the BFF Postgres.""" + with patch('laborious.activities.model_import.create_engine') as create_engine_mock: + activity = ModelImport( + status_config=STATUS_CONFIG, + import_config=IMPORT_CONFIG, + minio_repository=MagicMock(), + mlflow_repository=MagicMock(), + mongo_repository=AsyncMock(), + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + + create_engine_mock.assert_not_called() + assert activity._status_engine is None + + +def test_first_access_creates_the_engine_and_later_accesses_reuse_it(): + with patch( + 'laborious.activities.model_import.create_engine', return_value=FakeStatusEngine() + ) as create_engine_mock: + activity = ModelImport( + status_config=STATUS_CONFIG, + import_config=IMPORT_CONFIG, + minio_repository=MagicMock(), + mlflow_repository=MagicMock(), + mongo_repository=AsyncMock(), + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + + first = activity.status_engine + second = activity.status_engine + + create_engine_mock.assert_called_once() + assert first is second + + +@mark.asyncio +async def test_close_never_accessed_the_engine_does_not_create_one(import_activity: ModelImport): + """`close()` must not force the lazy engine into existence just to dispose it.""" + import_activity._status_engine = None + + with patch('laborious.activities.model_import.create_engine') as create_engine_mock: + import_activity.close() + + create_engine_mock.assert_not_called() diff --git a/tests/laborious/utils/bundle/test_reader.py b/tests/laborious/utils/bundle/test_reader.py index e912ad5..da22105 100644 --- a/tests/laborious/utils/bundle/test_reader.py +++ b/tests/laborious/utils/bundle/test_reader.py @@ -437,7 +437,7 @@ def test_gate_5_wraps_an_os_error_during_unpacking(work_dir, reader, monkeypatch def _raise_os_error(*_args, **_kwargs): raise OSError('disk full') - monkeypatch.setattr('laborious.utils.bundle.reader.shutil.copyfileobj', _raise_os_error) + monkeypatch.setattr(zipfile.ZipExtFile, 'read', _raise_os_error) with pytest.raises(BundleGateError, match='could not be unpacked') as error: _open(reader, work_dir, payload) @@ -463,6 +463,52 @@ def test_extract_member_rechecks_containment_on_the_resolved_path(work_dir, read assert error.value.gate == 5 +def test_extract_member_rejects_a_stream_larger_than_its_declared_size(work_dir, reader, monkeypatch): + """A decompressed stream longer than the entry's declared `file_size` is rejected mid-write, + instead of trusting the central-directory metadata gate 4 already checked in aggregate. + + A single crafted DEFLATE chunk can make `ZipExtFile.read()` return far more bytes than + `file_size` declares before its own EOF/CRC bookkeeping catches up, so the stand-in here + forces exactly that response shape rather than fighting zipfile's own CRC validation.""" + root = work_dir / 'root' + root.mkdir() + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, 'w') as archive: + archive.writestr('payload.bin', b'x' * 10) + + call_count = {'n': 0} + + def _oversized_read(self, *_args, **_kwargs): + call_count['n'] += 1 + return b'x' * 1000 if call_count['n'] == 1 else b'' + + monkeypatch.setattr(zipfile.ZipExtFile, 'read', _oversized_read) + + with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive: + entry = archive.infolist()[0] + with pytest.raises(BundleGateError, match='more than its declared') as error: + reader._extract_member(archive, entry, root, root.resolve()) + + assert error.value.gate == 5 + + +def test_extract_member_writes_a_well_formed_member_in_full(work_dir, reader): + """A member whose decompressed size matches its declared `file_size` is extracted byte for + byte, exactly as before the byte-counting check was added.""" + root = work_dir / 'root' + root.mkdir() + payload = b'y' * 100 + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, 'w') as archive: + archive.writestr('payload.bin', payload) + + with zipfile.ZipFile(io.BytesIO(buffer.getvalue())) as archive: + entry = archive.infolist()[0] + reader._extract_member(archive, entry, root, root.resolve()) + + assert (root / 'payload.bin').read_bytes() == payload + + # ------------------------------------------------------------------- gate 6 diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index a6ea506..883da62 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -557,12 +557,15 @@ async def test_download_model_invalid_model_type(mlflow_repository): 'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')] ) async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type, tmp_path): + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) result = await mlflow_repository.download_model('test_model', {}, model_type[0], 'pyfunc', True) - mlflow_repository.dowload_artifacts.assert_called_once_with('test_model', {}, model_type[1]) + mlflow_repository.dowload_artifacts.assert_called_once_with( + 'test_model', {}, model_type[1], run_id='run-123' + ) mlflow.pyfunc.load_model.assert_called_once_with( mlflow_repository.dowload_artifacts.return_value @@ -576,13 +579,16 @@ async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type @pytest.mark.asyncio async def test_download_model_predict(mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.load_predict_model = AsyncMock(return_value=('model', 'local_dir')) mlflow_repository.load_transform_model = AsyncMock() result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', False) - mlflow_repository.load_predict_model.assert_called_once_with('test_model', {}, 'pyfunc') + mlflow_repository.load_predict_model.assert_called_once_with( + 'test_model', {}, 'pyfunc', run_id='run-123' + ) mlflow_repository.load_transform_model.assert_not_called() assert result == ('model', 'local_dir') @@ -590,6 +596,7 @@ async def test_download_model_predict(mlflow_repository): @pytest.mark.asyncio async def test_download_model_transform(mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.load_predict_model = AsyncMock() mlflow_repository.load_transform_model = AsyncMock(return_value=('model', 'local_dir')) @@ -597,7 +604,9 @@ async def test_download_model_transform(mlflow_repository): result = await mlflow_repository.download_model('test_model', {}, 'transform', 'pyfunc', False) mlflow_repository.load_predict_model.assert_not_called() - mlflow_repository.load_transform_model.assert_called_once_with('test_model', {}, 'pyfunc') + mlflow_repository.load_transform_model.assert_called_once_with( + 'test_model', {}, 'pyfunc', run_id='run-123' + ) assert result == ('model', 'local_dir') @@ -616,10 +625,11 @@ async def test_download_model_prepares_the_code_dir_first( calls.append('code') return 'code-artifact-dir' - async def record_load(*_a): + async def record_load(*_a, **_k): calls.append('load') return 'model', None + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(side_effect=record_code_dir) loader = AsyncMock(side_effect=record_load) mlflow_repository.load_predict_model = loader @@ -628,7 +638,7 @@ async def test_download_model_prepares_the_code_dir_first( result = await mlflow_repository.download_model('test_model', {}, model_type, 'sklearn', False) mlflow_repository.download_artifact_code_dir.assert_called_once_with( - 'test_model', artifact_name, 'sklearn', {} + 'test_model', artifact_name, 'sklearn', {}, run_id='run-123' ) assert calls == ['code', 'load'] assert result == ('model', 'code-artifact-dir') @@ -637,6 +647,7 @@ async def test_download_model_prepares_the_code_dir_first( @pytest.mark.asyncio async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlflow_repository): """With a `code` directory downloaded, the wrapper branch returns it over its own download.""" + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value='code-artifact-dir') mlflow_repository.dowload_artifacts = AsyncMock(return_value='wrapper-dir') @@ -647,6 +658,42 @@ async def test_download_model_wrapper_returns_the_code_artifact_dir(mlflow, mlfl assert artifact_dir == 'code-artifact-dir' +@pytest.mark.asyncio +async def test_download_model_resolves_run_id_once_non_wrapper(mlflow_repository): + """The `run_id` `download_model` resolves once is the same one every downstream call uses.""" + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) + mlflow_repository.load_predict_model = AsyncMock(return_value=('model', None)) + + await mlflow_repository.download_model('test_model', {}, 'predict', 'sklearn', False) + + mlflow_repository.get_model_run_id.assert_called_once_with(model_name='test_model', stage='Production') + mlflow_repository.download_artifact_code_dir.assert_called_once_with( + 'test_model', 'prediction_model', 'sklearn', {}, run_id='run-123' + ) + mlflow_repository.load_predict_model.assert_called_once_with( + 'test_model', {}, 'sklearn', run_id='run-123' + ) + + +@pytest.mark.asyncio +async def test_download_model_resolves_run_id_once_wrapper(mlflow, mlflow_repository, tmp_path): + """Same single resolution, for the `load_wrapper=True` branch.""" + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) + mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) + + await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True) + + mlflow_repository.get_model_run_id.assert_called_once_with(model_name='test_model', stage='Production') + mlflow_repository.download_artifact_code_dir.assert_called_once_with( + 'test_model', 'prediction_model', 'pyfunc', {}, run_id='run-123' + ) + mlflow_repository.dowload_artifacts.assert_called_once_with( + 'test_model', {}, 'prediction_model', run_id='run-123' + ) + + def test_detect_and_parse_datetime_index_empty(mlflow_repository): input_data = DataFrame() response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) @@ -1965,6 +2012,7 @@ async def test_load_transform_model_flavor_failure_propagates( @patch(JOBLIB_PATH) @pytest.mark.asyncio async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.dowload_artifacts = AsyncMock() loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom') @@ -1983,6 +2031,7 @@ async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, ml async def test_download_model_wrapper_failure_emits_error_metric( joblib_mock, mlflow, mlflow_repository ): + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) mlflow_repository.dowload_artifacts = AsyncMock() loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom') @@ -2503,13 +2552,14 @@ async def test_download_model_delegates_every_flavor_to_the_loaders( mlflow_repository, model_type, loader ): """`download_model` owns no flavor branch: joblib delegates like the other two.""" + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') mlflow_repository.download_artifact_code_dir = AsyncMock(return_value=None) delegate = AsyncMock(return_value=('model', 'local-dir')) setattr(mlflow_repository, loader, delegate) result = await mlflow_repository.download_model('test_model', {}, model_type, 'joblib') - delegate.assert_called_once_with('test_model', {}, 'joblib') + delegate.assert_called_once_with('test_model', {}, 'joblib', run_id='run-123') assert result == ('model', 'local-dir') diff --git a/tests/laborious/workflows/test_import_model.py b/tests/laborious/workflows/test_import_model.py index 1e68237..1bc0266 100644 --- a/tests/laborious/workflows/test_import_model.py +++ b/tests/laborious/workflows/test_import_model.py @@ -12,7 +12,7 @@ workflow is: a sequence of activity invocations with policies attached. import asyncio from datetime import timedelta from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from pytest import fixture, mark, raises from sientia_do.temporal.policies import retry_policy @@ -725,6 +725,95 @@ async def test_reporting_the_report_is_where_it_stops(import_workflow: ImportMod assert result['run_id'] == RUN_ID +# --------------------------------------------------- cancellation is never swallowed + + +def cancelling_activities(): + """Patch the workflow module's `workflow` object so every activity call raises cancellation.""" + patcher = patch('laborious.workflows.import_model.workflow') + workflow_mock = patcher.start() + workflow_mock.execute_activity_method = AsyncMock(side_effect=asyncio.CancelledError()) + return patcher + + +BASE = {'metadata': {}, 'import_run_id': 77, 'workflow_id': WORKFLOW_ID} + + +@mark.asyncio +async def test_a_cancelled_hint_propagates_instead_of_being_reported(import_workflow: ImportModel): + patcher = cancelling_activities() + try: + with raises(asyncio.CancelledError): + await import_workflow._record_hint( + BASE, Activities.record_import_status, {'step': ImportStep.DOWNLOAD.value} + ) + finally: + patcher.stop() + + +@mark.asyncio +async def test_a_cancelled_terminal_write_propagates_instead_of_being_reported( + import_workflow: ImportModel, +): + patcher = cancelling_activities() + try: + with raises(asyncio.CancelledError): + await import_workflow._record_terminal( + BASE, + succeeded=True, + step=ImportStep.MODEL_DOCUMENT.value, + gate=None, + source={}, + cleanup_failed=False, + ) + finally: + patcher.stop() + + +@mark.asyncio +async def test_a_cancelled_pipeline_failure_report_propagates(import_workflow: ImportModel): + patcher = cancelling_activities() + try: + with raises(asyncio.CancelledError): + await import_workflow._report_failure( + BASE, RuntimeError('boom'), step=ImportStep.REGISTRATION.value, gate=None + ) + finally: + patcher.stop() + + +@mark.asyncio +async def test_a_cancelled_status_write_failure_report_propagates(import_workflow: ImportModel): + patcher = cancelling_activities() + try: + with raises(asyncio.CancelledError): + await import_workflow._report_status_write_failure(BASE, RuntimeError('db down')) + finally: + patcher.stop() + + +@mark.asyncio +async def test_a_cancelled_retention_check_propagates(import_workflow: ImportModel): + patcher = cancelling_activities() + try: + with raises(asyncio.CancelledError): + await import_workflow._ensure_retention(BASE, valid_input()) + finally: + patcher.stop() + + +@mark.asyncio +async def test_a_cancelled_cleanup_propagates_instead_of_returning_not_cleaned( + import_workflow: ImportModel, +): + patcher = cancelling_activities() + try: + with raises(asyncio.CancelledError): + await import_workflow._cleanup(BASE) + finally: + patcher.stop() + + # --------------------------------------------------------------- the failure report