Fixes and more compatibility with Mamba 2.0

Netflix · Dec 10, 2024 · aadce54 · aadce54
1 parent 6c9b173
commit aadce54
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 9 deletions.
diff --git a/metaflow_extensions/netflix_ext/cmd/debug/debug_stub_generator.py b/metaflow_extensions/netflix_ext/cmd/debug/debug_stub_generator.py
@@ -21,7 +21,12 @@ def __init__(self, task_pathspec):
         self.run_id = self.task.parent.parent.id
         self.flow_name = self.task.parent.parent.parent.id
         self.is_new_conda_step = self.is_new_conda_step()
-        self.workflow_dag = self.task["_graph_info"].data
+
+        parameters_task = Step(
+            "%s/_parameters" % self.task.parent.pathspec, _namespace_check=False
+        ).task
+
+        self.workflow_dag = parameters_task["_graph_info"].data
         self.file_name = self.workflow_dag["file"]
         self._dag_structure = self.get_dag_structure(self.workflow_dag["steps"])
         self.step_type = self.get_step_type(self.step_name)

diff --git a/metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py b/metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py
@@ -36,13 +36,13 @@
     ResolvedEnvironment,
     env_type_for_deps,
 )
-from metaflow_extensions.nflx.plugins.new_conda.envsresolver import EnvsResolver
+from metaflow_extensions.netflix_ext.plugins.conda.envsresolver import EnvsResolver
 
-from metaflow_extensions.nflx.plugins.new_conda.parsers import (
+from metaflow_extensions.netflix_ext.plugins.conda.parsers import (
     parse_req_value,
     parse_yml_value,
 )
-from metaflow_extensions.nflx.plugins.new_conda.utils import (
+from metaflow_extensions.netflix_ext.plugins.conda.utils import (
     AliasType,
     arch_id,
     channel_or_url,

diff --git a/metaflow_extensions/netflix_ext/plugins/conda/conda.py b/metaflow_extensions/netflix_ext/plugins/conda/conda.py
@@ -267,17 +267,36 @@ def call_conda(
             raise InvalidEnvironmentException("Binary '%s' unknown" % binary)
         try:
             env = {"CONDA_JSON": "True"}
+            initial_command = args[0] if args else None
             if self._conda_executable_type == "mamba":
                 env.update({"MAMBA_NO_BANNER": "1", "MAMBA_JSON": "True"})
-            if addl_env:
-                env.update(addl_env)
-
             if (
-                args
-                and args[0] not in ("package", "info")
+                self._mode == "local"
+                and CONDA_LOCAL_PATH
                 and (self.is_non_conda_exec or binary == "micromamba")
             ):
+                # In this case, we need to prepend some options to the arguments
+                # to ensure that it uses CONDA_LOCAL_PATH and not the system one
+                args = [
+                    "--no-env",
+                    "--rc-file",
+                    os.path.join(CONDA_LOCAL_PATH, ".mambarc"),
+                    "-r",
+                    CONDA_LOCAL_PATH,
+                ] + args
+                if initial_command and initial_command not in ("package", "info"):
+                    args.append("--json")
+            elif (
+                (self._conda_executable_type == "micromamba" or binary == "micromamba")
+                and initial_command
+                and initial_command not in ("package", "info")
+            ):
+                # This is in the remote case.
                 args.extend(["-r", self.root_prefix, "--json"])
+
+            if addl_env:
+                env.update(addl_env)
+
             debug.conda_exec("Conda call: %s" % str([self._bins[binary]] + args))
             return cast(
                 bytes,

diff --git a/metaflow_extensions/netflix_ext/plugins/conda/utils.py b/metaflow_extensions/netflix_ext/plugins/conda/utils.py
@@ -213,6 +213,10 @@ def get_sys_packages(
     result = {
         k: v for k, v in result.items() if k in CONDA_SYS_DEPENDENCIES and v is not None
     }
+    if "__cuda" in result and not result["__cuda"].endswith("=0"):
+        # Make things consistent so users can specify 12.2 and it is like 12.2=0
+        # (otherwise it sometimes causes re-resolution)
+        result["__cuda"] = result["__cuda"] + "=0"
     return result