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

mlrun/common/schemas/common.py

@@ -14,10 +14,12 @@
 
 import typing
 
-import pydantic
+import pydantic.v1
 
+import mlrun.errors
 
-class ImageBuilder(pydantic.BaseModel):
+
+class ImageBuilder(pydantic.v1.BaseModel):
     functionSourceCode: typing.Optional[str] = None  # noqa: N815
     codeEntryType: typing.Optional[str] = None  # noqa: N815
     codeEntryAttributes: typing.Optional[str] = None  # noqa: N815
@@ -40,4 +42,52 @@ class ImageBuilder(pydantic.BaseModel):
     source_code_target_dir: typing.Optional[str] = None
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
+
+
+class LabelsModel(pydantic.v1.BaseModel):
+    """
+    This class accepts either a dictionary, a list, or a string for filtering by labels.
+
+    :param labels:
+        - If a dictionary is provided, it should be in the format {'label_name': 'value'}.
+          The values can also be `None`, which will result in the format 'label_name' (without a value).
+          This will be converted to a list of strings in the format 'label_name=value'.
+        - If a list is provided, all items must be strings. Each string can either
+          be a simple label name (e.g., 'label1') or a key-value pair in the format
+          'label=value'.
+        - If a string is provided, it should be a comma-separated list of labels
+          (e.g., 'label1,label2').
+        - If no labels are specified, the default is an empty list.
+    """
+
+    labels: typing.Optional[
+        typing.Union[str, dict[str, typing.Optional[str]], list[str]]
+    ]
+
+    @pydantic.v1.validator("labels")
+    @classmethod
+    def validate(cls, labels) -> list[str]:
+        if labels is None:
+            return []
+
+        # If labels is a string, split it by commas
+        if isinstance(labels, str):
+            return [label.strip() for label in labels.split(",") if label.strip()]
+
+        if isinstance(labels, list):
+            if not all(isinstance(item, str) for item in labels):
+                raise mlrun.errors.MLRunValueError(
+                    "All items in the list must be strings."
+                )
+            return labels
+
+        if isinstance(labels, dict):
+            return [
+                f"{key}={value}" if value is not None else key
+                for key, value in labels.items()
+            ]
+
+        raise mlrun.errors.MLRunValueError(
+            "Invalid labels format. Must be a string, dictionary of strings, or a list of strings."
+        )

mlrun/common/schemas/constants.py

@@ -133,6 +133,21 @@ class RunPartitionByField(mlrun.common.types.StrEnum):
             )
 
 
+class ArtifactPartitionByField(mlrun.common.types.StrEnum):
+    name = "name"  # Supported for artifacts objects
+    project_and_name = "project_and_name"  # Supported for artifacts objects
+
+    def to_partition_by_db_field(self, db_cls):
+        if self.value == ArtifactPartitionByField.name:
+            return db_cls.key
+        elif self.value == ArtifactPartitionByField.project_and_name:
+            return db_cls.project, db_cls.key
+        else:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Unknown group by field: {self.value}"
+            )
+
+
 class SortField(mlrun.common.types.StrEnum):
     created = "created"
     updated = "updated"

mlrun/common/schemas/datastore_profile.py

@@ -15,7 +15,7 @@
 
 import typing
 
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 
 
 class DatastoreProfile(BaseModel):

mlrun/common/schemas/feature_store.py

@@ -14,7 +14,7 @@
 #
 from typing import Optional
 
-import pydantic
+import pydantic.v1
 
 from .auth import AuthorizationResourceTypes, Credentials
 from .object import (
@@ -27,7 +27,7 @@ from .object import (
 )
 
 
-class FeatureStoreBaseModel(pydantic.BaseModel):
+class FeatureStoreBaseModel(pydantic.v1.BaseModel):
     """
     Intermediate base class, in order to override pydantic's configuration, as per
     https://docs.pydantic.dev/1.10/usage/model_config/#change-behaviour-globally
@@ -43,7 +43,7 @@ class Feature(FeatureStoreBaseModel):
     labels: Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class Entity(FeatureStoreBaseModel):
@@ -52,17 +52,17 @@ class Entity(FeatureStoreBaseModel):
     labels: Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class FeatureSetSpec(ObjectSpec):
     entities: list[Entity] = []
     features: list[Feature] = []
-    engine: Optional[str] = pydantic.Field(default="storey")
+    engine: Optional[str] = pydantic.v1.Field(default="storey")
 
 
 class FeatureSet(FeatureStoreBaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.feature_set, const=True)
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.feature_set, const=True)
     metadata: ObjectMetadata
     spec: FeatureSetSpec
     status: ObjectStatus
@@ -155,7 +155,7 @@ class EntitiesOutput(FeatureStoreBaseModel):
 
 
 class FeatureVector(FeatureStoreBaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.feature_vector, const=True)
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.feature_vector, const=True)
     metadata: ObjectMetadata
     spec: ObjectSpec
     status: ObjectStatus
@@ -183,7 +183,7 @@ class DataSource(FeatureStoreBaseModel):
     path: str
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class DataTarget(FeatureStoreBaseModel):
@@ -192,7 +192,7 @@ class DataTarget(FeatureStoreBaseModel):
     path: Optional[str]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class FeatureSetIngestInput(FeatureStoreBaseModel):

mlrun/common/schemas/frontend_spec.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
@@ -43,20 +43,20 @@ class NuclioStreamsFeatureFlag(mlrun.common.types.StrEnum):
     disabled = "disabled"
 
 
-class FeatureFlags(pydantic.BaseModel):
+class FeatureFlags(pydantic.v1.BaseModel):
     project_membership: ProjectMembershipFeatureFlag
     authentication: AuthenticationFeatureFlag
     nuclio_streams: NuclioStreamsFeatureFlag
     preemption_nodes: PreemptionNodesFeatureFlag
 
 
-class ArtifactLimits(pydantic.BaseModel):
+class ArtifactLimits(pydantic.v1.BaseModel):
     max_chunk_size: int
     max_preview_size: int
     max_download_size: int
 
 
-class FrontendSpec(pydantic.BaseModel):
+class FrontendSpec(pydantic.v1.BaseModel):
     jobs_dashboard_url: typing.Optional[str]
     model_monitoring_dashboard_url: typing.Optional[str]
     abortable_function_kinds: list[str] = []

mlrun/common/schemas/function.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
@@ -90,42 +90,42 @@ class SecurityContextEnrichmentModes(mlrun.common.types.StrEnum):
     disabled = "disabled"
 
 
-class ImagePullSecret(pydantic.BaseModel):
+class ImagePullSecret(pydantic.v1.BaseModel):
     default: typing.Optional[str]
 
 
-class Pipelines(pydantic.BaseModel):
+class Pipelines(pydantic.v1.BaseModel):
     kfp_pod_user_unix_id: typing.Optional[int]
 
 
-class SecurityContext(pydantic.BaseModel):
+class SecurityContext(pydantic.v1.BaseModel):
     default: typing.Optional[str]
     enrichment_mode: typing.Optional[SecurityContextEnrichmentModes]
     enrichment_group_id: typing.Optional[int]
     pipelines: typing.Optional[Pipelines]
 
 
-class ServiceAccount(pydantic.BaseModel):
+class ServiceAccount(pydantic.v1.BaseModel):
     default: typing.Optional[str]
 
 
-class StateThresholds(pydantic.BaseModel):
+class StateThresholds(pydantic.v1.BaseModel):
     default: typing.Optional[dict[str, str]]
 
 
-class FunctionSpec(pydantic.BaseModel):
+class FunctionSpec(pydantic.v1.BaseModel):
     image_pull_secret: typing.Optional[ImagePullSecret]
     security_context: typing.Optional[SecurityContext]
     service_account: typing.Optional[ServiceAccount]
     state_thresholds: typing.Optional[StateThresholds]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class Function(pydantic.BaseModel):
+class Function(pydantic.v1.BaseModel):
     spec: typing.Optional[FunctionSpec]
     application: typing.Optional[dict[str, typing.Any]]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow

mlrun/common/schemas/hub.py

@@ -15,7 +15,7 @@
 from datetime import datetime, timezone
 from typing import Optional
 
-from pydantic import BaseModel, Extra, Field
+from pydantic.v1 import BaseModel, Extra, Field
 
 import mlrun.common.types
 import mlrun.config

mlrun/common/schemas/k8s.py

@@ -14,18 +14,18 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
 
-class ResourceSpec(pydantic.BaseModel):
+class ResourceSpec(pydantic.v1.BaseModel):
     cpu: typing.Optional[str]
     memory: typing.Optional[str]
     gpu: typing.Optional[str]
 
 
-class Resources(pydantic.BaseModel):
+class Resources(pydantic.v1.BaseModel):
     requests: ResourceSpec = ResourceSpec()
     limits: ResourceSpec = ResourceSpec()
 

mlrun/common/schemas/memory_reports.py

@@ -15,14 +15,14 @@
 
 import typing
 
-import pydantic
+import pydantic.v1
 
 
-class MostCommonObjectTypesReport(pydantic.BaseModel):
+class MostCommonObjectTypesReport(pydantic.v1.BaseModel):
     object_types: list[tuple[str, int]]
 
 
-class ObjectTypeReport(pydantic.BaseModel):
+class ObjectTypeReport(pydantic.v1.BaseModel):
     object_type: str
     sample_size: int
     start_index: typing.Optional[int]

mlrun/common/schemas/model_monitoring/__init__.py

@@ -14,6 +14,7 @@
 
 from .constants import (
     V3IO_MODEL_MONITORING_DB,
+    ApplicationEvent,
     ControllerPolicy,
     DriftStatus,
     EndpointType,
@@ -26,6 +27,7 @@ from .constants import (
     FunctionURI,
     MetricData,
     ModelEndpointMonitoringMetricType,
+    ModelEndpointSchema,
     ModelEndpointTarget,
     ModelEndpointTargetSchemas,
     ModelMonitoringMode,
@@ -36,7 +38,6 @@ from .constants import (
     ResultData,
     ResultKindApp,
     ResultStatusApp,
-    SchedulingKeys,
     SpecialApps,
     TDEngineSuperTables,
     TSDBTarget,

mlrun/common/schemas/model_monitoring/constants.py

@@ -29,6 +29,48 @@ class MonitoringStrEnum(StrEnum):
         return list(map(lambda c: c.value, cls))
 
 
+class ModelEndpointSchema(MonitoringStrEnum):
+    # metadata
+    UID = "uid"
+    PROJECT = "project"
+    ENDPOINT_TYPE = "endpoint_type"
+    NAME = "name"
+    CREATED = "created"
+    UPDATED = "updated"
+    LABELS = "labels"
+
+    # spec
+    FUNCTION_NAME = "function_name"
+    FUNCTION_TAG = "function_tag"
+    FUNCTION_UID = "function_uid"
+    MODEL_NAME = "model_name"
+    MODEL_DB_KEY = "model_db_key"
+    MODEL_TAG = "model_tag"
+    MODEL_CLASS = "model_class"
+    MODEL_UID = "model_uid"
+    FEATURE_NAMES = "feature_names"
+    LABEL_NAMES = "label_names"
+    FEATURE_STATS = "feature_stats"
+    MONITORING_FEATURE_SET_URI = "monitoring_feature_set_uri"
+    CHILDREN = "children"
+    CHILDREN_UIDS = "children_uids"
+    FUNCTION_URI = "function_uri"
+    MODEL_URI = "model_uri"
+
+    # status
+    STATE = "state"
+    MONITORING_MODE = "monitoring_mode"
+    FIRST_REQUEST = "first_request"
+
+    # status - operative
+    LAST_REQUEST = "last_request"
+    RESULT_STATUS = "result_status"
+    AVG_LATENCY = "avg_latency"
+    ERROR_COUNT = "error_count"
+    CURRENT_STATS = "current_stats"
+    DRIFT_MEASURES = "drift_measures"
+
+
 class EventFieldType:
     FUNCTION_URI = "function_uri"
     FUNCTION = "function"
@@ -40,6 +82,7 @@ class EventFieldType:
     TIMESTAMP = "timestamp"
     # `endpoint_id` is deprecated as a field in the model endpoint schema since 1.3.1, replaced by `uid`.
     ENDPOINT_ID = "endpoint_id"
+    ENDPOINT_NAME = "endpoint_name"
     UID = "uid"
     ENDPOINT_TYPE = "endpoint_type"
     REQUEST_ID = "request_id"
@@ -55,6 +98,8 @@ class EventFieldType:
     NAMED_PREDICTIONS = "named_predictions"
     ERROR_COUNT = "error_count"
     MODEL_ERROR = "model_error"
+    ERROR_TYPE = "error_type"
+    INFER_ERROR = "infer_error"
     ENTITIES = "entities"
     FIRST_REQUEST = "first_request"
     LAST_REQUEST = "last_request"
@@ -106,21 +151,24 @@ class ApplicationEvent:
     START_INFER_TIME = "start_infer_time"
     END_INFER_TIME = "end_infer_time"
     ENDPOINT_ID = "endpoint_id"
+    ENDPOINT_NAME = "endpoint_name"
     OUTPUT_STREAM_URI = "output_stream_uri"
 
 
 class WriterEvent(MonitoringStrEnum):
+    ENDPOINT_NAME = "endpoint_name"
     APPLICATION_NAME = "application_name"
     ENDPOINT_ID = "endpoint_id"
     START_INFER_TIME = "start_infer_time"
     END_INFER_TIME = "end_infer_time"
-    EVENT_KIND = "event_kind"  # metric or result
+    EVENT_KIND = "event_kind"  # metric or result or stats
     DATA = "data"
 
 
 class WriterEventKind(MonitoringStrEnum):
     METRIC = "metric"
     RESULT = "result"
+    STATS = "stats"
 
 
 class MetricData(MonitoringStrEnum):
@@ -134,7 +182,17 @@ class ResultData(MonitoringStrEnum):
     RESULT_KIND = "result_kind"
     RESULT_STATUS = "result_status"
     RESULT_EXTRA_DATA = "result_extra_data"
+
+
+class StatsData(MonitoringStrEnum):
+    STATS_NAME = "stats_name"
+    STATS = "stats"
+    TIMESTAMP = "timestamp"
+
+
+class StatsKind(MonitoringStrEnum):
     CURRENT_STATS = "current_stats"
+    DRIFT_MEASURES = "drift_measures"
 
 
 class EventLiveStats:
@@ -169,7 +227,6 @@ class TSDBTarget(MonitoringStrEnum):
 
 
 class ProjectSecretKeys:
-    ENDPOINT_STORE_CONNECTION = "MODEL_MONITORING_ENDPOINT_STORE_CONNECTION"
     ACCESS_KEY = "MODEL_MONITORING_ACCESS_KEY"
     STREAM_PATH = "STREAM_PATH"
     TSDB_CONNECTION = "TSDB_CONNECTION"
@@ -177,7 +234,6 @@ class ProjectSecretKeys:
     @classmethod
     def mandatory_secrets(cls):
         return [
-            cls.ENDPOINT_STORE_CONNECTION,
             cls.STREAM_PATH,
             cls.TSDB_CONNECTION,
         ]
@@ -194,13 +250,6 @@ class ModelMonitoringStoreKinds:
     EVENTS = "events"
 
 
-class SchedulingKeys:
-    LAST_ANALYZED = "last_analyzed"
-    ENDPOINT_ID = "endpoint_id"
-    APPLICATION_NAME = "application_name"
-    UID = "uid"
-
-
 class FileTargetKind:
     ENDPOINTS = "endpoints"
     EVENTS = "events"
@@ -209,14 +258,13 @@ class FileTargetKind:
     PARQUET = "parquet"
     APPS_PARQUET = "apps_parquet"
     LOG_STREAM = "log_stream"
-    APP_RESULTS = "app_results"
-    APP_METRICS = "app_metrics"
     MONITORING_SCHEDULES = "monitoring_schedules"
     MONITORING_APPLICATION = "monitoring_application"
     ERRORS = "errors"
+    STATS = "stats"
 
 
-class ModelMonitoringMode(str, Enum):
+class ModelMonitoringMode(StrEnum):
     enabled = "enabled"
     disabled = "disabled"
 
@@ -225,6 +273,11 @@ class EndpointType(IntEnum):
     NODE_EP = 1  # end point that is not a child of a router
     ROUTER = 2  # endpoint that is router
     LEAF_EP = 3  # end point that is a child of a router
+    BATCH_EP = 4  # endpoint that is representing an offline batch endpoint
+
+    @classmethod
+    def top_level_list(cls):
+        return [cls.NODE_EP, cls.ROUTER, cls.BATCH_EP]
 
 
 class MonitoringFunctionNames(MonitoringStrEnum):
@@ -244,6 +297,7 @@ class TDEngineSuperTables(MonitoringStrEnum):
     APP_RESULTS = "app_results"
     METRICS = "metrics"
     PREDICTIONS = "predictions"
+    ERRORS = "errors"
 
 
 @dataclass
@@ -364,7 +418,6 @@ class SpecialApps:
 
 _RESERVED_FUNCTION_NAMES = MonitoringFunctionNames.list() + [SpecialApps.MLRUN_INFRA]
 
-
 V3IO_MODEL_MONITORING_DB = "v3io"
 
 

mlrun/common/schemas/model_monitoring/grafana.py

@@ -14,7 +14,7 @@
 
 from typing import Optional, Union
 
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 
 import mlrun.common.types