mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0rc2__py3-none-any.whl

mlrun/datastore/targets.py

@@ -396,7 +396,7 @@ class BaseStoreTarget(DataTargetBase):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         columns=None,
         partitioned: bool = False,
@@ -405,8 +405,8 @@ class BaseStoreTarget(DataTargetBase):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        schema: Optional[dict[str, Any]] = None,
         credentials_prefix=None,
     ):
         super().__init__(
@@ -834,16 +834,16 @@ class ParquetTarget(BaseStoreTarget):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         columns=None,
-        partitioned: bool = None,
+        partitioned: Optional[bool] = None,
         key_bucketing_number: Optional[int] = None,
         partition_cols: Optional[list[str]] = None,
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = 10000,
         flush_after_seconds: Optional[int] = 900,
-        storage_options: dict[str, str] = None,
+        storage_options: Optional[dict[str, str]] = None,
     ):
         self.path = path
         if partitioned is None:
@@ -1136,8 +1136,7 @@ class CSVTarget(BaseStoreTarget):
         import pyspark.sql.functions as funcs
 
         for col_name, col_type in df.dtypes:
-            # covers TimestampType and TimestampNTZType, which was added in PySpark 3.4.0
-            if col_type.startswith("timestamp"):
+            if col_type == "timestamp":
                 # df.write.csv saves timestamps with millisecond precision, but we want microsecond precision
                 # for compatibility with storey.
                 df = df.withColumn(
@@ -1200,7 +1199,7 @@ class SnowflakeTarget(BaseStoreTarget):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         columns=None,
         partitioned: bool = False,
@@ -1209,15 +1208,15 @@ class SnowflakeTarget(BaseStoreTarget):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        schema: Optional[dict[str, Any]] = None,
         credentials_prefix=None,
-        url: str = None,
-        user: str = None,
-        db_schema: str = None,
-        database: str = None,
-        warehouse: str = None,
-        table_name: str = None,
+        url: Optional[str] = None,
+        user: Optional[str] = None,
+        db_schema: Optional[str] = None,
+        database: Optional[str] = None,
+        warehouse: Optional[str] = None,
+        table_name: Optional[str] = None,
     ):
         attributes = attributes or {}
         if url:
@@ -1904,7 +1903,7 @@ class SQLTarget(BaseStoreTarget):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         partitioned: bool = False,
         key_bucketing_number: Optional[int] = None,
@@ -1912,16 +1911,16 @@ class SQLTarget(BaseStoreTarget):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        db_url: str = None,
-        table_name: str = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        db_url: Optional[str] = None,
+        table_name: Optional[str] = None,
+        schema: Optional[dict[str, Any]] = None,
         primary_key_column: str = "",
         if_exists: str = "append",
         create_table: bool = False,
         # create_according_to_data: bool = False,
         varchar_len: int = 50,
-        parse_dates: list[str] = None,
+        parse_dates: Optional[list[str]] = None,
     ):
         """
         Write to SqlDB as output target for a flow.

mlrun/datastore/utils.py

@@ -26,7 +26,7 @@ import mlrun.datastore
 
 
 def parse_kafka_url(
-    url: str, brokers: typing.Union[list, str] = None
+    url: str, brokers: typing.Optional[typing.Union[list, str]] = None
 ) -> tuple[str, list]:
     """Generating Kafka topic and adjusting a list of bootstrap servers.
 
@@ -71,7 +71,7 @@ def upload_tarball(source_dir, target, secrets=None):
 
 def filter_df_start_end_time(
     df: typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]],
-    time_column: str = None,
+    time_column: typing.Optional[str] = None,
     start_time: pd.Timestamp = None,
     end_time: pd.Timestamp = None,
 ) -> typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]]:

mlrun/datastore/v3io.py

@@ -14,6 +14,7 @@
 
 import time
 from datetime import datetime
+from typing import Optional
 
 import fsspec
 import v3io
@@ -33,7 +34,9 @@ V3IO_DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024 * 10
 
 
 class V3ioStore(DataStore):
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self.endpoint = self.endpoint or mlrun.mlconf.v3io_api
 

mlrun/datastore/wasbfs/fs.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 from urllib.parse import urlparse
 
 from fsspec import AbstractFileSystem
@@ -22,23 +23,23 @@ class WasbFS(AbstractFileSystem):
 
     def __init__(
         self,
-        account_name: str = None,
-        account_key: str = None,
-        connection_string: str = None,
-        credential: str = None,
-        sas_token: str = None,
+        account_name: Optional[str] = None,
+        account_key: Optional[str] = None,
+        connection_string: Optional[str] = None,
+        credential: Optional[str] = None,
+        sas_token: Optional[str] = None,
         request_session=None,
-        socket_timeout: int = None,
-        blocksize: int = None,
-        client_id: str = None,
-        client_secret: str = None,
-        tenant_id: str = None,
+        socket_timeout: Optional[int] = None,
+        blocksize: Optional[int] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        tenant_id: Optional[str] = None,
         anon: bool = True,
-        location_mode: str = None,
+        location_mode: Optional[str] = None,
         loop=None,
         asynchronous: bool = False,
         default_fill_cache: bool = True,
-        default_cache_type: str = None,
+        default_cache_type: Optional[str] = None,
         **kwargs,
     ):
         from adlfs import AzureBlobFileSystem

mlrun/db/base.py

@@ -14,7 +14,7 @@
 
 import datetime
 from abc import ABC, abstractmethod
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 from deprecated import deprecated
 
@@ -23,6 +23,7 @@ import mlrun.common
 import mlrun.common.formatters
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.model_endpoints as mm_endpoints
 import mlrun.model_monitoring
 
 
@@ -73,7 +74,7 @@ class RunDBInterface(ABC):
         name: Optional[str] = None,
         uid: Optional[Union[str, list[str]]] = None,
         project: Optional[str] = None,
-        labels: Optional[Union[str, list[str]]] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         state: Optional[
             mlrun.common.runtimes.constants.RunStates
         ] = None,  # Backward compatibility
@@ -81,10 +82,10 @@ class RunDBInterface(ABC):
         sort: bool = True,
         last: int = 0,
         iter: bool = False,
-        start_time_from: datetime.datetime = None,
-        start_time_to: datetime.datetime = None,
-        last_update_time_from: datetime.datetime = None,
-        last_update_time_to: datetime.datetime = None,
+        start_time_from: Optional[datetime.datetime] = None,
+        start_time_to: Optional[datetime.datetime] = None,
+        last_update_time_from: Optional[datetime.datetime] = None,
+        last_update_time_to: Optional[datetime.datetime] = None,
         partition_by: Union[mlrun.common.schemas.RunPartitionByField, str] = None,
         rows_per_partition: int = 1,
         partition_sort_by: Union[mlrun.common.schemas.SortField, str] = None,
@@ -96,12 +97,30 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    @abstractmethod
+    def paginated_list_runs(
+        self,
+        *args,
+        page: Optional[int] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+        **kwargs,
+    ):
+        pass
+
     @abstractmethod
     def del_run(self, uid, project="", iter=0):
         pass
 
     @abstractmethod
-    def del_runs(self, name="", project="", labels=None, state="", days_ago=0):
+    def del_runs(
+        self,
+        name: str = "",
+        project: str = "",
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
+        state: Optional[mlrun.common.runtimes.constants.RunStates] = None,
+        days_ago: int = 0,
+    ):
         pass
 
     @abstractmethod
@@ -126,19 +145,30 @@ class RunDBInterface(ABC):
     @abstractmethod
     def list_artifacts(
         self,
-        name="",
-        project="",
-        tag="",
-        labels=None,
+        name: Optional[str] = "",
+        project: Optional[str] = "",
+        tag: Optional[str] = "",
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         since=None,
         until=None,
-        iter: int = None,
+        iter: Optional[int] = None,
         best_iteration: bool = False,
-        kind: str = None,
+        kind: Optional[str] = None,
         category: Union[str, mlrun.common.schemas.ArtifactCategories] = None,
-        tree: str = None,
+        tree: Optional[str] = None,
         format_: mlrun.common.formatters.ArtifactFormat = mlrun.common.formatters.ArtifactFormat.full,
-        limit: int = None,
+        limit: Optional[int] = None,
+    ):
+        pass
+
+    @abstractmethod
+    def paginated_list_artifacts(
+        self,
+        *args,
+        page: Optional[int] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+        **kwargs,
     ):
         pass
 
@@ -153,13 +183,19 @@ class RunDBInterface(ABC):
         deletion_strategy: mlrun.common.schemas.artifact.ArtifactsDeletionStrategies = (
             mlrun.common.schemas.artifact.ArtifactsDeletionStrategies.metadata_only
         ),
-        secrets: dict = None,
+        secrets: Optional[dict] = None,
         iter=None,
     ):
         pass
 
     @abstractmethod
-    def del_artifacts(self, name="", project="", tag="", labels=None):
+    def del_artifacts(
+        self,
+        name: Optional[str] = "",
+        project: Optional[str] = "",
+        tag: Optional[str] = "",
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
+    ):
         pass
 
     @abstractmethod
@@ -176,7 +212,26 @@ class RunDBInterface(ABC):
 
     @abstractmethod
     def list_functions(
-        self, name=None, project="", tag="", labels=None, since=None, until=None
+        self,
+        name: Optional[str] = None,
+        project: Optional[str] = None,
+        tag: Optional[str] = None,
+        kind: Optional[str] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
+        format_: mlrun.common.formatters.FunctionFormat = mlrun.common.formatters.FunctionFormat.full,
+        since: Optional[datetime.datetime] = None,
+        until: Optional[datetime.datetime] = None,
+    ):
+        pass
+
+    @abstractmethod
+    def paginated_list_functions(
+        self,
+        *args,
+        page: Optional[int] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+        **kwargs,
     ):
         pass
 
@@ -254,6 +309,14 @@ class RunDBInterface(ABC):
             kind="artifact", identifiers=artifact_identifiers
         )
 
+    def get_model_endpoint_monitoring_metrics(
+        self,
+        project: str,
+        endpoint_id: str,
+        type: Literal["results", "metrics", "all"] = "all",
+    ) -> list[mm_endpoints.ModelEndpointMonitoringMetric]:
+        pass
+
     @abstractmethod
     def delete_project(
         self,
@@ -289,9 +352,9 @@ class RunDBInterface(ABC):
     @abstractmethod
     def list_projects(
         self,
-        owner: str = None,
+        owner: Optional[str] = None,
         format_: mlrun.common.formatters.ProjectFormat = mlrun.common.formatters.ProjectFormat.name_only,
-        labels: list[str] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         state: mlrun.common.schemas.ProjectState = None,
     ) -> mlrun.common.schemas.ProjectsOutput:
         pass
@@ -319,7 +382,11 @@ class RunDBInterface(ABC):
 
     @abstractmethod
     def get_feature_set(
-        self, name: str, project: str = "", tag: str = None, uid: str = None
+        self,
+        name: str,
+        project: str = "",
+        tag: Optional[str] = None,
+        uid: Optional[str] = None,
     ) -> dict:
         pass
 
@@ -333,10 +400,10 @@ class RunDBInterface(ABC):
     def list_features(
         self,
         project: str,
-        name: str = None,
-        tag: str = None,
-        entities: list[str] = None,
-        labels: list[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,
     ) -> mlrun.common.schemas.FeaturesOutput:
         pass
 
@@ -344,10 +411,10 @@ class RunDBInterface(ABC):
     def list_features_v2(
         self,
         project: str,
-        name: str = None,
-        tag: str = None,
-        entities: list[str] = None,
-        labels: list[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,
     ) -> mlrun.common.schemas.FeaturesOutputV2:
         pass
 
@@ -361,9 +428,9 @@ class RunDBInterface(ABC):
     def list_entities(
         self,
         project: str,
-        name: str = None,
-        tag: str = None,
-        labels: list[str] = None,
+        name: Optional[str] = None,
+        tag: Optional[str] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
     ) -> mlrun.common.schemas.EntitiesOutput:
         pass
 
@@ -371,9 +438,9 @@ class RunDBInterface(ABC):
     def list_entities_v2(
         self,
         project: str,
-        name: str = None,
-        tag: str = None,
-        labels: list[str] = None,
+        name: Optional[str] = None,
+        tag: Optional[str] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
     ) -> mlrun.common.schemas.EntitiesOutputV2:
         pass
 
@@ -381,12 +448,12 @@ class RunDBInterface(ABC):
     def list_feature_sets(
         self,
         project: str = "",
-        name: str = None,
-        tag: str = None,
-        state: str = None,
-        entities: list[str] = None,
-        features: list[str] = None,
-        labels: list[str] = None,
+        name: Optional[str] = None,
+        tag: Optional[str] = None,
+        state: Optional[str] = None,
+        entities: Optional[list[str]] = None,
+        features: Optional[list[str]] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         partition_by: Union[
             mlrun.common.schemas.FeatureStorePartitionByField, str
         ] = None,
@@ -442,7 +509,11 @@ class RunDBInterface(ABC):
 
     @abstractmethod
     def get_feature_vector(
-        self, name: str, project: str = "", tag: str = None, uid: str = None
+        self,
+        name: str,
+        project: str = "",
+        tag: Optional[str] = None,
+        uid: Optional[str] = None,
     ) -> dict:
         pass
 
@@ -450,10 +521,10 @@ class RunDBInterface(ABC):
     def list_feature_vectors(
         self,
         project: str = "",
-        name: str = None,
-        tag: str = None,
-        state: str = None,
-        labels: list[str] = None,
+        name: Optional[str] = None,
+        tag: Optional[str] = None,
+        state: Optional[str] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         partition_by: Union[
             mlrun.common.schemas.FeatureStorePartitionByField, str
         ] = None,
@@ -499,12 +570,12 @@ class RunDBInterface(ABC):
     def get_pipeline(
         self,
         run_id: str,
-        namespace: str = None,
+        namespace: Optional[str] = None,
         timeout: int = 30,
         format_: Union[
             str, mlrun.common.formatters.PipelineFormat
         ] = mlrun.common.formatters.PipelineFormat.summary,
-        project: str = None,
+        project: Optional[str] = None,
     ):
         pass
 
@@ -512,14 +583,14 @@ class RunDBInterface(ABC):
     def list_pipelines(
         self,
         project: str,
-        namespace: str = None,
+        namespace: Optional[str] = None,
         sort_by: str = "",
         page_token: str = "",
         filter_: str = "",
         format_: Union[
             str, mlrun.common.formatters.PipelineFormat
         ] = mlrun.common.formatters.PipelineFormat.metadata_only,
-        page_size: int = None,
+        page_size: Optional[int] = None,
     ) -> mlrun.common.schemas.PipelinesOutput:
         pass
 
@@ -530,7 +601,7 @@ class RunDBInterface(ABC):
         provider: Union[
             str, mlrun.common.schemas.SecretProviderName
         ] = mlrun.common.schemas.SecretProviderName.kubernetes,
-        secrets: dict = None,
+        secrets: Optional[dict] = None,
     ):
         pass
 
@@ -542,7 +613,7 @@ class RunDBInterface(ABC):
         provider: Union[
             str, mlrun.common.schemas.SecretProviderName
         ] = mlrun.common.schemas.SecretProviderName.kubernetes,
-        secrets: list[str] = None,
+        secrets: Optional[list[str]] = None,
     ) -> mlrun.common.schemas.SecretsData:
         pass
 
@@ -553,7 +624,7 @@ class RunDBInterface(ABC):
         provider: Union[
             str, mlrun.common.schemas.SecretProviderName
         ] = mlrun.common.schemas.SecretProviderName.kubernetes,
-        token: str = None,
+        token: Optional[str] = None,
     ) -> mlrun.common.schemas.SecretKeysData:
         pass
 
@@ -564,7 +635,7 @@ class RunDBInterface(ABC):
         provider: Union[
             str, mlrun.common.schemas.SecretProviderName
         ] = mlrun.common.schemas.SecretProviderName.kubernetes,
-        secrets: list[str] = None,
+        secrets: Optional[list[str]] = None,
     ):
         pass
 
@@ -575,7 +646,7 @@ class RunDBInterface(ABC):
         provider: Union[
             str, mlrun.common.schemas.SecretProviderName
         ] = mlrun.common.schemas.SecretProviderName.vault,
-        secrets: dict = None,
+        secrets: Optional[dict] = None,
     ):
         pass
 
@@ -602,11 +673,13 @@ class RunDBInterface(ABC):
         project: str,
         model: Optional[str] = None,
         function: Optional[str] = None,
-        labels: list[str] = None,
+        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         start: str = "now-1h",
         end: str = "now",
         metrics: Optional[list[str]] = None,
-    ):
+        top_level: bool = False,
+        uids: Optional[list[str]] = None,
+    ) -> list[mlrun.model_monitoring.model_endpoint.ModelEndpoint]:
         pass
 
     @abstractmethod
@@ -665,8 +738,8 @@ class RunDBInterface(ABC):
     def get_hub_catalog(
         self,
         source_name: str,
-        version: str = None,
-        tag: str = None,
+        version: Optional[str] = None,
+        tag: Optional[str] = None,
         force_refresh: bool = False,
     ):
         pass
@@ -676,7 +749,7 @@ class RunDBInterface(ABC):
         self,
         source_name: str,
         item_name: str,
-        version: str = None,
+        version: Optional[str] = None,
         tag: str = "latest",
         force_refresh: bool = False,
     ):
@@ -744,6 +817,7 @@ class RunDBInterface(ABC):
         alert_name: str,
         alert_data: Union[dict, mlrun.alerts.alert.AlertConfig],
         project="",
+        force_reset: bool = False,
     ):
         pass
 
@@ -771,6 +845,35 @@ class RunDBInterface(ABC):
     def list_alert_templates(self):
         pass
 
+    @abstractmethod
+    def list_alert_activations(
+        self,
+        project: Optional[str] = None,
+        name: Optional[str] = None,
+        since: Optional[datetime.datetime] = None,
+        until: Optional[datetime.datetime] = None,
+        entity: Optional[str] = None,
+        severity: Optional[
+            list[Union[mlrun.common.schemas.alert.AlertSeverity, str]]
+        ] = None,
+        entity_kind: Optional[
+            Union[mlrun.common.schemas.alert.EventEntityKind, str]
+        ] = None,
+        event_kind: Optional[Union[mlrun.common.schemas.alert.EventKind, str]] = None,
+    ):
+        pass
+
+    @abstractmethod
+    def paginated_list_alert_activations(
+        self,
+        *args,
+        page: Optional[int] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+        **kwargs,
+    ):
+        pass
+
     @abstractmethod
     def get_builder_status(
         self,
@@ -779,6 +882,7 @@ class RunDBInterface(ABC):
         logs: bool = True,
         last_log_timestamp: float = 0.0,
         verbose: bool = False,
+        events_offset: int = 0,
     ):
         pass
 
@@ -805,7 +909,7 @@ class RunDBInterface(ABC):
         self,
         notification_objects: list[mlrun.model.Notification],
         run_uid: str,
-        project: str = None,
+        project: Optional[str] = None,
         mask_params: bool = True,
     ):
         pass
@@ -859,7 +963,9 @@ class RunDBInterface(ABC):
 
     @abstractmethod
     def start_function(
-        self, func_url: str = None, function: "mlrun.runtimes.BaseRuntime" = None
+        self,
+        func_url: Optional[str] = None,
+        function: "mlrun.runtimes.BaseRuntime" = None,
     ):
         pass
 
@@ -878,7 +984,7 @@ class RunDBInterface(ABC):
         source: Optional[str] = None,
         run_name: Optional[str] = None,
         namespace: Optional[str] = None,
-        notifications: list["mlrun.model.Notification"] = None,
+        notifications: Optional[list["mlrun.model.Notification"]] = None,
     ) -> "mlrun.common.schemas.WorkflowResponse":
         pass
 
@@ -911,7 +1017,7 @@ class RunDBInterface(ABC):
         delete_stream_function: bool = False,
         delete_histogram_data_drift_app: bool = True,
         delete_user_applications: bool = False,
-        user_application_list: list[str] = None,
+        user_application_list: Optional[list[str]] = None,
     ) -> bool:
         pass