diff --git a/packages/kumo-core/src/kumo_core/train/formatting.py b/packages/kumo-core/src/kumo_core/train/formatting.py index 53ce12c3a..0d5585d73 100644 --- a/packages/kumo-core/src/kumo_core/train/formatting.py +++ b/packages/kumo-core/src/kumo_core/train/formatting.py @@ -31,6 +31,15 @@ class LinkPredictionTable(Protocol): _ANCHOR_TIME_COL = "ANCHOR_TIMESTAMP" +_CLASS_TASKS = frozenset({ + TaskType.MULTICLASS_CLASSIFICATION, + TaskType.MULTILABEL_CLASSIFICATION, +}) +_LONG_FORMAT_TASKS = _CLASS_TASKS | { + TaskType.MULTILABEL_RANKING, + TaskType.TEMPORAL_LINK_PREDICTION, + TaskType.STATIC_LINK_PREDICTION, +} def _get_col_names( @@ -67,6 +76,23 @@ def _get_col_names( return result +def prediction_output_columns(task_type: TaskType) -> list[str]: + r"""The prediction columns a scored row carries after its entity column, + for a target named ``TARGET`` and a time column named ``TIMESTAMP``.""" + if task_type in _CLASS_TASKS: + return ['PREDICTED', 'CLASS', 'SCORE'] + if task_type.is_ranking: + return ['CLASS', 'SCORE'] + if task_type == TaskType.FORECASTING: + return ['TIMESTAMP', 'TARGET_PRED'] + if task_type == TaskType.BINARY_CLASSIFICATION: + return [ + 'TARGET_PRED', + *_get_col_names(pd.Index([]), task_type, is_wide_format=True) + ] + return ['TARGET_PRED'] + + def _format_prediction_sanity_checks( target_encoder: TargetEncoder, task_type: TaskType, @@ -284,13 +310,7 @@ def format_prediction( 'metadata. Rename the time column before forecasting.') format: Literal['wide', 'long'] - if task_type in { - TaskType.MULTICLASS_CLASSIFICATION, - TaskType.MULTILABEL_CLASSIFICATION, - TaskType.MULTILABEL_RANKING, - TaskType.TEMPORAL_LINK_PREDICTION, - TaskType.STATIC_LINK_PREDICTION, - }: + if task_type in _LONG_FORMAT_TASKS: format = 'long' else: assert top_k is None, f"`top_k` is not supported for {task_type}." diff --git a/packages/kumo-core/tests/test_prediction_output_columns.py b/packages/kumo-core/tests/test_prediction_output_columns.py new file mode 100644 index 000000000..eee36d837 --- /dev/null +++ b/packages/kumo-core/tests/test_prediction_output_columns.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import pytest +from kumo_core.train.formatting import prediction_output_columns +from kumo_core.types import TaskType + + +@pytest.mark.parametrize('task_type, columns', [ + (TaskType.BINARY_CLASSIFICATION, + ['TARGET_PRED', 'False_PROB', 'True_PROB']), + (TaskType.REGRESSION, ['TARGET_PRED']), + (TaskType.FORECASTING, ['TIMESTAMP', 'TARGET_PRED']), + (TaskType.MULTICLASS_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']), + (TaskType.MULTILABEL_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']), + (TaskType.MULTILABEL_RANKING, ['CLASS', 'SCORE']), + (TaskType.TEMPORAL_LINK_PREDICTION, ['CLASS', 'SCORE']), + (TaskType.STATIC_LINK_PREDICTION, ['CLASS', 'SCORE']), +]) +def test_prediction_output_columns(task_type, columns) -> None: + assert prediction_output_columns(task_type) == columns diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 3d7df6802..fd57ca84d 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -203,9 +203,50 @@ names: `start(record)` with the fit run, `on_trial(record, event)` with a trial attempt's run after each of its events, and `finish(record, error)` with the fit run's final record (`None` when tracking stopped early). A tracker that raises is not called again for that fit and its run is ended as failed. -Optional members: `prepare()`, `tags`, and the best-effort hooks `log_model`, +Optional members: `prepare()`, `tags`, `log_model(result, snapshot, registry)` (see MLflow models below), and the best-effort hooks `log_prediction_run` and `log_monitoring`. Params are the resolved model plan (fit) or searched hyperparameters and `attempt` (trial). `RunRecord.inputs` maps each source table to its `connector`, `location`, `version` and `version_kind`. A tracker's `tags=` are added as given, except `kumo.*` names and param names. + +### MLflow models + +With `registry=ModelRegistryLocation(name, alias=...)` and an `MlflowTracking` tracker, every fit in the block logs its model on the fit run as an MLflow pyfunc model and registers it as a new version of `name` (Unity Catalog with `MLFLOW_REGISTRY_URI=databricks-uc`, SageMaker managed MLflow, or any MLflow registry). One alias is set on the new version. A bare `ModelRegistryLocation` without such a tracker, a fixed `version`, or an invalid alias raises when `kumo.MLOps` is built. + +```python +from kumo.engine import ModelRegistryLocation +from kumo.engine.integrations.mlflow import MlflowTracking, load_model + +cfg = kumo.MLOps("s3://bucket/kumo", + tracking=MlflowTracking(experiment_name="/Shared/users_spend"), + registry=ModelRegistryLocation("main.ml.users_spend", alias="champion")) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +result.registered_model # PublishedModel(name='main.ml.users_spend', version='4', uri='models:/main.ml.users_spend/4') + +trainer = load_model("models:/main.ml.users_spend@champion") # the Trainer back +model = mlflow.pyfunc.load_model("models:/main.ml.users_spend@champion") +scores = model.predict(entities, params={"snapshot": "s3://bucket/kumo/snapshots//"}) +``` + +- `fit` waits for the upload and registration before it returns, for up to an hour, with a 600 s MLflow HTTP timeout during the upload (unless `MLFLOW_HTTP_REQUEST_TIMEOUT` is set). If they fail, the tracking thread stops or the hour passes, `fit` warns and still returns its result with `registered_model=None`, the run ends `FAILED`, and a model that finishes uploading later is not registered. +- `MlflowTracking(log_model=True)` logs the model without registering it. `log_model(trainer, snapshot=...)` logs a fitted or loaded model in any run, with the bundle's source tables as run inputs. +- The model's metadata is `kumo_bundle_id`; its `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version. +- The pyfunc model loads a snapshot on its first `predict` with that URI and keeps it for the life of the process: each serving replica, and each Spark executor of `mlflow.pyfunc.spark_udf`, downloads the snapshot once and keeps every snapshot it has loaded (no eviction), so its disk must hold them. With `serving_snapshot=` (on `log_model` or `MlflowTracking`), each replica downloads that snapshot when it starts, and `predict` uses it when no `snapshot` param is given. `params={"population": True}` scores every entity and returns the whole frame in memory; for large populations use `Trainer.predict` with a `PredictionOutputConfig`. Each `predict` call works in its own scratch directory and removes it before returning. +- The `mlflow` extra needs an MLflow 3.13 or later client: DiskGraph pins pyarrow 24, which older MLflow releases exclude. SageMaker managed MLflow servers run MLflow 3.0; a 3.17 client against a 3.0 server handles logging, registration, aliases and `load_model`. +- On Databricks, the parent folder of an MLflow experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`, which the `mlflow` extra does not install. + +#### Which registry + +Pick one per registered model, since the two store different formats. Both record only the `kumo_bundle_id` pointer; lineage stays in the bundle's `_provenance.json`. Both accept the same aliases (`kumoapi.plugin.check_model_alias`). + +| | Connector registry | MLflow model | +|---|---|---| +| Register | `kumo.MLOps(registry=ModelOutputConfig(connector, location=ModelRegistryLocation(name, alias=...)))` | `kumo.MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(name, alias=...))` | +| Load | `Trainer.load_from_connector(connector, location=...)` | `mlflow.pyfunc.load_model(uri)` or `load_model(uri)` | +| Promote | `set_model_alias(connector, location, alias)` | MLflow's `set_registered_model_alias` | +| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow | +| SageMaker | not available | SageMaker managed MLflow (`MLFLOW_TRACKING_URI` = tracking server ARN) | +| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | +| Plain VM | not available (store under the artifact root instead) | any MLflow tracking server | diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 20831531a..cb33e39b7 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -8,6 +8,8 @@ from dataclasses import dataclass, field from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union +from kumoapi.plugin import ModelRegistryLocation, check_model_alias + from kumo.engine._manifest import KINDS @@ -79,6 +81,17 @@ def __post_init__(self) -> None: if unknown: raise ValueError(f'Unknown export kinds {unknown}; choose from ' f'{list(KINDS)}.') + if isinstance(self.registry, ModelRegistryLocation): + if not callable(getattr(self.tracking, 'log_model', None)): + raise ValueError( + 'registry=ModelRegistryLocation(...) registers through ' + 'the MLflow tracker: pass tracking=MlflowTracking(...), ' + 'or registry=ModelOutputConfig(connector, location=...).') + if self.registry.version is not None: + raise ValueError('MLflow assigns registered model versions; ' + 'leave ModelRegistryLocation.version unset.') + if self.registry.alias is not None: + check_model_alias(self.registry.alias) object.__setattr__(self, 'root', root) object.__setattr__(self, 'exports', kinds) diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index 7381e5cb5..8a73ce5fa 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -35,11 +35,15 @@ import hashlib import uuid +from contextlib import contextmanager +from contextvars import ContextVar from pathlib import Path -from typing import Union +from typing import Iterator, Optional, Union _DEFAULT_ROOT = Path.home() / ".cache" / "kumo-engine" _root = _DEFAULT_ROOT +_scoped_root: ContextVar[Optional[Path]] = ContextVar('kumo_scratch_root', + default=None) def set_scratch_dir(path: Union[str, Path]) -> None: @@ -65,7 +69,18 @@ def set_scratch_dir(path: Union[str, Path]) -> None: def get_scratch_dir() -> Path: - return _root + return _scoped_root.get() or _root + + +@contextmanager +def scoped_scratch_dir(path: Path) -> Iterator[None]: + r"""Nests this thread's (context's) new scratch directories under + ``path`` instead of the configured root.""" + token = _scoped_root.set(path) + try: + yield + finally: + _scoped_root.reset(token) def _spool_dir() -> str: diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index c79aa608b..aee04c13d 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -32,6 +32,9 @@ from kumo.engine._events import Event, _finite_floats if TYPE_CHECKING: + from kumoapi.plugin import ModelRegistryLocation, PublishedModel + + from kumo.engine.graph import GraphSnapshot from kumo.engine.trainer.trainer import TrainingResult logger = logging.getLogger(__name__) @@ -42,6 +45,7 @@ # never stalls training; these bound how long the fit waits for it. QUEUE_SIZE = 1000 FINISH_SECONDS = 30.0 +MODEL_SECONDS = 3600.0 @dataclass(frozen=True) @@ -219,8 +223,12 @@ class Tracking(Protocol): the event that updated it, and `finish` the fit run's final record (`None` if tracking stopped early, which should end the run as failed). Optional members: ``prepare()``, called first on the fit's own thread; - ``tags`` (extra fit-run tags); and the best-effort hooks - ``log_model(result, snapshot)``, ``log_prediction_run(record)`` and + ``tags`` (extra fit-run tags); ``log_model(result, snapshot, registry)``, + called after a fit that succeeded and before `finish`, which the fit waits + for while the tracking thread is alive, up to an hour (an error, a stopped + thread or the time limit ends the run as failed), and which returns the + version it registered at ``registry``, a ``ModelRegistryLocation``; and the + best-effort hooks ``log_prediction_run(record)`` and ``log_monitoring(record)``. """ def start(self, record: RunRecord) -> None: @@ -286,6 +294,56 @@ def on_event(self, event: Event) -> None: self._submit(lambda: self._call( lambda tracker: tracker.on_trial(record, event))) + def log_model( + self, + result: 'TrainingResult', + snapshot: 'GraphSnapshot', + registry: Optional['ModelRegistryLocation'] = None, + ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]: + r"""Runs the tracker's ``log_model`` on the tracking thread and waits + for it; returns the registered version and the error it raised.""" + if (self._worker is None + or not hasattr(self._tracker, 'log_model')): + return None, None + done = threading.Event() + failed: List[BaseException] = [] + registered: List[Any] = [] + + def log() -> None: + try: + tracker = self._tracker + if tracker is not None: + registered.append( + getattr(tracker, 'log_model')(result, snapshot, + registry)) + except BaseException as error: + failed.append(error) + finally: + done.set() + + worker = self._worker + deadline = time.monotonic() + MODEL_SECONDS + try: + self._jobs.put(log, timeout=MODEL_SECONDS) + except queue.Full: + pass + while not done.wait(1.0): + if worker.is_alive() and time.monotonic() < deadline: + continue + failed.append( + TimeoutError(f'it did not finish within {MODEL_SECONDS:.0f}s') + if worker.is_alive() else + RuntimeError('the tracking thread stopped')) + tracker, self._tracker = self._tracker, None + if tracker is not None: + _abandon(tracker) + break + if failed: + warnings.warn(f'Could not log the model: {failed[0]}', + UserWarning, stacklevel=3) + return None, failed[0] + return (registered or [None])[0], None + def finish(self, result: Optional['TrainingResult'], error: Optional[BaseException]) -> None: worker = self._worker diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py new file mode 100644 index 000000000..0e2993a0a --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import tempfile +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Tuple + +import pandas as pd + +from kumo.engine.connector._extras import import_provider +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + +_pyfunc = import_provider('mlflow.pyfunc', 'mlflow') +_schema = import_provider('mlflow.types.schema', 'mlflow') +_signature = import_provider('mlflow.models.signature', 'mlflow') + +_DTYPES = { + 'PREDICTED': 'boolean', + 'SCORE': 'double', + 'TIMESTAMP': 'datetime', + 'False_PROB': 'double', + 'True_PROB': 'double', +} + + +class KumoModel(_pyfunc.PythonModel): + def load_context(self, context: Any) -> None: + from kumo.engine.trainer.trainer import Trainer + + self._trainer = Trainer.load(context.artifacts['bundle']) + self._columns = list(context.model_config['columns']) + self._snapshots: Dict[str, GraphSnapshot] = {} + self._default = context.model_config.get('snapshot') + if self._default: + self._snapshot(self._default) + + def predict(self, context: Any, model_input: pd.DataFrame, + params: Optional[Dict[str, Any]] = None) -> pd.DataFrame: + params = params or {} + if len(model_input) == 0 and not params.get('population'): + return pd.DataFrame(columns=self._columns) + uri = params.get('snapshot') or self._default + if not uri: + raise ValueError("Pass params={'snapshot': }.") + return predict(self._trainer, self._snapshot(uri), model_input, + params) + + def _snapshot(self, uri: str) -> GraphSnapshot: + if uri not in self._snapshots: + connector, path = snapshot_location(uri) + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + return self._snapshots[uri] + + +def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + if len(rows) == 0: + return trainer.predict( + snapshot, prediction_time=None + if time is None else time.to_pydatetime(), **options).frame + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + return trainer.predict( + snapshot, + entities=FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'], **options).frame + + +def snapshot_location(uri: str) -> Tuple[Connector, str]: + from kumo.engine.connector import connector_for_root + + parent, slash, name = uri.rstrip('/').rpartition('/') + return connector_for_root(parent or slash or '.'), name + + +def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: + from kumo_core.train.formatting import prediction_output_columns + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + bound = PredictiveQuery.from_snapshot(snapshot, query) + parsed = bound._parsed + entity = snapshot._native_graph.table(parsed.entity_table) + entity_type = _entity_type(str(entity.column(entity.primary_key).dtype)) + task_type = TaskType(parsed.task_type) + dtypes = { + **_DTYPES, 'TARGET_PRED': + 'boolean' if task_type == TaskType.BINARY_CLASSIFICATION else 'double' + } + inputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + _schema.ColSpec('datetime', 'TIMESTAMP', required=False), + ]) + outputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + *(_schema.ColSpec(dtypes.get(name) or _schema.AnyType(), name) + for name in prediction_output_columns(task_type)), + ]) + params = _schema.ParamSchema([ + _schema.ParamSpec('snapshot', 'string', None), + _schema.ParamSpec('prediction_time', 'string', None), + _schema.ParamSpec('binary_classification_threshold', 'double', None), + _schema.ParamSpec('num_classes_to_return', 'long', None), + _schema.ParamSpec('population', 'boolean', False), + ]) + example = pd.DataFrame({ + 'ENTITY': { + 'long': [0, 1], + 'double': [0.0, 1.0], + }.get(entity_type, ['entity-0', 'entity-1']) + }) + return _signature.ModelSignature(inputs, outputs, params), example + + +def _entity_type(dtype: str) -> str: + if dtype.startswith('int') or dtype == 'byte': + return 'long' + if dtype.startswith('float'): + return 'double' + return 'string' diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py new file mode 100644 index 000000000..03cedae08 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_model import KumoModel + +import_provider('mlflow.models', 'mlflow').set_model(KumoModel()) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index e799aa68f..3d4c45901 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""MLflow tracking for `Trainer.fit` (``pip install 'kumo-engine[mlflow]'``). -""" +r"""MLflow tracking and models for `Trainer.fit` +(``pip install 'kumo-engine[mlflow]'``).""" import hashlib import json @@ -13,12 +13,33 @@ import time import warnings from contextlib import ExitStack, contextmanager -from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Set, Tuple +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + FrozenSet, + Iterator, + Mapping, + Optional, + Set, + Tuple, +) + +from kumoapi.plugin import ModelRegistryLocation, PublishedModel from kumo.engine._events import Event from kumo.engine._tracking import DROPPED_PARAMS, RunRecord +from kumo.engine._version import __version__ from kumo.engine.connector._extras import import_provider +if TYPE_CHECKING: + from kumo_core.embedding_models import EmbeddingModelStore + + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer.trainer import Trainer, TrainingResult + logger = logging.getLogger(__name__) _MAX_PARAM_KEY = 250 @@ -30,6 +51,8 @@ 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', } +# Model upload only, and only where the limits above apply. +_UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} _PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS', 'ADLS') @@ -48,6 +71,12 @@ class MlflowTracking: log_system_metrics: sample CPU, memory, disk, network and (with ``nvidia-ml-py``) GPU usage on the fit run. tags: extra tags for the fit run. + log_model: log the fitted model on the fit run (see `log_model`) + before the fit returns. ``kumo.MLOps(registry= + ModelRegistryLocation(...))`` also registers it, which implies + this. + serving_snapshot: exported snapshot URI the served model loads at + startup (see `log_model`). ``run_id`` is the MLflow id of the latest fit run, once it has started. """ @@ -58,12 +87,16 @@ def __init__( run_name: Optional[str] = None, log_system_metrics: bool = True, tags: Optional[Mapping[str, str]] = None, + log_model: bool = False, + serving_snapshot: Optional[str] = None, ) -> None: self._mlflow = import_provider('mlflow', 'mlflow') self.experiment_name = experiment_name self.run_name = run_name self.log_system_metrics = log_system_metrics self.tags = dict(tags or {}) + self.serving_snapshot = serving_snapshot + self.logs_model = log_model self._outer: Optional[str] = None self.run_id: Optional[str] = None self._reset() @@ -74,6 +107,8 @@ def _reset(self) -> None: self._experiment_id: Optional[str] = None self._children: Dict[Tuple[int, int], str] = {} self._artifacts: Set[str] = set() + self._model_id: Optional[str] = None + self._http_owned: FrozenSet[str] = frozenset() self._guard = ExitStack() def prepare(self) -> None: @@ -83,7 +118,9 @@ def prepare(self) -> None: def start(self, record: RunRecord) -> None: self._reset() self.run_id = None - self._guard.enter_context(_limited_http()) + self._http_owned = frozenset(name for name in _HTTP_LIMITS + if name not in os.environ) + self._guard.enter_context(_limited_http(self._http_owned)) try: autologging = import_provider('mlflow.utils.autologging_utils', 'mlflow') @@ -114,11 +151,7 @@ def start(self, record: RunRecord) -> None: self._experiment_id = run.info.experiment_id self._log_params(self._parent, record.params) self._log_artifacts(record.artifacts) - try: - self._log_inputs(record.inputs) - except Exception: - logger.warning('Could not log the fit\'s source tables to MLflow', - exc_info=True) + _log_inputs(mlflow, self._client, self._parent, record.inputs) def on_trial(self, record: RunRecord, event: Event) -> None: key = (event.trial or 0, int(event.params['attempt'])) @@ -139,12 +172,42 @@ def on_trial(self, record: RunRecord, event: Event) -> None: self._client.set_terminated(run_id, 'FINISHED' if succeeded else 'FAILED') + def log_model( + self, + result: 'TrainingResult', + snapshot: 'GraphSnapshot', + registry: Optional[ModelRegistryLocation] = None, + ) -> Optional[PublishedModel]: + run_id = self._parent + if not (self.logs_model or registry) or run_id is None: + return None + upload = { + name: value + for name, value in _UPLOAD_HTTP.items() if name in self._http_owned + } + os.environ.update(upload) + try: + info = _log_model(self._mlflow, result, snapshot, name='model', + run_id=run_id, + serving_snapshot=self.serving_snapshot) + finally: + if self._parent == run_id: + os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) + if self._parent != run_id: + logger.warning('The fit run %s ended before its model was ' + 'logged; the model is not registered', run_id) + return None + self._model_id = info.model_id + return (None if registry is None else _register( + self._mlflow, info, registry)) + def finish(self, record: Optional[RunRecord], error: Optional[BaseException]) -> None: status = 'FAILED' try: if record is not None and self._parent is not None: - self._log_metrics(self._parent, record.metrics) + self._log_metrics(self._parent, record.metrics, + model_id=self._model_id) self._log_tags(self._parent, record.tags) self._log_artifacts(record.artifacts) if error is None: @@ -168,25 +231,6 @@ def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: self._client.log_text(self._parent, text, name) self._artifacts.add(name) - def _log_inputs(self, - inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: - meta = import_provider('mlflow.data.meta_dataset', 'mlflow') - entities = self._mlflow.entities - datasets = [] - for name, entry in inputs.items(): - source = _dataset_source(entry) - if source is None: - continue - digest = hashlib.sha256( - json.dumps(dict(entry), sort_keys=True).encode()).hexdigest() - dataset = meta.MetaDataset(source, name=name, digest=digest[:8]) - datasets.append( - entities.DatasetInput( - dataset._to_mlflow_entity(), - tags=[entities.InputTag('mlflow.data.context', 'train')])) - if datasets: - self._client.log_inputs(self._parent, datasets=datasets) - def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: entities = self._mlflow.entities self._client.log_batch( @@ -207,18 +251,167 @@ def _log_params(self, run_id: str, params: Mapping[str, str]) -> None: ], tags=tags) def _log_metrics(self, run_id: Optional[str], - metrics: Mapping[str, float], step: int = 0) -> None: + metrics: Mapping[str, float], step: int = 0, + model_id: Optional[str] = None) -> None: if run_id is None or not metrics: return entities = self._mlflow.entities timestamp = int(time.time() * 1000) self._client.log_batch( run_id, metrics=[ - entities.Metric(key, float(value), timestamp, step) + entities.Metric(key, float(value), timestamp, step, + model_id=model_id) for key, value in metrics.items() ]) +def log_model( + trainer: 'Trainer', + *, + snapshot: 'GraphSnapshot', + name: str = 'model', + run_id: Optional[str] = None, + serving_snapshot: Optional[str] = None, +) -> Any: + r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or + the active run, else a new one) with its source tables as run inputs; + returns its `ModelInfo`. Scores with + ``predict(entities, params={'snapshot': uri})``; with + ``serving_snapshot`` (an exported snapshot URI), each serving replica + downloads that snapshot once at startup and scores against it when no + ``snapshot`` param is given.""" + from kumo.engine._tracking import fit_inputs + + mlflow = import_provider('mlflow', 'mlflow') + result = trainer.result + if result is None: + raise ValueError('log_model requires a fitted or loaded Trainer') + with _in_run(mlflow, run_id) as active: + _log_inputs(mlflow, mlflow.MlflowClient(), active, + fit_inputs(result.manifest)) + info = _log_model(mlflow, result, snapshot, name=name, run_id=active, + serving_snapshot=serving_snapshot) + if result.metrics: + mlflow.log_metrics(dict(result.metrics), run_id=info.run_id, + model_id=info.model_id) + return info + + +def load_model( + model_uri: str, + embedding_models: Optional['EmbeddingModelStore'] = None, +) -> 'Trainer': + r"""Restore the `Trainer` of a model `log_model` logged.""" + from kumo.engine._scratch import new_run_id, trainer_scratch_dir + from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle + + mlflow = import_provider('mlflow', 'mlflow') + destination = trainer_scratch_dir(new_run_id()) / 'mlflow-model' + destination.mkdir(parents=True) + local = Path( + mlflow.artifacts.download_artifacts(artifact_uri=model_uri, + dst_path=str(destination))) + model = mlflow.models.Model.load(str(local)) + bundle = model.flavors['python_function']['artifacts']['bundle']['path'] + trainer = Trainer.load(local / bundle, embedding_models) + warn_if_other_bundle(model.metadata, trainer) + return trainer + + +def _log_model(mlflow: Any, result: 'TrainingResult', + snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, + serving_snapshot: Optional[str]) -> Any: + from kumo_core.artifacts import PROVENANCE_FILE + + from kumo.engine.integrations import _mlflow_model, _mlflow_model_code + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + if snapshot is None: + raise ValueError('log_model needs the snapshot the model trained on') + spec = read_training_spec(result.bundle_dir) + signature, example = _mlflow_model.signature(snapshot, + spec.predictive_query) + provenance = result.bundle_dir / PROVENANCE_FILE + model_logger = logging.getLogger('mlflow.models.model') + model_logger.addFilter(_no_serving_check) + try: + with _in_run(mlflow, run_id): + return mlflow.pyfunc.log_model( + name=name, + python_model=_mlflow_model_code.__file__, + model_config={ + 'columns': signature.outputs.input_names(), + 'snapshot': serving_snapshot, + }, + artifacts={ + 'bundle': str(result.bundle_dir), + **({'provenance': str(provenance)} + if provenance.is_file() else {}), + }, + signature=signature, + input_example=example, + metadata=registry_pointers(result), + pip_requirements=[f'kumo-engine[s3,gcs]=={__version__}'], + ) + finally: + model_logger.removeFilter(_no_serving_check) + + +def _log_inputs(mlflow: Any, client: Any, run_id: str, + inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: + try: + meta = import_provider('mlflow.data.meta_dataset', 'mlflow') + entities = mlflow.entities + datasets = [] + for name, entry in inputs.items(): + source = _dataset_source(entry) + if source is None: + continue + digest = hashlib.sha256( + json.dumps(dict(entry), sort_keys=True).encode()).hexdigest() + dataset = meta.MetaDataset(source, name=name, digest=digest[:8]) + datasets.append( + entities.DatasetInput( + dataset._to_mlflow_entity(), + tags=[entities.InputTag('mlflow.data.context', 'train')])) + if datasets: + client.log_inputs(run_id, datasets=datasets) + except Exception: + logger.warning('Could not log the fit\'s source tables to MLflow', + exc_info=True) + + +def _register(mlflow: Any, info: Any, + location: ModelRegistryLocation) -> PublishedModel: + name = location.model_name + version = str(mlflow.register_model(info.model_uri, name).version) + client = mlflow.MlflowClient() + if location.comment: + client.update_model_version(name, version, + description=location.comment) + if location.alias: + client.set_registered_model_alias(name, location.alias, version) + return PublishedModel(name=name, version=version, + uri=f'models:/{name}/{version}') + + +def _no_serving_check(record: logging.LogRecord) -> bool: + # MLflow's input-example check predicts without a snapshot; no switch. + return not record.getMessage().startswith( + 'Failed to validate serving input example') + + +@contextmanager +def _in_run(mlflow: Any, run_id: Optional[str]) -> Iterator[str]: + active = mlflow.active_run() + if active is not None and run_id in (None, active.info.run_id): + yield active.info.run_id + return + with mlflow.start_run(run_id=run_id, nested=active is not None) as run: + yield run.info.run_id + + def _dataset_source(entry: Mapping[str, Optional[str]]) -> Any: connector = (entry.get('connector') or '').upper() location = entry.get('location') @@ -269,11 +462,10 @@ def _quietly(function: Callable[..., Any], *args: Any) -> None: @contextmanager -def _limited_http() -> Iterator[None]: - unset = [name for name in _HTTP_LIMITS if name not in os.environ] - os.environ.update({name: _HTTP_LIMITS[name] for name in unset}) +def _limited_http(names: FrozenSet[str]) -> Iterator[None]: + os.environ.update({name: _HTTP_LIMITS[name] for name in names}) try: yield finally: - for name in unset: + for name in names: os.environ.pop(name, None) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 992fd91d0..16d268244 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -340,6 +340,11 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: check_model_alias(export_config.location.alias) +def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: + registry = None if settings is None else settings.registry + return registry if isinstance(registry, ModelRegistryLocation) else None + + def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: registry = None if settings is None else settings.registry return registry if isinstance(registry, ModelOutputConfig) else None @@ -352,8 +357,7 @@ def _publish(training_result: TrainingResult, name: str(getattr(spec, name)) for name in ('spec_version', 'predictive_query', 'snapshot_id') } - if training_result.bundle_id: - metadata[BUNDLE_ID_KEY] = training_result.bundle_id + metadata.update(registry_pointers(training_result)) location = export_config.location alias = location.alias if location else None if location and alias: @@ -577,7 +581,12 @@ def fit( events.emit('fit_ended', metrics={ **result.metrics, 'seconds': time.monotonic() - fit_start }, params={'status': 'succeeded'}) - session.finish(result, None) + registered, error = session.log_model(result, snapshot, + _mlflow_registry(settings)) + session.finish(result, error) + if registered is not None: + result = replace(result, registered_model=registered) + self._training_result = result if registry is not None: result = self._register(result, registry) return result @@ -1196,13 +1205,7 @@ def load_from_connector( location=address.location) trainer = cls._load_bundle(bundle_dir, embedding_models, accelerator=accelerator) - recorded = (getattr(published, 'metadata', None) - or {}).get(BUNDLE_ID_KEY) - if recorded and trainer.bundle_id and recorded != trainer.bundle_id: - warnings.warn( - f"The registry records bundle {recorded} for this version, " - f"but the fetched bundle is {trainer.bundle_id}.", UserWarning, - stacklevel=2) + warn_if_other_bundle(getattr(published, 'metadata', None), trainer) return trainer @classmethod @@ -1242,3 +1245,17 @@ def set_model_alias(connector: Connector, location: ModelRegistryLocation, raise NotImplementedError( f'{type(connector).__name__} cannot set model aliases') sink.set_model_alias(cast(SourceConnector, connector), location, alias) + + +def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], + trainer: Trainer) -> None: + recorded = (metadata or {}).get(BUNDLE_ID_KEY) + if recorded and trainer.bundle_id and recorded != trainer.bundle_id: + warnings.warn( + f"The registry records bundle {recorded} for this version, but " + f"the fetched bundle is {trainer.bundle_id}.", UserWarning, + stacklevel=3) + + +def registry_pointers(result: TrainingResult) -> Dict[str, str]: + return {BUNDLE_ID_KEY: result.bundle_id} if result.bundle_id else {} diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py new file mode 100644 index 000000000..bb41ff18b --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import warnings +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumoapi.model_plan import ModelPlan +from kumoapi.plugin import PublishedModel +from test_mlflow_tracking import ( # noqa: F401 + _runs, + table_and_snapshot, + tracking_uri, +) + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + +mlflow = pytest.importorskip('mlflow') + + +def _plan(): + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + return plan + + +def _sorted(frame): + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, + tmp_path, caplog, + monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking, load_model + + table, snapshot = table_and_snapshot + plan = _plan() + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking(log_system_metrics=False), + registry=kumo.ModelRegistryLocation('m', + alias='champion')) + trainer = kumo.Trainer(plan) + result = trainer.fit(table, snapshot=snapshot, mlops=cfg) + + assert 'MLflow tracking failed' not in caplog.text + parent, = [ + run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags + ] + model, = mlflow.search_logged_models(output_format='list') + assert model.source_run_id == parent.info.run_id + metrics = {metric.key: metric.value for metric in model.metrics} + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') + assert version.source == model.model_uri + assert {k: v for k, v in parent.data.tags.items() + if k.startswith('kumo.')} == { + 'kumo.kind': 'fit', + 'kumo.run_id': parent.data.tags['kumo.run_id'], + 'kumo.bundle_id': result.id, + } + assert result.registered_model == PublishedModel( + name='m', version=str(version.version), + uri=f'models:/m/{version.version}') + assert trainer.result is result + + info = mlflow.models.get_model_info('models:/m@champion') + assert info.metadata == {'kumo_bundle_id': result.id} + provenance = json.loads(mlflow.artifacts.load_text( + f'{model.artifact_location}/artifacts/{PROVENANCE_FILE}')) + assert provenance == result.manifest.to_dict() + assert [(c['name'], c['type']) for c in info.signature.inputs.to_dict() + ] == [('ENTITY', 'long'), ('TIMESTAMP', 'datetime')] + assert [(c['name'], c['type']) for c in info.signature.outputs.to_dict() + ] == [('ENTITY', 'long'), ('TARGET_PRED', 'double')] + + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), + 'snapshots')) + served = mlflow.pyfunc.load_model('models:/m@champion').predict( + pd.DataFrame({'ENTITY': range(0, 400, 7)}), + params={'snapshot': str(tmp_path / 'exports' / location)}) + assert list(served) == ['ENTITY', 'TARGET_PRED'] + assert len(served) == len(range(0, 400, 7)) + assert load_model('models:/m@champion').bundle_id == result.id + + original = mlflow.models.Model.load + + def other_bundle(path): + model = original(path) + model.metadata['kumo_bundle_id'] = 'other' + return model + + with pytest.MonkeyPatch.context() as patch, pytest.warns( + UserWarning, match='records bundle other'): + patch.setattr(mlflow.models.Model, 'load', other_bundle) + load_model('models:/m@champion') + + +def test_logged_model_scores_the_query_population_and_timed_rows( + fitted, tracking_uri, tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import log_model + + trainer, table, snapshot = fitted + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), + 'snapshots')) + + info = log_model(trainer, snapshot=snapshot, + serving_snapshot=str(tmp_path / 'exports' / location)) + + assert mlflow.get_run(info.run_id).info.status == 'FINISHED' + requirements = mlflow.pyfunc.get_model_dependencies(info.model_uri) + assert f'kumo-engine[gcs,s3]=={kumo.__version__}' in open( + requirements).read().splitlines() + assert {metric.key + for metric in mlflow.get_logged_model(info.model_id).metrics + } == set(trainer.result.metrics) + assert [c['name'] for c in info.signature.outputs.to_dict() + ] == ['ENTITY', 'TARGET_PRED', 'False_PROB', 'True_PROB'] + pyfunc = mlflow.pyfunc.load_model(info.model_uri) + params = { + 'snapshot': str(tmp_path / 'exports' / location), + 'prediction_time': str(pd.Timestamp(table.timestamps[0])), + 'binary_classification_threshold': 0.5, + } + expected = _sorted( + trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame) + empty = pd.DataFrame({'ENTITY': pd.Series([], dtype='int64')}) + tables = get_scratch_dir() / 'prediction_tables' + before = set(tables.glob('*')) + predict = kumo.Trainer.predict + + def predict_beside_another_call(self, *args, **kwargs): + (tables / 'other-call').mkdir(parents=True, exist_ok=True) + return predict(self, *args, **kwargs) + + monkeypatch.setattr(kumo.Trainer, 'predict', predict_beside_another_call) + + assert pyfunc.predict(empty, params=params).empty + served = pyfunc.unwrap_python_model()._snapshots + assert list(served) == [params['snapshot']] + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(empty, params={ + **params, 'snapshot': None, 'population': True + })), expected) + for rows, population in ((empty, True), + (pd.DataFrame({'ENTITY': range(50)}), False)): + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(rows, + params={**params, 'population': population})), + expected) + scratch = pyfunc.unwrap_python_model()._trainer.scratch_dir + assert list(scratch.iterdir()) == [] + assert set(tables.glob('*')) == before | {tables / 'other-call'} + + +@pytest.mark.parametrize('tracking,registry', [ + (None, kumo.ModelRegistryLocation('m')), + (SimpleNamespace(start=None), kumo.ModelRegistryLocation('m')), + ('mlflow', kumo.ModelRegistryLocation('m', alias='SageMakerLifecycle-p')), + ('mlflow', kumo.ModelRegistryLocation('m', version='3')), +]) +def test_invalid_mlflow_registry_is_refused(tracking_uri, tmp_path, tracking, + registry): + from kumo.engine.integrations.mlflow import MlflowTracking + + if tracking == 'mlflow': + tracking = MlflowTracking() + with pytest.raises(ValueError): + kumo.MLOps(str(tmp_path), tracking=tracking, registry=registry) + + +def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( + monkeypatch): + import time + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) + + class Tracker: + def __init__(self, fail): + self.fail, self.calls = fail, [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot, registry): + time.sleep(0.2) + self.calls.append('log_model') + if self.fail: + raise RuntimeError('upload failed') + + def finish(self, record, error): + self.calls.append(('finish', error)) + + result = SimpleNamespace(metrics={}, bundle_dir='missing', bundle_id='b') + for fail in (False, True): + tracker = Tracker(fail) + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + _, error = session.log_model(result, None) + session.finish(result, error) + session._worker.join(5) + assert tracker.calls[:2] == ['start', 'log_model'] + assert tracker.calls[2] == ('finish', error) + assert (error is not None) is fail + assert any('upload failed' in str(w.message) for w in caught) is fail + + +@pytest.mark.parametrize('stall', ['slow upload', 'dead thread']) +def test_model_logging_wait_is_bounded(monkeypatch, stall): + import threading + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) + release = threading.Event() + + class Tracker: + calls = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot, registry): + release.wait(30) + + def finish(self, record, error): + self.calls.append((record, error)) + + tracker = Tracker() + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + if stall == 'dead thread': + session._submit(lambda: (_ for _ in ()).throw(SystemExit())) + session._worker.join(5) + with pytest.warns(UserWarning, match='Could not log the model'): + _, error = session.log_model(SimpleNamespace(), None) + release.set() + session.finish(None, error) + + assert type(error) is (TimeoutError + if stall == 'slow upload' else RuntimeError) + assert tracker.calls == [(None, None)] + + +def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( + tracking_uri, monkeypatch, caplog): + import os + + from test_tracking_contract import CONTEXT + + from kumo.engine._tracking import TrackingSession + from kumo.engine.integrations import mlflow as integration + + for name in integration._HTTP_LIMITS: + monkeypatch.delenv(name, raising=False) + tracker = integration.MlflowTracking(log_system_metrics=False) + location = kumo.ModelRegistryLocation('m') + seen = [] + + def log_model(mlflow, result, snapshot, **kwargs): + seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) + if len(seen) == 2: + tracker.finish(None, None) + return SimpleNamespace(model_id=None, model_uri='runs:/x/model') + + monkeypatch.setattr(integration, '_log_model', log_model) + monkeypatch.setattr(integration, '_register', lambda *args: 'registered') + for _ in range(2): + session = TrackingSession(tracker) + session.start(lambda: CONTEXT) + registered, error = session.log_model(SimpleNamespace(), None, + location) + assert error is None + during = os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT') + session.finish(None, None) + session._worker.join(5) + assert registered == ('registered' if len(seen) == 1 else None) + assert during == ('30' if len(seen) == 1 else None) + + assert seen == ['600', '600'] + assert 'MLFLOW_HTTP_REQUEST_TIMEOUT' not in os.environ + assert 'the model is not registered' in caplog.text