more changes

metaflow/plugins/__init__.py

@@ -63,7 +63,10 @@
 ]
 
 # Add environments here
-ENVIRONMENTS_DESC = [("conda", ".pypi.conda_environment.CondaEnvironment")]
+ENVIRONMENTS_DESC = [
+    ("conda", ".pypi.conda_environment.CondaEnvironment"),
+    ("pypi", ".pypi.pypi_environment.PyPIEnvironment"),
+]
 
 # Add metadata providers here
 METADATA_PROVIDERS_DESC = [

metaflow/plugins/pypi/conda_decorator.py

@@ -38,7 +38,7 @@ class CondaStepDecorator(StepDecorator):
         "packages": {},
         "libraries": {},  # Deprecated! Use packages going forward
         "python": None,
-        # TODO: Add support for disabled
+        "disabled": None,
     }
     # To define conda channels for the whole solve, users can specify
     # CONDA_CHANNELS in their environment. For pinning specific packages to specific
@@ -55,10 +55,52 @@ def __init__(self, attributes=None, statically_defined=False):
         del self.attributes["libraries"]
 
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        # The init_environment hook for Environment creates the relevant virtual
+        # environments. The step_init hook sets up the relevant state for that hook to
+        # do it's magic.
+
+        self.flow = flow
+        self.step = step
+        self.environment = environment
+        self.datastore = flow_datastore
+
+        # Support flow-level decorator.
+        if "conda_base" in self.flow._flow_decorators:
+            super_attributes = self.flow._flow_decorators["conda_base"][0].attributes
+            self.attributes["packages"] = {
+                **super_attributes["packages"],
+                **self.attributes["packages"],
+            }
+            self.attributes["python"] = (
+                self.attributes["python"] or super_attributes["python"]
+            )
+            self.attributes["disabled"] = (
+                self.attributes["disabled"]
+                if self.attributes["disabled"] is not None
+                else super_attributes["disabled"]
+            )
+        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"]
+
         # @conda uses a conda environment to create a virtual environment.
         # The conda environment can be created through micromamba.
         _supported_virtual_envs = ["conda"]
 
+        # To placate people who don't want to see a shred of conda in UX, we symlink
+        # --environment=pypi to --environment=conda
+        _supported_virtual_envs.extend(["pypi"])
+
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda
         # implementation.
@@ -67,7 +109,7 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 "@%s decorator requires %s"
                 % (
                     self.name,
-                    "or ".join(
+                    " or ".join(
                         ["--environment=%s" % env for env in _supported_virtual_envs]
                     ),
                 )
@@ -76,37 +118,13 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         # At this point, the list of 32 bit instance types is shrinking quite rapidly.
         # We can worry about supporting them when there is a need.
 
-        # The init_environment hook for Environment creates the relevant virtual
-        # environments. The step_init hook sets up the relevant state for that hook to
-        # do it's magic.
-
-        self.flow = flow
-        self.step = step
-        self.environment = environment
-        self.datastore = flow_datastore
-
         # 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))
 
-        # Support flow-level decorator
-        if "conda_base" in self.flow._flow_decorators:
-            super_attributes = self.flow._flow_decorators["conda_base"][0].attributes
-            self.attributes["packages"] = {
-                **super_attributes["packages"],
-                **self.attributes["packages"],
-            }
-            self.attributes["python"] = (
-                self.attributes["python"] or super_attributes["python"]
-            )
-
-        # Set Python interpreter to user's Python if necessary.
-        if not self.attributes["python"]:
-            self.attributes["python"] = platform.python_version()  # CPython!
-
     def runtime_init(self, flow, graph, package, run_id):
         # Create a symlink to metaflow installed outside the virtual environment.
         self.metaflow_dir = tempfile.TemporaryDirectory(dir="/tmp")
@@ -241,7 +259,7 @@ class CondaFlowDecorator(FlowDecorator):
         "packages": {},
         "libraries": {},  # Deprecated! Use packages going forward.
         "python": None,
-        # TODO: Add support for disabled
+        "disabled": None,
         # TODO: Support `@conda(python='3.10')` before shipping!!
     }
 
@@ -264,6 +282,10 @@ def flow_init(
         # The conda environment can be created through micromamba.
         _supported_virtual_envs = ["conda"]
 
+        # To placate people who don't want to see a shred of conda in UX, we symlink
+        # --environment=pypi to --environment=conda
+        _supported_virtual_envs.extend(["pypi"])
+
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda
         # implementation.
@@ -272,7 +294,7 @@ def flow_init(
                 "@%s decorator requires %s"
                 % (
                     self.name,
-                    "or ".join(
+                    " or ".join(
                         ["--environment=%s" % env for env in _supported_virtual_envs]
                     ),
                 )

metaflow/plugins/pypi/conda_environment.py

@@ -162,10 +162,14 @@ def cache(storage, results, type_):
 
     def executable(self, step_name, default=None):
         step = next(step for step in self.flow if step.name == step_name)
-        id_ = self.get_environment(step)["id_"]
-        # bootstrap.py is responsible for ensuring the validity of this executable.
-        # -s is important! Can otherwise leak packages to other environments.
-        return os.path.join(id_, "bin/python -s")
+        id_ = self.get_environment(step).get("id_")
+        if id_:
+            # bootstrap.py is responsible for ensuring the validity of this executable.
+            # -s is important! Can otherwise leak packages to other environments.
+            return os.path.join(id_, "bin/python -s")
+        else:
+            # for @conda/@pypi(disabled=True).
+            return super().executable(step_name, default)
 
     def interpreter(self, step_name):
         step = next(step for step in self.flow if step.name == step_name)
@@ -180,6 +184,9 @@ def get_environment(self, step):
             # @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 {}
         # TODO: Support dependencies for `--metadata`.
         # TODO: Introduce support for `--telemetry` as a follow up.
         # Certain packages are required for metaflow runtime to function correctly.
@@ -282,13 +289,17 @@ def add_to_package(self):
     def bootstrap_commands(self, step_name, datastore_type):
         # Bootstrap conda and execution environment for step
         step = next(step for step in self.flow if step.name == step_name)
-        id_ = self.get_environment(step)["id_"]
-        return [
-            "echo 'Bootstrapping virtual environment...'",
-            'python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
-            % (self.flow.name, id_, self.datastore_type),
-            "echo 'Environment bootstrapped.'",
-        ]
+        id_ = self.get_environment(step).get("id_")
+        if id_:
+            return [
+                "echo 'Bootstrapping virtual environment...'",
+                'python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
+                % (self.flow.name, id_, self.datastore_type),
+                "echo 'Environment bootstrapped.'",
+            ]
+        else:
+            # for @conda/@pypi(disabled=True).
+            return super().bootstrap_commands(step_name, datastore_type)
 
     # TODO: Make this an instance variable once local_root is part of the object
     #       constructor.

metaflow/plugins/pypi/pypi_decorator.py

@@ -39,6 +39,10 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
 
         _supported_virtual_envs = ["conda"]  # , "venv"]
 
+        # To placate people who don't want to see a shred of conda in UX, we symlink
+        # --environment=pypi to --environment=conda
+        _supported_virtual_envs.extend(["pypi"])
+
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @pypi
         # implementation.
@@ -47,7 +51,7 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 "@%s decorator requires %s"
                 % (
                     self.name,
-                    "or ".join(
+                    " or ".join(
                         ["--environment=%s" % env for env in _supported_virtual_envs]
                     ),
                 )
@@ -96,3 +100,25 @@ def flow_init(
         from metaflow import decorators
 
         decorators._attach_decorators(flow, ["pypi"])
+
+        # @pypi uses a conda environment to create a virtual environment.
+        # The conda environment can be created through micromamba.
+        _supported_virtual_envs = ["conda"]
+
+        # To placate people who don't want to see a shred of conda in UX, we symlink
+        # --environment=pypi to --environment=conda
+        _supported_virtual_envs.extend(["pypi"])
+
+        # The --environment= requirement ensures that valid virtual environments are
+        # created for every step to execute it, greatly simplifying the @conda
+        # implementation.
+        if environment.TYPE not in _supported_virtual_envs:
+            raise InvalidEnvironmentException(
+                "@%s decorator requires %s"
+                % (
+                    self.name,
+                    " or ".join(
+                        ["--environment=%s" % env for env in _supported_virtual_envs]
+                    ),
+                )
+            )

metaflow/plugins/pypi/pypi_environment.py

@@ -0,0 +1,7 @@
+from .conda_environment import CondaEnvironment
+
+
+# To placate people who don't want to see a shred of conda in UX, we symlink
+# --environment=pypi to --environment=conda
+class PyPIEnvironment(CondaEnvironment):
+    TYPE = "pypi"