fix node-selector init in compute_pool case, and fix typo in compute …

…pool node-selector. add ability to skip passing some attributes to step CLI
Netflix · Aug 19, 2024 · b245b20 · b245b20
1 parent 7b64927
commit b245b20
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/metaflow/plugins/kubernetes/kubernetes_decorator.py b/metaflow/plugins/kubernetes/kubernetes_decorator.py
@@ -158,10 +158,12 @@ def __init__(self, attributes=None, statically_defined=False):
                 self.attributes["node_selector"].split(",")
             )
         if self.attributes["compute_pool"]:
+            if self.attributes["node_selector"] is None:
+                self.attributes["node_selector"] = {}
             self.attributes["node_selector"].update(
-                {"outerbounds.co/compute_pool": self.attributes["compute_pool"]}
+                {"outerbounds.co/compute-pool": self.attributes["compute_pool"]}
             )
-            del self.attributes["compute_pool"]
+            # del self.attributes["compute_pool"]
 
         if self.attributes["tolerations"]:
             try:
@@ -379,9 +381,13 @@ def runtime_step_cli(
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
 
+            # skip certain keys as CLI arguments
+            _skip_keys = ["compute_pool"]
             # --namespace is used to specify Metaflow namespace (a different
             # concept from k8s namespace).
             for k, v in self.attributes.items():
+                if k in _skip_keys:
+                    continue
                 if k == "namespace":
                     cli_args.command_options["k8s_namespace"] = v
                 elif k in {"node_selector"} and v: