mlrun 1.7.0rc5__py3-none-any.whl → 1.7.2__py3-none-any.whl

mlrun/runtimes/local.py

@@ -33,6 +33,7 @@ from sys import executable
 from nuclio import Event
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 from mlrun.lists import RunList
 
 from ..errors import err_to_str
@@ -57,7 +58,9 @@ class ParallelRunner:
 
         return TrackerManager()
 
-    def _get_handler(self, handler, context):
+    def _get_handler(
+        self, handler: str, context: MLClientCtx, embed_in_sys: bool = True
+    ):
         return handler
 
     def _get_dask_client(self, options):
@@ -85,7 +88,7 @@ class ParallelRunner:
         handler = runobj.spec.handler
         self._force_handler(handler)
         set_paths(self.spec.pythonpath)
-        handler = self._get_handler(handler, execution)
+        handler = self._get_handler(handler, execution, embed_in_sys=False)
 
         client, function_name = self._get_dask_client(generator.options)
         parallel_runs = generator.options.parallel_runs or 4
@@ -142,7 +145,10 @@ class ParallelRunner:
         if function_name and generator.options.teardown_dask:
             logger.info("Tearing down the dask cluster..")
             mlrun.get_run_db().delete_runtime_resources(
-                kind="dask", object_id=function_name, force=True
+                project=self.metadata.project,
+                kind=mlrun.runtimes.RuntimeKinds.dask,
+                object_id=function_name,
+                force=True,
             )
 
         return results
@@ -223,12 +229,14 @@ class LocalRuntime(BaseRuntime, ParallelRunner):
     def is_deployed(self):
         return True
 
-    def _get_handler(self, handler, context):
+    def _get_handler(
+        self, handler: str, context: MLClientCtx, embed_in_sys: bool = True
+    ):
         command = self.spec.command
         if not command and self.spec.build.functionSourceCode:
             # if the code is embedded in the function object extract or find it
             command, _ = mlrun.run.load_func_code(self)
-        return load_module(command, handler, context)
+        return load_module(command, handler, context, embed_in_sys=embed_in_sys)
 
     def _pre_run(self, runobj: RunObject, execution: MLClientCtx):
         workdir = self.spec.workdir
@@ -257,7 +265,8 @@ class LocalRuntime(BaseRuntime, ParallelRunner):
             set_paths(os.path.realpath("."))
 
         if (
-            runobj.metadata.labels.get("kind") == RemoteSparkRuntime.kind
+            runobj.metadata.labels.get(mlrun_constants.MLRunInternalLabels.kind)
+            == RemoteSparkRuntime.kind
             and environ["MLRUN_SPARK_CLIENT_IGZ_SPARK"] == "true"
         ):
             from mlrun.runtimes.remotesparkjob import igz_spark_pre_hook
@@ -370,8 +379,20 @@ class LocalRuntime(BaseRuntime, ParallelRunner):
             return run_obj_dict
 
 
-def load_module(file_name, handler, context):
-    """Load module from file name"""
+def load_module(
+    file_name: str,
+    handler: str,
+    context: MLClientCtx,
+    embed_in_sys: bool = True,
+):
+    """
+    Load module from filename
+    :param file_name:       The module path to load
+    :param handler:         The callable to load
+    :param context:         Execution context
+    :param embed_in_sys:    Embed the file-named module in sys.modules. This is not persistent with remote
+                            environments and therefore can effect pickling.
+    """
     module = None
     if file_name:
         path = Path(file_name)
@@ -382,13 +403,21 @@ def load_module(file_name, handler, context):
         if spec is None:
             raise RunError(f"Cannot import from {file_name!r}")
         module = imputil.module_from_spec(spec)
+        if embed_in_sys:
+            sys.modules[mod_name] = module
         spec.loader.exec_module(module)
 
     class_args = {}
     if context:
         class_args = copy(context._parameters.get("_init_args", {}))
 
-    return get_handler_extended(handler, context, class_args, namespaces=module)
+    return get_handler_extended(
+        handler,
+        context,
+        class_args,
+        namespaces=module,
+        reload_modules=context._reset_on_run,
+    )
 
 
 def run_exec(cmd, args, env=None, cwd=None):
@@ -493,7 +522,7 @@ def exec_from_params(handler, runobj: RunObject, context: MLClientCtx, cwd=None)
                 logger.warning("Run was aborted", err=err_to_str(exc))
                 # Run was aborted, the state run state is updated by the abort job, no need to commit again
                 context.set_state(
-                    mlrun.runtimes.constants.RunStates.aborted, commit=False
+                    mlrun.common.runtimes.constants.RunStates.aborted, commit=False
                 )
                 commit = False
             except Exception as exc:

mlrun/runtimes/mpijob/__init__.py

@@ -21,28 +21,8 @@ from mlrun.config import config
 from .. import MPIJobCRDVersions
 from .abstract import AbstractMPIJobRuntime
 from .v1 import MpiRuntimeV1
-from .v1alpha1 import MpiRuntimeV1Alpha1
 
 
 def _resolve_mpijob_crd_version():
     # config is expected to get enriched from the API through the client-spec
     return config.mpijob_crd_version or MPIJobCRDVersions.default()
-
-
-class MpiRuntimeContainer(containers.DeclarativeContainer):
-    resolver = providers.Callable(
-        _resolve_mpijob_crd_version,
-    )
-
-    selector = providers.Selector(
-        resolver,
-        v1=providers.Object(MpiRuntimeV1),
-        v1alpha1=providers.Object(MpiRuntimeV1Alpha1),
-    )
-
-    # An empty selector to be overriden by the API
-    handler_selector = providers.Selector(
-        resolver,
-        v1=providers.Object(None),
-        v1alpha1=providers.Object(None),
-    )

mlrun/runtimes/mpijob/abstract.py

@@ -223,14 +223,14 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
         ```
         # Define the wanted MPI arguments
         mpi_args = []
-        mpi_args.append('-x')
-        mpi_args.append('NCCL_DEBUG=INFO')
-        mpi_args.append('-x')
-        mpi_args.append('NCCL_SOCKET_NTHREADS=2')
-        mpi_args.append('-x')
-        mpi_args.append('NCCL_NSOCKS_PERTHREAD=8')
-        mpi_args.append('-x')
-        mpi_args.append('NCCL_MIN_NCHANNELS=4')
+        mpi_args.append("-x")
+        mpi_args.append("NCCL_DEBUG=INFO")
+        mpi_args.append("-x")
+        mpi_args.append("NCCL_SOCKET_NTHREADS=2")
+        mpi_args.append("-x")
+        mpi_args.append("NCCL_NSOCKS_PERTHREAD=8")
+        mpi_args.append("-x")
+        mpi_args.append("NCCL_MIN_NCHANNELS=4")
 
         # Set the MPI arguments in the function
         fn.set_mpi_args(mpi_args)

mlrun/runtimes/mpijob/v1.py

@@ -11,7 +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.
-from mlrun.runtimes.constants import MPIJobCRDVersions, MPIJobV1CleanPodPolicies
+from mlrun.common.runtimes.constants import MPIJobCRDVersions, MPIJobV1CleanPodPolicies
 from mlrun.runtimes.mpijob.abstract import AbstractMPIJobRuntime, MPIResourceSpec