mlrun 1.10.0rc11__py3-none-any.whl → 1.10.0rc13__py3-none-any.whl

mlrun/datastore/datastore_profile.py

@@ -456,6 +456,36 @@ class DatastoreProfileTDEngine(DatastoreProfile):
         )
 
 
+class OpenAIProfile(DatastoreProfile):
+    type: str = pydantic.v1.Field("openai")
+    _private_attributes = "api_key"
+    api_key: typing.Optional[str] = None
+    organization: typing.Optional[str] = None
+    project: typing.Optional[str] = None
+    base_url: typing.Optional[str] = None
+    timeout: typing.Optional[float] = None
+    max_retries: typing.Optional[int] = None
+
+    def secrets(self) -> dict:
+        res = {}
+        if self.api_key:
+            res["OPENAI_API_KEY"] = self.api_key
+        if self.organization:
+            res["OPENAI_ORG_ID"] = self.organization
+        if self.project:
+            res["OPENAI_PROJECT_ID"] = self.project
+        if self.base_url:
+            res["OPENAI_BASE_URL"] = self.base_url
+        if self.timeout:
+            res["OPENAI_TIMEOUT"] = self.timeout
+        if self.max_retries:
+            res["OPENAI_MAX_RETRIES"] = self.max_retries
+        return res
+
+    def url(self, subpath):
+        return f"{self.type}://{subpath.lstrip('/')}"
+
+
 _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "v3io": DatastoreProfileV3io,
     "s3": DatastoreProfileS3,
@@ -469,6 +499,7 @@ _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "hdfs": DatastoreProfileHdfs,
     "taosws": DatastoreProfileTDEngine,
     "config": ConfigProfile,
+    "openai": OpenAIProfile,
 }
 
 

mlrun/datastore/dbfs_store.py

@@ -104,7 +104,7 @@ class DBFSStore(DataStore):
             token=self._get_secret_or_env("DATABRICKS_TOKEN"),
             instance=self._get_secret_or_env("DATABRICKS_HOST"),
         )
-        return self._sanitize_storage_options(res)
+        return self._sanitize_options(res)
 
     def _verify_filesystem_and_key(self, key: str):
         if not self.filesystem:

mlrun/datastore/google_cloud_storage.py

@@ -105,12 +105,12 @@ class GoogleCloudStorageStore(DataStore):
             except json.JSONDecodeError:
                 # If it's not json, handle it as a filename
                 token = credentials
-            return self._sanitize_storage_options(dict(token=token))
+            return self._sanitize_options(dict(token=token))
         else:
             logger.info(
                 "No GCS credentials available - auth will rely on auto-discovery of credentials"
             )
-            return self._sanitize_storage_options(None)
+            return self._sanitize_options(None)
 
     def get_storage_options(self):
         return self.storage_options

mlrun/datastore/model_provider/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 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.

mlrun/datastore/model_provider/model_provider.py

@@ -0,0 +1,160 @@
+# Copyright 2025 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.
+from collections.abc import Awaitable
+from typing import Callable, Optional, TypeVar, Union
+
+import mlrun.errors
+from mlrun.datastore.remote_client import (
+    BaseRemoteClient,
+)
+
+T = TypeVar("T")
+
+
+class ModelProvider(BaseRemoteClient):
+    """
+    The ModelProvider class is an abstract base for integrating with external
+    model providers, primarily generative AI (GenAI) services.
+
+    Designed to be subclassed, it defines a consistent interface and shared
+    functionality for tasks such as text generation, embeddings, and invoking
+    fine-tuned models. Subclasses should implement provider-specific logic,
+    including SDK client initialization, model invocation, and custom operations.
+
+    Key Features:
+    - Establishes a consistent, reusable client management for model provider integrations.
+    - Simplifies GenAI service integration by abstracting common operations.
+    - Reduces duplication through shared components for common tasks.
+    - Holds default invocation parameters (e.g., temperature, max_tokens) to avoid boilerplate
+    code and promote consistency.
+    """
+
+    support_async = False
+
+    def __init__(
+        self,
+        parent,
+        kind,
+        name,
+        endpoint="",
+        secrets: Optional[dict] = None,
+        default_invoke_kwargs: Optional[dict] = None,
+    ):
+        super().__init__(
+            parent=parent, name=name, kind=kind, endpoint=endpoint, secrets=secrets
+        )
+        self.default_invoke_kwargs = default_invoke_kwargs or {}
+        self._client = None
+        self._default_operation = None
+        self._async_client = None
+        self._default_async_operation = None
+
+    def load_client(self) -> None:
+        """
+        Initializes the SDK client for the model provider with the given keyword arguments
+        and assigns it to an instance attribute (e.g., self._client).
+
+        Subclasses should override this method to:
+        - Create and configure the provider-specific client instance.
+        - Assign the client instance to self._client.
+        - Define a default operation callable (e.g., a method to invoke model completions)
+        and assign it to self._default_operation.
+        """
+
+        raise NotImplementedError("load_client method is not implemented")
+
+    def invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Optional[Union[str, T]]:
+        """
+        Invokes a generative AI model with the provided messages and additional parameters.
+        This method is designed to be a flexible interface for interacting with various
+        generative AI backends (e.g., OpenAI, Hugging Face, etc.). It allows users to send
+        a list of messages (following a standardized format) and receive a response. The
+        response can be returned as plain text or in its full structured format, depending
+        on the `as_str` parameter.
+
+        :param messages:    A list of dictionaries representing the conversation history or input messages.
+                            Each dictionary should follow the format::
+                            {"role": "system"| "user" | "assistant" ..., "content": "Message content as a string"}
+                            Example:
+
+                            .. code-block:: json
+
+                                [
+                                    {"role": "system", "content": "You are a helpful assistant."},
+                                    {"role": "user", "content": "What is the capital of France?"}
+                                ]
+
+                            This format is consistent across all backends. Defaults to None if no messages
+                            are provided.
+
+        :param as_str:      A boolean flag indicating whether to return the response as a plain string.
+                            - If True, the function extracts and returns the main content of the first
+                            response.
+                            - If False, the function returns the full response object,
+                            which may include additional metadata or multiple response options.
+                            Defaults to False.
+
+        :param invoke_kwargs:
+                            Additional keyword arguments to be passed to the underlying model API call.
+                            These can include parameters such as temperature, max tokens, etc.,
+                            depending on the capabilities of the specific backend being used.
+
+        :return:
+                            - If `as_str` is True: Returns the main content of the first response as a string.
+                            - If `as_str` is False: Returns the full response object.
+
+        """
+        raise NotImplementedError("invoke method is not implemented")
+
+    def customized_invoke(
+        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
+    ) -> Optional[T]:
+        raise NotImplementedError("customized_invoke method is not implemented")
+
+    @property
+    def client(self):
+        return self._client
+
+    @property
+    def model(self):
+        return None
+
+    def get_invoke_kwargs(self, invoke_kwargs):
+        kwargs = self.default_invoke_kwargs.copy()
+        kwargs.update(invoke_kwargs)
+        return kwargs
+
+    @property
+    def async_client(self):
+        if not self.support_async:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"{self.__class__.__name__} does not support async operations"
+            )
+        return self._async_client
+
+    async def async_customized_invoke(self, **kwargs):
+        raise NotImplementedError("async_customized_invoke is not implemented")
+
+    async def async_invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Awaitable[str]:
+        raise NotImplementedError("async_invoke is not implemented")

mlrun/datastore/model_provider/openai_provider.py

@@ -0,0 +1,144 @@
+# Copyright 2025 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.
+
+from typing import Callable, Optional, TypeVar, Union
+
+import mlrun
+from mlrun.datastore.model_provider.model_provider import ModelProvider
+
+T = TypeVar("T")
+
+
+class OpenAIProvider(ModelProvider):
+    """
+    OpenAIProvider is a wrapper around the OpenAI SDK that provides an interface
+    for interacting with OpenAI's generative AI services.
+
+    It supports both synchronous and asynchronous operations, allowing flexible
+    integration into various workflows.
+
+    This class extends the ModelProvider base class and implements OpenAI-specific
+    functionality, including client initialization, model invocation, and custom
+    operations tailored to the OpenAI API.
+    """
+
+    def __init__(
+        self,
+        parent,
+        schema,
+        name,
+        endpoint="",
+        secrets: Optional[dict] = None,
+        default_invoke_kwargs: Optional[dict] = None,
+    ):
+        endpoint = endpoint or mlrun.mlconf.model_providers.openai_default_model
+        if schema != "openai":
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "OpenAIProvider supports only 'openai' as the provider kind."
+            )
+        super().__init__(
+            parent=parent,
+            kind=schema,
+            name=name,
+            endpoint=endpoint,
+            secrets=secrets,
+            default_invoke_kwargs=default_invoke_kwargs,
+        )
+        self.options = self.get_client_options()
+        self.load_client()
+
+    @classmethod
+    def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
+        if endpoint and subpath:
+            endpoint = endpoint + subpath
+            #  in openai there is no usage of subpath variable. if the model contains "/", it is part of the model name.
+            subpath = ""
+        return endpoint, subpath
+
+    @property
+    def model(self):
+        return self.endpoint
+
+    def load_client(self) -> None:
+        """
+        Initializes the OpenAI SDK client using the provided options.
+
+        This method imports the `OpenAI` class from the `openai` package, instantiates
+        a client with the given keyword arguments (`self.options`), and assigns it to
+        `self._client`.
+
+        It also sets the default operation to `self.client.chat.completions.create`, which is
+        typically used for invoking chat-based model completions.
+
+        Raises:
+            ImportError: If the `openai` package is not installed.
+        """
+        try:
+            from openai import OpenAI  # noqa
+
+            self._client = OpenAI(**self.options)
+            self._default_operation = self.client.chat.completions.create
+        except ImportError as exc:
+            raise ImportError("openai package is not installed") from exc
+
+    def get_client_options(self):
+        res = dict(
+            api_key=self._get_secret_or_env("OPENAI_API_KEY"),
+            organization=self._get_secret_or_env("OPENAI_ORG_ID"),
+            project=self._get_secret_or_env("OPENAI_PROJECT_ID"),
+            base_url=self._get_secret_or_env("OPENAI_BASE_URL"),
+            timeout=self._get_secret_or_env("OPENAI_TIMEOUT"),
+            max_retries=self._get_secret_or_env("OPENAI_MAX_RETRIES"),
+        )
+        return self._sanitize_options(res)
+
+    def customized_invoke(
+        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
+    ) -> Optional[T]:
+        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
+        if operation:
+            return operation(**invoke_kwargs, model=self.model)
+        else:
+            return self._default_operation(**invoke_kwargs, model=self.model)
+
+    def invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Optional[Union[str, T]]:
+        """
+        OpenAI-specific implementation of `ModelProvider.invoke`.
+        Invokes an OpenAI model operation using the sync client.
+        For full details, see `ModelProvider.invoke`.
+
+        :param messages:    Same as ModelProvider.invoke.
+
+        :param as_str: bool
+                            If `True`, returns only the main content of the first response
+                            (`response.choices[0].message.content`).
+                            If `False`, returns the full response object, whose type depends on
+                            the specific OpenAI SDK operation used (e.g., chat completion, completion, etc.).
+
+        :param invoke_kwargs:
+                            Same as ModelProvider.invoke.
+
+        """
+        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
+        response = self._default_operation(
+            model=self.endpoint, messages=messages, **invoke_kwargs
+        )
+        if as_str:
+            return response.choices[0].message.content
+        return response

mlrun/datastore/remote_client.py

@@ -0,0 +1,65 @@
+# Copyright 2025 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.
+
+from typing import Optional
+
+import mlrun
+
+
+class BaseRemoteClient:
+    """
+    The BaseRemoteClient class serves as a foundational component for managing
+    secrets and configurations.
+    It is designed to be extended by subclasses that interact with external services,
+     such as file systems (e.g., Datastore) or model providers (e.g., ModelProvider).
+
+    This class is intended to provide shared functionality and should not be
+    used directly. Instead, create a subclass to implement logic specific to
+    your use case, such as interactions with S3 storage or invoking model providers like OpenAI.
+    """
+
+    def __init__(self, parent, kind, name, endpoint="", secrets: Optional[dict] = None):
+        self._parent = parent
+        self.kind = kind
+        self.name = name
+        self.endpoint = endpoint
+        self._secrets = secrets or {}
+        self.secret_pfx = ""
+
+    def _get_secret_or_env(self, key, default=None):
+        # Project-secrets are mounted as env variables whose name can be retrieved from SecretsStore
+        return mlrun.get_secret_or_env(
+            key, secret_provider=self._get_secret, default=default
+        )
+
+    def _get_parent_secret(self, key):
+        return self._parent.secret(self.secret_pfx + key)
+
+    def _get_secret(self, key: str, default=None):
+        return self._secrets.get(key, default) or self._get_parent_secret(key)
+
+    @property
+    def url(self):
+        return f"{self.kind}://{self.endpoint}"
+
+    @staticmethod
+    def _sanitize_options(options):
+        if not options:
+            return {}
+        options = {k: v for k, v in options.items() if v is not None and v != ""}
+        return options
+
+    @classmethod
+    def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
+        return endpoint, subpath

mlrun/datastore/s3.py

@@ -186,7 +186,7 @@ class S3Store(DataStore):
         if profile:
             storage_options["profile"] = profile
 
-        return self._sanitize_storage_options(storage_options)
+        return self._sanitize_options(storage_options)
 
     @property
     def spark_url(self):

mlrun/datastore/storeytargets.py

@@ -46,7 +46,7 @@ def get_url_and_storage_options(path, external_storage_options=None):
         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)
+    return url, DataStore._sanitize_options(storage_options)
 
 
 class TDEngineStoreyTarget(storey.TDEngineTarget):

mlrun/datastore/utils.py

@@ -311,3 +311,25 @@ class KafkaParameters:
             valid_keys.update(ref_dict.keys())
         # Return a new dictionary with only valid keys
         return {k: v for k, v in input_dict.items() if k in valid_keys}
+
+
+def parse_url(url):
+    if url and url.startswith("v3io://") and not url.startswith("v3io:///"):
+        url = url.replace("v3io://", "v3io:///", 1)
+    parsed_url = urlparse(url)
+    schema = parsed_url.scheme.lower()
+    endpoint = parsed_url.hostname
+    if endpoint:
+        # HACK - urlparse returns the hostname after in lower case - we want the original case:
+        # the hostname is a substring of the netloc, in which it's the original case, so we find the indexes of the
+        # hostname in the netloc and take it from there
+        lower_hostname = parsed_url.hostname
+        netloc = str(parsed_url.netloc)
+        lower_netloc = netloc.lower()
+        hostname_index_in_netloc = lower_netloc.index(str(lower_hostname))
+        endpoint = netloc[
+            hostname_index_in_netloc : hostname_index_in_netloc + len(lower_hostname)
+        ]
+    if parsed_url.port:
+        endpoint += f":{parsed_url.port}"
+    return schema, endpoint, parsed_url

mlrun/datastore/v3io.py

@@ -97,7 +97,7 @@ class V3ioStore(DataStore):
             v3io_access_key=self._get_secret_or_env("V3IO_ACCESS_KEY"),
             v3io_api=mlrun.mlconf.v3io_api,
         )
-        return self._sanitize_storage_options(res)
+        return self._sanitize_options(res)
 
     def _upload(
         self,

mlrun/db/base.py

@@ -44,7 +44,7 @@ class RunDBInterface(ABC):
         pass
 
     @abstractmethod
-    def get_log(self, uid, project="", offset=0, size=0):
+    def get_log(self, uid, project="", offset=0, size=0, attempt=None):
         pass
 
     @abstractmethod

mlrun/db/httpdb.py

@@ -608,7 +608,7 @@ class HTTPRunDB(RunDBInterface):
         error = f"store log {project}/{uid}"
         self.api_call("POST", path, error, params, body)
 
-    def get_log(self, uid, project="", offset=0, size=None):
+    def get_log(self, uid, project="", offset=0, size=None, attempt=None):
         """Retrieve 1 MB data of log.
 
         :param uid: Log unique ID
@@ -616,6 +616,8 @@ class HTTPRunDB(RunDBInterface):
         :param offset: Retrieve partial log, get up to ``size`` bytes starting at offset ``offset``
             from beginning of log (must be >= 0)
         :param size: If set to ``-1`` will retrieve and print all data to end of the log by chunks of 1MB each.
+        :param attempt: For retriable runs, the attempt number to retrieve the log for.
+            1 is the initial attempt.
         :returns: The following objects:
 
             - state - The state of the runtime object which generates this log, if it exists. In case no known state
@@ -636,6 +638,8 @@ class HTTPRunDB(RunDBInterface):
             return state, offset
 
         params = {"offset": offset, "size": size}
+        if attempt:
+            params["attempt"] = attempt
         path = self._path_of("logs", project, uid)
         error = f"get log {project}/{uid}"
         resp = self.api_call("GET", path, error, params=params)
@@ -658,7 +662,7 @@ class HTTPRunDB(RunDBInterface):
         resp = self.api_call("GET", path, error)
         return resp.json()["size"]
 
-    def watch_log(self, uid, project="", watch=True, offset=0):
+    def watch_log(self, uid, project="", watch=True, offset=0, attempt=None):
         """Retrieve logs of a running process by chunks of 1MB, and watch the progress of the execution until it
         completes. This method will print out the logs and continue to periodically poll for, and print,
         new logs as long as the state of the runtime which generates this log is either ``pending`` or ``running``.
@@ -668,10 +672,11 @@ class HTTPRunDB(RunDBInterface):
         :param watch: If set to ``True`` will continue tracking the log as described above. Otherwise this function
             is practically equivalent to the :py:func:`~get_log` function.
         :param offset: Minimal offset in the log to watch.
+        :param attempt: For retriable runs, the attempt number to retrieve the log for. 1 is the initial attempt.
         :returns: The final state of the log being watched and the final offset.
         """
 
-        state, text = self.get_log(uid, project, offset=offset)
+        state, text = self.get_log(uid, project, offset=offset, attempt=attempt)
         if text:
             print(text.decode(errors=mlrun.mlconf.httpdb.logs.decode.errors))
         nil_resp = 0
@@ -687,7 +692,7 @@ class HTTPRunDB(RunDBInterface):
                         mlrun.mlconf.httpdb.logs.pull_logs_backoff_no_logs_default_interval
                     )
                 )
-            state, text = self.get_log(uid, project, offset=offset)
+            state, text = self.get_log(uid, project, offset=offset, attempt=attempt)
             if text:
                 nil_resp = 0
                 print(

mlrun/db/nopdb.py

@@ -63,7 +63,7 @@ class NopDB(RunDBInterface):
     def store_log(self, uid, project="", body=None, append=False):
         pass
 
-    def get_log(self, uid, project="", offset=0, size=0):
+    def get_log(self, uid, project="", offset=0, size=0, attempt=None):
         pass
 
     def store_run(self, struct, uid, project="", iter=0):