mlrun 1.6.4rc8__py3-none-any.whl → 1.7.0__py3-none-any.whl

mlrun/serving/v2_serving.py

@@ -15,12 +15,13 @@
 import threading
 import time
 import traceback
-from typing import Dict, 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
@@ -29,8 +30,6 @@ from .utils import StepToDict, _extract_input_data, _update_result_body
 
 
 class V2ModelServer(StepToDict):
-    """base model serving class (v2), using similar API to KFServing v2 and Triton"""
-
     def __init__(
         self,
         context=None,
@@ -64,11 +63,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()}
@@ -103,7 +102,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()
@@ -149,7 +148,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
             )
@@ -177,9 +176,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
         ----------
@@ -221,6 +220,8 @@ class V2ModelServer(StepToDict):
 
     def _pre_event_processing_actions(self, event, event_body, op):
         self._check_readiness(event)
+        if "_dict" in op:
+            event_body = self._inputs_to_list(event_body)
         request = self.preprocess(event_body, op)
         return self.validate(request, op)
 
@@ -237,7 +238,12 @@ class V2ModelServer(StepToDict):
         if not op and event.method != "GET":
             op = "infer"
 
-        if op == "predict" or op == "infer":
+        if (
+            op == "predict"
+            or op == "infer"
+            or op == "infer_dict"
+            or op == "predict_dict"
+        ):
             # predict operation
             request = self._pre_event_processing_actions(event, event_body, op)
             try:
@@ -252,6 +258,7 @@ class V2ModelServer(StepToDict):
                 "id": event_id,
                 "model_name": self.name,
                 "outputs": outputs,
+                "timestamp": start.isoformat(sep=" ", timespec="microseconds"),
             }
             if self.version:
                 response["model_version"] = self.version
@@ -329,6 +336,7 @@ class V2ModelServer(StepToDict):
             else:
                 track_request = {"id": event_id, "inputs": inputs or []}
                 track_response = {"outputs": outputs or []}
+                # TODO : check dict/list
                 self._model_logger.push(start, track_request, track_response, op)
         event.body = _update_result_body(self._result_path, original_body, response)
         return event
@@ -362,22 +370,61 @@ class V2ModelServer(StepToDict):
 
         return request
 
-    def preprocess(self, request: Dict, operation) -> Dict:
+    def preprocess(self, request: dict, operation) -> dict:
         """preprocess the event body before validate and action"""
         return request
 
-    def postprocess(self, request: Dict) -> Dict:
+    def postprocess(self, request: dict) -> dict:
         """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:
+    def explain(self, request: dict) -> dict:
         """model explain operation"""
         raise NotImplementedError()
 
+    def _inputs_to_list(self, request: dict) -> dict:
+        """
+        Convert the inputs from list of dictionary / dictionary to list of lists / list
+        where the internal list order is according to the ArtifactModel inputs.
+
+        :param request: event
+        :return: evnet body converting the inputs to be list of lists
+        """
+        if self.model_spec and self.model_spec.inputs:
+            input_order = [feature.name for feature in self.model_spec.inputs]
+        else:
+            raise mlrun.MLRunInvalidArgumentError(
+                "In order to use predict_dict or infer_dict operation you have to provide `model_path` "
+                "to the model server and to load it by `load()` function"
+            )
+        inputs = request.get("inputs")
+        try:
+            if isinstance(inputs, list) and all(
+                isinstance(item, dict) for item in inputs
+            ):
+                new_inputs = [
+                    [input_dict[key] for key in input_order] for input_dict in inputs
+                ]
+            elif isinstance(inputs, dict):
+                new_inputs = [inputs[key] for key in input_order]
+            else:
+                raise mlrun.MLRunInvalidArgumentError(
+                    "When using predict_dict or infer_dict operation the inputs must be "
+                    "of type `list[dict]` or `dict`"
+                )
+        except KeyError:
+            raise mlrun.MLRunInvalidArgumentError(
+                f"Input dictionary don't contain all the necessary input keys : {input_order}"
+            )
+        request["inputs"] = new_inputs
+        return request
+
 
 class _ModelLogPusher:
     def __init__(self, model, context, output_stream=None):
@@ -481,11 +528,17 @@ 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:
+    if model.model_path and model.model_path.startswith("store://"):
+        # Enrich the model server with the model artifact metadata
+        model.get_model()
+        if not model.version:
+            # Enrich the model version with the model artifact tag
+            model.version = model.model_spec.tag
+        model.labels = model.model_spec.labels
         versioned_model_name = f"{model.name}:{model.version}"
     else:
         versioned_model_name = f"{model.name}:latest"
@@ -495,48 +548,74 @@ def _init_endpoint_record(
         function_uri=graph_server.function_uri, versioned_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=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/tracker_manager.py

@@ -14,7 +14,7 @@
 
 import importlib
 import inspect
-from typing import List, Union
+from typing import Union
 
 import mlrun.errors
 from mlrun.config import config as mlconf
@@ -28,7 +28,7 @@ _TRACKERS = ["mlflow"]
 
 # A list for the available trackers during runtime. It will be setup at the beginning of the run by the function
 # `_collect_available_trackers`:
-_AVAILABLE_TRACKERS: List[Tracker] = None
+_AVAILABLE_TRACKERS: list[Tracker] = None
 
 
 class TrackerManager(metaclass=Singleton):
@@ -41,7 +41,7 @@ class TrackerManager(metaclass=Singleton):
         """
         Initialize a new empty tracker manager.
         """
-        self._trackers: List[Tracker] = []
+        self._trackers: list[Tracker] = []
 
         # Check general config for tracking usage, if false we return an empty manager
         if mlconf.external_platform_tracking.enabled:

mlrun/track/trackers/mlflow_tracker.py

@@ -15,7 +15,6 @@ import os
 import pathlib
 import tempfile
 import zipfile
-from typing import List
 
 import mlflow
 import mlflow.entities
@@ -443,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:
@@ -526,7 +530,7 @@ class MLFlowTracker(Tracker):
         )
 
     @staticmethod
-    def _schema_to_feature(schema: mlflow.types.Schema) -> List[Feature]:
+    def _schema_to_feature(schema: mlflow.types.Schema) -> list[Feature]:
         """
         Cast MLFlow schema to MLRun features.
 

mlrun/utils/async_http.py

@@ -16,7 +16,7 @@
 import asyncio
 import logging
 import typing
-from typing import List, Optional
+from typing import Optional
 
 import aiohttp
 import aiohttp.http_exceptions
@@ -38,12 +38,10 @@ class AsyncClientWithRetry(RetryClient):
         self,
         max_retries: int = config.http_retry_defaults.max_retries,
         retry_backoff_factor: float = config.http_retry_defaults.backoff_factor,
-        retry_on_status_codes: typing.List[
-            int
-        ] = config.http_retry_defaults.status_codes,
+        retry_on_status_codes: list[int] = config.http_retry_defaults.status_codes,
         retry_on_exception: bool = True,
         raise_for_status: bool = True,
-        blacklisted_methods: typing.Optional[typing.List[str]] = None,
+        blacklisted_methods: typing.Optional[list[str]] = None,
         logger: logging.Logger = None,
         *args,
         **kwargs,
@@ -82,7 +80,7 @@ class AsyncClientWithRetry(RetryClient):
 
     def _make_requests(
         self,
-        params_list: List[RequestParams],
+        params_list: list[RequestParams],
         retry_options: Optional[RetryOptionsBase] = None,
         raise_for_status: Optional[bool] = None,
     ) -> "_CustomRequestContext":
@@ -117,7 +115,7 @@ class ExponentialRetryOverride(ExponentialRetry):
     def __init__(
         self,
         retry_on_exception: bool,
-        blacklisted_methods: typing.List[str],
+        blacklisted_methods: list[str],
         *args,
         **kwargs,
     ):
@@ -239,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,

mlrun/utils/azure_vault.py

@@ -63,7 +63,7 @@ class AzureVaultStore:
             mlconf.secret_stores.azure_vault.secret_path + "/" + file_name
         )
         if os.path.isfile(full_path):
-            with open(full_path, "r") as secret_file:
+            with open(full_path) as secret_file:
                 contents = secret_file.read()
             return contents
         return None

mlrun/utils/clones.py

@@ -18,7 +18,6 @@ import tarfile
 import tempfile
 import zipfile
 from os import path, remove
-from typing import Tuple
 from urllib.parse import urlparse
 
 from git import Repo
@@ -91,7 +90,7 @@ def get_repo_url(repo):
     return url
 
 
-def add_credentials_git_remote_url(url: str, secrets=None) -> Tuple[str, bool]:
+def add_credentials_git_remote_url(url: str, secrets=None) -> tuple[str, bool]:
     """Enrich a Git remote URL with credential related secrets, if any are available
     If no secrets are supplied, or if the secrets are insufficient, the original URL is returned
     Besides the URL, this function also returns a bool indicating if any enrichment was done

mlrun/utils/condition_evaluator.py

@@ -19,7 +19,7 @@ from mlrun.utils import logger
 
 
 def evaluate_condition_in_separate_process(
-    condition: str, context: typing.Dict[str, typing.Any], timeout: int = 5
+    condition: str, context: dict[str, typing.Any], timeout: int = 5
 ):
     if not condition:
         return True
@@ -44,13 +44,13 @@ def evaluate_condition_in_separate_process(
 
 
 def _evaluate_condition_wrapper(
-    connection, condition: str, context: typing.Dict[str, typing.Any]
+    connection, condition: str, context: dict[str, typing.Any]
 ):
     connection.send(_evaluate_condition(condition, context))
     return connection.close()
 
 
-def _evaluate_condition(condition: str, context: typing.Dict[str, typing.Any]):
+def _evaluate_condition(condition: str, context: dict[str, typing.Any]):
     import jinja2.sandbox
 
     jinja_env = jinja2.sandbox.SandboxedEnvironment()

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
 
@@ -19,7 +20,7 @@ from sqlalchemy.orm import class_mapper
 
 
 class BaseModel:
-    def to_dict(self, exclude=None):
+    def to_dict(self, exclude=None, strip: bool = False):
         """
         NOTE - this function (currently) does not handle serializing relationships
         """
@@ -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
@@ -44,10 +55,17 @@ class HasStruct(BaseModel):
     def struct(self, value):
         self.body = pickle.dumps(value)
 
-    def to_dict(self, exclude=None):
+    def to_dict(self, exclude=None, strip: bool = False):
         """
         NOTE - this function (currently) does not handle serializing relationships
         """
         exclude = exclude or []
         exclude.append("body")
-        return super().to_dict(exclude)
+        return super().to_dict(exclude, strip=strip)
+
+    @abc.abstractmethod
+    def get_identifier_string(self):
+        """
+        This method must be implemented by any subclass.
+        """
+        pass