mlrun 1.7.0rc5__py3-none-any.whl → 1.7.2__py3-none-any.whl

mlrun/serving/v2_serving.py

@@ -15,12 +15,13 @@
 import threading
 import time
 import traceback
-from typing import Union
+from typing import Optional, Union
 
-import mlrun.common.model_monitoring
+import mlrun.artifacts
+import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring
-from mlrun.artifacts import ModelArtifact  # noqa: F401
-from mlrun.config import config
+import mlrun.model_monitoring
+from mlrun.errors import err_to_str
 from mlrun.utils import logger, now_date
 
 from ..common.helpers import parse_versioned_object_uri
@@ -38,6 +39,7 @@ class V2ModelServer(StepToDict):
         protocol=None,
         input_path: str = None,
         result_path: str = None,
+        shard_by_endpoint: Optional[bool] = None,
         **kwargs,
     ):
         """base model serving class (v2), using similar API to KFServing v2 and Triton
@@ -62,11 +64,11 @@ class V2ModelServer(StepToDict):
             class MyClass(V2ModelServer):
                 def load(self):
                     # load and initialize the model and/or other elements
-                    model_file, extra_data = self.get_model(suffix='.pkl')
+                    model_file, extra_data = self.get_model(suffix=".pkl")
                     self.model = load(open(model_file, "rb"))
 
                 def predict(self, request):
-                    events = np.array(request['inputs'])
+                    events = np.array(request["inputs"])
                     dmatrix = xgb.DMatrix(events)
                     result: xgb.DMatrix = self.model.predict(dmatrix)
                     return {"outputs": result.tolist()}
@@ -90,6 +92,8 @@ class V2ModelServer(StepToDict):
                               this require that the event body will behave like a dict, example:
                               event: {"x": 5} , result_path="resp" means the returned response will be written
                               to event["y"] resulting in {"x": 5, "resp": <result>}
+        :param shard_by_endpoint: whether to use the endpoint as the partition/sharding key when writing to model
+                                  monitoring stream. Defaults to True.
         :param kwargs:     extra arguments (can be accessed using self.get_param(key))
         """
         self.name = name
@@ -101,7 +105,7 @@ class V2ModelServer(StepToDict):
         self.error = ""
         self.protocol = protocol or "v2"
         self.model_path = model_path
-        self.model_spec: mlrun.artifacts.ModelArtifact = None
+        self.model_spec: Optional[mlrun.artifacts.ModelArtifact] = None
         self._input_path = input_path
         self._result_path = result_path
         self._kwargs = kwargs  # for to_dict()
@@ -118,7 +122,9 @@ class V2ModelServer(StepToDict):
         if model:
             self.model = model
             self.ready = True
+        self._versioned_model_name = None
         self.model_endpoint_uid = None
+        self.shard_by_endpoint = shard_by_endpoint
 
     def _load_and_update_state(self):
         try:
@@ -147,7 +153,7 @@ class V2ModelServer(StepToDict):
             logger.warn("GraphServer not initialized for VotingEnsemble instance")
             return
 
-        if not self.context.is_mock or self.context.server.track_models:
+        if not self.context.is_mock or self.context.monitoring_mock:
             self.model_endpoint_uid = _init_endpoint_record(
                 graph_server=server, model=self
             )
@@ -175,9 +181,9 @@ class V2ModelServer(StepToDict):
         ::
 
             def load(self):
-                model_file, extra_data = self.get_model(suffix='.pkl')
+                model_file, extra_data = self.get_model(suffix=".pkl")
                 self.model = load(open(model_file, "rb"))
-                categories = extra_data['categories'].as_df()
+                categories = extra_data["categories"].as_df()
 
         Parameters
         ----------
@@ -224,6 +230,23 @@ class V2ModelServer(StepToDict):
         request = self.preprocess(event_body, op)
         return self.validate(request, op)
 
+    @property
+    def versioned_model_name(self):
+        if self._versioned_model_name:
+            return self._versioned_model_name
+
+        # Generating version model value based on the model name and model version
+        if self.model_path and self.model_path.startswith("store://"):
+            # Enrich the model server with the model artifact metadata
+            self.get_model()
+            if not self.version:
+                # Enrich the model version with the model artifact tag
+                self.version = self.model_spec.tag
+            self.labels = self.model_spec.labels
+        version = self.version or "latest"
+        self._versioned_model_name = f"{self.name}:{version}"
+        return self._versioned_model_name
+
     def do_event(self, event, *args, **kwargs):
         """main model event handler method"""
         start = now_date()
@@ -231,6 +254,11 @@ class V2ModelServer(StepToDict):
         event_body = _extract_input_data(self._input_path, event.body)
         event_id = event.id
         op = event.path.strip("/")
+
+        partition_key = (
+            self.model_endpoint_uid if self.shard_by_endpoint is not False else None
+        )
+
         if event_body and isinstance(event_body, dict):
             op = op or event_body.get("operation")
             event_id = event_body.get("id", event_id)
@@ -250,13 +278,20 @@ class V2ModelServer(StepToDict):
             except Exception as exc:
                 request["id"] = event_id
                 if self._model_logger:
-                    self._model_logger.push(start, request, op=op, error=exc)
+                    self._model_logger.push(
+                        start,
+                        request,
+                        op=op,
+                        error=exc,
+                        partition_key=partition_key,
+                    )
                 raise exc
 
             response = {
                 "id": event_id,
                 "model_name": self.name,
                 "outputs": outputs,
+                "timestamp": start.isoformat(sep=" ", timespec="microseconds"),
             }
             if self.version:
                 response["model_version"] = self.version
@@ -286,7 +321,7 @@ class V2ModelServer(StepToDict):
             setattr(event, "terminated", True)
             event_body = {
                 "name": self.name,
-                "version": self.version,
+                "version": self.version or "",
                 "inputs": [],
                 "outputs": [],
             }
@@ -306,7 +341,13 @@ class V2ModelServer(StepToDict):
             except Exception as exc:
                 request["id"] = event_id
                 if self._model_logger:
-                    self._model_logger.push(start, request, op=op, error=exc)
+                    self._model_logger.push(
+                        start,
+                        request,
+                        op=op,
+                        error=exc,
+                        partition_key=partition_key,
+                    )
                 raise exc
 
             response = {
@@ -330,11 +371,20 @@ class V2ModelServer(StepToDict):
         if self._model_logger:
             inputs, outputs = self.logged_results(request, response, op)
             if inputs is None and outputs is None:
-                self._model_logger.push(start, request, response, op)
+                self._model_logger.push(
+                    start, request, response, op, partition_key=partition_key
+                )
             else:
                 track_request = {"id": event_id, "inputs": inputs or []}
                 track_response = {"outputs": outputs or []}
-                self._model_logger.push(start, track_request, track_response, op)
+                # TODO : check dict/list
+                self._model_logger.push(
+                    start,
+                    track_request,
+                    track_response,
+                    op,
+                    partition_key=partition_key,
+                )
         event.body = _update_result_body(self._result_path, original_body, response)
         return event
 
@@ -375,8 +425,10 @@ class V2ModelServer(StepToDict):
         """postprocess, before returning response"""
         return request
 
-    def predict(self, request: dict) -> dict:
-        """model prediction operation"""
+    def predict(self, request: dict) -> list:
+        """model prediction operation
+        :return: list with the model prediction results (can be multi-port) or list of lists for multiple predictions
+        """
         raise NotImplementedError()
 
     def explain(self, request: dict) -> dict:
@@ -449,7 +501,7 @@ class _ModelLogPusher:
             base_data["labels"] = self.model.labels
         return base_data
 
-    def push(self, start, request, resp=None, op=None, error=None):
+    def push(self, start, request, resp=None, op=None, error=None, partition_key=None):
         start_str = start.isoformat(sep=" ", timespec="microseconds")
         if error:
             data = self.base_data()
@@ -460,7 +512,7 @@ class _ModelLogPusher:
             if self.verbose:
                 message = f"{message}\n{traceback.format_exc()}"
             data["error"] = message
-            self.output_stream.push([data])
+            self.output_stream.push([data], partition_key=partition_key)
             return
 
         self._sample_iter = (self._sample_iter + 1) % self.stream_sample
@@ -486,7 +538,7 @@ class _ModelLogPusher:
                         "metrics",
                     ]
                     data["values"] = self._batch
-                    self.output_stream.push([data])
+                    self.output_stream.push([data], partition_key=partition_key)
             else:
                 data = self.base_data()
                 data["request"] = request
@@ -496,7 +548,7 @@ class _ModelLogPusher:
                 data["microsec"] = microsec
                 if getattr(self.model, "metrics", None):
                     data["metrics"] = self.model.metrics
-                self.output_stream.push([data])
+                self.output_stream.push([data], partition_key=partition_key)
 
 
 def _init_endpoint_record(
@@ -523,62 +575,83 @@ def _init_endpoint_record(
             graph_server.function_uri
         )
     except Exception as e:
-        logger.error("Failed to parse function URI", exc=e)
+        logger.error("Failed to parse function URI", exc=err_to_str(e))
         return None
 
-    # Generating version model value based on the model name and model version
-    if model.version:
-        versioned_model_name = f"{model.name}:{model.version}"
-    else:
-        versioned_model_name = f"{model.name}:latest"
-
     # Generating model endpoint ID based on function uri and model version
     uid = mlrun.common.model_monitoring.create_model_endpoint_uid(
-        function_uri=graph_server.function_uri, versioned_model=versioned_model_name
+        function_uri=graph_server.function_uri,
+        versioned_model=model.versioned_model_name,
     ).uid
 
-    # If model endpoint object was found in DB, skip the creation process.
     try:
-        mlrun.get_run_db().get_model_endpoint(project=project, endpoint_id=uid)
-
+        model_ep = mlrun.get_run_db().get_model_endpoint(
+            project=project, endpoint_id=uid
+        )
     except mlrun.errors.MLRunNotFoundError:
+        model_ep = None
+    except mlrun.errors.MLRunBadRequestError as err:
+        logger.info(
+            "Cannot get the model endpoints store", err=mlrun.errors.err_to_str(err)
+        )
+        return
+
+    if model.context.server.track_models and not model_ep:
         logger.info("Creating a new model endpoint record", endpoint_id=uid)
+        model_endpoint = mlrun.common.schemas.ModelEndpoint(
+            metadata=mlrun.common.schemas.ModelEndpointMetadata(
+                project=project, labels=model.labels, uid=uid
+            ),
+            spec=mlrun.common.schemas.ModelEndpointSpec(
+                function_uri=graph_server.function_uri,
+                model=model.versioned_model_name,
+                model_class=model.__class__.__name__,
+                model_uri=model.model_path,
+                stream_path=model.context.stream.stream_uri,
+                active=True,
+                monitoring_mode=mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled,
+            ),
+            status=mlrun.common.schemas.ModelEndpointStatus(
+                endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP
+            ),
+        )
 
-        try:
-            model_endpoint = mlrun.common.schemas.ModelEndpoint(
-                metadata=mlrun.common.schemas.ModelEndpointMetadata(
-                    project=project, labels=model.labels, uid=uid
-                ),
-                spec=mlrun.common.schemas.ModelEndpointSpec(
-                    function_uri=graph_server.function_uri,
-                    model=versioned_model_name,
-                    model_class=model.__class__.__name__,
-                    model_uri=model.model_path,
-                    stream_path=config.model_endpoint_monitoring.store_prefixes.default.format(
-                        project=project, kind="stream"
-                    ),
-                    active=True,
-                    monitoring_mode=mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
-                    if model.context.server.track_models
-                    else mlrun.common.schemas.model_monitoring.ModelMonitoringMode.disabled,
-                ),
-                status=mlrun.common.schemas.ModelEndpointStatus(
-                    endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP
-                ),
-            )
+        db = mlrun.get_run_db()
+        db.create_model_endpoint(
+            project=project,
+            endpoint_id=uid,
+            model_endpoint=model_endpoint.dict(),
+        )
 
+    elif model_ep:
+        attributes = {}
+        old_model_uri = model_ep.spec.model_uri
+        mlrun.model_monitoring.helpers.enrich_model_endpoint_with_model_uri(
+            model_endpoint=model_ep,
+            model_obj=model.model_spec,
+        )
+        if model_ep.spec.model_uri != old_model_uri:
+            attributes["model_uri"] = model_ep.spec.model_uri
+        if (
+            model_ep.spec.monitoring_mode
+            == mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
+        ) != model.context.server.track_models:
+            attributes["monitoring_mode"] = (
+                mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
+                if model.context.server.track_models
+                else mlrun.common.schemas.model_monitoring.ModelMonitoringMode.disabled
+            )
+        if attributes:
             db = mlrun.get_run_db()
-
-            db.create_model_endpoint(
+            db.patch_model_endpoint(
                 project=project,
                 endpoint_id=uid,
-                model_endpoint=model_endpoint.dict(),
+                attributes=attributes,
+            )
+            logger.info(
+                "Updating model endpoint attributes",
+                attributes=attributes,
+                endpoint_id=uid,
             )
-
-        except Exception as e:
-            logger.error("Failed to create endpoint record", exc=e)
-
-    except Exception as e:
-        logger.error("Failed to retrieve model endpoint object", exc=e)
 
     return uid

mlrun/track/tracker.py

@@ -31,8 +31,9 @@ class Tracker(ABC):
     * Offline: Manually importing models and artifacts into an MLRun project using the `import_x` methods.
     """
 
+    @staticmethod
     @abstractmethod
-    def is_enabled(self) -> bool:
+    def is_enabled() -> bool:
         """
         Checks if tracker is enabled.
 

mlrun/track/trackers/mlflow_tracker.py

@@ -442,6 +442,11 @@ class MLFlowTracker(Tracker):
         # Prepare the archive path:
         model_uri = pathlib.Path(model_uri)
         archive_path = pathlib.Path(tmp_path) / f"{model_uri.stem}.zip"
+        if not os.path.exists(model_uri):
+            local_path = mlflow.artifacts.download_artifacts(
+                artifact_uri=str(model_uri)
+            )
+            model_uri = pathlib.Path(local_path)
 
         # TODO add progress bar for the case of large files
         # Zip the artifact:

mlrun/utils/async_http.py

@@ -24,7 +24,7 @@ from aiohttp_retry import ExponentialRetry, RequestParams, RetryClient, RetryOpt
 from aiohttp_retry.client import _RequestContext
 
 from mlrun.config import config
-from mlrun.errors import err_to_str
+from mlrun.errors import err_to_str, raise_for_status
 
 from .helpers import logger as mlrun_logger
 
@@ -46,12 +46,21 @@ class AsyncClientWithRetry(RetryClient):
         *args,
         **kwargs,
     ):
+        # do not retry on PUT / PATCH as they might have side effects (not truly idempotent)
+        blacklisted_methods = (
+            blacklisted_methods
+            if blacklisted_methods is not None
+            else [
+                "POST",
+                "PUT",
+                "PATCH",
+            ]
+        )
         super().__init__(
             *args,
             retry_options=ExponentialRetryOverride(
                 retry_on_exception=retry_on_exception,
-                # do not retry on PUT / PATCH as they might have side effects (not truly idempotent)
-                blacklisted_methods=blacklisted_methods or ["POST", "PUT", "PATCH"],
+                blacklisted_methods=blacklisted_methods,
                 attempts=max_retries,
                 statuses=retry_on_status_codes,
                 factor=retry_backoff_factor,
@@ -63,6 +72,12 @@ class AsyncClientWithRetry(RetryClient):
             **kwargs,
         )
 
+    def methods_blacklist_update_required(self, new_blacklist: str):
+        self._retry_options: ExponentialRetryOverride
+        return set(self._retry_options.blacklisted_methods).difference(
+            set(new_blacklist)
+        )
+
     def _make_requests(
         self,
         params_list: list[RequestParams],
@@ -173,7 +188,7 @@ class _CustomRequestContext(_RequestContext):
                 last_attempt = current_attempt == self._retry_options.attempts
                 if self._is_status_code_ok(response.status) or last_attempt:
                     if self._raise_for_status:
-                        response.raise_for_status()
+                        raise_for_status(response)
 
                     self._response = response
                     return response
@@ -222,7 +237,7 @@ class _CustomRequestContext(_RequestContext):
                 retry_wait = self._retry_options.get_timeout(
                     attempt=current_attempt, response=None
                 )
-                self._logger.debug(
+                self._logger.warning(
                     "Request failed on retryable exception, retrying",
                     retry_wait_secs=retry_wait,
                     method=params.method,
@@ -275,6 +290,11 @@ class _CustomRequestContext(_RequestContext):
                 if isinstance(exc.os_error, exc_type):
                     return
         if exc.__cause__:
-            return self.verify_exception_type(exc.__cause__)
+            # If the cause exception is retriable, return, otherwise, raise the original exception
+            try:
+                self.verify_exception_type(exc.__cause__)
+            except Exception:
+                raise exc
+            return
         else:
             raise exc

mlrun/utils/db.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import abc
 import pickle
 from datetime import datetime
 
@@ -28,12 +29,22 @@ class BaseModel:
         columns = [column.key for column in mapper.columns if column.key not in exclude]
 
         def get_key_value(c):
+            # all (never say never) DB classes have "object" defined as "full_object"
+            if c == "object":
+                c = "full_object"
             if isinstance(getattr(self, c), datetime):
                 return c, getattr(self, c).isoformat()
             return c, getattr(self, c)
 
         return dict(map(get_key_value, columns))
 
+    @abc.abstractmethod
+    def get_identifier_string(self):
+        """
+        This method must be implemented by any subclass.
+        """
+        pass
+
 
 class HasStruct(BaseModel):
     @property
@@ -51,3 +62,10 @@ class HasStruct(BaseModel):
         exclude = exclude or []
         exclude.append("body")
         return super().to_dict(exclude, strip=strip)
+
+    @abc.abstractmethod
+    def get_identifier_string(self):
+        """
+        This method must be implemented by any subclass.
+        """
+        pass