mlrun 1.10.0rc24__py3-none-any.whl → 1.10.0rc26__py3-none-any.whl

mlrun/artifacts/llm_prompt.py

@@ -62,12 +62,19 @@ class LLMPromptArtifactSpec(ArtifactSpec):
             parent_uri=model_artifact.uri
             if isinstance(model_artifact, model_art.ModelArtifact)
             else model_artifact,
+            format=kwargs.pop("format", "") or "json",
             **kwargs,
         )
 
         self.prompt_template = prompt_template
         self.prompt_legend = prompt_legend
-        self.model_configuration = model_configuration
+        if model_configuration is not None and not isinstance(
+            model_configuration, dict
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "LLMPromptArtifact model_configuration must be a dictionary or None"
+            )
+        self.model_configuration = model_configuration or {}
         self.description = description
         self._model_artifact = (
             model_artifact

mlrun/common/model_monitoring/helpers.py

@@ -14,6 +14,7 @@
 
 import sys
 import typing
+from datetime import datetime
 
 import mlrun.common
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
@@ -24,6 +25,7 @@ BinCounts = typing.NewType("BinCounts", list[int])
 BinEdges = typing.NewType("BinEdges", list[float])
 
 _MAX_FLOAT = sys.float_info.max
+logger = mlrun.utils.create_logger(level="info", name="mm_helpers")
 
 
 def parse_model_endpoint_project_prefix(path: str, project_name: str):
@@ -87,3 +89,87 @@ def pad_features_hist(feature_stats: FeatureStats) -> None:
     for feature in feature_stats.values():
         if hist_key in feature:
             pad_hist(Histogram(feature[hist_key]))
+
+
+def get_model_endpoints_creation_task_status(
+    server,
+) -> tuple[
+    mlrun.common.schemas.BackgroundTaskState,
+    typing.Optional[datetime],
+    typing.Optional[set[str]],
+]:
+    background_task = None
+    background_task_state = mlrun.common.schemas.BackgroundTaskState.running
+    background_task_check_timestamp = None
+    model_endpoint_uids = None
+    try:
+        background_task = mlrun.get_run_db().get_project_background_task(
+            server.project, server.model_endpoint_creation_task_name
+        )
+        background_task_check_timestamp = mlrun.utils.now_date()
+        log_background_task_state(
+            server, background_task.status.state, background_task_check_timestamp
+        )
+        background_task_state = background_task.status.state
+    except mlrun.errors.MLRunNotFoundError:
+        logger.warning(
+            "Model endpoint creation task not found listing model endpoints",
+            project=server.project,
+            task_name=server.model_endpoint_creation_task_name,
+        )
+    if background_task is None:
+        model_endpoints = mlrun.get_run_db().list_model_endpoints(
+            project=server.project,
+            function_name=server.function_name,
+            function_tag=server.function_tag,
+            tsdb_metrics=False,
+        )
+        if model_endpoints:
+            model_endpoint_uids = {
+                endpoint.metadata.uid for endpoint in model_endpoints.endpoints
+            }
+            logger.info(
+                "Model endpoints found after background task not found, model monitoring will monitor "
+                "events",
+                project=server.project,
+                function_name=server.function_name,
+                function_tag=server.function_tag,
+                uids=model_endpoint_uids,
+            )
+            background_task_state = mlrun.common.schemas.BackgroundTaskState.succeeded
+        else:
+            logger.warning(
+                "Model endpoints not found after background task not found, model monitoring will not "
+                "monitor events",
+                project=server.project,
+                function_name=server.function_name,
+                function_tag=server.function_tag,
+            )
+            background_task_state = mlrun.common.schemas.BackgroundTaskState.failed
+    return background_task_state, background_task_check_timestamp, model_endpoint_uids
+
+
+def log_background_task_state(
+    server,
+    background_task_state: mlrun.common.schemas.BackgroundTaskState,
+    background_task_check_timestamp: typing.Optional[datetime],
+):
+    logger.info(
+        "Checking model endpoint creation task status",
+        task_name=server.model_endpoint_creation_task_name,
+    )
+    if (
+        background_task_state
+        in mlrun.common.schemas.BackgroundTaskState.terminal_states()
+    ):
+        logger.info(
+            f"Model endpoint creation task completed with state {background_task_state}"
+        )
+    else:  # in progress
+        logger.info(
+            f"Model endpoint creation task is still in progress with the current state: "
+            f"{background_task_state}. Events will not be monitored for the next "
+            f"{mlrun.mlconf.model_endpoint_monitoring.model_endpoint_creation_check_period} seconds",
+            function_name=server.function.name,
+            background_task_check_timestamp=background_task_check_timestamp.isoformat(),
+        )

mlrun/common/schemas/hub.py

@@ -15,6 +15,7 @@
 from datetime import datetime, timezone
 from typing import Optional
 
+import deepdiff
 from pydantic.v1 import BaseModel, Extra, Field
 
 import mlrun.common.types
@@ -36,9 +37,9 @@ class HubObjectMetadata(BaseModel):
         extra = Extra.allow
 
 
-# Currently only functions are supported. Will add more in the future.
 class HubSourceType(mlrun.common.types.StrEnum):
     functions = "functions"
+    modules = "modules"
 
 
 # Sources-related objects
@@ -46,7 +47,6 @@ class HubSourceSpec(ObjectSpec):
     path: str  # URL to base directory, should include schema (s3://, etc...)
     channel: str
     credentials: Optional[dict] = {}
-    object_type: HubSourceType = Field(HubSourceType.functions, const=True)
 
 
 class HubSource(BaseModel):
@@ -55,11 +55,11 @@ class HubSource(BaseModel):
     spec: HubSourceSpec
     status: Optional[ObjectStatus] = ObjectStatus(state="created")
 
-    def get_full_uri(self, relative_path):
-        return f"{self.spec.path}/{self.spec.object_type}/{self.spec.channel}/{relative_path}"
+    def get_full_uri(self, relative_path, object_type):
+        return f"{self.spec.path}/{object_type}/{self.spec.channel}/{relative_path}"
 
-    def get_catalog_uri(self):
-        return self.get_full_uri(mlrun.mlconf.hub.catalog_filename)
+    def get_catalog_uri(self, object_type):
+        return self.get_full_uri(mlrun.mlconf.hub.catalog_filename, object_type)
 
     @classmethod
     def generate_default_source(cls):
@@ -78,11 +78,23 @@ class HubSource(BaseModel):
             spec=HubSourceSpec(
                 path=mlrun.mlconf.hub.default_source.url,
                 channel=mlrun.mlconf.hub.default_source.channel,
-                object_type=HubSourceType(mlrun.mlconf.hub.default_source.object_type),
             ),
             status=ObjectStatus(state="created"),
         )
 
+    def diff(self, another_source: "HubSource") -> dict:
+        """
+        Compare this HubSource with another one.
+        Returns a dict of differences (metadata, spec, status).
+        """
+        exclude_paths = [
+            "root['metadata']['updated']",
+            "root['metadata']['created']",
+        ]
+        return deepdiff.DeepDiff(
+            self.dict(), another_source.dict(), exclude_paths=exclude_paths
+        )
+
 
 last_source_index = -1
 
@@ -94,21 +106,16 @@ class IndexedHubSource(BaseModel):
 
 # Item-related objects
 class HubItemMetadata(HubObjectMetadata):
-    source: HubSourceType = Field(HubSourceType.functions, const=True)
+    source: HubSourceType = HubSourceType.functions
     version: str
     tag: Optional[str]
 
     def get_relative_path(self) -> str:
-        if self.source == HubSourceType.functions:
-            # This is needed since the hub deployment script modifies the paths to use _ instead of -.
-            modified_name = self.name.replace("-", "_")
-            # Prefer using the tag if exists. Otherwise, use version.
-            version = self.tag or self.version
-            return f"{modified_name}/{version}/"
-        else:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"Bad source for hub item - {self.source}"
-            )
+        # This is needed since the hub deployment script modifies the paths to use _ instead of -.
+        modified_name = self.name.replace("-", "_")
+        # Prefer using the tag if exists. Otherwise, use version.
+        version = self.tag or self.version
+        return f"{modified_name}/{version}/"
 
 
 class HubItemSpec(ObjectSpec):

mlrun/common/schemas/model_monitoring/constants.py

@@ -331,6 +331,7 @@ class EndpointType(IntEnum):
 class EndpointMode(IntEnum):
     REAL_TIME = 0
     BATCH = 1
+    BATCH_LEGACY = 2  # legacy batch mode, used for endpoints created through the batch inference job
 
 
 class MonitoringFunctionNames(MonitoringStrEnum):

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -119,7 +119,7 @@ class ModelEndpointMetadata(ObjectMetadata, ModelEndpointParser):
     project: constr(regex=PROJECT_PATTERN)
     endpoint_type: EndpointType = EndpointType.NODE_EP
     uid: Optional[constr(regex=MODEL_ENDPOINT_ID_PATTERN)]
-    mode: EndpointMode = EndpointMode.REAL_TIME
+    mode: Optional[EndpointMode] = None
 
     @classmethod
     def mutable_fields(cls):
@@ -131,6 +131,15 @@ class ModelEndpointMetadata(ObjectMetadata, ModelEndpointParser):
             return str(v)
         return v
 
+    @validator("mode", pre=True, always=True)
+    def _set_mode_based_on_endpoint_type(cls, v, values):  # noqa: N805
+        if v is None:
+            if values.get("endpoint_type") == EndpointType.BATCH_EP:
+                return EndpointMode.BATCH_LEGACY
+            else:
+                return EndpointMode.REAL_TIME
+        return v
+
 
 class ModelEndpointSpec(ObjectSpec, ModelEndpointParser):
     model_class: Optional[str] = ""

mlrun/config.py

@@ -718,7 +718,6 @@ default_config = {
             "name": "default",
             "description": "MLRun global function hub",
             "url": "https://mlrun.github.io/marketplace",
-            "object_type": "functions",
             "channel": "master",
         },
     },
@@ -1000,9 +999,9 @@ class Config:
         )
 
     @staticmethod
-    def get_default_hub_source() -> str:
+    def get_default_hub_source_url_prefix(object_type) -> str:
         default_source = config.hub.default_source
-        return f"{default_source.url}/{default_source.object_type}/{default_source.channel}/"
+        return f"{default_source.url}/{object_type}/{default_source.channel}/"
 
     @staticmethod
     def decode_base64_config_and_load_to_object(

mlrun/datastore/__init__.py

@@ -43,7 +43,7 @@ import storey
 
 import mlrun.datastore.wasbfs
 from mlrun.datastore.datastore_profile import (
-    DatastoreProfileKafkaSource,
+    DatastoreProfileKafkaStream,
     DatastoreProfileKafkaTarget,
     DatastoreProfileV3io,
 )
@@ -123,7 +123,7 @@ def get_stream_pusher(stream_path: str, **kwargs):
         )
         if isinstance(
             datastore_profile,
-            (DatastoreProfileKafkaSource, DatastoreProfileKafkaTarget),
+            (DatastoreProfileKafkaStream, DatastoreProfileKafkaTarget),
         ):
             attributes = datastore_profile.attributes()
             brokers = attributes.pop("brokers", None)

mlrun/datastore/azure_blob.py

@@ -229,18 +229,25 @@ class AzureBlobStore(DataStore):
         st = self.storage_options
         service = "blob"
         primary_url = None
-        if st.get("connection_string"):
+
+        # Parse connection string (fills account_name/account_key or SAS)
+        connection_string = st.get("connection_string")
+        if connection_string:
             primary_url, _, parsed_credential = parse_connection_str(
-                st.get("connection_string"), credential=None, service=service
+                connection_string, credential=None, service=service
             )
-            for key in ["account_name", "account_key"]:
-                parsed_value = parsed_credential.get(key)
-                if parsed_value:
+
+            if isinstance(parsed_credential, str):
+                # SharedAccessSignature as raw string
+                parsed_credential = {"sas_token": parsed_credential}
+
+            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:
                         if key == "account_name":
                             raise mlrun.errors.MLRunInvalidArgumentError(
-                                f"Storage option for '{key}' is '{st[key]}',\
-                                    which does not match corresponding connection string '{parsed_value}'"
+                                f"Storage option for '{key}' is '{st[key]}', "
+                                f"which does not match corresponding connection string '{parsed_value}'"
                             )
                         else:
                             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -249,6 +256,7 @@ class AzureBlobStore(DataStore):
                     st[key] = parsed_value
 
         account_name = st.get("account_name")
+        # Derive host (prefer connection string primary URL)
         if primary_url:
             if primary_url.startswith("http://"):
                 primary_url = primary_url[len("http://") :]
@@ -258,48 +266,63 @@ class AzureBlobStore(DataStore):
         elif account_name:
             host = f"{account_name}.{service}.core.windows.net"
         else:
+            # nothing to configure yet
             return res
 
-        if "account_key" in st:
+        host = host.rstrip("/")
+
+        # Account key (optional; WASB supports it)
+        if "account_key" in st and st["account_key"]:
             res[f"spark.hadoop.fs.azure.account.key.{host}"] = st["account_key"]
 
-        if "client_secret" in st or "client_id" in st or "tenant_id" in st:
-            res[f"spark.hadoop.fs.azure.account.auth.type.{host}"] = "OAuth"
-            res[f"spark.hadoop.fs.azure.account.oauth.provider.type.{host}"] = (
-                "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider"
-            )
-            if "client_id" in st:
-                res[f"spark.hadoop.fs.azure.account.oauth2.client.id.{host}"] = st[
-                    "client_id"
-                ]
-            if "client_secret" in st:
-                res[f"spark.hadoop.fs.azure.account.oauth2.client.secret.{host}"] = st[
-                    "client_secret"
-                ]
-            if "tenant_id" in st:
-                tenant_id = st["tenant_id"]
-                res[f"spark.hadoop.fs.azure.account.oauth2.client.endpoint.{host}"] = (
-                    f"https://login.microsoftonline.com/{tenant_id}/oauth2/token"
-                )
+        # --- 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"):
+                # fs.azure.sas.<container>.<account>.blob.core.windows.net = <sas>
+                res[f"spark.hadoop.fs.azure.sas.{container}.{host}"] = sas
 
-        if "sas_token" in st:
-            res[f"spark.hadoop.fs.azure.account.auth.type.{host}"] = "SAS"
-            res[f"spark.hadoop.fs.azure.sas.token.provider.type.{host}"] = (
-                "org.apache.hadoop.fs.azurebfs.sas.FixedSASTokenProvider"
-            )
-            res[f"spark.hadoop.fs.azure.sas.fixed.token.{host}"] = st["sas_token"]
+            else:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Container name is required for WASB SAS. "
+                    "Set self.endpoint or storage_options['container']."
+                )
         return res
 
     @property
     def spark_url(self):
-        spark_options = self.get_spark_options()
-        url = f"wasbs://{self.endpoint}"
-        prefix = "spark.hadoop.fs.azure.account.key."
-        if spark_options:
-            for key in spark_options:
-                if key.startswith(prefix):
-                    account_key = key[len(prefix) :]
-                    if not url.endswith(account_key):
-                        url += f"@{account_key}"
-                    break
-        return url
+        # Build: wasbs://<container>@<host>
+        st = self.storage_options
+        service = "blob"
+
+        container = getattr(self, "endpoint", None) or st.get("container")
+        if not container:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Container is required to build the WASB URL "
+                "(self.endpoint or storage_options['container'])."
+            )
+
+        # Prefer host from connection string; else synthesize from account_name
+        host = None
+        account_name = st.get("account_name")
+        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("/")
+
+        if not host and account_name:
+            host = f"{account_name}.{service}.core.windows.net"
+
+        if not host:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "account_name is required (or provide a connection_string) to build the WASB URL."
+            )
+
+        return f"wasbs://{container}@{host}"

mlrun/datastore/datastore_profile.py

@@ -19,6 +19,7 @@ import typing
 from urllib.parse import ParseResult, urlparse
 
 import pydantic.v1
+from deprecated import deprecated
 from mergedeep import merge
 
 import mlrun
@@ -138,6 +139,15 @@ class ConfigProfile(DatastoreProfile):
         return res
 
 
+# TODO: Remove in 1.12.0
+@deprecated(
+    version="1.10.0",
+    reason=(
+        "This class is deprecated from mlrun 1.10.0, and will be removed in 1.12.0. "
+        "Use `DatastoreProfileKafkaStream` instead."
+    ),
+    category=FutureWarning,
+)
 class DatastoreProfileKafkaTarget(DatastoreProfile):
     type: str = pydantic.v1.Field("kafka_target")
     _private_attributes = "kwargs_private"
@@ -158,8 +168,8 @@ class DatastoreProfileKafkaTarget(DatastoreProfile):
         return attributes
 
 
-class DatastoreProfileKafkaSource(DatastoreProfile):
-    type: str = pydantic.v1.Field("kafka_source")
+class DatastoreProfileKafkaStream(DatastoreProfile):
+    type: str = pydantic.v1.Field("kafka_stream")
     _private_attributes = ("kwargs_private", "sasl_user", "sasl_pass")
     brokers: typing.Union[str, list[str]]
     topics: typing.Union[str, list[str]]
@@ -198,6 +208,19 @@ class DatastoreProfileKafkaSource(DatastoreProfile):
         return attributes
 
 
+# TODO: Remove in 1.12.0
+@deprecated(
+    version="1.10.0",
+    reason=(
+        "This class is deprecated from mlrun 1.10.0, and will be removed in 1.12.0. "
+        "Use `DatastoreProfileKafkaStream` instead."
+    ),
+    category=FutureWarning,
+)
+class DatastoreProfileKafkaSource(DatastoreProfileKafkaStream):
+    type: str = pydantic.v1.Field("kafka_source")
+
+
 class DatastoreProfileV3io(DatastoreProfile):
     type: str = pydantic.v1.Field("v3io")
     v3io_access_key: typing.Optional[str] = None
@@ -232,7 +255,7 @@ class DatastoreProfileS3(DatastoreProfile):
         if self.secret_key:
             res["AWS_SECRET_ACCESS_KEY"] = self.secret_key
         if self.endpoint_url:
-            res["S3_ENDPOINT_URL"] = self.endpoint_url
+            res["AWS_ENDPOINT_URL_S3"] = self.endpoint_url
         if self.force_non_anonymous:
             res["S3_NON_ANONYMOUS"] = self.force_non_anonymous
         if self.profile_name:
@@ -333,7 +356,9 @@ class DatastoreProfileGCS(DatastoreProfile):
             #  in gcs the path after schema is starts with bucket, wherefore it should not start with "/".
             subpath = subpath[1:]
         if self.bucket:
-            return f"gcs://{self.bucket}/{subpath}"
+            return (
+                f"gcs://{self.bucket}/{subpath}" if subpath else f"gcs://{self.bucket}"
+            )
         else:
             return f"gcs://{subpath}"
 
@@ -370,7 +395,11 @@ class DatastoreProfileAzureBlob(DatastoreProfile):
             #  in azure the path after schema is starts with container, wherefore it should not start with "/".
             subpath = subpath[1:]
         if self.container:
-            return f"az://{self.container}/{subpath}"
+            return (
+                f"az://{self.container}/{subpath}"
+                if subpath
+                else f"az://{self.container}"
+            )
         else:
             return f"az://{subpath}"
 
@@ -518,6 +547,7 @@ _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "basic": DatastoreProfileBasic,
     "kafka_target": DatastoreProfileKafkaTarget,
     "kafka_source": DatastoreProfileKafkaSource,
+    "kafka_stream": DatastoreProfileKafkaStream,
     "dbfs": DatastoreProfileDBFS,
     "gcs": DatastoreProfileGCS,
     "az": DatastoreProfileAzureBlob,