fix issue

metaflow/plugins/pypi/conda_decorator.py

@@ -79,19 +79,21 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 if self.attributes["disabled"] is not None
                 else super_attributes["disabled"]
             )
+
+        # Set default for `disabled` argument.
         if not self.attributes["disabled"]:
             self.attributes["disabled"] = False
         # Set Python interpreter to user's Python if necessary.
         if not self.attributes["python"]:
             self.attributes["python"] = platform.python_version()  # CPython!
 
         # Take care of `disabled` argument.
-        if self.attributes["disabled"]:
-            _step = next(step for step in self.flow if step.name == self.step)
-            _step.decorators[:] = [
-                deco for deco in _step.decorators if deco.name not in ["conda", "pypi"]
-            ]
-        del self.attributes["disabled"]
+        # if self.attributes["disabled"]:
+        #     _step = next(step for step in self.flow if step.name == self.step)
+        #     # _step.decorators[:] = [
+        #     #     deco for deco in _step.decorators if deco.name not in ["conda", "pypi"]
+        #     # ]
+        # del self.attributes["disabled"]
 
         # @conda uses a conda environment to create a virtual environment.
         # The conda environment can be created through micromamba.
@@ -126,6 +128,8 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         environment.set_local_root(LocalStorage.get_datastore_root_from_config(logger))
 
     def runtime_init(self, flow, graph, package, run_id):
+        if self.attributes["disabled"]:
+            return
         # Create a symlink to metaflow installed outside the virtual environment.
         self.metaflow_dir = tempfile.TemporaryDirectory(dir="/tmp")
         os.symlink(
@@ -169,6 +173,8 @@ def runtime_init(self, flow, graph, package, run_id):
     def runtime_task_created(
         self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
     ):
+        if self.attributes["disabled"]:
+            return
         self.interpreter = (
             self.environment.interpreter(self.step)
             if not any(
@@ -194,6 +200,8 @@ def task_pre_step(
         ubf_context,
         inputs,
     ):
+        if self.attributes["disabled"]:
+            return
         # Add Python interpreter's parent to the path to ensure that any non-pythonic
         # dependencies in the virtual environment are visible to the user code.
         # sys.executable points to the Python interpreter in the virtual environment
@@ -212,6 +220,8 @@ def task_pre_step(
     def runtime_step_cli(
         self, cli_args, retry_count, max_user_code_retries, ubf_context
     ):
+        if self.attributes["disabled"]:
+            return
         # TODO: Support unbounded for-each
         # TODO: Check what happens when PYTHONPATH is defined via @environment
         # Ensure local installation of Metaflow is visible to user code
@@ -229,6 +239,8 @@ def runtime_step_cli(
             cli_args.entrypoint[0] = self.interpreter
 
     def runtime_finished(self, exception):
+        if self.attributes["disabled"]:
+            return
         self.metaflow_dir.cleanup()
 
 

metaflow/plugins/pypi/conda_environment.py

@@ -183,10 +183,14 @@ def get_environment(self, step):
         for decorator in step.decorators:
             # @conda decorator is guaranteed to exist thanks to self.decospecs
             if decorator.name in ["conda", "pypi"]:
-                environment[decorator.name] = dict(decorator.attributes)
-        # environment can be empty for @conda/@pypi(disabled=True)
-        if not environment:
-            return {}
+                # handle @conda/@pypi(disabled=True)
+                if decorator.attributes["disabled"]:
+                    return {}
+                environment[decorator.name] = {
+                    k: decorator.attributes[k]
+                    for k in decorator.attributes
+                    if k != "disabled"
+                }
         # TODO: Support dependencies for `--metadata`.
         # TODO: Introduce support for `--telemetry` as a follow up.
         # Certain packages are required for metaflow runtime to function correctly.

metaflow/plugins/pypi/pypi_decorator.py

@@ -48,13 +48,14 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 else super_attributes["disabled"]
             )
 
-        # Take care of `disabled` argument.
-        if self.attributes["disabled"]:
-            _step = next(step for step in self.flow if step.name == self.step)
-            _step.decorators[:] = [
-                deco for deco in _step.decorators if deco.name not in ["conda", "pypi"]
-            ]
-        del self.attributes["disabled"]
+        # # Take care of `disabled` argument.
+        # if self.attributes["disabled"]:
+        #     _step = next(step for step in self.flow if step.name == self.step)
+        #     # _step.decorators[:] = [
+        #     #     deco for deco in _step.decorators if deco.name not in ["conda", "pypi"]
+        #     # ]
+
+        # del self.attributes["disabled"]
 
         # At the moment, @pypi uses a conda environment as a virtual environment. This
         # is to ensure that we can have a dedicated Python interpreter within the
@@ -88,13 +89,6 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 )
             )
 
-        # TODO: This code snippet can be done away with by altering the constructor of
-        #       MetaflowEnvironment. A good first-task exercise.
-        # Avoid circular import
-        from metaflow.plugins.datastores.local_storage import LocalStorage
-
-        environment.set_local_root(LocalStorage.get_datastore_root_from_config(logger))
-
 
 class PyPIFlowDecorator(FlowDecorator):
     """