mlrun 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/runtimes/function_reference.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.

mlrun/runtimes/generators.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.

mlrun/runtimes/kubejob.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -14,15 +14,15 @@
 
 import os
 import time
+import warnings
 
 from kubernetes import client
 from kubernetes.client.rest import ApiException
 
-import mlrun.api.schemas
+import mlrun.common.schemas
 import mlrun.errors
 from mlrun.runtimes.base import BaseRuntimeHandler
 
-from ..builder import build_runtime
 from ..db import RunDBError
 from ..errors import err_to_str
 from ..kfpops import build_op
@@ -30,7 +30,7 @@ from ..model import RunObject
 from ..utils import get_in, logger
 from .base import RunError, RuntimeClassMode
 from .pod import KubeResource, kube_resource_spec_to_pod_spec
-from .utils import AsyncLogWriter
+from .utils import get_k8s
 
 
 class KubejobRuntime(KubeResource):
@@ -44,12 +44,13 @@ class KubejobRuntime(KubeResource):
         if self.spec.image:
             return True
 
-        if self._is_remote_api():
-            db = self._get_db()
-            try:
-                db.get_builder_status(self, logs=False)
-            except Exception:
-                pass
+        db = self._get_db()
+        try:
+            # getting builder status enriches the runtime when it needs to be fetched from the API,
+            # otherwise it's a no-op
+            db.get_builder_status(self, logs=False)
+        except Exception:
+            pass
 
         if self.spec.image:
             return True
@@ -58,27 +59,30 @@ class KubejobRuntime(KubeResource):
         return False
 
     def with_source_archive(
-        self, source, workdir=None, handler=None, pull_at_runtime=True
+        self, source, workdir=None, handler=None, pull_at_runtime=True, target_dir=None
     ):
         """load the code from git/tar/zip archive at runtime or build
 
-        :param source:     valid path to git, zip, or tar file, e.g.
-                           git://github.com/mlrun/something.git
-                           http://some/url/file.zip
-        :param handler: default function handler
-        :param workdir: working dir relative to the archive root or absolute (e.g. './subdir')
+        :param source:          valid absolute path or URL to git, zip, or tar file, e.g.
+                                git://github.com/mlrun/something.git
+                                http://some/url/file.zip
+                                note path source must exist on the image or exist locally when run is local
+                                (it is recommended to use 'workdir' when source is a filepath instead)
+        :param handler:         default function handler
+        :param workdir:         working dir relative to the archive root (e.g. './subdir') or absolute to the image root
         :param pull_at_runtime: load the archive into the container at job runtime vs on build/deploy
+        :param target_dir:      target dir on runtime pod or repo clone / archive extraction
         """
-        if source.endswith(".zip") and not pull_at_runtime:
-            logger.warn(
-                "zip files are not natively extracted by docker, use tar.gz for faster loading during build"
-            )
+        mlrun.utils.helpers.validate_builder_source(source, pull_at_runtime, workdir)
 
         self.spec.build.source = source
         if handler:
             self.spec.default_handler = handler
         if workdir:
             self.spec.workdir = workdir
+        if target_dir:
+            self.spec.clone_target_dir = target_dir
+
         self.spec.build.load_source_on_run = pull_at_runtime
         if (
             self.spec.build.base_image
@@ -86,7 +90,7 @@ class KubejobRuntime(KubeResource):
             and pull_at_runtime
             and not self.spec.image
         ):
-            # if we load source from repo and dont need a full build use the base_image as the image
+            # if we load source from repo and don't need a full build use the base_image as the image
             self.spec.image = self.spec.build.base_image
         elif not pull_at_runtime:
             # clear the image so build will not be skipped
@@ -106,7 +110,9 @@ class KubejobRuntime(KubeResource):
         auto_build=None,
         requirements=None,
         overwrite=False,
-        verify_base_image=True,
+        verify_base_image=False,
+        prepare_image_for_deploy=True,
+        requirements_file=None,
     ):
         """specify builder configuration for the deploy operation
 
@@ -121,43 +127,43 @@ class KubejobRuntime(KubeResource):
         :param with_mlrun: add the current mlrun package to the container build
         :param auto_build: when set to True and the function require build it will be built on the first
                            function run, use only if you dont plan on changing the build config between runs
-        :param requirements: requirements.txt file to install or list of packages to install
+        :param requirements: a list of packages to install
+        :param requirements_file: requirements file to install
         :param overwrite:  overwrite existing build configuration
 
            * False: the new params are merged with the existing (currently merge is applied to requirements and
              commands)
            * True: the existing params are replaced by the new ones
-        :param verify_base_image: verify the base image is set
+        :param verify_base_image:           verify that the base image is configured
+                                            (deprecated, use prepare_image_for_deploy)
+        :param prepare_image_for_deploy:    prepare the image/base_image spec for deployment
         """
-        if image:
-            self.spec.build.image = image
-        if base_image:
-            self.spec.build.base_image = base_image
-        # if overwrite and requirements or commands passed, clear the existing commands
-        # (requirements are added to the commands parameter)
-        if (requirements or commands) and overwrite:
-            self.spec.build.commands = None
-        if requirements:
-            self.with_requirements(
-                requirements, overwrite=False, verify_base_image=False
-            )
-        if commands:
-            self.with_commands(commands, overwrite=False, verify_base_image=False)
-        if extra:
-            self.spec.build.extra = extra
-        if secret is not None:
-            self.spec.build.secret = secret
-        if source:
-            self.spec.build.source = source
-        if load_source_on_run:
-            self.spec.build.load_source_on_run = load_source_on_run
-        if with_mlrun is not None:
-            self.spec.build.with_mlrun = with_mlrun
-        if auto_build:
-            self.spec.build.auto_build = auto_build
-
-        if verify_base_image:
-            self.verify_base_image()
+
+        image = mlrun.utils.helpers.remove_image_protocol_prefix(image)
+        self.spec.build.build_config(
+            image,
+            base_image,
+            commands,
+            secret,
+            source,
+            extra,
+            load_source_on_run,
+            with_mlrun,
+            auto_build,
+            requirements,
+            requirements_file,
+            overwrite,
+        )
+
+        if verify_base_image or prepare_image_for_deploy:
+            if verify_base_image:
+                # TODO: remove verify_base_image in 1.6.0
+                warnings.warn(
+                    "verify_base_image is deprecated in 1.4.0 and will be removed in 1.6.0, "
+                    "use prepare_image_for_deploy",
+                    category=FutureWarning,
+                )
+            self.prepare_image_for_deploy()
 
     def deploy(
         self,
@@ -194,7 +200,13 @@ class KubejobRuntime(KubeResource):
                     or "/mlrun/" in build.base_image
                 )
 
-        if not build.source and not build.commands and not build.extra and with_mlrun:
+        if (
+            not build.source
+            and not build.commands
+            and not build.requirements
+            and not build.extra
+            and with_mlrun
+        ):
             logger.info(
                 "running build to add mlrun package, set "
                 "with_mlrun=False to skip if its already in the image"
@@ -209,6 +221,7 @@ class KubejobRuntime(KubeResource):
         if is_kfp:
             watch = True
 
+        ready = False
         if self._is_remote_api():
             db = self._get_db()
             data = db.remote_builder(
@@ -223,6 +236,8 @@ class KubejobRuntime(KubeResource):
             self.spec.build.base_image = self.spec.build.base_image or get_in(
                 data, "data.spec.build.base_image"
             )
+            # get the clone target dir in case it was enriched due to loading source
+            self.spec.clone_target_dir = get_in(data, "data.spec.clone_target_dir")
             ready = data.get("ready", False)
             if not ready:
                 logger.info(
@@ -232,17 +247,6 @@ class KubejobRuntime(KubeResource):
                 state = self._build_watch(watch, show_on_failure=show_on_failure)
                 ready = state == "ready"
                 self.status.state = state
-        else:
-            self.save(versioned=False)
-            ready = build_runtime(
-                mlrun.api.schemas.AuthInfo(),
-                self,
-                with_mlrun,
-                mlrun_version_specifier,
-                skip_deployed,
-                watch,
-            )
-            self.save(versioned=False)
 
         if watch and not ready:
             raise mlrun.errors.MLRunRuntimeError("Deploy failed")
@@ -279,36 +283,6 @@ class KubejobRuntime(KubeResource):
         print()
         return self.status.state
 
-    def builder_status(self, watch=True, logs=True):
-        if self._is_remote_api():
-            return self._build_watch(watch, logs)
-
-        else:
-            pod = self.status.build_pod
-            if not self.status.state == "ready" and pod:
-                k8s = self._get_k8s()
-                status = k8s.get_pod_status(pod)
-                if logs:
-                    if watch:
-                        status = k8s.watch(pod)
-                    else:
-                        resp = k8s.logs(pod)
-                        if resp:
-                            print(resp.encode())
-
-                if status == "succeeded":
-                    self.status.build_pod = None
-                    self.status.state = "ready"
-                    logger.info("build completed successfully")
-                    return "ready"
-                if status in ["failed", "error"]:
-                    self.status.state = status
-                    logger.error(f" build {status}, watch the build pod logs: {pod}")
-                    return status
-
-                logger.info(f"builder status is: {status}, wait for it to complete")
-            return None
-
     def deploy_step(
         self,
         image=None,
@@ -335,23 +309,14 @@ class KubejobRuntime(KubeResource):
         )
 
     def _run(self, runobj: RunObject, execution):
-
         command, args, extra_env = self._get_cmd_args(runobj)
 
         if runobj.metadata.iteration:
             self.store_run(runobj)
-        k8s = self._get_k8s()
         new_meta = self._get_meta(runobj)
 
         self._add_secrets_to_spec_before_running(runobj)
-        workdir = self.spec.workdir
-        if workdir:
-            if self.spec.build.source and self.spec.build.load_source_on_run:
-                # workdir will be set AFTER the clone
-                workdir = None
-            elif not workdir.startswith("/"):
-                # relative path mapped to real path in the job pod
-                workdir = os.path.join("/mlrun", workdir)
+        workdir = self._resolve_workdir()
 
         pod_spec = func_to_pod(
             self.full_image_path(
@@ -368,23 +333,41 @@ class KubejobRuntime(KubeResource):
         )
         pod = client.V1Pod(metadata=new_meta, spec=pod_spec)
         try:
-            pod_name, namespace = k8s.create_pod(pod)
+            pod_name, namespace = get_k8s().create_pod(pod)
         except ApiException as exc:
             raise RunError(err_to_str(exc))
 
-        if pod_name and self.kfp:
-            writer = AsyncLogWriter(self._db_conn, runobj)
-            status = k8s.watch(pod_name, namespace, writer=writer)
-
-            if status in ["failed", "error"]:
-                raise RunError(f"pod exited with {status}, check logs")
-        else:
-            txt = f"Job is running in the background, pod: {pod_name}"
-            logger.info(txt)
-            runobj.status.status_text = txt
+        txt = f"Job is running in the background, pod: {pod_name}"
+        logger.info(txt)
+        runobj.status.status_text = txt
 
         return None
 
+    def _resolve_workdir(self):
+        """
+        The workdir is relative to the source root, if the source is not loaded on run then the workdir
+        is relative to the clone target dir (where the source was copied to).
+        Otherwise, if the source is loaded on run, the workdir is resolved on the run as well.
+        If the workdir is absolute, keep it as is.
+        """
+        workdir = self.spec.workdir
+        if self.spec.build.source and self.spec.build.load_source_on_run:
+            # workdir will be set AFTER the clone which is done in the pre-run of local runtime
+            return None
+
+        if workdir and os.path.isabs(workdir):
+            return workdir
+
+        if self.spec.clone_target_dir:
+            workdir = workdir or ""
+            if workdir.startswith("./"):
+                # TODO: use 'removeprefix' when we drop python 3.7 support
+                # workdir.removeprefix("./")
+                workdir = workdir[2:]
+            return os.path.join(self.spec.clone_target_dir, workdir)
+
+        return workdir
+
 
 def func_to_pod(image, runtime, extra_env, command, args, workdir):
     container = client.V1Container(

mlrun/runtimes/local.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -41,7 +41,7 @@ from ..execution import MLClientCtx
 from ..model import RunObject
 from ..utils import get_handler_extended, get_in, logger, set_paths
 from ..utils.clones import extract_source
-from .base import BaseRuntime, FunctionSpec, spec_fields
+from .base import BaseRuntime
 from .kubejob import KubejobRuntime
 from .remotesparkjob import RemoteSparkRuntime
 from .utils import RunError, global_context, log_std
@@ -172,48 +172,10 @@ class HandlerRuntime(BaseRuntime, ParallelRunner):
         return context.to_dict()
 
 
-class LocalFunctionSpec(FunctionSpec):
-    _dict_fields = spec_fields + ["clone_target_dir"]
-
-    def __init__(
-        self,
-        command=None,
-        args=None,
-        mode=None,
-        default_handler=None,
-        pythonpath=None,
-        entry_points=None,
-        description=None,
-        workdir=None,
-        build=None,
-        clone_target_dir=None,
-    ):
-        super().__init__(
-            command=command,
-            args=args,
-            mode=mode,
-            build=build,
-            entry_points=entry_points,
-            description=description,
-            workdir=workdir,
-            default_handler=default_handler,
-            pythonpath=pythonpath,
-        )
-        self.clone_target_dir = clone_target_dir
-
-
 class LocalRuntime(BaseRuntime, ParallelRunner):
     kind = "local"
     _is_remote = False
 
-    @property
-    def spec(self) -> LocalFunctionSpec:
-        return self._spec
-
-    @spec.setter
-    def spec(self, spec):
-        self._spec = self._verify_dict(spec, "spec", LocalFunctionSpec)
-
     def to_job(self, image=""):
         struct = self.to_dict()
         obj = KubejobRuntime.from_dict(struct)
@@ -224,12 +186,12 @@ class LocalRuntime(BaseRuntime, ParallelRunner):
     def with_source_archive(self, source, workdir=None, handler=None, target_dir=None):
         """load the code from git/tar/zip archive at runtime or build
 
-        :param source:     valid path to git, zip, or tar file, e.g.
-                           git://github.com/mlrun/something.git
-                           http://some/url/file.zip
-        :param handler: default function handler
-        :param workdir: working dir relative to the archive root or absolute (e.g. './subdir')
-        :param target_dir: local target dir for repo clone (by default its <current-dir>/code)
+        :param source:      valid path to git, zip, or tar file, e.g.
+                            git://github.com/mlrun/something.git
+                            http://some/url/file.zip
+        :param handler:     default function handler
+        :param workdir:     working dir relative to the archive root (e.g. './subdir') or absolute
+        :param target_dir:  local target dir for repo clone (by default its <current-dir>/code)
         """
         self.spec.build.source = source
         self.spec.build.load_source_on_run = True
@@ -255,6 +217,8 @@ class LocalRuntime(BaseRuntime, ParallelRunner):
         execution._current_workdir = workdir
         execution._old_workdir = None
 
+        # _is_run_local is set when the user specifies local=True in run()
+        # in this case we don't want to extract the source code and contaminate the user's local dir
         if self.spec.build.source and not hasattr(self, "_is_run_local"):
             target_dir = extract_source(
                 self.spec.build.source,
@@ -456,7 +420,7 @@ def exec_from_params(handler, runobj: RunObject, context: MLClientCtx, cwd=None)
     if runobj.spec.verbose:
         logger.set_logger_level("DEBUG")
 
-    # Prepare the inputs type hints (user may pass type hints as part of the inputs keys):
+    # Prepare the inputs type hints (user may pass type hints as part of the input keys):
     runobj.spec.extract_type_hints_from_inputs()
     # Read the keyword arguments to pass to the function (combining params and inputs from the run spec):
     kwargs = get_func_arg(handler, runobj, context)
@@ -471,20 +435,23 @@ def exec_from_params(handler, runobj: RunObject, context: MLClientCtx, cwd=None)
             if cwd:
                 os.chdir(cwd)
             # Apply the MLRun handler decorator for parsing inputs using type hints and logging outputs using log hints
-            # (Expected behavior: inputs are being parsed when they have type hints in code or given by user.
-            # outputs are logged only if log hints are provided by the user):
-            val = mlrun.handler(
-                inputs=(
-                    runobj.spec.inputs_type_hints
-                    if runobj.spec.inputs_type_hints
-                    else True  # True will use type hints if provided in user's code.
-                ),
-                outputs=(
-                    runobj.spec.returns
-                    if runobj.spec.returns
-                    else None  # None will turn off outputs logging.
-                ),
-            )(handler)(**kwargs)
+            # (Expected behavior: inputs are being parsed when they have type hints in code or given by user. Outputs
+            # are logged only if log hints are provided by the user):
+            if mlrun.mlconf.packagers.enabled:
+                val = mlrun.handler(
+                    inputs=(
+                        runobj.spec.inputs_type_hints
+                        if runobj.spec.inputs_type_hints
+                        else True  # True will use type hints if provided in user's code.
+                    ),
+                    outputs=(
+                        runobj.spec.returns
+                        if runobj.spec.returns
+                        else None  # None will turn off outputs logging.
+                    ),
+                )(handler)(**kwargs)
+            else:
+                val = handler(**kwargs)
             context.set_state("completed", commit=False)
         except Exception as exc:
             err = err_to_str(exc)
@@ -511,6 +478,18 @@ def get_func_arg(handler, runobj: RunObject, context: MLClientCtx, is_nuclio=Fal
     kwargs = {}
     args = inspect.signature(handler).parameters
 
+    def _get_input_value(input_key: str):
+        input_obj = context.get_input(input_key, inputs[input_key])
+        # If there is no type hint annotation but there is a default value and its type is string, point the data
+        # item to local downloaded file path (`local()` returns the downloaded temp path string):
+        if (
+            args[input_key].annotation is inspect.Parameter.empty
+            and type(args[input_key].default) is str
+        ):
+            return input_obj.local()
+        else:
+            return input_obj
+
     for key in args.keys():
         if key == "context":
             kwargs[key] = context
@@ -519,9 +498,23 @@ def get_func_arg(handler, runobj: RunObject, context: MLClientCtx, is_nuclio=Fal
         elif key in params:
             kwargs[key] = copy(params[key])
         elif key in inputs:
-            obj = context.get_input(key, inputs[key])
-            if type(args[key].default) is str or args[key].annotation == str:
-                kwargs[key] = obj.local()
-            else:
-                kwargs[key] = context.get_input(key, inputs[key])
+            kwargs[key] = _get_input_value(key)
+
+    list_of_params = list(args.values())
+    if len(list_of_params) == 0:
+        return kwargs
+
+    # get the last parameter, as **kwargs can only be last in the function's parameters list
+    last_param = list_of_params[-1]
+    # VAR_KEYWORD meaning : A dict of keyword arguments that aren’t bound to any other parameter.
+    # This corresponds to a **kwargs parameter in a Python function definition.
+    if last_param.kind == last_param.VAR_KEYWORD:
+        # if handler has **kwargs, pass all parameters provided by the user to the handler which were not already set
+        # as part of the previous loop which handled all parameters which were explicitly defined in the handler
+        for key in params:
+            if key not in kwargs:
+                kwargs[key] = copy(params[key])
+        for key in inputs:
+            if key not in kwargs:
+                kwargs[key] = _get_input_value(key)
     return kwargs

mlrun/runtimes/mpijob/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.

mlrun/runtimes/mpijob/abstract.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -24,7 +24,7 @@ from mlrun.execution import MLClientCtx
 from mlrun.model import RunObject
 from mlrun.runtimes.kubejob import KubejobRuntime
 from mlrun.runtimes.pod import KubeResourceSpec
-from mlrun.runtimes.utils import RunError
+from mlrun.runtimes.utils import RunError, get_k8s
 from mlrun.utils import get_in, logger
 
 
@@ -60,6 +60,7 @@ class MPIResourceSpec(KubeResourceSpec):
         tolerations=None,
         preemption_mode=None,
         security_context=None,
+        clone_target_dir=None,
     ):
         super().__init__(
             command=command,
@@ -88,6 +89,7 @@ class MPIResourceSpec(KubeResourceSpec):
             tolerations=tolerations,
             preemption_mode=preemption_mode,
             security_context=security_context,
+            clone_target_dir=clone_target_dir,
         )
         self.mpi_args = mpi_args or [
             "-x",
@@ -189,10 +191,9 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
     def _submit_mpijob(self, job, namespace=None):
         mpi_group, mpi_version, mpi_plural = self._get_crd_info()
 
-        k8s = self._get_k8s()
-        namespace = k8s.resolve_namespace(namespace)
+        namespace = get_k8s().resolve_namespace(namespace)
         try:
-            resp = k8s.crdapi.create_namespaced_custom_object(
+            resp = get_k8s().crdapi.create_namespaced_custom_object(
                 mpi_group,
                 mpi_version,
                 namespace=namespace,
@@ -208,7 +209,7 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
 
     def delete_job(self, name, namespace=None):
         mpi_group, mpi_version, mpi_plural = self._get_crd_info()
-        k8s = self._get_k8s()
+        k8s = get_k8s()
         namespace = k8s.resolve_namespace(namespace)
         try:
             # delete the mpi job
@@ -223,11 +224,10 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
 
     def list_jobs(self, namespace=None, selector="", show=True):
         mpi_group, mpi_version, mpi_plural = self._get_crd_info()
-        k8s = self._get_k8s()
-        namespace = k8s.resolve_namespace(namespace)
+        namespace = get_k8s().resolve_namespace(namespace)
         items = []
         try:
-            resp = k8s.crdapi.list_namespaced_custom_object(
+            resp = get_k8s().crdapi.list_namespaced_custom_object(
                 mpi_group,
                 mpi_version,
                 namespace,
@@ -247,10 +247,9 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
 
     def get_job(self, name, namespace=None):
         mpi_group, mpi_version, mpi_plural = self._get_crd_info()
-        k8s = self._get_k8s()
-        namespace = k8s.resolve_namespace(namespace)
+        namespace = get_k8s().resolve_namespace(namespace)
         try:
-            resp = k8s.crdapi.get_namespaced_custom_object(
+            resp = get_k8s().crdapi.get_namespaced_custom_object(
                 mpi_group, mpi_version, namespace, mpi_plural, name
             )
         except client.exceptions.ApiException as exc:
@@ -259,12 +258,11 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
         return resp
 
     def get_pods(self, name=None, namespace=None, launcher=False):
-        k8s = self._get_k8s()
-        namespace = k8s.resolve_namespace(namespace)
+        namespace = get_k8s().resolve_namespace(namespace)
 
         selector = self._generate_pods_selector(name, launcher)
 
-        pods = k8s.list_pods(selector=selector, namespace=namespace)
+        pods = get_k8s().list_pods(selector=selector, namespace=namespace)
         if pods:
             return {p.metadata.name: p.status.phase for p in pods}
 

mlrun/runtimes/mpijob/v1.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -62,6 +62,7 @@ class MPIV1ResourceSpec(MPIResourceSpec):
         tolerations=None,
         preemption_mode=None,
         security_context=None,
+        clone_target_dir=None,
     ):
         super().__init__(
             command=command,
@@ -91,6 +92,7 @@ class MPIV1ResourceSpec(MPIResourceSpec):
             tolerations=tolerations,
             preemption_mode=preemption_mode,
             security_context=security_context,
+            clone_target_dir=clone_target_dir,
         )
         self.clean_pod_policy = clean_pod_policy or MPIJobV1CleanPodPolicies.default()