mlrun 1.10.0rc18__py3-none-any.whl → 1.10.0rc20__py3-none-any.whl

mlrun/datastore/model_provider/openai_provider.py

@@ -13,13 +13,19 @@
 # limitations under the License.
 import inspect
 from collections.abc import Awaitable
-from typing import Callable, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 import mlrun
-from mlrun.datastore.model_provider.model_provider import ModelProvider
+from mlrun.datastore.model_provider.model_provider import (
+    InvokeResponseFormat,
+    ModelProvider,
+    UsageResponseKeys,
+)
 from mlrun.datastore.utils import accepts_param
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from openai._models import BaseModel  # noqa
+    from openai.types.chat.chat_completion import ChatCompletion
 
 
 class OpenAIProvider(ModelProvider):
@@ -36,6 +42,7 @@ class OpenAIProvider(ModelProvider):
     """
 
     support_async = True
+    response_class = None
 
     def __init__(
         self,
@@ -62,6 +69,27 @@ class OpenAIProvider(ModelProvider):
         self.options = self.get_client_options()
         self.load_client()
 
+    @classmethod
+    def _import_response_class(cls) -> None:
+        if not cls.response_class:
+            try:
+                from openai.types.chat.chat_completion import ChatCompletion
+            except ImportError as exc:
+                raise ImportError("openai package is not installed") from exc
+            cls.response_class = ChatCompletion
+
+    @staticmethod
+    def _extract_string_output(response: "ChatCompletion") -> str:
+        """
+        Extracts the first generated string from Hugging Face pipeline output,
+        regardless of whether it's plain text-generation or chat-style output.
+        """
+        if len(response.choices) != 1:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "OpenAIProvider: extracting string from response is only supported for single-response outputs"
+            )
+        return response.choices[0].message.content
+
     @classmethod
     def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
         if endpoint and subpath:
@@ -101,8 +129,8 @@ class OpenAIProvider(ModelProvider):
         return self._sanitize_options(res)
 
     def custom_invoke(
-        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
-    ) -> Optional[T]:
+        self, operation: Optional[Callable] = None, **invoke_kwargs
+    ) -> Union["ChatCompletion", "BaseModel"]:
         """
         OpenAI-specific implementation of `ModelProvider.custom_invoke`.
 
@@ -139,9 +167,9 @@ class OpenAIProvider(ModelProvider):
 
     async def async_custom_invoke(
         self,
-        operation: Optional[Callable[..., Awaitable[T]]] = None,
+        operation: Optional[Callable[..., Awaitable[Any]]] = None,
         **invoke_kwargs,
-    ) -> Optional[T]:
+    ) -> Union["ChatCompletion", "BaseModel"]:
         """
         OpenAI-specific implementation of `ModelProvider.async_custom_invoke`.
 
@@ -178,60 +206,90 @@ class OpenAIProvider(ModelProvider):
                 **invoke_kwargs, **model_kwargs
             )
 
+    def _response_handler(
+        self,
+        response: "ChatCompletion",
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **kwargs,
+    ) -> ["ChatCompletion", str, dict[str, Any]]:
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
+            str_response = self._extract_string_output(response)
+            if invoke_response_format == InvokeResponseFormat.STRING:
+                return str_response
+            if invoke_response_format == InvokeResponseFormat.USAGE:
+                stats = response.to_dict()["usage"]
+                response = {
+                    UsageResponseKeys.ANSWER: str_response,
+                    UsageResponseKeys.USAGE: stats,
+                }
+        return response
+
     def invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Optional[Union[str, T]]:
+    ) -> Union[dict[str, Any], str, "ChatCompletion"]:
         """
         OpenAI-specific implementation of `ModelProvider.invoke`.
-        Invokes an OpenAI model operation using the sync client.
+        Invokes an OpenAI model operation using the synchronous client.
         For full details, see `ModelProvider.invoke`.
 
-        :param messages:    Same as ModelProvider.invoke.
+        :param messages:
+            Same as `ModelProvider.invoke`.
 
-        :param as_str: bool
-                            If `True`, returns only the main content of the first response
-                            (`response.choices[0].message.content`).
-                            If `False`, returns the full response object, whose type depends on
-                            the specific OpenAI SDK operation used (e.g., chat completion, completion, etc.).
+        :param invoke_response_format: InvokeResponseFormat
+            Specifies the format of the returned response. Options:
+
+            - "string": Returns only the generated text content, taken from a single response.
+            - "stats": Combines the generated text with metadata (e.g., token usage), returning a dictionary:
+
+              .. code-block:: json
+                 {
+                     "answer": "<generated_text>",
+                     "stats": <ChatCompletion>.to_dict()["usage"]
+                 }
+
+            - "full": Returns the full OpenAI `ChatCompletion` object.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.invoke.
-        :return:            Same as ModelProvider.invoke.
+            Additional keyword arguments passed to the OpenAI client. Same as in `ModelProvider.invoke`.
 
+        :return:
+            A string, dictionary, or `ChatCompletion` object, depending on `invoke_response_format`.
         """
+
         response = self.custom_invoke(messages=messages, **invoke_kwargs)
-        if as_str:
-            return response.choices[0].message.content
-        return response
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )
 
     async def async_invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format=InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> str:
+    ) -> Union[str, "ChatCompletion", dict]:
         """
         OpenAI-specific implementation of `ModelProvider.async_invoke`.
         Invokes an OpenAI model operation using the async client.
-        For full details, see `ModelProvider.async_invoke`.
+        For full details, see `ModelProvider.async_invoke` and `OpenAIProvider.invoke`.
 
-        :param messages:    Same as ModelProvider.async_invoke.
+        :param messages:    Same as `OpenAIProvider.invoke`.
 
-        :param as_str: bool
-                            If `True`, returns only the main content of the first response
-                            (`response.choices[0].message.content`).
-                            If `False`, returns the full awaited response object, whose type depends on
-                            the specific OpenAI SDK operation used (e.g., chat completion, completion, etc.).
+        :param invoke_response_format: InvokeResponseFormat
+                            Same as `OpenAIProvider.invoke`.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.async_invoke.
-        :returns            Same as ModelProvider.async_invoke.
+                            Same as `OpenAIProvider.invoke`.
+        :returns            Same as `ModelProvider.async_invoke`.
 
         """
         response = await self.async_custom_invoke(messages=messages, **invoke_kwargs)
-        if as_str:
-            return response.choices[0].message.content
-        return response
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )

mlrun/db/base.py

@@ -16,8 +16,6 @@ import datetime
 from abc import ABC, abstractmethod
 from typing import Literal, Optional, Union
 
-from deprecated import deprecated
-
 import mlrun.alerts
 import mlrun.common
 import mlrun.common.formatters
@@ -445,23 +443,6 @@ class RunDBInterface(ABC):
     ) -> dict:
         pass
 
-    # TODO: remove in 1.10.0
-    @deprecated(
-        version="1.7.0",
-        reason="'list_features' will be removed in 1.10.0, use 'list_features_v2' instead",
-        category=FutureWarning,
-    )
-    @abstractmethod
-    def list_features(
-        self,
-        project: str,
-        name: Optional[str] = None,
-        tag: Optional[str] = None,
-        entities: Optional[list[str]] = None,
-        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
-    ) -> mlrun.common.schemas.FeaturesOutput:
-        pass
-
     @abstractmethod
     def list_features_v2(
         self,

mlrun/db/httpdb.py

@@ -24,6 +24,7 @@ from datetime import datetime, timedelta
 from os import environ, path, remove
 from typing import Literal, Optional, Union
 from urllib.parse import urlparse
+from uuid import UUID
 
 import pydantic.v1
 import requests
@@ -2554,50 +2555,6 @@ class HTTPRunDB(RunDBInterface):
         resp = self.api_call("GET", path, error_message)
         return FeatureSet.from_dict(resp.json())
 
-    def list_features(
-        self,
-        project: Optional[str] = None,
-        name: Optional[str] = None,
-        tag: Optional[str] = None,
-        entities: Optional[list[str]] = None,
-        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
-    ) -> list[dict]:
-        """List feature-sets which contain specific features. This function may return multiple versions of the same
-        feature-set if a specific tag is not requested. Note that the various filters of this function actually
-        refer to the feature-set object containing the features, not to the features themselves.
-
-        :param project: Project which contains these features.
-        :param name: Name of the feature to look for. The name is used in a like query, and is not case-sensitive. For
-            example, looking for ``feat`` will return features which are named ``MyFeature`` as well as ``defeat``.
-        :param tag: Return feature-sets which contain the features looked for, and are tagged with the specific tag.
-        :param entities: Return only feature-sets which contain an entity whose name is contained in this list.
-        :param labels: Filter feature-sets by label key-value pairs or key existence. This can be provided as:
-            - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
-            or `{"label": None}` to check for key existence.
-            - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
-            or just `"label"` for key existence.
-            - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
-            the specified key-value pairs or key existence.
-        :returns: A list of mapping from feature to a digest of the feature-set, which contains the feature-set
-            meta-data. Multiple entries may be returned for any specific feature due to multiple tags or versions
-            of the feature-set.
-        """
-
-        project = project or config.active_project
-        labels = self._parse_labels(labels)
-        params = {
-            "name": name,
-            "tag": tag,
-            "entity": entities or [],
-            "label": labels,
-        }
-
-        path = f"projects/{project}/features"
-
-        error_message = f"Failed listing features, project: {project}, query: {params}"
-        resp = self.api_call("GET", path, error_message, params=params)
-        return resp.json()["features"]
-
     def list_features_v2(
         self,
         project: Optional[str] = None,
@@ -3834,8 +3791,8 @@ class HTTPRunDB(RunDBInterface):
                                 If tsdb_metrics=False, this parameter will be ignored and no tsdb metrics
                                 will be included.
         :param top_level:       Whether to return only top level model endpoints.
-        :param mode:            Specifies the mode of the model endpoint. Can be "real-time", "batch", or both if set
-                                to None.
+        :param mode:            Specifies the mode of the model endpoint. Can be "real-time" (0), "batch" (1), or
+                                both if set to None.
         :param uids:            A list of unique ids to filter by.
         :param latest_only:     Whether to return only the latest model endpoint version.
         :return:                A list of model endpoints.
@@ -3968,6 +3925,13 @@ class HTTPRunDB(RunDBInterface):
             raise MLRunInvalidArgumentError(
                 "Either endpoint_uid or function_name and function_tag must be provided"
             )
+        if uid:
+            try:
+                UUID(uid)
+            except (ValueError, TypeError):
+                raise MLRunInvalidArgumentError(
+                    "endpoint_id must be a valid UUID string"
+                )
 
     def update_model_monitoring_controller(
         self,

mlrun/db/nopdb.py

@@ -376,16 +376,6 @@ class NopDB(RunDBInterface):
     ) -> dict:
         pass
 
-    def list_features(
-        self,
-        project: str,
-        name: Optional[str] = None,
-        tag: Optional[str] = None,
-        entities: Optional[list[str]] = None,
-        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
-    ) -> mlrun.common.schemas.FeaturesOutput:
-        pass
-
     def list_features_v2(
         self,
         project: str,

mlrun/launcher/base.py

@@ -157,6 +157,19 @@ class BaseLauncher(abc.ABC):
         ]:
             mlrun.utils.helpers.warn_on_deprecated_image(image)
 
+        # Raise an error if retry is configured for a runtime that doesn't support retries.
+        # For local runs, we intentionally skip this validation and allow the run to proceed, since they are typically
+        # used for debugging purposes, and in such cases we avoid blocking their execution.
+        if (
+            not mlrun.runtimes.RuntimeKinds.is_local_runtime(runtime.kind)
+            and run.spec.retry.count
+            and runtime.kind not in mlrun.runtimes.RuntimeKinds.retriable_runtimes()
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Retry is not supported for {runtime.kind} runtime, supported runtimes are: "
+                f"{mlrun.runtimes.RuntimeKinds.retriable_runtimes()}"
+            )
+
     @staticmethod
     def _validate_output_path(
         runtime: "mlrun.runtimes.BaseRuntime",
@@ -268,12 +281,6 @@ class BaseLauncher(abc.ABC):
 
         run.metadata.name = mlrun.utils.normalize_name(
             name=name or run.metadata.name or def_name,
-            # if name or runspec.metadata.name are set then it means that is user defined name and we want to warn the
-            # user that the passed name needs to be set without underscore, if its not user defined but rather enriched
-            # from the handler(function) name then we replace the underscore without warning the user.
-            # most of the time handlers will have `_` in the handler name (python convention is to separate function
-            # words with `_`), therefore we don't want to be noisy when normalizing the run name
-            verbose=bool(name or run.metadata.name),
         )
         mlrun.utils.verify_field_regex(
             "run.metadata.name", run.metadata.name, mlrun.utils.regex.run_name

mlrun/model_monitoring/api.py

@@ -18,6 +18,7 @@ from datetime import datetime
 
 import numpy as np
 import pandas as pd
+from deprecated import deprecated
 
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.datastore.base
@@ -45,6 +46,14 @@ DatasetType = typing.Union[
 ]
 
 
+# TODO: Remove this in 1.12.0
+@deprecated(
+    version="1.10.0",
+    reason="This function is deprecated and will be removed in 1.12. You can generate a model endpoint by either "
+    "deploying a monitored serving function as a real-time service or running it as an offline job. "
+    "To retrieve model endpoints, use `project.list_model_endpoints()`",
+    category=FutureWarning,
+)
 def get_or_create_model_endpoint(
     project: str,
     model_endpoint_name: str,
@@ -67,8 +76,8 @@ def get_or_create_model_endpoint(
     :param model_endpoint_name:      If a new model endpoint is created, the model endpoint name will be presented
                                      under this endpoint (applicable only to new endpoint_id).
     :param model_path:               The model store path (applicable only to new endpoint_id).
-    :param endpoint_id:              Model endpoint unique ID. If not exist in DB, will generate a new record based
-                                     on the provided `endpoint_id`.
+    :param endpoint_id:              Model endpoint unique ID. If not exist in DB, will generate a new record with a
+                                     newly generated ID.
     :param function_name:            If a new model endpoint is created, use this function name.
     :param function_tag:             If a new model endpoint is created, use this function tag.
     :param context:                  MLRun context. If `function_name` not provided, use the context to generate the
@@ -91,25 +100,26 @@ def get_or_create_model_endpoint(
         function_name = FunctionURI.from_string(
             context.to_dict()["spec"]["function"]
         ).function
-    try:
-        model_endpoint = db_session.get_model_endpoint(
-            project=project,
-            name=model_endpoint_name,
-            endpoint_id=endpoint_id,
-            function_name=function_name,
-            function_tag=function_tag or "latest",
-            feature_analysis=feature_analysis,
-        )
-        # If other fields provided, validate that they are correspond to the existing model endpoint data
-        _model_endpoint_validations(
-            model_endpoint=model_endpoint,
-            model_path=model_path,
-            sample_set_statistics=sample_set_statistics,
-        )
+    if endpoint_id or function_name:
+        try:
+            model_endpoint = db_session.get_model_endpoint(
+                project=project,
+                name=model_endpoint_name,
+                endpoint_id=endpoint_id,
+                function_name=function_name,
+                function_tag=function_tag or "latest",
+                feature_analysis=feature_analysis,
+            )
+            # If other fields provided, validate that they are correspond to the existing model endpoint data
+            _model_endpoint_validations(
+                model_endpoint=model_endpoint,
+                model_path=model_path,
+                sample_set_statistics=sample_set_statistics,
+            )
 
-    except (mlrun.errors.MLRunNotFoundError, mlrun.errors.MLRunInvalidArgumentError):
-        # Create a new model endpoint with the provided details
-        pass
+        except mlrun.errors.MLRunNotFoundError:
+            # Create a new model endpoint with the provided details
+            pass
     if not model_endpoint:
         model_endpoint = _generate_model_endpoint(
             project=project,
@@ -123,6 +133,13 @@ def get_or_create_model_endpoint(
     return model_endpoint
 
 
+# TODO: Remove this in 1.12.0
+@deprecated(
+    version="1.10.0",
+    reason="This function is deprecated and will be removed in 1.12. "
+    "Instead, run a monitored serving function as a job",
+    category=FutureWarning,
+)
 def record_results(
     project: str,
     model_path: str,
@@ -144,8 +161,8 @@ def record_results(
     :param model_path:               The model Store path.
     :param model_endpoint_name:      If a new model endpoint is generated, the model endpoint name will be presented
                                      under this endpoint.
-    :param endpoint_id:              Model endpoint unique ID. If not exist in DB, will generate a new record based
-                                     on the provided `endpoint_id`.
+    :param endpoint_id:              Model endpoint unique ID. If not exist in DB, will generate a new record with a
+                                     newly generated ID.
     :param function_name:            If a new model endpoint is created, use this function name for generating the
                                      function URI.
     :param context:                  MLRun context. Note that the context is required generating the model endpoint.
@@ -236,6 +253,7 @@ def _model_endpoint_validations(
             key=model_obj.key,
             iter=model_obj.iter,
             tree=model_obj.tree,
+            uid=model_obj.uid,
         )
 
         # Enrich the uri schema with the store prefix
@@ -325,12 +343,15 @@ def _generate_model_endpoint(
 
     :return `mlrun.common.schemas.ModelEndpoint` object.
     """
+
     current_time = datetime_now()
     model_endpoint = mlrun.common.schemas.ModelEndpoint(
         metadata=mlrun.common.schemas.ModelEndpointMetadata(
             project=project,
             name=model_endpoint_name,
             endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.BATCH_EP,
+            # Due to backwards compatibility, old batch model endpoint will be analyzed as real time endpoint
+            mode=mlrun.common.schemas.model_monitoring.EndpointMode.REAL_TIME,
         ),
         spec=mlrun.common.schemas.ModelEndpointSpec(
             function_name=function_name or "function",

mlrun/model_monitoring/applications/base.py

@@ -647,7 +647,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
             else:
                 class_name = handler_to_class.split(".")[-1].split("::")[0]
 
-            job_name = mlrun.utils.normalize_name(class_name, verbose=False)
+            job_name = mlrun.utils.normalize_name(class_name)
 
         if not mm_constants.APP_NAME_REGEX.fullmatch(job_name):
             raise mlrun.errors.MLRunValueError(