mlrun 1.7.1rc10__py3-none-any.whl → 1.8.0rc11__py3-none-any.whl

mlrun/run.py

@@ -30,15 +30,15 @@ from typing import Optional, Union
 
 import nuclio
 import yaml
-from mlrun_pipelines.common.models import RunStatuses
-from mlrun_pipelines.common.ops import format_summary_from_kfp_run, show_kfp_run
-from mlrun_pipelines.utils import get_client
 
 import mlrun.common.constants as mlrun_constants
 import mlrun.common.formatters
 import mlrun.common.schemas
 import mlrun.errors
 import mlrun.utils.helpers
+from mlrun_pipelines.common.models import RunStatuses
+from mlrun_pipelines.common.ops import format_summary_from_kfp_run, show_kfp_run
+from mlrun_pipelines.utils import get_client
 
 from .common.helpers import parse_versioned_object_uri
 from .config import config as mlconf
@@ -73,6 +73,9 @@ from .utils import (
     update_in,
 )
 
+if typing.TYPE_CHECKING:
+    from mlrun.datastore import DataItem
+
 
 def function_to_module(code="", workdir=None, secrets=None, silent=False):
     """Load code, notebook or mlrun function as .py module
@@ -200,7 +203,7 @@ def get_or_create_ctx(
     event=None,
     spec: Optional[dict] = None,
     with_env: bool = True,
-    rundb: str = "",
+    rundb: Union[str, "mlrun.db.RunDBInterface"] = "",
     project: str = "",
     upload_artifacts: bool = False,
     labels: Optional[dict] = None,
@@ -306,7 +309,7 @@ def get_or_create_ctx(
     out = rundb or mlconf.dbpath or environ.get("MLRUN_DBPATH")
     if out:
         autocommit = True
-        logger.info(f"logging run results to: {out}")
+        logger.info(f"Logging run results to: {out}")
 
     newspec["metadata"]["project"] = (
         newspec["metadata"].get("project") or project or mlconf.default_project
@@ -434,7 +437,7 @@ def new_function(
     mode: Optional[str] = None,
     handler: Optional[str] = None,
     source: Optional[str] = None,
-    requirements: Union[str, list[str]] = None,
+    requirements: Optional[Union[str, list[str]]] = None,
     kfp: Optional[bool] = None,
     requirements_file: str = "",
 ):
@@ -906,13 +909,53 @@ def _run_pipeline(
     return pipeline_run_id
 
 
+def retry_pipeline(
+    run_id: str,
+    project: Optional[str] = None,
+    namespace: Optional[str] = None,
+) -> str:
+    """Retry a pipeline run.
+
+    This function retries a previously executed pipeline run using the specified run ID. If the run is not in a
+    retryable state, a new run is created as a clone of the original run.
+
+    :param run_id: ID of the pipeline run to retry.
+    :param project: Optional; name of the project associated with the pipeline run.
+    :param namespace: Optional; Kubernetes namespace to use if not the default.
+
+    :returns: ID of the retried pipeline run or the ID of a cloned run if the original run is not retryable.
+    :raises ValueError: If access to the remote API service is not available.
+    """
+    mldb = mlrun.db.get_run_db()
+    if mldb.kind != "http":
+        raise ValueError(
+            "Retrying a pipeline requires access to remote API service. "
+            "Please set the dbpath URL."
+        )
+
+    pipeline_run_id = mldb.retry_pipeline(
+        run_id=run_id,
+        project=project,
+        namespace=namespace,
+    )
+    if pipeline_run_id == run_id:
+        logger.info(
+            f"Retried pipeline run ID={pipeline_run_id}, check UI for progress."
+        )
+    else:
+        logger.info(
+            f"Copy of pipeline {run_id} was retried as run ID={pipeline_run_id}, check UI for progress."
+        )
+    return pipeline_run_id
+
+
 def wait_for_pipeline_completion(
     run_id,
     timeout=60 * 60,
-    expected_statuses: list[str] = None,
+    expected_statuses: Optional[list[str]] = None,
     namespace=None,
     remote=True,
-    project: str = None,
+    project: Optional[str] = None,
 ):
     """Wait for Pipeline status, timeout in sec
 
@@ -1004,7 +1047,7 @@ def get_pipeline(
     format_: Union[
         str, mlrun.common.formatters.PipelineFormat
     ] = mlrun.common.formatters.PipelineFormat.summary,
-    project: str = None,
+    project: Optional[str] = None,
     remote: bool = True,
 ):
     """Get Pipeline status
@@ -1089,7 +1132,7 @@ def get_object(url, secrets=None, size=None, offset=0, db=None):
     return stores.object(url=url).get(size, offset)
 
 
-def get_dataitem(url, secrets=None, db=None) -> mlrun.datastore.DataItem:
+def get_dataitem(url, secrets=None, db=None) -> "DataItem":
     """get mlrun dataitem object (from path/url)"""
     stores = store_manager.set(secrets, db=db)
     return stores.object(url=url)

mlrun/runtimes/__init__.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
-
 __all__ = [
     "BaseRuntime",
     "KubejobRuntime",
@@ -60,7 +58,7 @@ from ..serving import MLModelServer, new_v1_model_server  # noqa isort: skip
 def new_model_server(
     name,
     model_class: str,
-    models: dict = None,
+    models: typing.Optional[dict] = None,
     filename="",
     protocol="",
     image="",

mlrun/runtimes/base.py

@@ -21,20 +21,21 @@ from os import environ
 from typing import Callable, Optional, Union
 
 import requests.exceptions
-from mlrun_pipelines.common.ops import mlrun_op
 from nuclio.build import mlrun_footer
 
 import mlrun.common.constants
 import mlrun.common.constants as mlrun_constants
+import mlrun.common.formatters
+import mlrun.common.runtimes
 import mlrun.common.schemas
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
-import mlrun.db
 import mlrun.errors
 import mlrun.launcher.factory
 import mlrun.utils.helpers
 import mlrun.utils.notifications
 import mlrun.utils.regex
 from mlrun.utils.helpers import generate_object_uri, verify_field_regex
+from mlrun_pipelines.common.ops import mlrun_op
 
 from ..config import config
 from ..datastore import store_manager
@@ -517,7 +518,7 @@ class BaseRuntime(ModelObj):
                 run_results = resp["status"].get("results", {})
                 if generator.eval_stop_condition(run_results):
                     logger.info(
-                        f"reached early stop condition ({generator.options.stop_condition}), stopping iterations!"
+                        f"Reached early stop condition ({generator.options.stop_condition}), stopping iterations!"
                     )
                     results.append(resp)
                     break
@@ -529,7 +530,7 @@ class BaseRuntime(ModelObj):
                 resp = self._update_run_state(task=task, err=error_string)
                 num_errors += 1
                 if num_errors > generator.max_errors:
-                    logger.error("too many errors, stopping iterations!")
+                    logger.error("Too many errors, stopping iterations!")
                     results.append(resp)
                     break
 
@@ -553,9 +554,9 @@ class BaseRuntime(ModelObj):
 
     def _update_run_state(
         self,
-        resp: dict = None,
+        resp: Optional[dict] = None,
         task: RunObject = None,
-        err: Union[Exception, str] = None,
+        err: Optional[Union[Exception, str]] = None,
         run_format: mlrun.common.formatters.RunFormat = mlrun.common.formatters.RunFormat.full,
     ) -> typing.Optional[dict]:
         """update the task state in the DB"""
@@ -644,7 +645,10 @@ class BaseRuntime(ModelObj):
         return updates
 
     def full_image_path(
-        self, image=None, client_version: str = None, client_python_version: str = None
+        self,
+        image=None,
+        client_version: Optional[str] = None,
+        client_python_version: Optional[str] = None,
     ):
         image = image or self.spec.image or ""
 
@@ -669,16 +673,16 @@ class BaseRuntime(ModelObj):
         handler=None,
         name: str = "",
         project: str = "",
-        params: dict = None,
+        params: Optional[dict] = None,
         hyperparams=None,
         selector="",
         hyper_param_options: HyperParamOptions = None,
-        inputs: dict = None,
-        outputs: list = None,
+        inputs: Optional[dict] = None,
+        outputs: Optional[list] = None,
         workdir: str = "",
         artifact_path: str = "",
         image: str = "",
-        labels: dict = None,
+        labels: Optional[dict] = None,
         use_db=True,
         verbose=None,
         scrape_metrics=False,

mlrun/runtimes/daskjob.py

@@ -363,7 +363,7 @@ class DaskCluster(KubejobRuntime):
         skip_deployed=False,
         is_kfp=False,
         mlrun_version_specifier=None,
-        builder_env: dict = None,
+        builder_env: Optional[dict] = None,
         show_on_failure: bool = False,
         force_build: bool = False,
     ):
@@ -406,9 +406,9 @@ class DaskCluster(KubejobRuntime):
 
     def with_scheduler_limits(
         self,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: Optional[str] = None,
+        cpu: Optional[str] = None,
+        gpus: Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -422,9 +422,9 @@ class DaskCluster(KubejobRuntime):
 
     def with_worker_limits(
         self,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: Optional[str] = None,
+        cpu: Optional[str] = None,
+        gpus: Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -442,7 +442,7 @@ class DaskCluster(KubejobRuntime):
         )
 
     def with_scheduler_requests(
-        self, mem: str = None, cpu: str = None, patch: bool = False
+        self, mem: Optional[str] = None, cpu: Optional[str] = None, patch: bool = False
     ):
         """
         set scheduler pod resources requests
@@ -451,7 +451,7 @@ class DaskCluster(KubejobRuntime):
         self.spec._verify_and_set_requests("scheduler_resources", mem, cpu, patch=patch)
 
     def with_worker_requests(
-        self, mem: str = None, cpu: str = None, patch: bool = False
+        self, mem: Optional[str] = None, cpu: Optional[str] = None, patch: bool = False
     ):
         """
         set worker pod resources requests

mlrun/runtimes/generators.py

@@ -15,6 +15,7 @@ import json
 import random
 import sys
 from copy import deepcopy
+from typing import Optional
 
 import pandas as pd
 
@@ -26,7 +27,7 @@ default_max_iterations = 10
 default_max_errors = 3
 
 
-def get_generator(spec: RunSpec, execution, param_file_secrets: dict = None):
+def get_generator(spec: RunSpec, execution, param_file_secrets: Optional[dict] = None):
     options = spec.hyper_param_options
     strategy = spec.strategy or options.strategy
     if not spec.is_hyper_job() or strategy == "custom":

mlrun/runtimes/kubejob.py

@@ -14,11 +14,10 @@
 import typing
 import warnings
 
-from mlrun_pipelines.common.ops import build_op
-
 import mlrun.common.schemas
 import mlrun.db
 import mlrun.errors
+from mlrun_pipelines.common.ops import build_op
 
 from ..model import RunObject
 from .pod import KubeResource
@@ -76,7 +75,7 @@ class KubejobRuntime(KubeResource):
         self,
         image="",
         base_image=None,
-        commands: list = None,
+        commands: typing.Optional[list] = None,
         secret=None,
         source=None,
         extra=None,
@@ -148,7 +147,7 @@ class KubejobRuntime(KubeResource):
         skip_deployed: bool = False,
         is_kfp: bool = False,
         mlrun_version_specifier: typing.Optional[bool] = None,
-        builder_env: dict = None,
+        builder_env: typing.Optional[dict] = None,
         show_on_failure: bool = False,
         force_build: bool = False,
     ) -> bool:
@@ -190,7 +189,7 @@ class KubejobRuntime(KubeResource):
         self,
         image=None,
         base_image=None,
-        commands: list = None,
+        commands: typing.Optional[list] = None,
         secret_name="",
         with_mlrun=True,
         skip_deployed=False,