mlrun 1.7.1rc10__py3-none-any.whl → 1.8.0rc11__py3-none-any.whl

mlrun/serving/states.py

@@ -25,7 +25,7 @@ import pathlib
 import traceback
 from copy import copy, deepcopy
 from inspect import getfullargspec, signature
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import storey.utils
 
@@ -104,7 +104,12 @@ class BaseStep(ModelObj):
     default_shape = "ellipse"
     _dict_fields = ["kind", "comment", "after", "on_error"]
 
-    def __init__(self, name: str = None, after: list = None, shape: str = None):
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        after: Optional[list] = None,
+        shape: Optional[str] = None,
+    ):
         self.name = name
         self._parent = None
         self.comment = None
@@ -154,14 +159,14 @@ class BaseStep(ModelObj):
 
     def error_handler(
         self,
-        name: str = None,
+        name: Optional[str] = None,
         class_name=None,
         handler=None,
         before=None,
         function=None,
-        full_event: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        full_event: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         **class_args,
     ):
         """set error handler on a step or the entire graph (to be executed on failure/raise)
@@ -297,13 +302,13 @@ class BaseStep(ModelObj):
     def to(
         self,
         class_name: Union[str, StepToDict] = None,
-        name: str = None,
-        handler: str = None,
-        graph_shape: str = None,
-        function: str = None,
-        full_event: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        name: Optional[str] = None,
+        handler: Optional[str] = None,
+        graph_shape: Optional[str] = None,
+        function: Optional[str] = None,
+        full_event: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         **class_args,
     ):
         """add a step right after this step and return the new step
@@ -404,16 +409,16 @@ class TaskStep(BaseStep):
 
     def __init__(
         self,
-        class_name: Union[str, type] = None,
-        class_args: dict = None,
-        handler: str = None,
-        name: str = None,
-        after: list = None,
-        full_event: bool = None,
-        function: str = None,
-        responder: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        class_name: Optional[Union[str, type]] = None,
+        class_args: Optional[dict] = None,
+        handler: Optional[str] = None,
+        name: Optional[str] = None,
+        after: Optional[list] = None,
+        full_event: Optional[bool] = None,
+        function: Optional[str] = None,
+        responder: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
     ):
         super().__init__(name, after)
         self.class_name = class_name
@@ -552,6 +557,8 @@ class TaskStep(BaseStep):
             self._object.post_init(mode)
             if hasattr(self._object, "model_endpoint_uid"):
                 self.endpoint_uid = self._object.model_endpoint_uid
+            if hasattr(self._object, "name"):
+                self.endpoint_name = self._object.name
 
     def respond(self):
         """mark this step as the responder.
@@ -607,16 +614,16 @@ class MonitoringApplicationStep(TaskStep):
 
     def __init__(
         self,
-        class_name: Union[str, type] = None,
-        class_args: dict = None,
-        handler: str = None,
-        name: str = None,
-        after: list = None,
-        full_event: bool = None,
-        function: str = None,
-        responder: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        class_name: Optional[Union[str, type]] = None,
+        class_args: Optional[dict] = None,
+        handler: Optional[str] = None,
+        name: Optional[str] = None,
+        after: Optional[list] = None,
+        full_event: Optional[bool] = None,
+        function: Optional[str] = None,
+        responder: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
     ):
         super().__init__(
             class_name=class_name,
@@ -641,16 +648,16 @@ class ErrorStep(TaskStep):
 
     def __init__(
         self,
-        class_name: Union[str, type] = None,
-        class_args: dict = None,
-        handler: str = None,
-        name: str = None,
-        after: list = None,
-        full_event: bool = None,
-        function: str = None,
-        responder: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        class_name: Optional[Union[str, type]] = None,
+        class_args: Optional[dict] = None,
+        handler: Optional[str] = None,
+        name: Optional[str] = None,
+        after: Optional[list] = None,
+        full_event: Optional[bool] = None,
+        function: Optional[str] = None,
+        responder: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
     ):
         super().__init__(
             class_name=class_name,
@@ -678,14 +685,14 @@ class RouterStep(TaskStep):
 
     def __init__(
         self,
-        class_name: Union[str, type] = None,
-        class_args: dict = None,
-        handler: str = None,
-        routes: list = None,
-        name: str = None,
-        function: str = None,
-        input_path: str = None,
-        result_path: str = None,
+        class_name: Optional[Union[str, type]] = None,
+        class_args: Optional[dict] = None,
+        handler: Optional[str] = None,
+        routes: Optional[list] = None,
+        name: Optional[str] = None,
+        function: Optional[str] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
     ):
         super().__init__(
             class_name,
@@ -813,12 +820,12 @@ class QueueStep(BaseStep):
 
     def __init__(
         self,
-        name: str = None,
-        path: str = None,
-        after: list = None,
-        shards: int = None,
-        retention_in_hours: int = None,
-        trigger_args: dict = None,
+        name: Optional[str] = None,
+        path: Optional[str] = None,
+        after: Optional[list] = None,
+        shards: Optional[int] = None,
+        retention_in_hours: Optional[int] = None,
+        trigger_args: Optional[dict] = None,
         **options,
     ):
         super().__init__(name, after)
@@ -850,13 +857,13 @@ class QueueStep(BaseStep):
     def to(
         self,
         class_name: Union[str, StepToDict] = None,
-        name: str = None,
-        handler: str = None,
-        graph_shape: str = None,
-        function: str = None,
-        full_event: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        name: Optional[str] = None,
+        handler: Optional[str] = None,
+        graph_shape: Optional[str] = None,
+        function: Optional[str] = None,
+        full_event: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         **class_args,
     ):
         if not function:
@@ -905,7 +912,7 @@ class FlowStep(BaseStep):
         self,
         name=None,
         steps=None,
-        after: list = None,
+        after: Optional[list] = None,
         engine=None,
         final_step=None,
     ):
@@ -948,9 +955,9 @@ class FlowStep(BaseStep):
         before=None,
         graph_shape=None,
         function=None,
-        full_event: bool = None,
-        input_path: str = None,
-        result_path: str = None,
+        full_event: Optional[bool] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         **class_args,
     ):
         """add task, queue or router step/class to the flow
@@ -1037,7 +1044,7 @@ class FlowStep(BaseStep):
         self._last_added = step
         return step
 
-    def clear_children(self, steps: list = None):
+    def clear_children(self, steps: Optional[list] = None):
         """remove some or all of the states, empty/None for all"""
         if not steps:
             steps = self._steps.keys()
@@ -1282,11 +1289,19 @@ class FlowStep(BaseStep):
         if self._controller:
             # async flow (using storey)
             event._awaitable_result = None
-            resp = self._controller.emit(
-                event, return_awaitable_result=self._wait_for_result
-            )
-            if self._wait_for_result and resp:
-                return resp.await_result()
+            if self.context.is_mock:
+                resp = self._controller.emit(
+                    event, return_awaitable_result=self._wait_for_result
+                )
+                if self._wait_for_result and resp:
+                    return resp.await_result()
+            else:
+                resp_awaitable = self._controller.emit(
+                    event, await_result=self._wait_for_result
+                )
+                if self._wait_for_result:
+                    return resp_awaitable
+                return self._await_and_return_id(resp_awaitable, event)
             event = copy(event)
             event.body = {"id": event.id}
             return event
@@ -1554,8 +1569,8 @@ def params_to_step(
     graph_shape=None,
     function=None,
     full_event=None,
-    input_path: str = None,
-    result_path: str = None,
+    input_path: Optional[str] = None,
+    result_path: Optional[str] = None,
     class_args=None,
 ):
     """return step object from provided params or classes/objects"""
@@ -1703,8 +1718,12 @@ def _init_async_objects(context, steps):
         is_explicit_ack_supported(context) and mlrun.mlconf.is_explicit_ack_enabled()
     )
 
-    # TODO: Change to AsyncEmitSource once we can drop support for nuclio<1.12.10
-    default_source = storey.SyncEmitSource(
+    if context.is_mock:
+        source_class = storey.SyncEmitSource
+    else:
+        source_class = storey.AsyncEmitSource
+
+    default_source = source_class(
         context=context,
         explicit_ack=explicit_ack,
         **source_args,

mlrun/serving/utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 import inspect
+from typing import Optional
 
 from mlrun.utils import get_in, update_in
 
@@ -55,7 +56,12 @@ class StepToDict:
         "kwargs",
     ]
 
-    def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
+    def to_dict(
+        self,
+        fields: Optional[list] = None,
+        exclude: Optional[list] = None,
+        strip: bool = False,
+    ):
         """convert the step object to a python dictionary"""
         fields = fields or getattr(self, "_dict_fields", None)
         if not fields:
@@ -105,5 +111,10 @@ class MonitoringApplicationToDict(StepToDict):
 class RouterToDict(StepToDict):
     _STEP_KIND = "router"
 
-    def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
+    def to_dict(
+        self,
+        fields: Optional[list] = None,
+        exclude: Optional[list] = None,
+        strip: bool = False,
+    ):
         return super().to_dict(exclude=["routes"], strip=strip)

mlrun/serving/v1_serving.py

@@ -18,6 +18,7 @@ import socket
 from copy import deepcopy
 from datetime import datetime
 from io import BytesIO
+from typing import Optional
 from urllib.request import urlopen
 
 import nuclio
@@ -33,7 +34,7 @@ serving_handler = "handler"
 def new_v1_model_server(
     name,
     model_class: str,
-    models: dict = None,
+    models: Optional[dict] = None,
     filename="",
     protocol="",
     image="",
@@ -68,7 +69,7 @@ def new_v1_model_server(
 
 
 class MLModelServer:
-    def __init__(self, name: str, model_dir: str = None, model=None):
+    def __init__(self, name: str, model_dir: Optional[str] = None, model=None):
         self.name = name
         self.ready = False
         self.model_dir = model_dir

mlrun/serving/v2_serving.py

@@ -21,10 +21,9 @@ import mlrun.artifacts
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring
 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
+from ..common.schemas.model_monitoring import ModelEndpointSchema
 from .server import GraphServer
 from .utils import StepToDict, _extract_input_data, _update_result_body
 
@@ -33,12 +32,12 @@ class V2ModelServer(StepToDict):
     def __init__(
         self,
         context=None,
-        name: str = None,
-        model_path: str = None,
+        name: Optional[str] = None,
+        model_path: Optional[str] = None,
         model=None,
         protocol=None,
-        input_path: str = None,
-        result_path: str = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         shard_by_endpoint: Optional[bool] = None,
         **kwargs,
     ):
@@ -110,12 +109,6 @@ class V2ModelServer(StepToDict):
         self._result_path = result_path
         self._kwargs = kwargs  # for to_dict()
         self._params = kwargs
-        self._model_logger = (
-            _ModelLogPusher(self, context)
-            if context and context.stream.enabled
-            else None
-        )
-
         self.metrics = {}
         self.labels = {}
         self.model = None
@@ -125,6 +118,7 @@ class V2ModelServer(StepToDict):
         self._versioned_model_name = None
         self.model_endpoint_uid = None
         self.shard_by_endpoint = shard_by_endpoint
+        self._model_logger = None
 
     def _load_and_update_state(self):
         try:
@@ -157,6 +151,11 @@ class V2ModelServer(StepToDict):
             self.model_endpoint_uid = _init_endpoint_record(
                 graph_server=server, model=self
             )
+        self._model_logger = (
+            _ModelLogPusher(self, self.context)
+            if self.context and self.context.stream.enabled
+            else None
+        )
 
     def get_param(self, key: str, default=None):
         """get param by key (specified in the model or the function)"""
@@ -198,13 +197,15 @@ class V2ModelServer(StepToDict):
             extra dataitems dictionary
 
         """
-        model_file, self.model_spec, extra_dataitems = mlrun.artifacts.get_model(
-            self.model_path, suffix
-        )
-        if self.model_spec and self.model_spec.parameters:
-            for key, value in self.model_spec.parameters.items():
-                self._params[key] = value
-        return model_file, extra_dataitems
+        if self.model_path:
+            model_file, self.model_spec, extra_dataitems = mlrun.artifacts.get_model(
+                self.model_path, suffix
+            )
+            if self.model_spec and self.model_spec.parameters:
+                for key, value in self.model_spec.parameters.items():
+                    self._params[key] = value
+            return model_file, extra_dataitems
+        return None, None
 
     def load(self):
         """model loading function, see also .get_model() method"""
@@ -474,7 +475,7 @@ class V2ModelServer(StepToDict):
 
 
 class _ModelLogPusher:
-    def __init__(self, model, context, output_stream=None):
+    def __init__(self, model: V2ModelServer, context, output_stream=None):
         self.model = model
         self.verbose = context.verbose
         self.hostname = context.stream.hostname
@@ -496,6 +497,7 @@ class _ModelLogPusher:
             "version": self.model.version,
             "host": self.hostname,
             "function_uri": self.function_uri,
+            "endpoint_id": self.model.model_endpoint_uid,
         }
         if getattr(self.model, "labels", None):
             base_data["labels"] = self.model.labels
@@ -567,26 +569,26 @@ def _init_endpoint_record(
     """
 
     logger.info("Initializing endpoint records")
-
-    # Generate required values for the model endpoint record
-    try:
-        # Getting project name from the function uri
-        project, uri, tag, hash_key = parse_versioned_object_uri(
-            graph_server.function_uri
-        )
-    except Exception as e:
-        logger.error("Failed to parse function URI", exc=err_to_str(e))
-        return None
-
-    # 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=model.versioned_model_name,
-    ).uid
-
+    if not model.model_spec:
+        model.get_model()
+    if model.model_spec:
+        model_name = model.model_spec.metadata.key
+        model_db_key = model.model_spec.spec.db_key
+        model_uid = model.model_spec.metadata.uid
+        model_tag = model.model_spec.tag
+        model_labels = model.model_spec.labels  # todo : check if we still need this
+    else:
+        model_name = None
+        model_db_key = None
+        model_uid = None
+        model_tag = None
+        model_labels = {}
     try:
         model_ep = mlrun.get_run_db().get_model_endpoint(
-            project=project, endpoint_id=uid
+            project=graph_server.project,
+            name=model.name,
+            function_name=graph_server.function_name,
+            function_tag=graph_server.function_tag or "latest",
         )
     except mlrun.errors.MLRunNotFoundError:
         model_ep = None
@@ -596,62 +598,93 @@ def _init_endpoint_record(
         )
         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(
+    function = mlrun.get_run_db().get_function(
+        name=graph_server.function_name,
+        project=graph_server.project,
+        tag=graph_server.function_tag or "latest",
+    )
+    function_uid = function.get("metadata", {}).get("uid")
+    if not model_ep and model.context.server.track_models:
+        logger.info(
+            "Creating a new model endpoint record",
+            name=model.name,
+            project=graph_server.project,
+            function_name=graph_server.function_name,
+            function_tag=graph_server.function_tag or "latest",
+            function_uid=function_uid,
+            model_name=model_name,
+            model_tag=model_tag,
+            model_db_key=model_db_key,
+            model_uid=model_uid,
+            model_class=model.__class__.__name__,
+        )
+        model_ep = mlrun.common.schemas.ModelEndpoint(
             metadata=mlrun.common.schemas.ModelEndpointMetadata(
-                project=project, labels=model.labels, uid=uid
+                project=graph_server.project,
+                labels=model_labels,
+                name=model.name,
+                endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP,
             ),
             spec=mlrun.common.schemas.ModelEndpointSpec(
-                function_uri=graph_server.function_uri,
-                model=model.versioned_model_name,
+                function_name=graph_server.function_name,
+                function_uid=function_uid,
+                function_tag=graph_server.function_tag or "latest",
+                model_name=model_name,
+                model_db_key=model_db_key,
+                model_uid=model_uid,
                 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,
+                model_tag=model_tag,
             ),
             status=mlrun.common.schemas.ModelEndpointStatus(
-                endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP
+                monitoring_mode=mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
+                if model.context.server.track_models
+                else mlrun.common.schemas.model_monitoring.ModelMonitoringMode.disabled,
             ),
         )
-
         db = mlrun.get_run_db()
-        db.create_model_endpoint(
-            project=project,
-            endpoint_id=uid,
-            model_endpoint=model_endpoint.dict(),
-        )
+        model_ep = db.create_model_endpoint(model_endpoint=model_ep)
 
     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 function_uid != model_ep.spec.function_uid:
+            attributes[ModelEndpointSchema.FUNCTION_UID] = function_uid
+        if model_name != model_ep.spec.model_name:
+            attributes[ModelEndpointSchema.MODEL_NAME] = model_name
+        if model_uid != model_ep.spec.model_uid:
+            attributes[ModelEndpointSchema.MODEL_UID] = model_uid
+        if model_tag != model_ep.spec.model_tag:
+            attributes[ModelEndpointSchema.MODEL_TAG] = model_tag
+        if model_db_key != model_ep.spec.model_db_key:
+            attributes[ModelEndpointSchema.MODEL_DB_KEY] = model_db_key
+        if model_labels != model_ep.metadata.labels:
+            attributes[ModelEndpointSchema.LABELS] = model_labels
+        if model.__class__.__name__ != model_ep.spec.model_class:
+            attributes[ModelEndpointSchema.MODEL_CLASS] = model.__class__.__name__
         if (
-            model_ep.spec.monitoring_mode
+            model_ep.status.monitoring_mode
             == mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
         ) != model.context.server.track_models:
-            attributes["monitoring_mode"] = (
+            attributes[ModelEndpointSchema.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.patch_model_endpoint(
-                project=project,
-                endpoint_id=uid,
-                attributes=attributes,
-            )
             logger.info(
                 "Updating model endpoint attributes",
                 attributes=attributes,
-                endpoint_id=uid,
+                project=model_ep.metadata.project,
+                name=model_ep.metadata.name,
+                function_name=model_ep.spec.function_name,
             )
+            db = mlrun.get_run_db()
+            model_ep = db.patch_model_endpoint(
+                project=model_ep.metadata.project,
+                name=model_ep.metadata.name,
+                endpoint_id=model_ep.metadata.uid,
+                attributes=attributes,
+            )
+    else:
+        return None
 
-    return uid
+    return model_ep.metadata.uid

mlrun/track/__init__.py

@@ -11,6 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from mlrun.track.tracker import Tracker
 from mlrun.track.tracker_manager import TrackerManager

mlrun/track/tracker.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from mlrun.artifacts import Artifact, ModelArtifact
 from mlrun.execution import MLClientCtx
@@ -63,7 +63,7 @@ class Tracker(ABC):
         project: MlrunProject,
         reference_id: Any,
         function_name: str,
-        handler: str = None,
+        handler: Optional[str] = None,
         **kwargs,
     ) -> RunObject:
         """