mlrun 1.7.0rc39__py3-none-any.whl → 1.7.0rc41__py3-none-any.whl

mlrun/common/constants.py

@@ -65,6 +65,9 @@ class MLRunInternalLabels:
     task_name = f"{MLRUN_LABEL_PREFIX}task-name"
     resource_name = f"{MLRUN_LABEL_PREFIX}resource_name"
     created = f"{MLRUN_LABEL_PREFIX}created"
+    producer_type = f"{MLRUN_LABEL_PREFIX}producer-type"
+    app_name = f"{MLRUN_LABEL_PREFIX}app-name"
+    endpoint_id = f"{MLRUN_LABEL_PREFIX}endpoint-id"
     host = "host"
     job_type = "job-type"
     kind = "kind"

mlrun/common/helpers.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.
-#
 
 
 def parse_versioned_object_uri(

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -21,7 +21,6 @@ from typing import Any, NamedTuple, Optional
 from pydantic import BaseModel, Field, validator
 from pydantic.main import Extra
 
-import mlrun.common.model_monitoring
 import mlrun.common.types
 
 from ..object import ObjectKind, ObjectSpec, ObjectStatus

mlrun/config.py

@@ -863,7 +863,7 @@ class Config:
                     f"Unable to decode {attribute_path}"
                 )
             parsed_attribute_value = json.loads(decoded_attribute_value)
-            if type(parsed_attribute_value) != expected_type:
+            if not isinstance(parsed_attribute_value, expected_type):
                 raise mlrun.errors.MLRunInvalidArgumentTypeError(
                     f"Expected type {expected_type}, got {type(parsed_attribute_value)}"
                 )

mlrun/data_types/to_pandas.py

@@ -21,7 +21,7 @@ import semver
 
 def _toPandas(spark_df):
     """
-    Modified version of spark DataFrame.toPandas() –
+    Modified version of spark DataFrame.toPandas() -
     https://github.com/apache/spark/blob/v3.2.3/python/pyspark/sql/pandas/conversion.py#L35
 
     The original code (which is only replaced in pyspark 3.5.0) fails with Pandas 2 installed, with the following error:
@@ -223,21 +223,21 @@ def _to_corrected_pandas_type(dt):
         TimestampType,
     )
 
-    if type(dt) == ByteType:
+    if isinstance(dt, ByteType):
         return np.int8
-    elif type(dt) == ShortType:
+    elif isinstance(dt, ShortType):
         return np.int16
-    elif type(dt) == IntegerType:
+    elif isinstance(dt, IntegerType):
         return np.int32
-    elif type(dt) == LongType:
+    elif isinstance(dt, LongType):
         return np.int64
-    elif type(dt) == FloatType:
+    elif isinstance(dt, FloatType):
         return np.float32
-    elif type(dt) == DoubleType:
+    elif isinstance(dt, DoubleType):
         return np.float64
-    elif type(dt) == BooleanType:
+    elif isinstance(dt, BooleanType):
         return bool
-    elif type(dt) == TimestampType:
+    elif isinstance(dt, TimestampType):
         return "datetime64[ns]"
     else:
         return None

mlrun/datastore/alibaba_oss.py

@@ -85,6 +85,7 @@ class OSSStore(DataStore):
         return oss.get_object(key).read()
 
     def put(self, key, data, append=False):
+        data, _ = self._prepare_put_data(data, append)
         bucket, key = self.get_bucket_and_key(key)
         oss = oss2.Bucket(self.auth, self.endpoint_url, bucket)
         oss.put_object(key, data)

mlrun/datastore/azure_blob.py

@@ -189,12 +189,7 @@ class AzureBlobStore(DataStore):
                 "Append mode not supported for Azure blob datastore"
             )
         remote_path = self._convert_key_to_remote_path(key)
-        if isinstance(data, bytes):
-            mode = "wb"
-        elif isinstance(data, str):
-            mode = "w"
-        else:
-            raise TypeError("Data type unknown.  Unable to put in Azure!")
+        data, mode = self._prepare_put_data(data, append)
         with self.filesystem.open(remote_path, mode) as f:
             f.write(data)
 

mlrun/datastore/base.py

@@ -157,6 +157,18 @@ class DataStore:
     def put(self, key, data, append=False):
         pass
 
+    def _prepare_put_data(self, data, append=False):
+        mode = "a" if append else "w"
+        if isinstance(data, bytearray):
+            data = bytes(data)
+
+        if isinstance(data, bytes):
+            return data, f"{mode}b"
+        elif isinstance(data, str):
+            return data, mode
+        else:
+            raise TypeError(f"Unable to put a value of type {type(self).__name__}")
+
     def stat(self, key):
         pass
 

mlrun/datastore/dbfs_store.py

@@ -130,11 +130,7 @@ class DBFSStore(DataStore):
                 "Append mode not supported for Databricks file system"
             )
         #  can not use append mode because it overrides data.
-        mode = "w"
-        if isinstance(data, bytes):
-            mode += "b"
-        elif not isinstance(data, str):
-            raise TypeError(f"Unknown data type {type(data)}")
+        data, mode = self._prepare_put_data(data, append)
         with self.filesystem.open(key, mode) as f:
             f.write(data)
 

mlrun/datastore/filestore.py

@@ -66,9 +66,7 @@ class FileStore(DataStore):
         dir_to_create = path.dirname(self._join(key))
         if dir_to_create:
             self._ensure_directory(dir_to_create)
-        mode = "a" if append else "w"
-        if isinstance(data, bytes):
-            mode = mode + "b"
+        data, mode = self._prepare_put_data(data, append)
         with open(self._join(key), mode) as fp:
             fp.write(data)
             fp.close()

mlrun/datastore/google_cloud_storage.py

@@ -131,15 +131,7 @@ class GoogleCloudStorageStore(DataStore):
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Append mode not supported for Google cloud storage datastore"
             )
-
-        if isinstance(data, bytes):
-            mode = "wb"
-        elif isinstance(data, str):
-            mode = "w"
-        else:
-            raise TypeError(
-                "Data type unknown.  Unable to put in Google cloud storage!"
-            )
+        data, mode = self._prepare_put_data(data, append)
         with self.filesystem.open(path, mode) as f:
             f.write(data)
 

mlrun/datastore/redis.py

@@ -126,6 +126,7 @@ class RedisStore(DataStore):
 
     def put(self, key, data, append=False):
         key = RedisStore.build_redis_key(key)
+        data, _ = self._prepare_put_data(data, append)
         if append:
             self.redis.append(key, data)
         else:

mlrun/datastore/s3.py

@@ -183,6 +183,7 @@ class S3Store(DataStore):
         return obj.get()["Body"].read()
 
     def put(self, key, data, append=False):
+        data, _ = self._prepare_put_data(data, append)
         bucket, key = self.get_bucket_and_key(key)
         self.s3.Object(bucket, key).put(Body=data)
 

mlrun/datastore/storeytargets.py

@@ -0,0 +1,147 @@
+# Copyright 2024 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 storey
+from mergedeep import merge
+from storey import V3ioDriver
+
+import mlrun
+import mlrun.model_monitoring.helpers
+from mlrun.datastore.base import DataStore
+
+from .utils import (
+    parse_kafka_url,
+)
+
+"""
+Storey targets expect storage_options, which may contain credentials.
+To avoid passing it openly within the graph, we use wrapper classes.
+"""
+
+
+def get_url_and_storage_options(path, external_storage_options=None):
+    store, resolved_store_path, url = mlrun.store_manager.get_or_create_store(path)
+    storage_options = store.get_storage_options()
+    if storage_options and external_storage_options:
+        # merge external storage options with the store's storage options. storage_options takes precedence
+        storage_options = merge(external_storage_options, storage_options)
+    else:
+        storage_options = storage_options or external_storage_options
+    return url, DataStore._sanitize_storage_options(storage_options)
+
+
+class TDEngineStoreyTarget(storey.TDEngineTarget):
+    def __init__(self, *args, **kwargs):
+        kwargs["url"] = mlrun.model_monitoring.helpers.get_tsdb_connection_string()
+        super().__init__(*args, **kwargs)
+
+
+class StoreyTargetUtils:
+    @staticmethod
+    def process_args_and_kwargs(args, kwargs):
+        args = list(args)
+        path = args[0] if args else kwargs.get("path")
+        external_storage_options = kwargs.get("storage_options")
+
+        url, storage_options = get_url_and_storage_options(
+            path, external_storage_options
+        )
+
+        if storage_options:
+            kwargs["storage_options"] = storage_options
+        if args:
+            args[0] = url
+        if "path" in kwargs:
+            kwargs["path"] = url
+        return args, kwargs
+
+
+class ParquetStoreyTarget(storey.ParquetTarget):
+    def __init__(self, *args, **kwargs):
+        args, kwargs = StoreyTargetUtils.process_args_and_kwargs(args, kwargs)
+        super().__init__(*args, **kwargs)
+
+
+class CSVStoreyTarget(storey.CSVTarget):
+    def __init__(self, *args, **kwargs):
+        args, kwargs = StoreyTargetUtils.process_args_and_kwargs(args, kwargs)
+        super().__init__(*args, **kwargs)
+
+
+class StreamStoreyTarget(storey.StreamTarget):
+    def __init__(self, *args, **kwargs):
+        args = list(args)
+
+        path = args[0] if args else kwargs.get("stream_path")
+        endpoint, storage_options = get_url_and_storage_options(path)
+
+        if not path:
+            raise mlrun.errors.MLRunInvalidArgumentError("StreamTarget requires a path")
+
+        access_key = storage_options.get("v3io_access_key")
+        storage = (
+            V3ioDriver(webapi=endpoint or mlrun.mlconf.v3io_api, access_key=access_key),
+        )
+
+        if storage_options:
+            kwargs["storage"] = storage
+        if args:
+            args[0] = endpoint
+        if "stream_path" in kwargs:
+            kwargs["stream_path"] = endpoint
+
+        super().__init__(*args, **kwargs)
+
+
+class KafkaStoreyTarget(storey.KafkaTarget):
+    def __init__(self, *args, **kwargs):
+        path = kwargs.pop("path")
+        attributes = kwargs.pop("attributes", None)
+        if path and path.startswith("ds://"):
+            datastore_profile = (
+                mlrun.datastore.datastore_profile.datastore_profile_read(path)
+            )
+            attributes = merge(attributes, datastore_profile.attributes())
+            brokers = attributes.pop(
+                "brokers", attributes.pop("bootstrap_servers", None)
+            )
+            topic = datastore_profile.topic
+        else:
+            brokers = attributes.pop(
+                "brokers", attributes.pop("bootstrap_servers", None)
+            )
+            topic, brokers = parse_kafka_url(path, brokers)
+
+        if not topic:
+            raise mlrun.errors.MLRunInvalidArgumentError("KafkaTarget requires a topic")
+        kwargs["brokers"] = brokers
+        kwargs["topic"] = topic
+        super().__init__(*args, **kwargs, **attributes)
+
+
+class NoSqlStoreyTarget(storey.NoSqlTarget):
+    pass
+
+
+class RedisNoSqlStoreyTarget(storey.NoSqlTarget):
+    def __init__(self, *args, **kwargs):
+        path = kwargs.pop("path")
+        endpoint, uri = mlrun.datastore.targets.RedisNoSqlTarget.get_server_endpoint(
+            path
+        )
+        kwargs["path"] = endpoint + "/" + uri
+        super().__init__(*args, **kwargs)
+
+
+class TSDBStoreyTarget(storey.TSDBTarget):
+    pass

mlrun/datastore/targets.py

@@ -47,7 +47,6 @@ from .spark_utils import spark_session_update_hadoop_options
 from .utils import (
     _generate_sql_query_with_time_filter,
     filter_df_start_end_time,
-    parse_kafka_url,
     select_columns_from_df,
 )
 
@@ -928,8 +927,9 @@ class ParquetTarget(BaseStoreTarget):
                 if time_unit == time_partitioning_granularity:
                     break
 
+        target_path = self.get_target_path()
         if not self.partitioned and not mlrun.utils.helpers.is_parquet_file(
-            self.get_target_path()
+            target_path
         ):
             partition_cols = []
 
@@ -937,25 +937,16 @@ class ParquetTarget(BaseStoreTarget):
         for key_column in key_columns:
             tuple_key_columns.append((key_column.name, key_column.value_type))
 
-        store, path_in_store, target_path = self._get_store_and_path()
-
-        storage_options = store.get_storage_options()
-        if storage_options and self.storage_options:
-            storage_options = merge(storage_options, self.storage_options)
-        else:
-            storage_options = storage_options or self.storage_options
-
         step = graph.add_step(
             name=self.name or "ParquetTarget",
             after=after,
             graph_shape="cylinder",
-            class_name="storey.ParquetTarget",
+            class_name="mlrun.datastore.storeytargets.ParquetStoreyTarget",
             path=target_path,
             columns=column_list,
             index_cols=tuple_key_columns,
             partition_cols=partition_cols,
             time_field=timestamp_key,
-            storage_options=storage_options,
             max_events=self.max_events,
             flush_after_seconds=self.flush_after_seconds,
             update_last_written=featureset_status.update_last_written_for_target,
@@ -1110,17 +1101,16 @@ class CSVTarget(BaseStoreTarget):
         column_list = self._get_column_list(
             features=features, timestamp_key=timestamp_key, key_columns=key_columns
         )
-        store, path_in_store, target_path = self._get_store_and_path()
+        target_path = self.get_target_path()
         graph.add_step(
             name=self.name or "CSVTarget",
             after=after,
             graph_shape="cylinder",
-            class_name="storey.CSVTarget",
+            class_name="mlrun.datastore.storeytargets.CSVStoreyTarget",
             path=target_path,
             columns=column_list,
             header=True,
             index_cols=key_columns,
-            storage_options=store.get_storage_options(),
             **self.attributes,
         )
 
@@ -1334,6 +1324,19 @@ class NoSqlBaseTarget(BaseStoreTarget):
         timestamp_key=None,
         featureset_status=None,
     ):
+        table, column_list = self._get_table_and_columns(features, key_columns)
+
+        graph.add_step(
+            name=self.name or self.writer_step_name,
+            after=after,
+            graph_shape="cylinder",
+            class_name="mlrun.datastore.storeytargets.NoSqlStoreyTarget",
+            columns=column_list,
+            table=table,
+            **self.attributes,
+        )
+
+    def _get_table_and_columns(self, features, key_columns):
         key_columns = list(key_columns.keys())
         table = self._resource.uri
         column_list = self._get_column_list(
@@ -1352,15 +1355,7 @@ class NoSqlBaseTarget(BaseStoreTarget):
                 col for col in column_list if col[0] not in aggregate_features
             ]
 
-        graph.add_step(
-            name=self.name or self.writer_step_name,
-            after=after,
-            graph_shape="cylinder",
-            class_name="storey.NoSqlTarget",
-            columns=column_list,
-            table=table,
-            **self.attributes,
-        )
+        return table, column_list
 
     def prepare_spark_df(self, df, key_columns, timestamp_key=None, spark_options=None):
         raise NotImplementedError()
@@ -1483,11 +1478,9 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
     support_spark = True
     writer_step_name = "RedisNoSqlTarget"
 
-    # Fetch server url from the RedisNoSqlTarget::__init__() 'path' parameter.
-    # If not set fetch it from 'mlrun.mlconf.redis.url' (MLRUN_REDIS__URL environment variable).
-    # Then look for username and password at REDIS_xxx secrets
-    def _get_server_endpoint(self):
-        endpoint, uri = parse_path(self.get_target_path())
+    @staticmethod
+    def get_server_endpoint(path):
+        endpoint, uri = parse_path(path)
         endpoint = endpoint or mlrun.mlconf.redis.url
         if endpoint.startswith("ds://"):
             datastore_profile = datastore_profile_read(endpoint)
@@ -1504,8 +1497,13 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
                 raise mlrun.errors.MLRunInvalidArgumentError(
                     "Provide Redis username and password only via secrets"
                 )
-            user = self._get_credential("REDIS_USER", "")
-            password = self._get_credential("REDIS_PASSWORD", "")
+            credentials_prefix = mlrun.get_secret_or_env(key="CREDENTIALS_PREFIX")
+            user = mlrun.get_secret_or_env(
+                "REDIS_USER", default="", prefix=credentials_prefix
+            )
+            password = mlrun.get_secret_or_env(
+                "REDIS_PASSWORD", default="", prefix=credentials_prefix
+            )
             host = parsed_endpoint.hostname
             port = parsed_endpoint.port if parsed_endpoint.port else "6379"
             scheme = parsed_endpoint.scheme
@@ -1519,7 +1517,7 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
         from storey import Table
         from storey.redis_driver import RedisDriver
 
-        endpoint, uri = self._get_server_endpoint()
+        endpoint, uri = self.get_server_endpoint(self.get_target_path())
 
         return Table(
             uri,
@@ -1528,7 +1526,7 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
         )
 
     def get_spark_options(self, key_column=None, timestamp_key=None, overwrite=True):
-        endpoint, uri = self._get_server_endpoint()
+        endpoint, uri = self.get_server_endpoint(self.get_target_path())
         parsed_endpoint = urlparse(endpoint)
         store, path_in_store, path = self._get_store_and_path()
         return {
@@ -1560,6 +1558,28 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
 
         return df
 
+    def add_writer_step(
+        self,
+        graph,
+        after,
+        features,
+        key_columns=None,
+        timestamp_key=None,
+        featureset_status=None,
+    ):
+        table, column_list = self._get_table_and_columns(features, key_columns)
+
+        graph.add_step(
+            path=self.get_target_path(),
+            name=self.name or self.writer_step_name,
+            after=after,
+            graph_shape="cylinder",
+            class_name="mlrun.datastore.storeytargets.RedisNoSqlStoreyTarget",
+            columns=column_list,
+            table=table,
+            **self.attributes,
+        )
+
 
 class StreamTarget(BaseStoreTarget):
     kind = TargetTypes.stream
@@ -1578,29 +1598,22 @@ class StreamTarget(BaseStoreTarget):
         timestamp_key=None,
         featureset_status=None,
     ):
-        from storey import V3ioDriver
-
         key_columns = list(key_columns.keys())
-        store, path_in_store, path = self._get_store_and_path()
-        if not path:
-            raise mlrun.errors.MLRunInvalidArgumentError("StreamTarget requires a path")
-        endpoint, uri = parse_path(path)
-        storage_options = store.get_storage_options()
-        access_key = storage_options.get("v3io_access_key")
+
         column_list = self._get_column_list(
             features=features, timestamp_key=timestamp_key, key_columns=key_columns
         )
+        stream_path = self.get_target_path()
+        if not stream_path:
+            raise mlrun.errors.MLRunInvalidArgumentError("StreamTarget requires a path")
 
         graph.add_step(
             name=self.name or "StreamTarget",
             after=after,
             graph_shape="cylinder",
-            class_name="storey.StreamTarget",
+            class_name="mlrun.datastore.storeytargets.StreamStoreyTarget",
             columns=column_list,
-            storage=V3ioDriver(
-                webapi=endpoint or mlrun.mlconf.v3io_api, access_key=access_key
-            ),
-            stream_path=uri,
+            stream_path=stream_path,
             **self.attributes,
         )
 
@@ -1676,34 +1689,19 @@ class KafkaTarget(BaseStoreTarget):
         column_list = self._get_column_list(
             features=features, timestamp_key=timestamp_key, key_columns=key_columns
         )
-        if self.path and self.path.startswith("ds://"):
-            datastore_profile = datastore_profile_read(self.path)
-            attributes = datastore_profile.attributes()
-            brokers = attributes.pop(
-                "brokers", attributes.pop("bootstrap_servers", None)
-            )
-            topic = datastore_profile.topic
-        else:
-            attributes = copy(self.attributes)
-            brokers = attributes.pop(
-                "brokers", attributes.pop("bootstrap_servers", None)
-            )
-            topic, brokers = parse_kafka_url(self.get_target_path(), brokers)
+        path = self.get_target_path()
 
-        if not topic:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "KafkaTarget requires a path (topic)"
-            )
+        if not path:
+            raise mlrun.errors.MLRunInvalidArgumentError("KafkaTarget requires a path")
 
         graph.add_step(
             name=self.name or "KafkaTarget",
             after=after,
             graph_shape="cylinder",
-            class_name="storey.KafkaTarget",
+            class_name="mlrun.datastore.storeytargets.KafkaStoreyTarget",
             columns=column_list,
-            topic=topic,
-            brokers=brokers,
-            **attributes,
+            path=path,
+            attributes=self.attributes,
         )
 
     def purge(self):
@@ -1740,7 +1738,7 @@ class TSDBTarget(BaseStoreTarget):
 
         graph.add_step(
             name=self.name or "TSDBTarget",
-            class_name="storey.TSDBTarget",
+            class_name="mlrun.datastore.storeytargets.TSDBStoreyTarget",
             after=after,
             graph_shape="cylinder",
             path=uri,
@@ -2029,7 +2027,7 @@ class SQLTarget(BaseStoreTarget):
             name=self.name or "SqlTarget",
             after=after,
             graph_shape="cylinder",
-            class_name="storey.NoSqlTarget",
+            class_name="mlrun.datastore.storeytargets.NoSqlStoreyTarget",
             columns=column_list,
             header=True,
             table=table,

mlrun/datastore/v3io.py

@@ -140,6 +140,7 @@ class V3ioStore(DataStore):
         max_chunk_size: int = V3IO_DEFAULT_UPLOAD_CHUNK_SIZE,
     ):
         """helper function for put method, allows for controlling max_chunk_size in testing"""
+        data, _ = self._prepare_put_data(data, append)
         container, path = split_path(self._join(key))
         buffer_size = len(data)  # in bytes
         buffer_offset = 0

mlrun/model_monitoring/api.py

@@ -147,8 +147,7 @@ def record_results(
                                      on the provided `endpoint_id`.
     :param function_name:            If a new model endpoint is created, use this function name for generating the
                                      function URI.
-    :param context:                  MLRun context. Note that the context is required for logging the artifacts
-                                     following the batch drift job.
+    :param context:                  MLRun context. Note that the context is required generating the model endpoint.
     :param infer_results_df:         DataFrame that will be stored under the model endpoint parquet target. Will be
                                      used for doing the drift analysis. Please make sure that the dataframe includes
                                      both feature names and label columns.