ob-metaflow 2.12.12.0__tar.gz → 2.12.15.1__tar.gz

{ob_metaflow-2.12.12.0/ob_metaflow.egg-info → ob_metaflow-2.12.15.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow
-Version: 2.12.12.0
+Version: 2.12.15.1
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Netflix, Outerbounds & the Metaflow Community
 Author-email: help@outerbounds.co
@@ -12,7 +12,7 @@ Requires-Dist: boto3
 Requires-Dist: pylint
 Requires-Dist: kubernetes
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.12.12.0; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.15.1; extra == "stubs"
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/cli.py

@@ -650,7 +650,7 @@ def resume(
             )
 
     if step_to_rerun is None:
-        clone_steps = set()
+        steps_to_rerun = set()
     else:
         # validate step name
         if step_to_rerun not in obj.graph.nodes:
@@ -660,7 +660,7 @@ def resume(
                     step_to_rerun, ",".join(list(obj.graph.nodes.keys()))
                 )
             )
-        clone_steps = {step_to_rerun}
+        steps_to_rerun = {step_to_rerun}
 
     if run_id:
         # Run-ids that are provided by the metadata service are always integers.
@@ -688,7 +688,7 @@ def resume(
         clone_run_id=origin_run_id,
         clone_only=clone_only,
         reentrant=reentrant,
-        clone_steps=clone_steps,
+        steps_to_rerun=steps_to_rerun,
         max_workers=max_workers,
         max_num_splits=max_num_splits,
         max_log_size=max_log_size * 1024 * 1024,

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/client/core.py

@@ -1027,6 +1027,8 @@ class MetaflowData(object):
         self._artifacts = dict((art.id, art) for art in artifacts)
 
     def __getattr__(self, name: str):
+        if name not in self._artifacts:
+            raise AttributeError(name)
         return self._artifacts[name].data
 
     def __contains__(self, var):

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/flowspec.py

@@ -17,6 +17,7 @@ from .exception import (
 )
 from .graph import FlowGraph
 from .unbounded_foreach import UnboundedForeachInput
+from .util import to_pod
 from .metaflow_config import INCLUDE_FOREACH_STACK, MAXIMUM_FOREACH_VALUE_CHARS
 
 # For Python 3 compatibility
@@ -201,7 +202,7 @@ class FlowSpec(metaclass=_FlowSpecMeta):
             "decorators": [
                 {
                     "name": deco.name,
-                    "attributes": deco.attributes,
+                    "attributes": to_pod(deco.attributes),
                     "statically_defined": deco.statically_defined,
                 }
                 for deco in flow_decorators(self)

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/graph.py

@@ -3,6 +3,9 @@ import ast
 import re
 
 
+from .util import to_pod
+
+
 def deindent_docstring(doc):
     if doc:
         # Find the indent to remove from the docstring. We consider the following possibilities:
@@ -72,7 +75,6 @@ class DAGNode(object):
         return "%s.%s" % (expr.value.id, expr.attr)
 
     def _parse(self, func_ast):
-
         self.num_args = len(func_ast.args.args)
         tail = func_ast.body[-1]
 
@@ -171,6 +173,8 @@ class FlowGraph(object):
         self.name = flow.__name__
         self.nodes = self._create_nodes(flow)
         self.doc = deindent_docstring(flow.__doc__)
+        # nodes sorted in topological order.
+        self.sorted_nodes = []
         self._traverse_graph()
         self._postprocess()
 
@@ -197,6 +201,7 @@ class FlowGraph(object):
 
     def _traverse_graph(self):
         def traverse(node, seen, split_parents):
+            self.sorted_nodes.append(node.name)
             if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
                 split_parents = split_parents + [node.name]
@@ -262,7 +267,6 @@ class FlowGraph(object):
         )
 
     def output_steps(self):
-
         steps_info = {}
         graph_structure = []
 
@@ -286,7 +290,7 @@ class FlowGraph(object):
                 "decorators": [
                     {
                         "name": deco.name,
-                        "attributes": deco.attributes,
+                        "attributes": to_pod(deco.attributes),
                         "statically_defined": deco.statically_defined,
                     }
                     for deco in node.decorators

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/metaflow_config.py

@@ -15,6 +15,17 @@ if sys.platform == "darwin":
 ## value, either set `METAFLOW_DEFAULT_DATASTORE` in your configuration file or set
 ## an environment variable called `METAFLOW_DEFAULT_DATASTORE`
 
+##
+# Constants (NOTE: these need to live before any from_conf)
+##
+
+# Path to the local directory to store artifacts for 'local' datastore.
+DATASTORE_LOCAL_DIR = ".metaflow"
+
+# Local configuration file (in .metaflow) containing overrides per-project
+LOCAL_CONFIG_FILE = "config.json"
+
+
 ###
 # Default configuration
 ###
@@ -42,8 +53,6 @@ USER = from_conf("USER")
 ###
 # Datastore configuration
 ###
-# Path to the local directory to store artifacts for 'local' datastore.
-DATASTORE_LOCAL_DIR = ".metaflow"
 DATASTORE_SYSROOT_LOCAL = from_conf("DATASTORE_SYSROOT_LOCAL")
 # S3 bucket and prefix to store artifacts for 's3' datastore.
 DATASTORE_SYSROOT_S3 = from_conf("DATASTORE_SYSROOT_S3")
@@ -247,7 +256,6 @@ DEFAULT_RUNTIME_LIMIT = from_conf("DEFAULT_RUNTIME_LIMIT", 5 * 24 * 60 * 60)
 # Organization customizations
 ###
 UI_URL = from_conf("UI_URL")
-PAGERDUTY_TEMPLATE_URL = from_conf("PAGERDUTY_TEMPLATE_URL")
 
 ###
 # Capture error logs from argo

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/metaflow_config_funcs.py

@@ -32,8 +32,40 @@ def init_config():
     return config
 
 
+def init_local_config():
+    # This function is heavily inspired from LocalStorage.get_datastore_root_from_config
+    # but simplifies certain things and also does not depend on DATASTORE_SYSROOT_LOCAL.
+    #
+    # In other words, since this config is meant to be local to a directory, it does not
+    # check in DATASTORE_SYSROOT_LOCAL but only up the current getcwd() path. This also
+    # prevents nasty circular dependencies :)
+
+    from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, LOCAL_CONFIG_FILE
+
+    current_path = os.getcwd()
+    check_dir = os.path.join(current_path, DATASTORE_LOCAL_DIR)
+    check_dir = os.path.realpath(check_dir)
+    while not os.path.isdir(check_dir):
+        new_path = os.path.dirname(current_path)
+        if new_path == current_path:  # No longer making upward progress
+            return {}
+        current_path = new_path
+        check_dir = os.path.join(current_path, DATASTORE_LOCAL_DIR)
+    path_to_config = os.path.join(check_dir, LOCAL_CONFIG_FILE)
+    # We found a directory to look for the config file in
+    if os.path.exists(path_to_config):
+        with open(path_to_config, encoding="utf-8") as f:
+            return json.load(f)
+    return {}
+
+
 # Initialize defaults required to setup environment variables.
-METAFLOW_CONFIG = init_config()
+# (initialized lazily in from_conf since init_local_config requires
+# some configuration values
+
+METAFLOW_CONFIG = None
+
+METAFLOW_LOCAL_CONFIG = None
 
 _all_configs = {}
 
@@ -51,7 +83,13 @@ def config_values(include=0):
 
 def from_conf(name, default=None, validate_fn=None):
     """
-    First try to pull value from environment, then from metaflow config JSON
+    Pull value from the environment or configuration.
+    Order is:
+    1. Environment (use any environment variable explicitly set by user)
+    2. Local config (use any value set in the local config file -- so stuff in
+       .metaflow/project.json for example)
+    3. Global config (use any value set in the global config file)
+    4. Default
 
     Prior to a value being returned, we will validate using validate_fn (if provided).
     Only non-None values are validated.
@@ -59,9 +97,19 @@ def from_conf(name, default=None, validate_fn=None):
     validate_fn should accept (name, value).
     If the value validates, return None, else raise an MetaflowException.
     """
+    global METAFLOW_CONFIG, METAFLOW_LOCAL_CONFIG
+
+    if METAFLOW_CONFIG is None:
+        METAFLOW_CONFIG = init_config()
+    if METAFLOW_LOCAL_CONFIG is None:
+        METAFLOW_LOCAL_CONFIG = init_local_config()
+
     is_default = True
     env_name = "METAFLOW_%s" % name
-    value = os.environ.get(env_name, METAFLOW_CONFIG.get(env_name, default))
+    value = os.environ.get(
+        env_name,
+        METAFLOW_LOCAL_CONFIG.get(env_name, METAFLOW_CONFIG.get(env_name, default)),
+    )
     if validate_fn and value is not None:
         validate_fn(env_name, value)
     if default is not None:

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/metaflow_environment.py

@@ -89,10 +89,16 @@ class MetaflowEnvironment(object):
         It should work silently if everything goes well.
         """
         if datastore_type == "s3":
-            return (
-                '%s -m awscli ${METAFLOW_S3_ENDPOINT_URL:+--endpoint-url=\\"${METAFLOW_S3_ENDPOINT_URL}\\"} '
-                + "s3 cp %s job.tar >/dev/null"
-            ) % (self._python(), code_package_url)
+            from .plugins.aws.aws_utils import parse_s3_full_path
+
+            bucket, s3_object = parse_s3_full_path(code_package_url)
+            # NOTE: the script quoting is extremely sensitive due to the way shlex.split operates and this being inserted
+            # into a quoted command elsewhere.
+            return "{python} -c '{script}'".format(
+                python=self._python(),
+                script='import boto3, os; boto3.client(\\"s3\\", endpoint_url=os.getenv(\\"METAFLOW_S3_ENDPOINT_URL\\")).download_file(\\"%s\\", \\"%s\\", \\"job.tar\\")'
+                % (bucket, s3_object),
+            )
         elif datastore_type == "azure":
             from .plugins.azure.azure_utils import parse_azure_full_path
 
@@ -119,25 +125,34 @@ class MetaflowEnvironment(object):
             )
 
     def _get_install_dependencies_cmd(self, datastore_type):
-        cmds = ["%s -m pip install requests -qqq" % self._python()]
-        if datastore_type == "s3":
-            cmds.append("%s -m pip install awscli boto3 -qqq" % self._python())
-        elif datastore_type == "azure":
-            cmds.append(
-                "%s -m pip install azure-identity azure-storage-blob azure-keyvault-secrets simple-azure-blob-downloader -qqq"
-                % self._python()
-            )
-        elif datastore_type == "gs":
-            cmds.append(
-                "%s -m pip install google-cloud-storage google-auth simple-gcp-object-downloader google-cloud-secret-manager -qqq"
-                % self._python()
-            )
-        else:
+        base_cmd = "{} -m pip install -qqq".format(self._python())
+
+        datastore_packages = {
+            "s3": ["boto3"],
+            "azure": [
+                "azure-identity",
+                "azure-storage-blob",
+                "azure-keyvault-secrets",
+                "simple-azure-blob-downloader",
+            ],
+            "gs": [
+                "google-cloud-storage",
+                "google-auth",
+                "simple-gcp-object-downloader",
+                "google-cloud-secret-manager",
+            ],
+        }
+
+        if datastore_type not in datastore_packages:
             raise NotImplementedError(
-                "We don't know how to generate an install dependencies cmd for datastore %s"
-                % datastore_type
+                "Unknown datastore type: {}".format(datastore_type)
             )
-        return " && ".join(cmds)
+
+        cmd = "{} {}".format(
+            base_cmd, " ".join(datastore_packages[datastore_type] + ["requests"])
+        )
+        # skip pip installs if we know that packages might already be available
+        return "if [ -z $METAFLOW_SKIP_INSTALL_DEPENDENCIES ]; then {}; fi".format(cmd)
 
     def get_package_commands(self, code_package_url, datastore_type):
         cmds = [

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/argo/argo_workflows.py

@@ -49,7 +49,6 @@ from metaflow.metaflow_config import (
     SERVICE_HEADERS,
     SERVICE_INTERNAL_URL,
     UI_URL,
-    PAGERDUTY_TEMPLATE_URL,
 )
 from metaflow.metaflow_config_funcs import config_values
 from metaflow.mflog import BASH_SAVE_LOGS, bash_capture_logs, export_mflog_env_vars
@@ -1659,6 +1658,9 @@ class ArgoWorkflows(object):
 
             # support Metaflow sandboxes
             env["METAFLOW_INIT_SCRIPT"] = KUBERNETES_SANDBOX_INIT_SCRIPT
+            env["METAFLOW_KUBERNETES_SANDBOX_INIT_SCRIPT"] = (
+                KUBERNETES_SANDBOX_INIT_SCRIPT
+            )
 
             # support for @secret
             env["METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE"] = DEFAULT_SECRETS_BACKEND_TYPE
@@ -1674,6 +1676,12 @@ class ArgoWorkflows(object):
             )
             env["METAFLOW_DATASTORE_SYSROOT_AZURE"] = DATASTORE_SYSROOT_AZURE
             env["METAFLOW_CARD_AZUREROOT"] = CARD_AZUREROOT
+            env["METAFLOW_ARGO_WORKFLOWS_KUBERNETES_SECRETS"] = (
+                ARGO_WORKFLOWS_KUBERNETES_SECRETS
+            )
+            env["METAFLOW_ARGO_WORKFLOWS_ENV_VARS_TO_SKIP"] = (
+                ARGO_WORKFLOWS_ENV_VARS_TO_SKIP
+            )
 
             # support for GCP
             env["METAFLOW_DATASTORE_SYSROOT_GS"] = DATASTORE_SYSROOT_GS
@@ -2379,18 +2387,10 @@ class ArgoWorkflows(object):
     def _pager_duty_notification_links(self):
         links = []
         if UI_URL:
-            if PAGERDUTY_TEMPLATE_URL:
-                pdproject = ''
-                pdbranch = ''
-                if getattr(current, "project_name", None):
-                    pdproject = current.project_name
-                    pdbranch = current.branch_name
-                href_val = PAGERDUTY_TEMPLATE_URL.format(pd_flow=self.flow.name, pd_namespace=KUBERNETES_NAMESPACE,pd_template=self.name,pd_project=pdproject, pd_branch=pdbranch)
-            else:
-                href_val = "%s/%s/%s"% (UI_URL.rstrip("/"), self.flow.name, "argo-{{workflow.name}}")
             links.append(
                 {
-                    "href": href_val,
+                    "href": "%s/%s/%s"
+                    % (UI_URL.rstrip("/"), self.flow.name, "argo-{{workflow.name}}"),
                     "text": "Metaflow UI",
                 }
             )

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/argo/argo_workflows_cli.py

@@ -36,6 +36,12 @@ from .argo_workflows import ArgoWorkflows
 
 VALID_NAME = re.compile(r"^[a-z0-9]([a-z0-9\.\-]*[a-z0-9])?$")
 
+unsupported_decorators = {
+    "snowpark": "Step *%s* is marked for execution on Snowpark with Argo Workflows which isn't currently supported.",
+    "slurm": "Step *%s* is marked for execution on Slurm with Argo Workflows which isn't currently supported.",
+    "nvidia": "Step *%s* is marked for execution on Nvidia with Argo Workflows which isn't currently supported.",
+}
+
 
 class IncorrectProductionToken(MetaflowException):
     headline = "Incorrect production token"
@@ -172,7 +178,7 @@ def argo_workflows(obj, name=None):
 )
 @click.option(
     "--enable-heartbeat-daemon/--no-enable-heartbeat-daemon",
-    default=False,
+    default=True,
     show_default=True,
     help="Use a daemon container to broadcast heartbeats.",
 )
@@ -186,7 +192,7 @@ def argo_workflows(obj, name=None):
 )
 @click.option(
     "--enable-error-msg-capture/--no-enable-error-msg-capture",
-    default=False,
+    default=True,
     show_default=True,
     help="Capture stack trace of first failed task in exit hook.",
 )
@@ -211,6 +217,11 @@ def create(
     deployer_attribute_file=None,
     enable_error_msg_capture=False,
 ):
+    for node in obj.graph:
+        for decorator, error_message in unsupported_decorators.items():
+            if any([d.name == decorator for d in node.decorators]):
+                raise MetaflowException(error_message % node.name)
+
     validate_tags(tags)
 
     if deployer_attribute_file:

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/argo/argo_workflows_deployer.py

@@ -226,7 +226,9 @@ def trigger(instance: DeployedFlow, **kwargs):
         )
 
         command_obj = instance.deployer.spm.get(pid)
-        content = handle_timeout(tfp_runner_attribute, command_obj)
+        content = handle_timeout(
+            tfp_runner_attribute, command_obj, instance.deployer.file_read_timeout
+        )
 
         if command_obj.process.returncode == 0:
             triggered_run = TriggeredRun(deployer=instance.deployer, content=content)

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/aws/aws_utils.py

@@ -5,6 +5,22 @@ from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import MAX_MEMORY_PER_TASK, MAX_CPU_PER_TASK
 
 
+def parse_s3_full_path(s3_uri):
+    from urllib.parse import urlparse
+
+    #  <scheme>://<netloc>/<path>;<params>?<query>#<fragment>
+    scheme, netloc, path, _, _, _ = urlparse(s3_uri)
+    assert scheme == "s3"
+    assert netloc is not None
+
+    bucket = netloc
+    path = path.lstrip("/").rstrip("/")
+    if path == "":
+        path = None
+
+    return bucket, path
+
+
 def get_ec2_instance_metadata():
     """
     Fetches the EC2 instance metadata through AWS instance metadata service

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/aws/batch/batch_decorator.py

@@ -123,6 +123,7 @@ class BatchDecorator(StepDecorator):
         "ephemeral_storage": None,
         "log_driver": None,
         "log_options": None,
+        "executable": None,
     }
     resource_defaults = {
         "cpu": "1",

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/aws/step_functions/step_functions_deployer.py

@@ -193,7 +193,9 @@ def trigger(instance: DeployedFlow, **kwargs):
         )
 
         command_obj = instance.deployer.spm.get(pid)
-        content = handle_timeout(tfp_runner_attribute, command_obj)
+        content = handle_timeout(
+            tfp_runner_attribute, command_obj, instance.deployer.file_read_timeout
+        )
 
         if command_obj.process.returncode == 0:
             triggered_run = TriggeredRun(deployer=instance.deployer, content=content)

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/kubernetes/kubernetes.py

@@ -17,6 +17,8 @@ from metaflow.metaflow_config import (
     ARGO_EVENTS_INTERNAL_WEBHOOK_URL,
     ARGO_EVENTS_SERVICE_ACCOUNT,
     ARGO_EVENTS_WEBHOOK_AUTH,
+    ARGO_WORKFLOWS_KUBERNETES_SECRETS,
+    ARGO_WORKFLOWS_ENV_VARS_TO_SKIP,
     AWS_SECRETS_MANAGER_DEFAULT_REGION,
     AZURE_KEY_VAULT_PREFIX,
     AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
@@ -280,6 +282,18 @@ class Kubernetes(object):
             .environment_variable(
                 "METAFLOW_INIT_SCRIPT", KUBERNETES_SANDBOX_INIT_SCRIPT
             )
+            .environment_variable(
+                "METAFLOW_KUBERNETES_SANDBOX_INIT_SCRIPT",
+                KUBERNETES_SANDBOX_INIT_SCRIPT,
+            )
+            .environment_variable(
+                "METAFLOW_ARGO_WORKFLOWS_KUBERNETES_SECRETS",
+                ARGO_WORKFLOWS_KUBERNETES_SECRETS,
+            )
+            .environment_variable(
+                "METAFLOW_ARGO_WORKFLOWS_ENV_VARS_TO_SKIP",
+                ARGO_WORKFLOWS_ENV_VARS_TO_SKIP,
+            )
             .environment_variable("METAFLOW_OTEL_ENDPOINT", OTEL_ENDPOINT)
             # Skip setting METAFLOW_DATASTORE_SYSROOT_LOCAL because metadata sync
             # between the local user instance and the remote Kubernetes pod
@@ -565,6 +579,18 @@ class Kubernetes(object):
             .environment_variable(
                 "METAFLOW_INIT_SCRIPT", KUBERNETES_SANDBOX_INIT_SCRIPT
             )
+            .environment_variable(
+                "METAFLOW_KUBERNETES_SANDBOX_INIT_SCRIPT",
+                KUBERNETES_SANDBOX_INIT_SCRIPT,
+            )
+            .environment_variable(
+                "METAFLOW_ARGO_WORKFLOWS_KUBERNETES_SECRETS",
+                ARGO_WORKFLOWS_KUBERNETES_SECRETS,
+            )
+            .environment_variable(
+                "METAFLOW_ARGO_WORKFLOWS_ENV_VARS_TO_SKIP",
+                ARGO_WORKFLOWS_ENV_VARS_TO_SKIP,
+            )
             .environment_variable("METAFLOW_OTEL_ENDPOINT", OTEL_ENDPOINT)
             # Skip setting METAFLOW_DATASTORE_SYSROOT_LOCAL because metadata sync
             # between the local user instance and the remote Kubernetes pod

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -99,6 +99,9 @@ class KubernetesDecorator(StepDecorator):
         Shared memory size (in MiB) required for this step
     port: int, optional
         Port number to specify in the Kubernetes job object
+    compute_pool : str, optional, default None
+        Compute pool to be used for for this step.
+        If not specified, any accessible compute pool within the perimeter is used.
     """
 
     name = "kubernetes"
@@ -123,6 +126,8 @@ class KubernetesDecorator(StepDecorator):
         "persistent_volume_claims": None,  # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
         "shared_memory": None,
         "port": None,
+        "compute_pool": None,
+        "executable": None,
     }
     package_url = None
     package_sha = None
@@ -155,6 +160,12 @@ class KubernetesDecorator(StepDecorator):
             self.attributes["node_selector"] = parse_kube_keyvalue_list(
                 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"]}
+            )
 
         if self.attributes["tolerations"]:
             try:
@@ -385,9 +396,13 @@ class KubernetesDecorator(StepDecorator):
             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:

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/pypi/conda_decorator.py

@@ -100,9 +100,9 @@ class CondaStepDecorator(StepDecorator):
         # --environment=pypi to --environment=conda
         _supported_virtual_envs.extend(["pypi"])
 
-        # TODO: Hardcoded for now to support Docker environment.
+        # TODO: Hardcoded for now to support the fast bakery environment.
         # We should introduce a more robust mechanism for appending supported environments, for example from within extensions.
-        _supported_virtual_envs.extend(["docker"])
+        _supported_virtual_envs.extend(["fast-bakery"])
 
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda
@@ -210,7 +210,7 @@ class CondaStepDecorator(StepDecorator):
         self.interpreter = (
             self.environment.interpreter(self.step)
             if not any(
-                decorator.name in ["batch", "kubernetes"]
+                decorator.name in ["batch", "kubernetes", "nvidia", "snowpark", "slurm"]
                 for decorator in next(
                     step for step in self.flow if step.name == self.step
                 ).decorators
@@ -344,9 +344,9 @@ class CondaFlowDecorator(FlowDecorator):
         # --environment=pypi to --environment=conda
         _supported_virtual_envs.extend(["pypi"])
 
-        # TODO: Hardcoded for now to support Docker environment.
+        # TODO: Hardcoded for now to support the fast bakery environment.
         # We should introduce a more robust mechanism for appending supported environments, for example from within extensions.
-        _supported_virtual_envs.extend(["docker"])
+        _supported_virtual_envs.extend(["fast-bakery"])
 
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/pypi/conda_environment.py

@@ -268,7 +268,7 @@ class CondaEnvironment(MetaflowEnvironment):
         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"]:
+            if decorator.name in ["batch", "kubernetes", "nvidia", "snowpark", "slurm"]:
                 # TODO: Support arm architectures
                 target_platform = "linux-64"
                 break
@@ -288,13 +288,10 @@ class CondaEnvironment(MetaflowEnvironment):
             # PyPI registries, the usage of environment variable `GOOGLE_APPLICATION_CREDENTIALS`
             # demands that `keyrings.google-artifactregistry-auth` has to be installed
             # and available in the underlying python environment.
-            
-            # commenting this out per https://outerboundsco.slack.com/archives/C040K733FND/p1719262399355449
-            # this should be a temporary workaround. Need to find a better fix
-            # if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
-            #     environment["conda"]["packages"][
-            #         "keyrings.google-artifactregistry-auth"
-            #     ] = ">=1.1.1"
+            if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
+                environment["conda"]["packages"][
+                    "keyrings.google-artifactregistry-auth"
+                ] = ">=1.1.1"
 
         # Z combinator for a recursive lambda
         deep_sort = (lambda f: f(f))(

{ob_metaflow-2.12.12.0 → ob_metaflow-2.12.15.1}/metaflow/plugins/pypi/pypi_decorator.py

@@ -24,6 +24,12 @@ class PyPIStepDecorator(StepDecorator):
     name = "pypi"
     defaults = {"packages": {}, "python": None, "disabled": None}  # wheels
 
+    def __init__(self, attributes=None, statically_defined=False):
+        self._user_defined_attributes = (
+            attributes.copy() if attributes is not None else {}
+        )
+        super().__init__(attributes, statically_defined)
+
     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
@@ -70,9 +76,9 @@ class PyPIStepDecorator(StepDecorator):
         # --environment=pypi to --environment=conda
         _supported_virtual_envs.extend(["pypi"])
 
-        # TODO: Hardcoded for now to support Docker environment.
+        # TODO: Hardcoded for now to support the fast bakery environment.
         # We should introduce a more robust mechanism for appending supported environments, for example from within extensions.
-        _supported_virtual_envs.extend(["docker"])
+        _supported_virtual_envs.extend(["fast-bakery"])
 
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @pypi
@@ -88,6 +94,9 @@ class PyPIStepDecorator(StepDecorator):
                 )
             )
 
+    def is_attribute_user_defined(self, name):
+        return name in self._user_defined_attributes
+
 
 class PyPIFlowDecorator(FlowDecorator):
     """
@@ -123,9 +132,9 @@ class PyPIFlowDecorator(FlowDecorator):
         # --environment=pypi to --environment=conda
         _supported_virtual_envs.extend(["pypi"])
 
-        # TODO: Hardcoded for now to support Docker environment.
+        # TODO: Hardcoded for now to support the fast bakery environment.
         # We should introduce a more robust mechanism for appending supported environments, for example from within extensions.
-        _supported_virtual_envs.extend(["docker"])
+        _supported_virtual_envs.extend(["fast-bakery"])
 
         # The --environment= requirement ensures that valid virtual environments are
         # created for every step to execute it, greatly simplifying the @conda