mlrun 1.10.0rc28__tar.gz → 1.10.0rc29__tar.gz

{mlrun-1.10.0rc28/mlrun.egg-info → mlrun-1.10.0rc29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlrun
-Version: 1.10.0rc28
+Version: 1.10.0rc29
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv
@@ -22,7 +22,7 @@ Requires-Python: >=3.9, <3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: urllib3>=1.26.20
-Requires-Dist: v3io-frames>=0.10.15
+Requires-Dist: v3io-frames>=0.10.16
 Requires-Dist: GitPython>=3.1.41,~=3.1
 Requires-Dist: aiohttp~=3.11
 Requires-Dist: aiohttp-retry~=2.9

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/__init__.py

@@ -37,6 +37,7 @@ from .datastore import DataItem, ModelProvider, store_manager
 from .db import get_run_db
 from .errors import MLRunInvalidArgumentError, MLRunNotFoundError
 from .execution import MLClientCtx
+from .hub import get_hub_module, import_module
 from .model import RunObject, RunTemplate, new_task
 from .package import ArtifactType, DefaultPackager, Packager, handler
 from .projects import (

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/common/schemas/hub.py

@@ -134,3 +134,8 @@ class HubCatalog(BaseModel):
     kind: ObjectKind = Field(ObjectKind.hub_catalog, const=True)
     channel: str
     catalog: list[HubItem]
+
+
+class HubModuleType(mlrun.common.types.StrEnum):
+    generic = "generic"
+    monitoring_app = "monitoring-app"

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/datastore/azure_blob.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Iguazio
+# Copyright 2025 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import time
 from pathlib import Path
 from typing import Optional
@@ -30,6 +31,40 @@ from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 
 class AzureBlobStore(DataStore):
+    """
+    Azure Blob Storage datastore implementation.
+
+    Supports multiple URL schemas: az://, wasbs://, wasb://
+
+    Supported Connection String Formats:
+    ====================================
+
+    1. Account Key (Standard):
+       "DefaultEndpointsProtocol=https;AccountName=<account>;AccountKey=<key>;EndpointSuffix=core.windows.net"
+
+    2. SAS Token:
+       "BlobEndpoint=https://<account>.blob.core.windows.net/;SharedAccessSignature=<sas_token>"
+
+    3. Minimal BlobEndpoint:
+       "BlobEndpoint=https://<account>.blob.core.windows.net/;AccountName=<account>;AccountKey=<key>"
+
+    4. Custom Domain:
+       "BlobEndpoint=https://<account>.mydomain.com/;AccountName=<account>;AccountKey=<key>"
+
+    5. China/Government Cloud:
+       "DefaultEndpointsProtocol=https;AccountName=<account>;AccountKey=<key>;EndpointSuffix=core.chinacloudapi.cn"
+
+    6. Full Service Endpoints with SAS:
+       "BlobEndpoint=https://<account>.blob.core.windows.net/;QueueEndpoint=...;SharedAccessSignature=<sas>"
+
+    Authentication Methods:
+    ======================
+    - Account Key (connection_string or storage_options)
+    - SAS Token (connection_string or storage_options)
+    - OAuth/Azure AD (storage_options: client_id, client_secret, tenant_id)
+
+    """
+
     using_bucket = True
     max_concurrency = 100
     max_blocksize = 1024 * 1024 * 4
@@ -40,6 +75,12 @@ class AzureBlobStore(DataStore):
     def __init__(
         self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
     ):
+        # Extract container from WASBS endpoint before calling super()
+        self._container_from_endpoint = None
+        if schema in ["wasbs", "wasb"] and endpoint and "@" in endpoint:
+            # Handle container@host format
+            self._container_from_endpoint, endpoint = endpoint.split("@", 1)
+
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self._service_client = None
         self._storage_options = None
@@ -67,6 +108,34 @@ class AzureBlobStore(DataStore):
                 or self._get_secret_or_env("AZURE_STORAGE_SAS_TOKEN"),
                 credential=self._get_secret_or_env("credential"),
             )
+            # Use container extracted from WASBS endpoint during initialization
+            if self._container_from_endpoint:
+                res["container"] = self._container_from_endpoint
+
+            # For az:// URLs, endpoint contains the container name
+            if not res.get("container") and self.kind in ["az"]:
+                if container := getattr(self, "endpoint", None):
+                    res["container"] = container
+
+            # Last resort: For wasbs:// without container, check if connection string has BlobEndpoint with container
+            if not res.get("container") and self.kind in ["wasbs", "wasb"]:
+                connection_string = res.get("connection_string")
+                if connection_string and "BlobEndpoint=" in connection_string:
+                    # Try to extract container from BlobEndpoint URL
+                    for part in connection_string.split(";"):
+                        if part.startswith("BlobEndpoint="):
+                            blob_endpoint = part.split("=", 1)[1]
+                            # Parse URL to get path component
+                            from urllib.parse import urlparse
+
+                            parsed = urlparse(blob_endpoint)
+                            if parsed.path and parsed.path.strip("/"):
+                                # Extract first path segment as container
+                                path_parts = parsed.path.strip("/").split("/")
+                                if path_parts[0]:
+                                    res["container"] = path_parts[0]
+                                    break
+
             self._storage_options = self._sanitize_options(res)
         return self._storage_options
 
@@ -243,10 +312,12 @@ class AzureBlobStore(DataStore):
 
             for key in ["account_name", "account_key", "sas_token"]:
                 if parsed_value := parsed_credential.get(key):
-                    if key in st and st[key] != parsed_value:
+                    # Only check for conflicts if storage options has a non-empty value for this key
+                    existing_value = st.get(key)
+                    if existing_value and existing_value != parsed_value:
                         if key == "account_name":
                             raise mlrun.errors.MLRunInvalidArgumentError(
-                                f"Storage option for '{key}' is '{st[key]}', "
+                                f"Storage option for '{key}' is '{existing_value}', "
                                 f"which does not match corresponding connection string '{parsed_value}'"
                             )
                         else:
@@ -262,7 +333,8 @@ class AzureBlobStore(DataStore):
                 primary_url = primary_url[len("http://") :]
             if primary_url.startswith("https://"):
                 primary_url = primary_url[len("https://") :]
-            host = primary_url
+            # Remove any path components from the host
+            host = primary_url.split("/")[0]
         elif account_name:
             host = f"{account_name}.{service}.core.windows.net"
         else:
@@ -278,7 +350,10 @@ class AzureBlobStore(DataStore):
         # --- WASB + SAS (container-scoped key; no provider classes needed) ---
         if "sas_token" in st and st["sas_token"]:
             sas = st["sas_token"].lstrip("?")
-            if container := getattr(self, "endpoint", None) or st.get("container"):
+
+            container = st.get("container")
+
+            if container:
                 # fs.azure.sas.<container>.<account>.blob.core.windows.net = <sas>
                 res[f"spark.hadoop.fs.azure.sas.{container}.{host}"] = sas
 
@@ -295,11 +370,12 @@ class AzureBlobStore(DataStore):
         st = self.storage_options
         service = "blob"
 
-        container = getattr(self, "endpoint", None) or st.get("container")
+        container = st.get("container")
+
         if not container:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                "Container is required to build the WASB URL "
-                "(self.endpoint or storage_options['container'])."
+                "Container name is required to build the WASB URL. "
+                "Set storage_options['container'] or use datastore profile with container specified."
             )
 
         # Prefer host from connection string; else synthesize from account_name
@@ -308,18 +384,23 @@ class AzureBlobStore(DataStore):
         connection_string = st.get("connection_string")
 
         if connection_string:
-            primary_url, _, _ = parse_connection_str(
-                connection_string, credential=None, service=service
-            )
-            if primary_url.startswith("http://"):
-                primary_url = primary_url[len("http://") :]
-            if primary_url.startswith("https://"):
-                primary_url = primary_url[len("https://") :]
-            host = primary_url.rstrip("/")
-
+            with contextlib.suppress(Exception):
+                primary_url, _, _ = parse_connection_str(
+                    connection_string, credential=None, service=service
+                )
+                if primary_url.startswith("http://"):
+                    primary_url = primary_url[len("http://") :]
+                if primary_url.startswith("https://"):
+                    primary_url = primary_url[len("https://") :]
+                # Remove any path components from the host
+                host = primary_url.split("/")[0].rstrip("/")
         if not host and account_name:
             host = f"{account_name}.{service}.core.windows.net"
 
+        # For wasbs:// URLs where endpoint is already the host
+        if not host and self.kind in ["wasbs", "wasb"] and hasattr(self, "endpoint"):
+            host = getattr(self, "endpoint", None)
+
         if not host:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "account_name is required (or provide a connection_string) to build the WASB URL."

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/datastore/utils.py

@@ -320,7 +320,13 @@ def parse_url(url):
     parsed_url = urlparse(url)
     schema = parsed_url.scheme.lower()
     endpoint = parsed_url.hostname
-    if endpoint:
+
+    # Special handling for WASBS URLs to preserve container information
+    if schema in ["wasbs", "wasb"] and parsed_url.netloc and "@" in parsed_url.netloc:
+        # For wasbs://container@host format, preserve the full netloc as endpoint
+        # This allows the datastore to extract container later
+        endpoint = parsed_url.netloc
+    elif endpoint:
         # HACK - urlparse returns the hostname after in lower case - we want the original case:
         # the hostname is a substring of the netloc, in which it's the original case, so we find the indexes of the
         # hostname in the netloc and take it from there
@@ -331,8 +337,8 @@ def parse_url(url):
         endpoint = netloc[
             hostname_index_in_netloc : hostname_index_in_netloc + len(lower_hostname)
         ]
-    if parsed_url.port:
-        endpoint += f":{parsed_url.port}"
+        if parsed_url.port:
+            endpoint += f":{parsed_url.port}"
     return schema, endpoint, parsed_url
 
 

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/db/base.py

@@ -774,6 +774,7 @@ class RunDBInterface(ABC):
         item_name: Optional[str] = None,
         tag: Optional[str] = None,
         version: Optional[str] = None,
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/db/httpdb.py

@@ -4310,6 +4310,7 @@ class HTTPRunDB(RunDBInterface):
         item_name: Optional[str] = None,
         tag: Optional[str] = None,
         version: Optional[str] = None,
+        item_type: HubSourceType = HubSourceType.functions,
     ) -> list[mlrun.common.schemas.hub.IndexedHubSource]:
         """
         List hub sources in the MLRun DB.
@@ -4317,6 +4318,7 @@ class HTTPRunDB(RunDBInterface):
         :param item_name:   Sources contain this item will be returned, If not provided all sources will be returned.
         :param tag:         Item tag to filter by, supported only if item name is provided.
         :param version:     Item version to filter by, supported only if item name is provided and tag is not.
+        :param item_type:   Item type to filter by, supported only if item name is provided.
 
         :returns: List of indexed hub sources.
         """
@@ -4324,6 +4326,7 @@ class HTTPRunDB(RunDBInterface):
         params = {}
         if item_name:
             params["item-name"] = normalize_name(item_name)
+            params["item-type"] = item_type
         if tag:
             params["tag"] = tag
         if version:

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/db/nopdb.py

@@ -673,6 +673,7 @@ class NopDB(RunDBInterface):
         item_name: Optional[str] = None,
         tag: Optional[str] = None,
         version: Optional[str] = None,
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 

mlrun-1.10.0rc29/mlrun/hub/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from .module import get_hub_module, import_module

mlrun-1.10.0rc29/mlrun/hub/module.py

@@ -0,0 +1,166 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Optional, Union
+
+import yaml
+
+import mlrun.common.types
+import mlrun.utils
+from mlrun.common.schemas.hub import HubModuleType, HubSourceType
+from mlrun.run import function_to_module, get_object
+from mlrun.utils import logger
+
+from ..model import ModelObj
+from ..utils import extend_hub_uri_if_needed
+
+
+class HubModule(ModelObj):
+    def __init__(
+        self,
+        name: str,
+        kind: Union[HubModuleType, str],
+        version: Optional[str] = None,
+        description: Optional[str] = None,
+        categories: Optional[list] = None,
+        requirements: Optional[list] = None,
+        local_path: Optional[str] = None,
+        filename: Optional[str] = None,
+        example: Optional[str] = None,
+        url: Optional[str] = None,
+        **kwargs,  # catch all for unused args
+    ):
+        self.name: str = name
+        self.version: str = version
+        self.kind: HubModuleType = kind
+        self.description: str = description or ""
+        self.categories: list = categories or []
+        self.requirements: list = requirements or []
+        self.local_path: str = local_path or ""
+        self.filename: str = filename or name + ".py"
+        self.example: str = example or ""
+        self.url: str = url or ""
+
+    def module(self):
+        """Import the module after downloading its fils to local_path"""
+        try:
+            return function_to_module(code=self.filename, workdir=self.local_path)
+        except FileNotFoundError:
+            searched_path = self.local_path or "./"
+            raise FileNotFoundError(
+                f"Module file {self.filename} not found in {searched_path}, try calling download_module_files() first"
+            )
+
+    def install_requirements(self) -> None:
+        """
+        Install pip-style requirements (e.g., ["pandas>=2.0.0", "requests==2.31.0"]).
+        """
+        for req in self.requirements:
+            logger.info(f"Installing {req} ...")
+            try:
+                subprocess.run(
+                    [sys.executable, "-m", "pip", "install", req], check=True, text=True
+                )
+                logger.info(f"Installed {req}")
+            except subprocess.CalledProcessError as e:
+                logger.error(f"Failed to install {req} (exit code {e.returncode})")
+
+    def download_module_files(self, local_path=None, secrets=None):
+        """
+        Download this hub module’s files (code file and, if available, an example notebook) to the target directory
+        specified by `local_path` (defaults to the current working directory).
+        This path will be used later to locate the code file when importing the module.
+        """
+        self.local_path = self.verify_directory(path=local_path)
+        source_url, _ = extend_hub_uri_if_needed(
+            uri=self.url, asset_type=HubSourceType.modules, file=self.filename
+        )
+        self._download_object(
+            obj_url=source_url, target_name=self.filename, secrets=secrets
+        )
+        if self.example:
+            example_url, _ = extend_hub_uri_if_needed(
+                uri=self.url, asset_type=HubSourceType.modules, file=self.example
+            )
+            self._download_object(
+                obj_url=example_url, target_name=self.example, secrets=secrets
+            )
+
+    def _download_object(self, obj_url, target_name, secrets=None):
+        data = get_object(url=obj_url, secrets=secrets)
+        target_dir = self.local_path if self.local_path is not None else os.getcwd()
+        target_filepath = os.path.join(target_dir, target_name)
+        with open(target_filepath, "wb") as f:
+            f.write(data)
+
+    @staticmethod
+    def verify_directory(path: str) -> Path:
+        """Validate that the given path is an existing directory."""
+        if path:
+            path = Path(path)
+            if not path.exists():
+                raise ValueError(f"Path does not exist: {path}")
+            if not path.is_dir():
+                raise ValueError(f"Path is not a directory: {path}")
+        return path
+
+
+def get_hub_module(
+    url="", download_files=True, secrets=None, local_path=None
+) -> HubModule:
+    """
+    Get a hub-module object containing metadata of the requested module.
+    :param url: Hub module url in the format "hub://[<source>/]<item-name>[:<tag>]"
+    :param download_files: When set to True, the module files (code file and example notebook) are downloaded
+    :param secrets: Optional, credentials dict for DB or URL (s3, v3io, ...)
+    :param local_path: Path to target directory for the module files. Ignored when download_files is set to False.
+                       Defaults to the current working directory.
+
+    :return: HubModule object
+    """
+    item_yaml_url, is_hub_uri = extend_hub_uri_if_needed(
+        uri=url, asset_type=HubSourceType.modules, file="item.yaml"
+    )
+    if not is_hub_uri:
+        raise mlrun.errors.MLRunInvalidArgumentError("Not a valid hub URL")
+    yaml_obj = get_object(url=item_yaml_url, secrets=secrets)
+    item_yaml = yaml.safe_load(yaml_obj)
+    spec = item_yaml.pop("spec", {})
+    hub_module = HubModule(**item_yaml, **spec, url=url)
+    if download_files:
+        hub_module.download_module_files(local_path=local_path, secrets=secrets)
+    return hub_module
+
+
+def import_module(url="", install_requirements=False, secrets=None, local_path=None):
+    """
+    Import a module from the hub to use directly.
+    :param url: hub module url in the format "hub://[<source>/]<item-name>[:<tag>]"
+    :param install_requirements: when set to True, the module's requirements are installed.
+    :param secrets: optional, credentials dict for DB or URL (s3, v3io, ...)
+    :param local_path: Path to target directory for the module files (code and example notebook).
+                       Defaults to the current working directory.
+
+    :return: the module
+    """
+    hub_module: HubModule = get_hub_module(
+        url=url, download_files=True, secrets=secrets, local_path=local_path
+    )
+    if install_requirements:
+        hub_module.install_requirements()
+    return hub_module.module()

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/projects/operations.py

@@ -177,7 +177,12 @@ def run_function(
                             This ensures latest code changes are executed. This argument must be used in
                             conjunction with the local=True argument.
     :param output_path:     path to store artifacts, when running in a workflow this will be set automatically
-    :param retry:           Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
+    :param retry:           Retry configuration for the run, can be a dict or an instance of
+                            :py:class:`~mlrun.model.Retry`.
+                            The `count` field in the `Retry` object specifies the number of retry attempts.
+                            If `count=0`, the run will not be retried.
+                            The `backoff` field specifies the retry backoff strategy between retry attempts.
+                            If not provided, no backoff is applied.
     :return: MLRun RunObject or PipelineNodeWrapper
     """
     if artifact_path:

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/projects/project.py

@@ -4104,7 +4104,12 @@ class MlrunProject(ModelObj):
                                 This ensures latest code changes are executed. This argument must be used in
                                 conjunction with the local=True argument.
         :param output_path:     path to store artifacts, when running in a workflow this will be set automatically
-        :param retry:           Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
+        :param retry:           Retry configuration for the run, can be a dict or an instance of
+                                :py:class:`~mlrun.model.Retry`.
+                                The `count` field in the `Retry` object specifies the number of retry attempts.
+                                If `count=0`, the run will not be retried.
+                                The `backoff` field specifies the retry backoff strategy between retry attempts.
+                                If not provided, no backoff is applied.
         :return: MLRun RunObject or PipelineNodeWrapper
         """
         if artifact_path:

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/runtimes/base.py

@@ -376,7 +376,12 @@ class BaseRuntime(ModelObj):
                              This ensures latest code changes are executed. This argument must be used in
                              conjunction with the local=True argument.
         :param output_path:    Default artifact output path.
-        :param retry:          Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
+        :param retry:          Retry configuration for the run, can be a dict or an instance of
+                               :py:class:`~mlrun.model.Retry`.
+                               The `count` field in the `Retry` object specifies the number of retry attempts.
+                               If `count=0`, the run will not be retried.
+                               The `backoff` field specifies the retry backoff strategy between retry attempts.
+                               If not provided, no backoff is applied.
         :return: Run context object (RunObject) with run metadata, results and status
         """
         if artifact_path or out_path:

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/runtimes/nuclio/application/application.py

@@ -698,6 +698,12 @@ class ApplicationRuntime(RemoteRuntime):
         """
         # create a function that includes only the reverse proxy, without the application
 
+        if not mlrun.get_current_project(silent=True):
+            raise mlrun.errors.MLRunMissingProjectError(
+                "An active project is required to run deploy_reverse_proxy_image(). "
+                "Use `mlrun.get_or_create_project()` or set an active project first."
+            )
+
         reverse_proxy_func = mlrun.run.new_function(
             name="reverse-proxy-temp", kind="remote"
         )

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/runtimes/nuclio/serving.py

@@ -14,6 +14,7 @@
 import json
 import os
 import warnings
+from base64 import b64decode
 from copy import deepcopy
 from typing import Optional, Union
 
@@ -678,6 +679,17 @@ class ServingRuntime(RemoteRuntime):
                         f"function {function} is used in steps and is not defined, "
                         "use the .add_child_function() to specify child function attributes"
                     )
+        if isinstance(self.spec.graph, RootFlowStep) and any(
+            isinstance(step_type, mlrun.serving.states.ModelRunnerStep)
+            for step_type in self.spec.graph.steps.values()
+        ):
+            # Add import for LLModel
+            decoded_code = b64decode(self.spec.build.functionSourceCode).decode("utf-8")
+            import_llmodel_code = "\nfrom mlrun.serving.states import LLModel\n"
+            if import_llmodel_code not in decoded_code:
+                decoded_code += import_llmodel_code
+            encoded_code = mlrun.utils.helpers.encode_user_code(decoded_code)
+            self.spec.build.functionSourceCode = encoded_code
 
         # Handle secret processing before handling child functions, since secrets are transferred to them
         if self.spec.secret_sources:

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/runtimes/pod.py

@@ -109,6 +109,7 @@ class KubeResourceSpec(FunctionSpec):
         "track_models",
         "parameters",
         "graph",
+        "filename",
     ]
     _default_fields_to_strip = FunctionSpec._default_fields_to_strip + [
         "volumes",

{mlrun-1.10.0rc28 → mlrun-1.10.0rc29}/mlrun/serving/states.py

@@ -1641,6 +1641,9 @@ class ModelRunnerStep(MonitoredStep):
         model_runner_step.add_model(..., model_class=MyModel(name="my_model"))
         graph.to(model_runner_step)
 
+    Note when ModelRunnerStep is used in a graph, MLRun automatically imports
+    the default language model class (LLModel) during function deployment.
+
     :param model_selector: ModelSelector instance whose select() method will be used to select models to run on each
       event. Optional. If not passed, all models will be run.
     :param raise_exception:  If True, an error will be raised when model selection fails or if one of the models raised
@@ -1829,7 +1832,9 @@ class ModelRunnerStep(MonitoredStep):
         Add a Model to this ModelRunner.
 
         :param endpoint_name:       str, will identify the model in the ModelRunnerStep, and assign model endpoint name
-        :param model_class:         Model class name
+        :param model_class:         Model class name. If LLModel is chosen
+                                    (either by name `LLModel` or by its full path, e.g. mlrun.serving.states.LLModel),
+                                    outputs will be overridden with UsageResponseKeys fields.
         :param execution_mechanism: Parallel execution mechanism to be used to execute this model. Must be one of:
             * "process_pool" – To run in a separate process from a process pool. This is appropriate for CPU or GPU
                 intensive tasks as they would otherwise block the main process by holding Python's Global Interpreter
@@ -1902,7 +1907,8 @@ class ModelRunnerStep(MonitoredStep):
                 "Cannot provide a model object as argument to `model_class` and also provide `model_parameters`."
             )
         if type(model_class) is LLModel or (
-            isinstance(model_class, str) and model_class == LLModel.__name__
+            isinstance(model_class, str)
+            and model_class.split(".")[-1] == LLModel.__name__
         ):
             if outputs:
                 warnings.warn(
@@ -2848,7 +2854,9 @@ class RootFlowStep(FlowStep):
         """
         Add a shared model to the graph, this model will be available to all the ModelRunners in the graph
         :param name:                Name of the shared model (should be unique in the graph)
-        :param model_class:         Model class name
+        :param model_class:         Model class name. If LLModel is chosen
+                                    (either by name `LLModel` or by its full path, e.g. mlrun.serving.states.LLModel),
+                                    outputs will be overridden with UsageResponseKeys fields.
         :param execution_mechanism: Parallel execution mechanism to be used to execute this model. Must be one of:
             * "process_pool" – To run in a separate process from a process pool. This is appropriate for CPU or GPU
                 intensive tasks as they would otherwise block the main process by holding Python's Global Interpreter
@@ -2892,7 +2900,8 @@ class RootFlowStep(FlowStep):
                 "Cannot provide a model object as argument to `model_class` and also provide `model_parameters`."
             )
         if type(model_class) is LLModel or (
-            isinstance(model_class, str) and model_class == LLModel.__name__
+            isinstance(model_class, str)
+            and model_class.split(".")[-1] == LLModel.__name__
         ):
             if outputs:
                 warnings.warn(