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

mlrun/artifacts/llm_prompt.py

@@ -61,7 +61,11 @@ class LLMPromptArtifactSpec(ArtifactSpec):
         self.prompt_legend = prompt_legend
         self.model_configuration = model_configuration
         self.description = description
-        self._model_artifact = None
+        self._model_artifact = (
+            model_artifact
+            if isinstance(model_artifact, model_art.ModelArtifact)
+            else None
+        )
 
     @property
     def model_uri(self):

mlrun/common/schemas/__init__.py

@@ -214,7 +214,7 @@ from .secret import (
     SecretsData,
     UserSecretCreationRequest,
 )
-from .serving import ModelRunnerStepData, MonitoringData
+from .serving import ModelRunnerStepData, ModelsData, MonitoringData
 from .tag import Tag, TagObjects
 from .workflow import (
     GetWorkflowResponse,

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -336,8 +336,8 @@ class ModelEndpointMonitoringMetricNoData(_ModelEndpointMonitoringMetricValuesBa
 
 class ApplicationBaseRecord(BaseModel):
     type: Literal["metric", "result"]
-    time: datetime
     value: float
+    time: Optional[datetime] = None
 
 
 class ApplicationResultRecord(ApplicationBaseRecord):

mlrun/common/schemas/serving.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
+
 from pydantic.v1 import BaseModel
 
 from mlrun.common.types import StrEnum
@@ -40,3 +42,8 @@ class MonitoringData(StrEnum):
     MODEL_PATH = "model_path"
     MODEL_ENDPOINT_UID = "model_endpoint_uid"
     MODEL_CLASS = "model_class"
+
+
+class ModelsData(enum.Enum):
+    MODEL_CLASS = 0
+    MODEL_PARAMETERS = 1

mlrun/config.py

@@ -125,6 +125,8 @@ default_config = {
                 "interval": "30",
                 # runs limit to fetch for retrying
                 "fetch_runs_limit": 1000,
+                # minutes until a run is considered stale and will be aborted
+                "staleness_threshold": 60 * 24 * 3,
             },
         },
         "projects": {

mlrun/datastore/model_provider/model_provider.py

@@ -12,7 +12,7 @@
 # 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
+from typing import Callable, Optional, TypeVar, Union
 
 import mlrun.errors
 from mlrun.datastore.remote_client import (
@@ -23,6 +23,23 @@ 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__(
@@ -44,9 +61,65 @@ class ModelProvider(BaseRemoteClient):
         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, prompt: Optional[str] = None, **invoke_kwargs) -> str:
+    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(
@@ -78,5 +151,10 @@ class ModelProvider(BaseRemoteClient):
     async def async_customized_invoke(self, **kwargs):
         raise NotImplementedError("async_customized_invoke is not implemented")
 
-    async def async_invoke(self, prompt: str, **invoke_kwargs) -> Awaitable[str]:
+    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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional, TypeVar
+from typing import Callable, Optional, TypeVar, Union
 
 import mlrun
 from mlrun.datastore.model_provider.model_provider import ModelProvider
@@ -21,6 +21,18 @@ 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,
@@ -59,6 +71,19 @@ class OpenAIProvider(ModelProvider):
         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
 
@@ -87,34 +112,33 @@ class OpenAIProvider(ModelProvider):
         else:
             return self._default_operation(**invoke_kwargs, model=self.model)
 
-    def _get_messages_parameter(
-        self, prompt: Optional[str] = None, **invoke_kwargs
-    ) -> (str, dict):
-        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
-        messages = invoke_kwargs.get("messages")
-        if messages:
-            if prompt:
-                raise mlrun.errors.MLRunInvalidArgumentError(
-                    "can not provide 'messages' and 'prompt' to invoke"
-                )
-        elif prompt:
-            messages = [
-                {
-                    "role": "user",
-                    "content": prompt,
-                },
-            ]
-        else:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "must provide 'messages' or 'prompt' to invoke"
-            )
-        return messages, invoke_kwargs
+    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`.
 
-    def invoke(self, prompt: Optional[str] = None, **invoke_kwargs) -> str:
-        messages, invoke_kwargs = self._get_messages_parameter(
-            prompt=prompt, **invoke_kwargs
-        )
+        :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
         )
-        return response.choices[0].message.content
+        if as_str:
+            return response.choices[0].message.content
+        return response

mlrun/datastore/remote_client.py

@@ -18,6 +18,17 @@ 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

mlrun/execution.py

@@ -961,6 +961,11 @@ class MLClientCtx:
         :returns: The logged `LLMPromptArtifact` object.
         """
 
+        if not prompt_string and not prompt_path:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Either 'prompt_string' or 'prompt_path' must be provided"
+            )
+
         llm_prompt = LLMPromptArtifact(
             key=key,
             project=self.project or "",

mlrun/model_monitoring/db/tsdb/v3io/v3io_connector.py

@@ -804,25 +804,45 @@ class V3IOTSDBConnector(TSDBConnector):
     @staticmethod
     def _get_sql_query(
         *,
-        endpoint_id: str,
         table_path: str,
+        endpoint_id: Optional[str] = None,
+        application_names: Optional[list[str]] = None,
         name: str = mm_schemas.ResultData.RESULT_NAME,
         metric_and_app_names: Optional[list[tuple[str, str]]] = None,
         columns: Optional[list[str]] = None,
+        group_by_columns: Optional[list[str]] = None,
     ) -> str:
         """Get the SQL query for the results/metrics table"""
+
+        if metric_and_app_names and not endpoint_id:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "If metric_and_app_names is provided, endpoint_id must also be provided"
+            )
+
+        if metric_and_app_names and application_names:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Cannot provide both metric_and_app_names and application_names"
+            )
+
         if columns:
             selection = ",".join(columns)
         else:
             selection = "*"
 
         with StringIO() as query:
-            query.write(
-                f"SELECT {selection} FROM '{table_path}' "
-                f"WHERE {mm_schemas.WriterEvent.ENDPOINT_ID}='{endpoint_id}'"
-            )
+            where_added = False
+            query.write(f"SELECT {selection} FROM '{table_path}'")
+            if endpoint_id:
+                query.write(
+                    f" WHERE {mm_schemas.WriterEvent.ENDPOINT_ID}='{endpoint_id}'"
+                )
+                where_added = True
             if metric_and_app_names:
-                query.write(" AND (")
+                if where_added:
+                    query.write(" AND (")
+                else:
+                    query.write(" WHERE (")
+                    where_added = True
 
                 for i, (app_name, result_name) in enumerate(metric_and_app_names):
                     sub_cond = (
@@ -835,6 +855,22 @@ class V3IOTSDBConnector(TSDBConnector):
 
                 query.write(")")
 
+            if application_names:
+                if where_added:
+                    query.write(" AND (")
+                else:
+                    query.write(" WHERE (")
+                for i, app_name in enumerate(application_names):
+                    sub_cond = f"{mm_schemas.WriterEvent.APPLICATION_NAME}='{app_name}'"
+                    if i != 0:  # not first sub condition
+                        query.write(" OR ")
+                    query.write(sub_cond)
+                query.write(")")
+
+            if group_by_columns:
+                query.write(" GROUP BY ")
+                query.write(",".join(group_by_columns))
+
             query.write(";")
             return query.getvalue()
 
@@ -1272,7 +1308,49 @@ class V3IOTSDBConnector(TSDBConnector):
         end: Optional[Union[datetime, str]] = None,
         application_names: Optional[Union[str, list[str]]] = None,
     ) -> dict[str, int]:
-        raise NotImplementedError
+        start, end = get_start_end(start=start, end=end, delta=timedelta(hours=24))
+        group_by_columns = [
+            mm_schemas.ApplicationEvent.APPLICATION_NAME,
+            mm_schemas.ApplicationEvent.ENDPOINT_ID,
+        ]
+
+        def get_application_endpoints_records(
+            record_type: Literal["metrics", "results"],
+        ):
+            if record_type == "results":
+                table_path = self.tables[mm_schemas.V3IOTSDBTables.APP_RESULTS]
+            else:
+                table_path = self.tables[mm_schemas.V3IOTSDBTables.METRICS]
+            sql_query = self._get_sql_query(
+                table_path=table_path,
+                columns=[mm_schemas.WriterEvent.START_INFER_TIME],
+                group_by_columns=group_by_columns,
+                application_names=application_names,
+            )
+            return self.frames_client.read(
+                backend=_TSDB_BE,
+                start=start,
+                end=end,
+                query=sql_query,
+            )
+
+        df_results = get_application_endpoints_records("results")
+        df_metrics = get_application_endpoints_records("metrics")
+
+        if df_results.empty and df_metrics.empty:
+            return {}
+
+        # Combine the two dataframes and count unique endpoints per application
+        combined_df = pd.concat([df_results, df_metrics], ignore_index=True)
+        if combined_df.empty:
+            return {}
+        combined_df.drop_duplicates(subset=group_by_columns, inplace=True)
+
+        grouped_df = combined_df.groupby(
+            mm_schemas.WriterEvent.APPLICATION_NAME
+        ).count()
+
+        return grouped_df[mm_schemas.WriterEvent.ENDPOINT_ID].to_dict()
 
     def calculate_latest_metrics(
         self,
@@ -1282,4 +1360,93 @@ class V3IOTSDBConnector(TSDBConnector):
     ) -> list[
         Union[mm_schemas.ApplicationResultRecord, mm_schemas.ApplicationMetricRecord]
     ]:
-        raise NotImplementedError
+        metric_list = []
+        start, end = get_start_end(start=start, end=end, delta=timedelta(hours=24))
+
+        # Get the latest results
+        def get_latest_metrics_records(
+            record_type: Literal["metrics", "results"],
+        ) -> pd.DataFrame:
+            group_by_columns = [mm_schemas.ApplicationEvent.APPLICATION_NAME]
+            if record_type == "results":
+                table_path = self.tables[mm_schemas.V3IOTSDBTables.APP_RESULTS]
+                columns = [
+                    f"last({mm_schemas.ResultData.RESULT_STATUS})",
+                    f"last({mm_schemas.ResultData.RESULT_VALUE})",
+                    f"last({mm_schemas.ResultData.RESULT_KIND})",
+                ]
+                group_by_columns += [
+                    mm_schemas.ResultData.RESULT_NAME,
+                ]
+            else:
+                table_path = self.tables[mm_schemas.V3IOTSDBTables.METRICS]
+                columns = [f"last({mm_schemas.MetricData.METRIC_VALUE})"]
+                group_by_columns += [
+                    mm_schemas.MetricData.METRIC_NAME,
+                ]
+            sql_query = self._get_sql_query(
+                table_path=table_path,
+                columns=columns,
+                group_by_columns=group_by_columns,
+                application_names=application_names,
+            )
+
+            return self.frames_client.read(
+                backend=_TSDB_BE,
+                start=start,
+                end=end,
+                query=sql_query,
+            )
+
+        df_results = get_latest_metrics_records("results")
+        df_metrics = get_latest_metrics_records("metrics")
+
+        if df_results.empty and df_metrics.empty:
+            return metric_list
+
+        # Convert the results DataFrame to a list of ApplicationResultRecord
+        def build_metric_objects() -> (
+            list[
+                Union[
+                    mm_schemas.ApplicationResultRecord,
+                    mm_schemas.ApplicationMetricRecord,
+                ]
+            ]
+        ):
+            metric_objects = []
+            if not df_results.empty:
+                df_results.rename(
+                    columns={
+                        f"last({mm_schemas.ResultData.RESULT_VALUE})": mm_schemas.ResultData.RESULT_VALUE,
+                        f"last({mm_schemas.ResultData.RESULT_STATUS})": mm_schemas.ResultData.RESULT_STATUS,
+                        f"last({mm_schemas.ResultData.RESULT_KIND})": mm_schemas.ResultData.RESULT_KIND,
+                    },
+                    inplace=True,
+                )
+                for _, row in df_results.iterrows():
+                    metric_objects.append(
+                        mm_schemas.ApplicationResultRecord(
+                            result_name=row[mm_schemas.ResultData.RESULT_NAME],
+                            kind=row[mm_schemas.ResultData.RESULT_KIND],
+                            status=row[mm_schemas.ResultData.RESULT_STATUS],
+                            value=row[mm_schemas.ResultData.RESULT_VALUE],
+                        )
+                    )
+            if not df_metrics.empty:
+                df_metrics.rename(
+                    columns={
+                        f"last({mm_schemas.MetricData.METRIC_VALUE})": mm_schemas.MetricData.METRIC_VALUE,
+                    },
+                    inplace=True,
+                )
+
+                for _, row in df_metrics.iterrows():
+                    metric_objects.append(
+                        mm_schemas.ApplicationMetricRecord(
+                            metric_name=row[mm_schemas.MetricData.METRIC_NAME],
+                            value=row[mm_schemas.MetricData.METRIC_VALUE],
+                        )
+                    )
+            return metric_objects
+
+        return build_metric_objects()

mlrun/projects/pipelines.py

@@ -1081,34 +1081,56 @@ def rerun_workflow(
     :param run_uid:      The run UID of the original workflow to retry.
     :param project_name: The project name.
     """
+    db = mlrun.get_run_db()
 
     try:
-        # TODO in followups: handle start and running notifications
-
-        # Retry the pipeline  - TODO: add submit-direct flag when created
-        db = mlrun.get_run_db()
+        # Invoke the KFP retry endpoint (direct-submit mode)
         new_pipeline_id = db.retry_pipeline(
-            run_uid, project_name, submit_mode=mlrun_constants.WorkflowSubmitMode.direct
+            run_id=run_uid,
+            project=project_name,
+            submit_mode=mlrun_constants.WorkflowSubmitMode.direct,
+        )
+        logger.info(
+            "KFP retry submitted",
+            new_pipeline_id=new_pipeline_id,
+            rerun_of_workflow=run_uid,
         )
 
-        # Store result for observability
-        context.set_label(
-            mlrun_constants.MLRunInternalLabels.workflow_id, new_pipeline_id
+    except mlrun.errors.MLRunHTTPError as http_exc:
+        logger.error(
+            "Failed calling KFP retry API",
+            run_id=run_uid,
+            error=err_to_str(http_exc),
         )
-        context.update_run()
+        raise
 
-        context.log_result("workflow_id", new_pipeline_id)
+    # Enqueue "running" notifications server-side for this RerunRunner run
+    db.push_run_notifications(context.uid, project_name)
 
-        # wait for pipeline completion so monitor will push terminal notifications
-        wait_for_pipeline_completion(
+    context.set_label(mlrun_constants.MLRunInternalLabels.workflow_id, new_pipeline_id)
+    context.update_run()
+
+    context.log_result("workflow_id", new_pipeline_id)
+
+    try:
+        pipeline = wait_for_pipeline_completion(
             new_pipeline_id,
             project=project_name,
         )
-
-    # Temporary exception
     except Exception as exc:
-        context.logger.error("Failed to rerun workflow", exc=err_to_str(exc))
-        raise
+        mlrun.utils.logger.error(
+            "Failed waiting for workflow completion",
+            rerun_pipeline_id=new_pipeline_id,
+            exc=err_to_str(exc),
+        )
+    else:
+        final_state = pipeline["run"]["status"]
+        context.log_result("workflow_state", final_state, commit=True)
+
+        if final_state != mlrun_pipelines.common.models.RunStatuses.succeeded:
+            raise mlrun.errors.MLRunRuntimeError(
+                f"Pipeline retry of {run_uid} finished in state={final_state}"
+            )
 
 
 def load_and_run(context, *args, **kwargs):
@@ -1201,13 +1223,13 @@ def load_and_run_workflow(
     start_notifications = [
         notification
         for notification in context.get_notifications(unmask_secret_params=True)
-        if "running" in notification.when
+        if mlrun.common.runtimes.constants.RunStates.running in notification.when
     ]
 
     # Prevent redundant notifications for run completion by ensuring that notifications are only triggered when the run
     # reaches the "running" state, as the server already handles the completion notifications.
     for notification in start_notifications:
-        notification.when = ["running"]
+        notification.when = [mlrun.common.runtimes.constants.RunStates.running]
 
     workflow_log_message = workflow_name or workflow_path
     context.logger.info(

mlrun/projects/project.py

@@ -1942,6 +1942,11 @@ class MlrunProject(ModelObj):
         :returns: The logged `LLMPromptArtifact` object.
         """
 
+        if not prompt_string and not prompt_path:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Either 'prompt_string' or 'prompt_path' must be provided"
+            )
+
         llm_prompt = LLMPromptArtifact(
             key=key,
             project=self.name,