From f9bbc2e9247784a4099a3addc5637225c7c4dccc Mon Sep 17 00:00:00 2001 From: David Pollack Date: Wed, 1 Feb 2023 12:54:45 +0100 Subject: [PATCH] override spec parser --- .../plugins/kubernetes/kubernetes_decorator.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/metaflow/plugins/kubernetes/kubernetes_decorator.py b/metaflow/plugins/kubernetes/kubernetes_decorator.py index 6112ee4b000..987a736de0e 100644 --- a/metaflow/plugins/kubernetes/kubernetes_decorator.py +++ b/metaflow/plugins/kubernetes/kubernetes_decorator.py @@ -400,6 +400,22 @@ def _save_package_once(cls, flow_datastore, package): [package.blob], len_hint=1 )[0] + @classmethod + def _parse_decorator_spec(cls, deco_spec: str): + if not deco_spec: + return cls() + + valid_options = "|".join(cls.defaults.keys()) + deco_spec_parts = [] + for part in re.split(f""",(?=[\s\w]+[{valid_options}]=)""", deco_spec): + name, val = part.split("=", 1) + if name in {"labels", "node_selector"}: + both = name == "node_selector" + val = json.dumps(cls.parse_kube_keyvalue_list(val.split(","), both)) + deco_spec_parts.append("=".join([name, val])) + deco_spec_parsed = ",".join(deco_spec_parts) + return super()._parse_decorator_spec(deco_spec_parsed) + @staticmethod def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True): try: