feature: enable extending pypi/conda execution decorator support (Net…

…flix#2098)

* add remote execution bookkeeping to step decorators by default. apply changes to kubernetes/batch decorators

* only introduce conda flags to specific remote decorators

* add back missing break

* rename flag to supports_conda_environment. add backward compatibility

* cleanup

* add note on backward compatibility and reintroduce all supported decorator names in the list.

metaflow/plugins/aws/batch/batch_decorator.py

@@ -134,6 +134,10 @@ class BatchDecorator(StepDecorator):
     package_sha = None
     run_time_limit = None
 
+    # Conda environment support
+    supports_conda_environment = True
+    target_platform = "linux-64"
+
     def __init__(self, attributes=None, statically_defined=False):
         super(BatchDecorator, self).__init__(attributes, statically_defined)
 

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -135,6 +135,10 @@ class KubernetesDecorator(StepDecorator):
     package_sha = None
     run_time_limit = None
 
+    # Conda environment support
+    supports_conda_environment = True
+    target_platform = "linux-64"
+
     def __init__(self, attributes=None, statically_defined=False):
         super(KubernetesDecorator, self).__init__(attributes, statically_defined)
 

metaflow/plugins/pypi/conda_environment.py

@@ -269,11 +269,19 @@ def get_environment(self, step):
         # Resolve `linux-64` Conda environments if @batch or @kubernetes are in play
         target_platform = conda_platform()
         for decorator in step.decorators:
-            # TODO: rather than relying on decorator names, rely on attributes
-            #       to make them extensible.
-            if decorator.name in ["batch", "kubernetes", "nvidia", "snowpark", "slurm"]:
+            # NOTE: Keep the list of supported decorator names for backward compatibility purposes.
+            # Older versions did not implement the 'support_conda_environment' attribute.
+            if getattr(
+                decorator, "supports_conda_environment", False
+            ) or decorator.name in [
+                "batch",
+                "kubernetes",
+                "nvidia",
+                "snowpark",
+                "slurm",
+            ]:
                 # TODO: Support arm architectures
-                target_platform = "linux-64"
+                target_platform = getattr(decorator, "target_platform", "linux-64")
                 break
 
         environment["conda"]["platforms"] = [target_platform]