mlrun 1.7.0rc26__py3-none-any.whl → 1.7.0rc31__py3-none-any.whl

mlrun/__main__.py

@@ -50,12 +50,12 @@ from .run import (
 from .runtimes import RemoteRuntime, RunError, RuntimeKinds, ServingRuntime
 from .secrets import SecretsStore
 from .utils import (
+    RunKeys,
     dict_to_yaml,
     get_in,
     is_relative_path,
     list2dict,
     logger,
-    run_keys,
     update_in,
 )
 from .utils.version import Version
@@ -380,15 +380,15 @@ def run(
     set_item(runobj.spec.hyper_param_options, hyper_param_strategy, "strategy")
     set_item(runobj.spec.hyper_param_options, selector, "selector")
 
-    set_item(runobj.spec, inputs, run_keys.inputs, list2dict(inputs))
+    set_item(runobj.spec, inputs, RunKeys.inputs, list2dict(inputs))
     set_item(
-        runobj.spec, returns, run_keys.returns, [py_eval(value) for value in returns]
+        runobj.spec, returns, RunKeys.returns, [py_eval(value) for value in returns]
     )
-    set_item(runobj.spec, in_path, run_keys.input_path)
-    set_item(runobj.spec, out_path, run_keys.output_path)
-    set_item(runobj.spec, outputs, run_keys.outputs, list(outputs))
+    set_item(runobj.spec, in_path, RunKeys.input_path)
+    set_item(runobj.spec, out_path, RunKeys.output_path)
+    set_item(runobj.spec, outputs, RunKeys.outputs, list(outputs))
     set_item(
-        runobj.spec, secrets, run_keys.secrets, line2keylist(secrets, "kind", "source")
+        runobj.spec, secrets, RunKeys.secrets, line2keylist(secrets, "kind", "source")
     )
     set_item(runobj.spec, verbose, "verbose")
     set_item(runobj.spec, scrape_metrics, "scrape_metrics")

mlrun/alerts/alert.py

@@ -26,7 +26,6 @@ class AlertConfig(ModelObj):
         "description",
         "summary",
         "severity",
-        "criteria",
         "reset_policy",
         "state",
     ]
@@ -34,6 +33,7 @@ class AlertConfig(ModelObj):
         "entities",
         "notifications",
         "trigger",
+        "criteria",
     ]
 
     def __init__(
@@ -104,6 +104,14 @@ class AlertConfig(ModelObj):
                     else self.trigger
                 )
             return None
+        if field_name == "criteria":
+            if self.criteria:
+                return (
+                    self.criteria.dict()
+                    if not isinstance(self.criteria, dict)
+                    else self.criteria
+                )
+            return None
         return super()._serialize_field(struct, field_name, strip)
 
     def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
@@ -137,6 +145,10 @@ class AlertConfig(ModelObj):
             trigger_obj = alert_objects.AlertTrigger.parse_obj(trigger_data)
             new_obj.trigger = trigger_obj
 
+        criteria_data = struct.get("criteria")
+        if criteria_data:
+            criteria_obj = alert_objects.AlertCriteria.parse_obj(criteria_data)
+            new_obj.criteria = criteria_obj
         return new_obj
 
     def with_notifications(self, notifications: list[alert_objects.AlertNotification]):

mlrun/artifacts/manager.py

@@ -100,6 +100,11 @@ class ArtifactProducer:
 
 def dict_to_artifact(struct: dict) -> Artifact:
     kind = struct.get("kind", "")
+
+    # TODO: remove this in 1.8.0
+    if mlrun.utils.is_legacy_artifact(struct):
+        return mlrun.artifacts.base.convert_legacy_artifact_to_new_format(struct)
+
     artifact_class = artifact_types[kind]
     return artifact_class.from_dict(struct)
 

mlrun/common/constants.py

@@ -11,7 +11,6 @@
 # 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.
-#
 
 IMAGE_NAME_ENRICH_REGISTRY_PREFIX = "."  # prefix for image name to enrich with registry
 MLRUN_SERVING_CONF = "serving-conf"
@@ -64,12 +63,13 @@ class MLRunInternalLabels:
     username = f"{MLRUN_LABEL_PREFIX}username"
     username_domain = f"{MLRUN_LABEL_PREFIX}username_domain"
     task_name = f"{MLRUN_LABEL_PREFIX}task-name"
+    resource_name = f"{MLRUN_LABEL_PREFIX}resource_name"
+    created = f"{MLRUN_LABEL_PREFIX}created"
     host = "host"
     job_type = "job-type"
     kind = "kind"
     component = "component"
-    resource_name = "resource_name"
-    created = "mlrun-created"
+    mlrun_type = "mlrun__type"
 
     owner = "owner"
     v3io_user = "v3io_user"

mlrun/common/formatters/artifact.py

@@ -37,6 +37,7 @@ class ArtifactFormat(ObjectFormat, mlrun.common.types.StrEnum):
                     "spec.db_key",
                     "spec.size",
                     "spec.framework",
+                    "spec.algorithm",
                     "spec.metrics",
                     "spec.target_path",
                 ]

mlrun/common/formatters/base.py

@@ -28,42 +28,42 @@ class ObjectFormat:
     full = "full"
 
     @staticmethod
-    def format_method(_format: str) -> typing.Optional[typing.Callable]:
+    def format_method(format_: str) -> typing.Optional[typing.Callable]:
         """
         Get the formatting method for the provided format.
         A `None` value signifies a pass-through formatting method (no formatting).
-        :param _format: The format as a string representation.
+        :param format_: The format as a string representation.
         :return: The formatting method.
         """
         return {
             ObjectFormat.full: None,
-        }[_format]
+        }[format_]
 
     @classmethod
     def format_obj(
         cls,
         obj: typing.Any,
-        _format: str,
+        format_: str,
         exclude_formats: typing.Optional[list[str]] = None,
     ) -> typing.Any:
         """
         Format the provided object based on the provided format.
         :param obj: The object to format.
-        :param _format: The format as a string representation.
+        :param format_: The format as a string representation.
         :param exclude_formats: A list of formats to exclude from the formatting process. If the provided format is in
                                 this list, an invalid format exception will be raised.
         """
         exclude_formats = exclude_formats or []
-        _format = _format or cls.full
+        format_ = format_ or cls.full
         invalid_format_exc = mlrun.errors.MLRunBadRequestError(
-            f"Provided format is not supported. format={_format}"
+            f"Provided format is not supported. format={format_}"
         )
 
-        if _format in exclude_formats:
+        if format_ in exclude_formats:
             raise invalid_format_exc
 
         try:
-            format_method = cls.format_method(_format)
+            format_method = cls.format_method(format_)
         except KeyError:
             raise invalid_format_exc
 

mlrun/common/schemas/alert.py

@@ -39,8 +39,8 @@ class EventKind(StrEnum):
     CONCEPT_DRIFT_SUSPECTED = "concept_drift_suspected"
     MODEL_PERFORMANCE_DETECTED = "model_performance_detected"
     MODEL_PERFORMANCE_SUSPECTED = "model_performance_suspected"
-    MODEL_SERVING_PERFORMANCE_DETECTED = "model_serving_performance_detected"
-    MODEL_SERVING_PERFORMANCE_SUSPECTED = "model_serving_performance_suspected"
+    SYSTEM_PERFORMANCE_DETECTED = "system_performance_detected"
+    SYSTEM_PERFORMANCE_SUSPECTED = "system_performance_suspected"
     MM_APP_ANOMALY_DETECTED = "mm_app_anomaly_detected"
     MM_APP_ANOMALY_SUSPECTED = "mm_app_anomaly_suspected"
     FAILED = "failed"
@@ -53,12 +53,8 @@ _event_kind_entity_map = {
     EventKind.CONCEPT_DRIFT_SUSPECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.MODEL_PERFORMANCE_DETECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.MODEL_PERFORMANCE_SUSPECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
-    EventKind.MODEL_SERVING_PERFORMANCE_DETECTED: [
-        EventEntityKind.MODEL_ENDPOINT_RESULT
-    ],
-    EventKind.MODEL_SERVING_PERFORMANCE_SUSPECTED: [
-        EventEntityKind.MODEL_ENDPOINT_RESULT
-    ],
+    EventKind.SYSTEM_PERFORMANCE_DETECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
+    EventKind.SYSTEM_PERFORMANCE_SUSPECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.MM_APP_ANOMALY_DETECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.MM_APP_ANOMALY_SUSPECTED: [EventEntityKind.MODEL_ENDPOINT_RESULT],
     EventKind.FAILED: [EventEntityKind.JOB],

mlrun/common/schemas/api_gateway.py

@@ -102,6 +102,13 @@ class APIGateway(_APIGatewayBaseModel):
             if upstream.nucliofunction.get("name")
         ]
 
+    def get_invoke_url(self):
+        return (
+            self.spec.host + self.spec.path
+            if self.spec.path and self.spec.host
+            else self.spec.host
+        )
+
     def enrich_mlrun_names(self):
         self._enrich_api_gateway_mlrun_name()
         self._enrich_mlrun_function_names()

mlrun/common/schemas/constants.py

@@ -120,10 +120,13 @@ class FeatureStorePartitionByField(mlrun.common.types.StrEnum):
 
 class RunPartitionByField(mlrun.common.types.StrEnum):
     name = "name"  # Supported for runs objects
+    project_and_name = "project_and_name"  # Supported for runs objects
 
     def to_partition_by_db_field(self, db_cls):
         if self.value == RunPartitionByField.name:
             return db_cls.name
+        elif self.value == RunPartitionByField.project_and_name:
+            return db_cls.project, db_cls.name
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"Unknown group by field: {self.value}"

mlrun/common/schemas/model_monitoring/__init__.py

@@ -25,6 +25,7 @@ from .constants import (
     FunctionURI,
     MetricData,
     ModelEndpointTarget,
+    ModelEndpointTargetSchemas,
     ModelMonitoringMode,
     ModelMonitoringStoreKinds,
     MonitoringFunctionNames,

mlrun/common/schemas/model_monitoring/constants.py

@@ -17,6 +17,7 @@ from dataclasses import dataclass
 from enum import Enum, IntEnum
 from typing import Optional
 
+import mlrun.common.constants
 import mlrun.common.helpers
 from mlrun.common.types import StrEnum
 
@@ -78,8 +79,6 @@ class EventFieldType:
     FEATURE_SET_URI = "monitoring_feature_set_uri"
     ALGORITHM = "algorithm"
     VALUE = "value"
-    DRIFT_DETECTED_THRESHOLD = "drift_detected_threshold"
-    POSSIBLE_DRIFT_THRESHOLD = "possible_drift_threshold"
     SAMPLE_PARQUET_PATH = "sample_parquet_path"
     TIME = "time"
     TABLE_COLUMN = "table_column"
@@ -158,19 +157,42 @@ class EventKeyMetrics:
     REAL_TIME = "real_time"
 
 
-class ModelEndpointTarget:
+class ModelEndpointTarget(MonitoringStrEnum):
     V3IO_NOSQL = "v3io-nosql"
     SQL = "sql"
 
 
+class StreamKind(MonitoringStrEnum):
+    V3IO_STREAM = "v3io_stream"
+    KAFKA = "kafka"
+
+
+class TSDBTarget(MonitoringStrEnum):
+    V3IO_TSDB = "v3io-tsdb"
+    TDEngine = "tdengine"
+    PROMETHEUS = "prometheus"
+
+
 class ProjectSecretKeys:
     ENDPOINT_STORE_CONNECTION = "MODEL_MONITORING_ENDPOINT_STORE_CONNECTION"
     ACCESS_KEY = "MODEL_MONITORING_ACCESS_KEY"
-    PIPELINES_ACCESS_KEY = "MODEL_MONITORING_PIPELINES_ACCESS_KEY"
-    KAFKA_BROKERS = "KAFKA_BROKERS"
     STREAM_PATH = "STREAM_PATH"
     TSDB_CONNECTION = "TSDB_CONNECTION"
 
+    @classmethod
+    def mandatory_secrets(cls):
+        return [
+            cls.ENDPOINT_STORE_CONNECTION,
+            cls.STREAM_PATH,
+            cls.TSDB_CONNECTION,
+        ]
+
+
+class ModelEndpointTargetSchemas(MonitoringStrEnum):
+    V3IO = "v3io"
+    MYSQL = "mysql"
+    SQLITE = "sqlite"
+
 
 class ModelMonitoringStoreKinds:
     ENDPOINTS = "endpoints"
@@ -318,7 +340,7 @@ class ResultKindApp(Enum):
     concept_drift = 1
     model_performance = 2
     system_performance = 3
-    custom = 4
+    mm_app_anomaly = 4
 
 
 class ResultStatusApp(IntEnum):
@@ -333,7 +355,7 @@ class ResultStatusApp(IntEnum):
 
 
 class ModelMonitoringAppLabel:
-    KEY = "mlrun__type"
+    KEY = mlrun.common.constants.MLRunInternalLabels.mlrun_type
     VAL = "mlrun__model-monitoring-application"
 
     def __str__(self) -> str:
@@ -344,12 +366,6 @@ class ControllerPolicy:
     BASE_PERIOD = "base_period"
 
 
-class TSDBTarget:
-    V3IO_TSDB = "v3io-tsdb"
-    TDEngine = "tdengine"
-    PROMETHEUS = "prometheus"
-
-
 class HistogramDataDriftApplicationConstants:
     NAME = "histogram-data-drift"
     GENERAL_RESULT_NAME = "general_drift"
@@ -362,3 +378,6 @@ class PredictionsQueryConstants:
 
 class SpecialApps:
     MLRUN_INFRA = "mlrun-infra"
+
+
+_RESERVED_FUNCTION_NAMES = MonitoringFunctionNames.list() + [SpecialApps.MLRUN_INFRA]

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -103,18 +103,6 @@ class ModelEndpointSpec(ObjectSpec):
             json_parse_values=json_parse_values,
         )
 
-    @validator("monitor_configuration")
-    @classmethod
-    def set_name(cls, monitor_configuration):
-        return monitor_configuration or {
-            EventFieldType.DRIFT_DETECTED_THRESHOLD: (
-                mlrun.mlconf.model_endpoint_monitoring.drift_thresholds.default.drift_detected
-            ),
-            EventFieldType.POSSIBLE_DRIFT_THRESHOLD: (
-                mlrun.mlconf.model_endpoint_monitoring.drift_thresholds.default.possible_drift
-            ),
-        }
-
     @validator("model_uri")
     @classmethod
     def validate_model_uri(cls, model_uri):

mlrun/common/schemas/project.py

@@ -114,18 +114,19 @@ class ProjectOwner(pydantic.BaseModel):
 
 class ProjectSummary(pydantic.BaseModel):
     name: str
-    files_count: int
-    feature_sets_count: int
-    models_count: int
-    runs_completed_recent_count: int
-    runs_failed_recent_count: int
-    runs_running_count: int
-    distinct_schedules_count: int
-    distinct_scheduled_jobs_pending_count: int
-    distinct_scheduled_pipelines_pending_count: int
+    files_count: int = 0
+    feature_sets_count: int = 0
+    models_count: int = 0
+    runs_completed_recent_count: int = 0
+    runs_failed_recent_count: int = 0
+    runs_running_count: int = 0
+    distinct_schedules_count: int = 0
+    distinct_scheduled_jobs_pending_count: int = 0
+    distinct_scheduled_pipelines_pending_count: int = 0
     pipelines_completed_recent_count: typing.Optional[int] = None
     pipelines_failed_recent_count: typing.Optional[int] = None
     pipelines_running_count: typing.Optional[int] = None
+    updated: typing.Optional[datetime.datetime] = None
 
 
 class IguazioProject(pydantic.BaseModel):

mlrun/common/schemas/schedule.py

@@ -96,7 +96,7 @@ class ScheduleUpdate(BaseModel):
     scheduled_object: Optional[Any]
     cron_trigger: Optional[Union[str, ScheduleCronTrigger]]
     desired_state: Optional[str]
-    labels: Optional[dict] = {}
+    labels: Optional[dict] = None
     concurrency_limit: Optional[int]
     credentials: Credentials = Credentials()
 

mlrun/config.py

@@ -52,6 +52,11 @@ default_config = {
     "kubernetes": {
         "kubeconfig_path": "",  # local path to kubeconfig file (for development purposes),
         # empty by default as the API already running inside k8s cluster
+        "pagination": {
+            # pagination config for interacting with k8s API
+            "list_pods_limit": 200,
+            "list_crd_objects_limit": 200,
+        },
     },
     "dbpath": "",  # db/api url
     # url to nuclio dashboard api (can be with user & token, e.g. https://username:password@dashboard-url.com)
@@ -64,11 +69,15 @@ default_config = {
     "api_base_version": "v1",
     "version": "",  # will be set to current version
     "images_tag": "",  # tag to use with mlrun images e.g. mlrun/mlrun (defaults to version)
-    "images_registry": "",  # registry to use with mlrun images e.g. quay.io/ (defaults to empty, for dockerhub)
+    # registry to use with mlrun images that start with "mlrun/" e.g. quay.io/ (defaults to empty, for dockerhub)
+    "images_registry": "",
+    # registry to use with non-mlrun images (don't start with "mlrun/") specified in 'images_to_enrich_registry'
+    # defaults to empty, for dockerhub
+    "vendor_images_registry": "",
     # comma separated list of images that are in the specified images_registry, and therefore will be enriched with this
     # registry when used. default to mlrun/* which means any image which is of the mlrun repository (mlrun/mlrun,
     # mlrun/ml-base, etc...)
-    "images_to_enrich_registry": "^mlrun/*",
+    "images_to_enrich_registry": "^mlrun/*,python:3.9",
     "kfp_url": "",
     "kfp_ttl": "14400",  # KFP ttl in sec, after that completed PODs will be deleted
     "kfp_image": "mlrun/mlrun",  # image to use for KFP runner (defaults to mlrun/mlrun)
@@ -104,7 +113,12 @@ default_config = {
             # max number of parallel abort run jobs in runs monitoring
             "concurrent_abort_stale_runs_workers": 10,
             "list_runs_time_period_in_days": 7,  # days
-        }
+        },
+        "projects": {
+            "summaries": {
+                "cache_interval": "30",
+            },
+        },
     },
     "crud": {
         "runs": {
@@ -250,7 +264,7 @@ default_config = {
             "remote": "mlrun/mlrun",
             "dask": "mlrun/ml-base",
             "mpijob": "mlrun/mlrun",
-            "application": "python:3.9-slim",
+            "application": "python:3.9",
         },
         # see enrich_function_preemption_spec for more info,
         # and mlrun.common.schemas.function.PreemptionModes for available options
@@ -265,6 +279,16 @@ default_config = {
                 "url": "",
                 "service": "mlrun-api-chief",
                 "port": 8080,
+                "feature_gates": {
+                    "scheduler": "enabled",
+                    "project_sync": "enabled",
+                    "cleanup": "enabled",
+                    "runs_monitoring": "enabled",
+                    "pagination_cache": "enabled",
+                    "project_summaries": "enabled",
+                    "start_logs": "enabled",
+                    "stop_logs": "enabled",
+                },
             },
             "worker": {
                 "sync_with_chief": {
@@ -433,7 +457,6 @@ default_config = {
             "followers": "",
             # This is used as the interval for the sync loop both when mlrun is leader and follower
             "periodic_sync_interval": "1 minute",
-            "counters_cache_ttl": "2 minutes",
             "project_owners_cache_ttl": "30 seconds",
             # access key to be used when the leader is iguazio and polling is done from it
             "iguazio_access_key": "",
@@ -504,13 +527,12 @@ default_config = {
     "model_endpoint_monitoring": {
         "serving_stream_args": {"shard_count": 1, "retention_period_hours": 24},
         "application_stream_args": {"shard_count": 1, "retention_period_hours": 24},
-        "drift_thresholds": {"default": {"possible_drift": 0.5, "drift_detected": 0.7}},
         # Store prefixes are used to handle model monitoring storing policies based on project and kind, such as events,
         # stream, and endpoints.
         "store_prefixes": {
             "default": "v3io:///users/pipelines/{project}/model-endpoints/{kind}",
             "user_space": "v3io:///projects/{project}/model-endpoints/{kind}",
-            "stream": "",
+            "stream": "",  # TODO: Delete in 1.9.0
             "monitoring_application": "v3io:///users/pipelines/{project}/monitoring-apps/",
         },
         # Offline storage path can be either relative or a full path. This path is used for general offline data
@@ -523,11 +545,12 @@ default_config = {
         "parquet_batching_max_events": 10_000,
         "parquet_batching_timeout_secs": timedelta(minutes=1).total_seconds(),
         # See mlrun.model_monitoring.db.stores.ObjectStoreFactory for available options
-        "store_type": "v3io-nosql",
+        "store_type": "v3io-nosql",  # TODO: Delete in 1.9.0
         "endpoint_store_connection": "",
         # See mlrun.model_monitoring.db.tsdb.ObjectTSDBFactory for available options
-        "tsdb_connector_type": "v3io-tsdb",
         "tsdb_connection": "",
+        # See mlrun.common.schemas.model_monitoring.constants.StreamKind for available options
+        "stream_connection": "",
     },
     "secret_stores": {
         # Use only in testing scenarios (such as integration tests) to avoid using k8s for secrets (will use in-memory
@@ -660,7 +683,9 @@ default_config = {
         "failed_runs_grace_period": 3600,
         "verbose": True,
         # the number of workers which will be used to trigger the start log collection
-        "concurrent_start_logs_workers": 15,
+        "concurrent_start_logs_workers": 50,
+        # the number of runs for which to start logs on api startup
+        "start_logs_startup_run_limit": 150,
         # the time in hours in which to start log collection from.
         # after upgrade, we might have runs which completed in the mean time or still in non-terminal state and
         # we want to collect their logs in the new log collection method (sidecar)
@@ -708,6 +733,8 @@ default_config = {
         # maximum number of alerts we allow to be configured.
         # user will get an error when exceeding this
         "max_allowed": 10000,
+        # maximum allowed value for count in criteria field inside AlertConfig
+        "max_criteria_count": 100,
     },
     "auth_with_client_id": {
         "enabled": False,
@@ -1118,7 +1145,6 @@ class Config:
             if store_prefix_dict.get(kind):
                 # Target exist in store prefix and has a valid string value
                 return store_prefix_dict[kind].format(project=project, **kwargs)
-
             if (
                 function_name
                 and function_name

mlrun/data_types/spark.py

@@ -20,10 +20,10 @@ import pytz
 from pyspark.sql.functions import to_utc_timestamp
 from pyspark.sql.types import BooleanType, DoubleType, TimestampType
 
+from mlrun.feature_store.retrieval.spark_merger import spark_df_to_pandas
 from mlrun.utils import logger
 
 from .data_types import InferOptions, spark_to_value_type
-from .to_pandas import toPandas
 
 try:
     import pyspark.sql.functions as funcs
@@ -75,7 +75,7 @@ def get_df_preview_spark(df, preview_lines=20):
     """capture preview data from spark df"""
     df = df.limit(preview_lines)
 
-    result_dict = toPandas(df).to_dict(orient="split")
+    result_dict = spark_df_to_pandas(df).to_dict(orient="split")
     return [result_dict["columns"], *result_dict["data"]]
 
 

mlrun/data_types/to_pandas.py

@@ -15,21 +15,11 @@
 import warnings
 from collections import Counter
 
-from pyspark.sql.types import (
-    BooleanType,
-    ByteType,
-    DoubleType,
-    FloatType,
-    IntegerType,
-    IntegralType,
-    LongType,
-    MapType,
-    ShortType,
-    TimestampType,
-)
-
-
-def toPandas(spark_df):
+import pandas as pd
+import semver
+
+
+def _toPandas(spark_df):
     """
     Modified version of spark DataFrame.toPandas() –
     https://github.com/apache/spark/blob/v3.2.3/python/pyspark/sql/pandas/conversion.py#L35
@@ -40,6 +30,12 @@ def toPandas(spark_df):
     This modification adds the missing unit to the dtype.
     """
     from pyspark.sql.dataframe import DataFrame
+    from pyspark.sql.types import (
+        BooleanType,
+        IntegralType,
+        MapType,
+        TimestampType,
+    )
 
     assert isinstance(spark_df, DataFrame)
 
@@ -48,7 +44,6 @@ def toPandas(spark_df):
     require_minimum_pandas_version()
 
     import numpy as np
-    import pandas as pd
 
     timezone = spark_df.sql_ctx._conf.sessionLocalTimeZone()
 
@@ -217,6 +212,16 @@ def toPandas(spark_df):
 
 def _to_corrected_pandas_type(dt):
     import numpy as np
+    from pyspark.sql.types import (
+        BooleanType,
+        ByteType,
+        DoubleType,
+        FloatType,
+        IntegerType,
+        LongType,
+        ShortType,
+        TimestampType,
+    )
 
     if type(dt) == ByteType:
         return np.int8
@@ -236,3 +241,30 @@ def _to_corrected_pandas_type(dt):
         return "datetime64[ns]"
     else:
         return None
+
+
+def spark_df_to_pandas(spark_df):
+    # as of pyspark 3.2.3, toPandas fails to convert timestamps unless we work around the issue
+    # when we upgrade pyspark, we should check whether this workaround is still necessary
+    # see https://stackoverflow.com/questions/76389694/transforming-pyspark-to-pandas-dataframe
+    if semver.parse(pd.__version__)["major"] >= 2:
+        import pyspark.sql.functions as pyspark_functions
+
+        type_conversion_dict = {}
+        for field in spark_df.schema.fields:
+            if str(field.dataType) == "TimestampType":
+                spark_df = spark_df.withColumn(
+                    field.name,
+                    pyspark_functions.date_format(
+                        pyspark_functions.to_timestamp(field.name),
+                        "yyyy-MM-dd'T'HH:mm:ss.SSSSSSSSS",
+                    ),
+                )
+                type_conversion_dict[field.name] = "datetime64[ns]"
+
+        df = _toPandas(spark_df)
+        if type_conversion_dict:
+            df = df.astype(type_conversion_dict)
+        return df
+    else:
+        return _toPandas(spark_df)

mlrun/datastore/__init__.py

@@ -117,6 +117,7 @@ def get_stream_pusher(stream_path: str, **kwargs):
         return OutputStream(stream_path, **kwargs)
     elif stream_path.startswith("v3io"):
         endpoint, stream_path = parse_path(stream_path)
+        endpoint = kwargs.pop("endpoint", None) or endpoint
         return OutputStream(stream_path, endpoint=endpoint, **kwargs)
     elif stream_path.startswith("dummy://"):
         return _DummyStream(**kwargs)

mlrun/datastore/azure_blob.py

@@ -208,6 +208,7 @@ class AzureBlobStore(DataStore):
             for key in spark_options:
                 if key.startswith(prefix):
                     account_key = key[len(prefix) :]
-                    url += f"@{account_key}"
+                    if not url.endswith(account_key):
+                        url += f"@{account_key}"
                     break
         return url