mlrun 1.7.2rc4__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/datastore/inmem.py

@@ -17,6 +17,7 @@ from io import BytesIO, StringIO
 import pandas as pd
 
 import mlrun
+import mlrun.utils.helpers
 
 from .base import DataStore, FileStats
 
@@ -35,7 +36,9 @@ class InMemoryStore(DataStore):
 
     def _get_item(self, key):
         if key not in self._items:
-            raise ValueError(f"item {key} not found in memory store")
+            raise mlrun.errors.MLRunNotFoundError(
+                f"item {key} not found in memory store"
+            )
         return self._items[key]
 
     def get(self, key, size=None, offset=0):

mlrun/datastore/redis.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
 
 import redis
@@ -30,7 +31,9 @@ class RedisStore(DataStore):
     - key and value sizes are limited to 512MB
     """
 
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         redis_default_port = "6379"
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self.headers = None

mlrun/datastore/s3.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import time
+from typing import Optional
 
 import boto3
 from boto3.s3.transfer import TransferConfig
@@ -20,13 +21,15 @@ from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
-from .base import DataStore, FileStats, get_range, make_datastore_schema_sanitizer
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 
 class S3Store(DataStore):
     using_bucket = True
 
-    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)
         # will be used in case user asks to assume a role and work through fsspec
         self._temp_credentials = None
@@ -105,6 +108,13 @@ class S3Store(DataStore):
                     "choose-signer.s3.*", disable_signing
                 )
 
+    @staticmethod
+    def get_range(size, offset):
+        byterange = f"bytes={offset}-"
+        if size:
+            byterange += str(offset + size - 1)
+        return byterange
+
     def get_spark_options(self):
         res = {}
         st = self.get_storage_options()
@@ -155,6 +165,7 @@ class S3Store(DataStore):
             key=access_key_id,
             secret=secret,
             token=token,
+            use_listings_cache=False,
         )
 
         if endpoint_url:
@@ -182,7 +193,7 @@ class S3Store(DataStore):
         bucket, key = self.get_bucket_and_key(key)
         obj = self.s3.Object(bucket, key)
         if size or offset:
-            return obj.get(Range=get_range(size, offset))["Body"].read()
+            return obj.get(Range=S3Store.get_range(size, offset))["Body"].read()
         return obj.get()["Body"].read()
 
     def put(self, key, data, append=False):

mlrun/datastore/sources.py

@@ -18,7 +18,7 @@ import warnings
 from base64 import b64encode
 from copy import copy
 from datetime import datetime
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import pandas as pd
 import semver
@@ -34,6 +34,7 @@ from mlrun.datastore.utils import transform_list_filters_to_tuple
 from mlrun.secrets import SecretsStore
 from mlrun.utils import logger
 
+from ..common.schemas.function import Function
 from ..model import DataSource
 from ..platforms.iguazio import parse_path
 from ..utils import get_class, is_explicit_ack_supported
@@ -181,10 +182,10 @@ class CSVSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        schedule: str = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         parse_dates: Union[None, int, str, list[int], list[str]] = None,
         **kwargs,
     ):
@@ -308,11 +309,11 @@ class ParquetSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
         end_time: Optional[Union[datetime, str]] = None,
         additional_filters: Optional[list[Union[tuple, list]]] = None,
@@ -392,7 +393,9 @@ class ParquetSource(BaseSourceDriver):
         )
 
     @classmethod
-    def from_dict(cls, struct=None, fields=None, deprecated_fields: dict = None):
+    def from_dict(
+        cls, struct=None, fields=None, deprecated_fields: Optional[dict] = None
+    ):
         new_obj = super().from_dict(
             struct=struct, fields=fields, deprecated_fields=deprecated_fields
         )
@@ -564,18 +567,18 @@ class BigQuerySource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        table: str = None,
-        max_results_for_table: int = None,
-        query: str = None,
-        materialization_dataset: str = None,
-        chunksize: int = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        table: Optional[str] = None,
+        max_results_for_table: Optional[int] = None,
+        query: Optional[str] = None,
+        materialization_dataset: Optional[str] = None,
+        chunksize: Optional[int] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time=None,
         end_time=None,
-        gcp_project: str = None,
-        spark_options: dict = None,
+        gcp_project: Optional[str] = None,
+        spark_options: Optional[dict] = None,
         **kwargs,
     ):
         if query and table:
@@ -776,27 +779,27 @@ class SnowflakeSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        key_field: str = None,
-        attributes: dict[str, object] = None,
-        time_field: str = None,
-        schedule: str = None,
+        key_field: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time=None,
         end_time=None,
-        query: str = None,
-        url: str = None,
-        user: str = None,
-        database: str = None,
-        schema: str = None,
-        db_schema: str = None,
-        warehouse: str = None,
+        query: Optional[str] = None,
+        url: Optional[str] = None,
+        user: Optional[str] = None,
+        database: Optional[str] = None,
+        schema: Optional[str] = None,
+        db_schema: Optional[str] = None,
+        warehouse: Optional[str] = None,
         **kwargs,
     ):
-        # TODO: Remove in 1.9.0
+        # TODO: Remove in 1.10.0
         if schema:
             warnings.warn(
-                "schema is deprecated in 1.7.0, and will be removed in 1.9.0, please use db_schema"
+                "schema is deprecated in 1.7.0, and will be removed in 1.10.0, please use db_schema"
             )
-        db_schema = db_schema or schema  # TODO: Remove in 1.9.0
+        db_schema = db_schema or schema  # TODO: Remove in 1.10.0
 
         attributes = attributes or {}
         if url:
@@ -850,9 +853,9 @@ class CustomSource(BaseSourceDriver):
 
     def __init__(
         self,
-        class_name: str = None,
+        class_name: Optional[str] = None,
         name: str = "",
-        schedule: str = None,
+        schedule: Optional[str] = None,
         **attributes,
     ):
         attributes = attributes or {}
@@ -930,12 +933,12 @@ class OnlineSource(BaseSourceDriver):
 
     def __init__(
         self,
-        name: str = None,
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        time_field: str = None,
-        workers: int = None,
+        name: Optional[str] = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        workers: Optional[int] = None,
     ):
         super().__init__(name, path, attributes, key_field, time_field)
         self.online = True
@@ -949,8 +952,7 @@ class OnlineSource(BaseSourceDriver):
             is_explicit_ack_supported(context)
             and mlrun.mlconf.is_explicit_ack_enabled()
         )
-        # TODO: Change to AsyncEmitSource once we can drop support for nuclio<1.12.10
-        src_class = storey.SyncEmitSource(
+        src_class = storey.AsyncEmitSource(
             context=context,
             key_field=self.key_field or key_field,
             full_event=True,
@@ -965,6 +967,21 @@ class OnlineSource(BaseSourceDriver):
             "This source type is not supported with ingestion service yet"
         )
 
+    @staticmethod
+    def set_explicit_ack_mode(function: Function, **extra_arguments) -> dict[str, Any]:
+        extra_arguments = extra_arguments or {}
+        engine = "sync"
+        if function.spec and hasattr(function.spec, "graph"):
+            engine = getattr(function.spec.graph, "engine", None) or engine
+        if mlrun.mlconf.is_explicit_ack_enabled() and engine == "async":
+            extra_arguments["explicit_ack_mode"] = extra_arguments.get(
+                "explicit_ack_mode", "explicitOnly"
+            )
+            extra_arguments["worker_allocation_mode"] = extra_arguments.get(
+                "worker_allocation_mode", "static"
+            )
+        return extra_arguments
+
 
 class HttpSource(OnlineSource):
     kind = "http"
@@ -986,7 +1003,7 @@ class StreamSource(OnlineSource):
         seek_to="earliest",
         shards=1,
         retention_in_hours=24,
-        extra_attributes: dict = None,
+        extra_attributes: Optional[dict] = None,
         **kwargs,
     ):
         """
@@ -1027,15 +1044,7 @@ class StreamSource(OnlineSource):
             raise_for_status=v3io.dataplane.RaiseForStatus.never,
         )
         res.raise_for_status([409, 204])
-
-        kwargs = {}
-        engine = "async"
-        if hasattr(function.spec, "graph") and function.spec.graph.engine:
-            engine = function.spec.graph.engine
-
-        if mlrun.mlconf.is_explicit_ack_enabled() and engine == "async":
-            kwargs["explicit_ack_mode"] = "explicitOnly"
-            kwargs["worker_allocation_mode"] = "static"
+        kwargs = self.set_explicit_ack_mode(function=function)
 
         function.add_v3io_stream_trigger(
             url,
@@ -1086,12 +1095,9 @@ class KafkaSource(OnlineSource):
         attributes["initial_offset"] = initial_offset
         if partitions is not None:
             attributes["partitions"] = partitions
-        sasl = attributes.pop("sasl", {})
-        if sasl_user and sasl_pass:
-            sasl["enabled"] = True
-            sasl["user"] = sasl_user
-            sasl["password"] = sasl_pass
-        if sasl:
+        if sasl := mlrun.datastore.utils.KafkaParameters(attributes).sasl(
+            usr=sasl_user, pwd=sasl_pass
+        ):
             attributes["sasl"] = sasl
         super().__init__(attributes=attributes, **kwargs)
 
@@ -1116,18 +1122,15 @@ class KafkaSource(OnlineSource):
         else:
             extra_attributes = copy(self.attributes)
         partitions = extra_attributes.pop("partitions", None)
-        explicit_ack_mode = None
-        engine = "async"
-        if hasattr(function.spec, "graph") and function.spec.graph.engine:
-            engine = function.spec.graph.engine
 
-        if mlrun.mlconf.is_explicit_ack_enabled() and engine == "async":
-            explicit_ack_mode = "explicitOnly"
-            extra_attributes["workerAllocationMode"] = extra_attributes.get(
-                "worker_allocation_mode", "static"
-            )
+        extra_attributes = self.set_explicit_ack_mode(function, **extra_attributes)
+        explicit_ack_mode = extra_attributes.get("explicit_ack_mode")
+        extra_attributes["workerAllocationMode"] = extra_attributes.get(
+            "worker_allocation_mode", "pool"
+        )
 
         trigger_kwargs = {}
+
         if "max_workers" in extra_attributes:
             trigger_kwargs = {"max_workers": extra_attributes.pop("max_workers")}
 
@@ -1168,7 +1171,7 @@ class KafkaSource(OnlineSource):
         self,
         num_partitions: int = 4,
         replication_factor: int = 1,
-        topics: list[str] = None,
+        topics: Optional[list[str]] = None,
     ):
         """
         Create Kafka topics with the specified number of partitions and replication factor.
@@ -1193,19 +1196,13 @@ class KafkaSource(OnlineSource):
         new_topics = [
             NewTopic(topic, num_partitions, replication_factor) for topic in topics
         ]
-        kafka_admin = KafkaAdminClient(
-            bootstrap_servers=brokers,
-            sasl_mechanism=self.attributes.get("sasl", {}).get("sasl_mechanism"),
-            sasl_plain_username=self.attributes.get("sasl", {}).get("username"),
-            sasl_plain_password=self.attributes.get("sasl", {}).get("password"),
-            sasl_kerberos_service_name=self.attributes.get("sasl", {}).get(
-                "sasl_kerberos_service_name", "kafka"
-            ),
-            sasl_kerberos_domain_name=self.attributes.get("sasl", {}).get(
-                "sasl_kerberos_domain_name"
-            ),
-            sasl_oauth_token_provider=self.attributes.get("sasl", {}).get("mechanism"),
-        )
+
+        kafka_admin_kwargs = {}
+        kafka_admin_kwargs = mlrun.datastore.utils.KafkaParameters(
+            self.attributes
+        ).admin()
+
+        kafka_admin = KafkaAdminClient(bootstrap_servers=brokers, **kafka_admin_kwargs)
         try:
             kafka_admin.create_topics(new_topics)
         finally:
@@ -1226,16 +1223,16 @@ class SQLSource(BaseSourceDriver):
     def __init__(
         self,
         name: str = "",
-        chunksize: int = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        chunksize: Optional[int] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
         end_time: Optional[Union[datetime, str]] = None,
-        db_url: str = None,
-        table_name: str = None,
-        spark_options: dict = None,
-        parse_dates: list[str] = None,
+        db_url: Optional[str] = None,
+        table_name: Optional[str] = None,
+        spark_options: Optional[dict] = None,
+        parse_dates: Optional[list[str]] = None,
         **kwargs,
     ):
         """

mlrun/datastore/store_resources.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-
 import mlrun
 import mlrun.artifacts
 from mlrun.config import config
@@ -165,11 +163,16 @@ def get_store_resource(
         return db.get_feature_vector(name, project, tag, uid)
 
     elif StorePrefix.is_artifact(kind):
-        project, key, iteration, tag, tree = parse_artifact_uri(
+        project, key, iteration, tag, tree, uid = parse_artifact_uri(
             uri, project or config.default_project
         )
         resource = db.read_artifact(
-            key, project=project, tag=tag, iter=iteration, tree=tree
+            key,
+            project=project,
+            tag=tag,
+            iter=iteration,
+            tree=tree,
+            uid=uid,
         )
         if resource.get("kind", "") == "link":
             # todo: support other link types (not just iter, move this to the db/api layer

mlrun/datastore/storeytargets.py

@@ -11,6 +11,8 @@
 # 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.
+from urllib.parse import urlparse
+
 import storey
 from mergedeep import merge
 from storey import V3ioDriver
@@ -18,6 +20,12 @@ from storey import V3ioDriver
 import mlrun
 import mlrun.model_monitoring.helpers
 from mlrun.datastore.base import DataStore
+from mlrun.datastore.datastore_profile import (
+    DatastoreProfileKafkaSource,
+    DatastoreProfileKafkaTarget,
+    DatastoreProfileTDEngine,
+    datastore_profile_read,
+)
 
 from ..platforms.iguazio import parse_path
 from .utils import (
@@ -42,9 +50,16 @@ def get_url_and_storage_options(path, external_storage_options=None):
 
 
 class TDEngineStoreyTarget(storey.TDEngineTarget):
-    def __init__(self, *args, **kwargs):
-        kwargs["url"] = mlrun.model_monitoring.helpers.get_tsdb_connection_string()
-        super().__init__(*args, **kwargs)
+    def __init__(self, *args, url: str, **kwargs):
+        if url.startswith("ds://"):
+            datastore_profile = datastore_profile_read(url)
+            if not isinstance(datastore_profile, DatastoreProfileTDEngine):
+                raise ValueError(
+                    f"Unexpected datastore profile type:{datastore_profile.type}."
+                    "Only DatastoreProfileTDEngine is supported"
+                )
+            url = datastore_profile.dsn()
+        super().__init__(*args, url=url, **kwargs)
 
 
 class StoreyTargetUtils:
@@ -69,7 +84,12 @@ class StoreyTargetUtils:
 
 class ParquetStoreyTarget(storey.ParquetTarget):
     def __init__(self, *args, **kwargs):
+        alt_key_name = kwargs.pop("alternative_v3io_access_key", None)
         args, kwargs = StoreyTargetUtils.process_args_and_kwargs(args, kwargs)
+        storage_options = kwargs.get("storage_options", {})
+        if storage_options and storage_options.get("v3io_access_key") and alt_key_name:
+            if alt_key := mlrun.get_secret_or_env(alt_key_name):
+                storage_options["v3io_access_key"] = alt_key
         super().__init__(*args, **kwargs)
 
 
@@ -89,17 +109,20 @@ class StreamStoreyTarget(storey.StreamTarget):
             raise mlrun.errors.MLRunInvalidArgumentError("StreamTarget requires a path")
 
         _, storage_options = get_url_and_storage_options(uri)
-        endpoint, path = parse_path(uri)
+        _, path = parse_path(uri)
 
         access_key = storage_options.get("v3io_access_key")
-        storage = V3ioDriver(
-            webapi=endpoint or mlrun.mlconf.v3io_api, access_key=access_key
-        )
+
+        if alt_key_name := kwargs.pop("alternative_v3io_access_key", None):
+            if alt_key := mlrun.get_secret_or_env(alt_key_name):
+                access_key = alt_key
+
+        storage = V3ioDriver(access_key=access_key)
 
         if storage_options:
             kwargs["storage"] = storage
         if args:
-            args[0] = endpoint
+            args[0] = path
         if "stream_path" in kwargs:
             kwargs["stream_path"] = path
 
@@ -108,17 +131,26 @@ class StreamStoreyTarget(storey.StreamTarget):
 
 class KafkaStoreyTarget(storey.KafkaTarget):
     def __init__(self, *args, **kwargs):
+        kwargs.pop("alternative_v3io_access_key", None)
         path = kwargs.pop("path")
-        attributes = kwargs.pop("attributes", None)
+        attributes = kwargs.pop("attributes", {})
         if path and path.startswith("ds://"):
-            datastore_profile = (
-                mlrun.datastore.datastore_profile.datastore_profile_read(path)
-            )
+            datastore_profile = datastore_profile_read(path)
+            if not isinstance(
+                datastore_profile,
+                (DatastoreProfileKafkaSource, DatastoreProfileKafkaTarget),
+            ):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Unsupported datastore profile type: {type(datastore_profile)}"
+                )
+
             attributes = merge(attributes, datastore_profile.attributes())
-            brokers = attributes.pop(
-                "brokers", attributes.pop("bootstrap_servers", None)
+            brokers = attributes.pop("brokers", None)
+            # Override the topic with the one in the url (if any)
+            parsed = urlparse(path)
+            topic = (
+                parsed.path.strip("/") if parsed.path else datastore_profile.get_topic()
             )
-            topic = datastore_profile.topic
         else:
             brokers = attributes.pop(
                 "brokers", attributes.pop("bootstrap_servers", None)
@@ -129,7 +161,10 @@ class KafkaStoreyTarget(storey.KafkaTarget):
             raise mlrun.errors.MLRunInvalidArgumentError("KafkaTarget requires a topic")
         kwargs["brokers"] = brokers
         kwargs["topic"] = topic
-        super().__init__(*args, **kwargs, **attributes)
+
+        attributes = mlrun.datastore.utils.KafkaParameters(attributes).producer()
+
+        super().__init__(*args, **kwargs, producer_options=attributes)
 
 
 class NoSqlStoreyTarget(storey.NoSqlTarget):