mlrun 1.10.0rc40__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/runtimes/pod.py

@@ -20,12 +20,14 @@ import typing
 import warnings
 from collections.abc import Iterable
 from enum import Enum
+from typing import Optional
 
 import dotenv
 import kubernetes.client as k8s_client
 from kubernetes.client import V1Volume, V1VolumeMount
 
 import mlrun.common.constants
+import mlrun.common.secrets
 import mlrun.errors
 import mlrun.runtimes.mounts
 import mlrun.utils.regex
@@ -708,19 +710,45 @@ class KubeResource(BaseRuntime):
     def spec(self, spec):
         self._spec = self._verify_dict(spec, "spec", KubeResourceSpec)
 
-    def set_env_from_secret(self, name, secret=None, secret_key=None):
-        """set pod environment var from secret"""
-        secret_key = secret_key or name
+    def set_env_from_secret(
+        self,
+        name: str,
+        secret: Optional[str] = None,
+        secret_key: Optional[str] = None,
+    ):
+        """
+        Set an environment variable from a Kubernetes Secret.
+        Client-side guard forbids MLRun internal auth/project secrets; no-op on API.
+        """
+        mlrun.common.secrets.validate_not_forbidden_secret(secret)
+        key = secret_key or name
         value_from = k8s_client.V1EnvVarSource(
-            secret_key_ref=k8s_client.V1SecretKeySelector(name=secret, key=secret_key)
+            secret_key_ref=k8s_client.V1SecretKeySelector(name=secret, key=key)
         )
-        return self._set_env(name, value_from=value_from)
+        return self._set_env(name=name, value_from=value_from)
 
-    def set_env(self, name, value=None, value_from=None):
-        """set pod environment var from value"""
-        if value is not None:
-            return self._set_env(name, value=str(value))
-        return self._set_env(name, value_from=value_from)
+    def set_env(
+        self,
+        name: str,
+        value: Optional[str] = None,
+        value_from: Optional[typing.Any] = None,
+    ):
+        """
+        Set an environment variable.
+        If value comes from a Secret, validate on client-side only.
+        """
+        if value_from is not None:
+            secret_name = self._extract_secret_name_from_value_from(
+                value_from=value_from
+            )
+            if secret_name:
+                mlrun.common.secrets.validate_not_forbidden_secret(secret_name)
+            return self._set_env(name=name, value_from=value_from)
+
+        # Plain literal value path
+        return self._set_env(
+            name=name, value=(str(value) if value is not None else None)
+        )
 
     def with_annotations(self, annotations: dict):
         """set a key/value annotations in the metadata of the pod"""
@@ -1366,6 +1394,27 @@ class KubeResource(BaseRuntime):
 
         return self.status.state
 
+    @staticmethod
+    def _extract_secret_name_from_value_from(
+        value_from: typing.Any,
+    ) -> Optional[str]:
+        """Extract secret name from a V1EnvVarSource or dict representation."""
+        if isinstance(value_from, k8s_client.V1EnvVarSource):
+            if value_from.secret_key_ref:
+                return value_from.secret_key_ref.name
+        elif isinstance(value_from, dict):
+            value_from = (
+                value_from.get("valueFrom")
+                or value_from.get("value_from")
+                or value_from
+            )
+            secret_key_ref = (value_from or {}).get("secretKeyRef") or (
+                value_from or {}
+            ).get("secret_key_ref")
+            if isinstance(secret_key_ref, dict):
+                return secret_key_ref.get("name")
+        return None
+
 
 def _resolve_if_type_sanitized(attribute_name, attribute):
     attribute_config = sanitized_attributes[attribute_name]

mlrun/secrets.py

@@ -12,9 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import os
 from ast import literal_eval
+from collections.abc import Callable
 from os import environ
-from typing import Callable, Optional, Union
+from typing import Optional, Union
+
+import mlrun.auth.utils
+import mlrun.utils.helpers
+from mlrun.config import is_running_as_api
 
 from .utils import AzureVaultStore, list2dict
 
@@ -48,6 +54,11 @@ class SecretsStore:
                 self._secrets[prefix + k] = str(v)
 
         elif kind == "file":
+            # Ensure files cannot be open from inside the API
+            if is_running_as_api():
+                raise RuntimeError(
+                    "add_source of kind 'file' is not allowed from the API"
+                )
             with open(source) as fp:
                 lines = fp.read().splitlines()
                 secrets_dict = list2dict(lines)
@@ -191,7 +202,7 @@ def get_secret_or_env(
         key = f"{prefix}_{key}"
 
     if secret_provider:
-        if isinstance(secret_provider, (dict, SecretsStore)):
+        if isinstance(secret_provider, dict | SecretsStore):
             secret_value = secret_provider.get(key)
         else:
             secret_value = secret_provider(key)
@@ -243,3 +254,36 @@ def _find_value_in_json_env_lists(
                     if value_in_entry:
                         return value_in_entry
     return None
+
+
+@mlrun.utils.iguazio_v4_only
+def sync_secret_tokens() -> None:
+    """
+    Synchronize local secret tokens with the backend. Doesn't sync when running from a runtime.
+
+    This function:
+      1. Reads the local token file (defaults to `mlrun.mlconf.auth_with_oauth_token.token_file` value).
+      2. Validates its content and converts validated tokens into `SecretToken` objects.
+      3. Uploads the tokens to the backend.
+      4. Logs a warning if any tokens were updated on the backend due to newer
+         expiration times found locally.
+    """
+
+    # Do not sync tokens from the file when using the offline token environment variable.
+    # The offline token from the env var takes precedence over the file.
+    # Using the env var is not the recommended approach, and tokens from the env var
+    # will not be saved as secrets in the backend.
+    if os.getenv("MLRUN_AUTH_OFFLINE_TOKEN") or mlrun.utils.is_running_in_runtime():
+        return
+
+    # The import is needed here to prevent a circular import, since this method is called from the mlrun.db connection.
+    from mlrun.db import get_run_db
+
+    secret_tokens = mlrun.auth.utils.load_and_prepare_secret_tokens(
+        auth_user_id=get_run_db().token_provider.authenticated_user_id
+    )
+
+    # The log_warning=False flag ensures the SDK doesn't log
+    # unnecessary warnings about local file updates, since
+    # this method reads from the file, not updates it.
+    get_run_db().store_secret_tokens(secret_tokens, log_warning=False)

mlrun/serving/__init__.py

@@ -27,6 +27,7 @@ __all__ = [
     "ModelRunner",
     "Model",
     "ModelSelector",
+    "ModelRunnerSelector",
     "MonitoredStep",
     "LLModel",
 ]
@@ -47,6 +48,7 @@ from .states import (
     ModelRunner,
     Model,
     ModelSelector,
+    ModelRunnerSelector,
     MonitoredStep,
     LLModel,
 )  # noqa

mlrun/serving/remote.py

@@ -168,7 +168,7 @@ class RemoteStep(storey.SendToHttp):
                 text = await resp.text()
                 raise RuntimeError(f"bad http response {resp.status}: {text}")
             return resp
-        except asyncio.TimeoutError as exc:
+        except TimeoutError as exc:
             logger.error(f"http request to {url} timed out in RemoteStep {self.name}")
             raise exc
 
@@ -241,7 +241,7 @@ class RemoteStep(storey.SendToHttp):
             headers[event_id_key] = event.id
         if method == "GET":
             body = None
-        elif body is not None and not isinstance(body, (str, bytes)):
+        elif body is not None and not isinstance(body, str | bytes):
             if self._body_function_handler:
                 body = self._body_function_handler(body)
             body = json.dumps(body)
@@ -253,7 +253,7 @@ class RemoteStep(storey.SendToHttp):
         if (
             self.return_json
             or headers.get("content-type", "").lower() == "application/json"
-        ) and isinstance(data, (str, bytes)):
+        ) and isinstance(data, str | bytes):
             data = json.loads(data)
         return data
 
@@ -390,7 +390,7 @@ class BatchHttpRequests(_ConcurrentJobExecution):
 
             if is_get:
                 body = None
-            elif body is not None and not isinstance(body, (str, bytes)):
+            elif body is not None and not isinstance(body, str | bytes):
                 if self._body_function_handler:
                     body = self._body_function_handler(body)
                 body = json.dumps(body)
@@ -458,7 +458,7 @@ class BatchHttpRequests(_ConcurrentJobExecution):
         if (
             self.return_json
             or headers.get("content-type", "").lower() == "application/json"
-        ) and isinstance(data, (str, bytes)):
+        ) and isinstance(data, str | bytes):
             data = json.loads(data)
         return data
 

mlrun/serving/routers.py

@@ -986,7 +986,7 @@ class VotingEnsemble(ParallelRun):
         List
             The model's predictions
         """
-        if isinstance(response, (list, numpy.ndarray)):
+        if isinstance(response, list | numpy.ndarray):
             return response
         try:
             self.format_response_with_col_name_flag = True
@@ -1123,7 +1123,7 @@ class EnrichmentModelRouter(ModelRouter):
 
     def preprocess(self, event):
         """Turn an entity identifier (source) to a Feature Vector"""
-        if isinstance(event.body, (str, bytes)):
+        if isinstance(event.body, str | bytes):
             event.body = json.loads(event.body)
         event.body["inputs"] = self._feature_service.get(
             event.body["inputs"], as_list=True
@@ -1275,7 +1275,7 @@ class EnrichmentVotingEnsemble(VotingEnsemble):
         """
         Turn an entity identifier (source) to a Feature Vector
         """
-        if isinstance(event.body, (str, bytes)):
+        if isinstance(event.body, str | bytes):
             event.body = json.loads(event.body)
         event.body["inputs"] = self._feature_service.get(
             event.body["inputs"], as_list=True

mlrun/serving/server.py

@@ -24,7 +24,7 @@ import socket
 import traceback
 import uuid
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from typing import Any, Optional, Union
 
 import pandas as pd
@@ -303,7 +303,7 @@ class GraphServer(ModelObj):
             if event_path_key in event.headers:
                 event.path = event.headers.get(event_path_key)
 
-        if isinstance(event.body, (str, bytes)) and (
+        if isinstance(event.body, str | bytes) and (
             not event.content_type or event.content_type in ["json", "application/json"]
         ):
             # assume it is json and try to load
@@ -348,7 +348,7 @@ class GraphServer(ModelObj):
         ):
             return body
 
-        if body and not isinstance(body, (str, bytes)):
+        if body and not isinstance(body, str | bytes):
             body = json.dumps(body)
             return context.Response(
                 body=body, content_type="application/json", status_code=200
@@ -363,8 +363,6 @@ class GraphServer(ModelObj):
 def add_error_raiser_step(
     graph: RootFlowStep, monitored_steps: dict[str, MonitoredStep]
 ) -> RootFlowStep:
-    monitored_steps_raisers = {}
-    user_steps = list(graph.steps.values())
     for monitored_step in monitored_steps.values():
         error_step = graph.add_step(
             class_name="mlrun.serving.states.ModelRunnerErrorRaiser",
@@ -379,21 +377,7 @@ def add_error_raiser_step(
         if monitored_step.responder:
             monitored_step.responder = False
             error_step.respond()
-        monitored_steps_raisers[monitored_step.name] = error_step.name
         error_step.on_error = monitored_step.on_error
-    if monitored_steps_raisers:
-        for step in user_steps:
-            if step.after:
-                if isinstance(step.after, list):
-                    for i in range(len(step.after)):
-                        if step.after[i] in monitored_steps_raisers:
-                            step.after[i] = monitored_steps_raisers[step.after[i]]
-                else:
-                    if (
-                        isinstance(step.after, str)
-                        and step.after in monitored_steps_raisers
-                    ):
-                        step.after = monitored_steps_raisers[step.after]
     return graph
 
 
@@ -649,7 +633,7 @@ async def async_execute_graph(
 
     if df.empty:
         context.logger.warn("Job terminated due to empty inputs (0 rows)")
-        return []
+        return
 
     track_models = spec.get("track_models")
 
@@ -676,7 +660,7 @@ async def async_execute_graph(
             start_time = end_time = df["timestamp"].iloc[0].isoformat()
     else:
         # end time will be set from clock time when the batch completes
-        start_time = datetime.now(tz=timezone.utc).isoformat()
+        start_time = datetime.now(tz=UTC).isoformat()
 
     server.graph = add_system_steps_to_graph(
         server.project,
@@ -756,7 +740,7 @@ async def async_execute_graph(
     server = GraphServer.from_dict(spec)
     server.init_states(None, namespace)
 
-    batch_completion_time = datetime.now(tz=timezone.utc).isoformat()
+    batch_completion_time = datetime.now(tz=UTC).isoformat()
 
     if not timestamp_column:
         end_time = batch_completion_time
@@ -779,30 +763,49 @@ async def async_execute_graph(
         model_endpoint_uids=model_endpoint_uids,
     )
 
-    # log the results as artifacts
-    num_of_meps_in_the_graph = len(server.graph.model_endpoints_names)
-    artifact_path = None
-    if (
-        "{{run.uid}}" not in context.artifact_path
-    ):  # TODO: delete when IG-22841 is resolved
-        artifact_path = "+/{{run.uid}}"  # will be concatenated to the context's path in extend_artifact_path
-    if num_of_meps_in_the_graph <= 1:
+    has_responder = False
+    for step in server.graph.steps.values():
+        if getattr(step, "responder", False):
+            has_responder = True
+            break
+
+    if has_responder:
+        # log the results as a dataset artifact
+        artifact_path = None
+        if (
+            "{{run.uid}}" not in context.artifact_path
+        ):  # TODO: delete when IG-22841 is resolved
+            artifact_path = "+/{{run.uid}}"  # will be concatenated to the context's path in extend_artifact_path
         context.log_dataset(
             "prediction", df=pd.DataFrame(responses), artifact_path=artifact_path
         )
-    else:
-        # turn this list of samples into a dict of lists, one per model endpoint
-        grouped = defaultdict(list)
-        for sample in responses:
-            for model_name, features in sample.items():
-                grouped[model_name].append(features)
-        # create a dataframe per model endpoint and log it
-        for model_name, features in grouped.items():
-            context.log_dataset(
-                f"prediction_{model_name}",
-                df=pd.DataFrame(features),
-                artifact_path=artifact_path,
-            )
+
+        # if we got responses that appear to be in the right format, try to log per-model datasets too
+        if (
+            responses
+            and responses[0]
+            and isinstance(responses[0], dict)
+            and isinstance(next(iter(responses[0].values())), dict | list)
+        ):
+            try:
+                # turn this list of samples into a dict of lists, one per model endpoint
+                grouped = defaultdict(list)
+                for sample in responses:
+                    for model_name, features in sample.items():
+                        grouped[model_name].append(features)
+                # create a dataframe per model endpoint and log it
+                for model_name, features in grouped.items():
+                    context.log_dataset(
+                        f"prediction_{model_name}",
+                        df=pd.DataFrame(features),
+                        artifact_path=artifact_path,
+                    )
+            except Exception as e:
+                context.logger.warning(
+                    "Failed to log per-model prediction datasets",
+                    error=err_to_str(e),
+                )
+
     context.log_result("num_rows", run_call_count)
 
 

mlrun/serving/serving_wrapper.py

@@ -11,6 +11,7 @@
 # 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 asyncio
 
 # serving runtime hooks, used in empty serving functions
 from mlrun.runtimes import nuclio_init_hook
@@ -20,5 +21,8 @@ def init_context(context):
     nuclio_init_hook(context, globals(), "serving_v2")
 
 
-def handler(context, event):
-    return context.mlrun_handler(context, event)
+async def handler(context, event):
+    result = context.mlrun_handler(context, event)
+    if asyncio.iscoroutine(result):
+        return await result
+    return result