ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/aws/batch/batch_cli.py

@@ -7,11 +7,12 @@ import traceback
 from metaflow import util
 from metaflow import R
 from metaflow.exception import CommandException, METAFLOW_EXIT_DISALLOW_RETRY
-from metaflow.metadata.util import sync_local_metadata_from_datastore
+from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.mflog import TASK_LOG_SOURCE
-
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from .batch import Batch, BatchKilledException
+from ..aws_utils import validate_aws_tag
 
 
 @click.group()
@@ -47,7 +48,7 @@ def _execute_cmd(func, flow_name, run_id, user, my_runs, echo):
     func(flow_name, run_id, user, echo)
 
 
-@batch.command(help="List unfinished AWS Batch tasks of this flow")
+@batch.command("list", help="List unfinished AWS Batch tasks of this flow")
 @click.option(
     "--my-runs",
     default=False,
@@ -61,7 +62,7 @@ def _execute_cmd(func, flow_name, run_id, user, my_runs, echo):
     help="List unfinished tasks corresponding to the run id.",
 )
 @click.pass_context
-def list(ctx, run_id, user, my_runs):
+def _list(ctx, run_id, user, my_runs):
     batch = Batch(ctx.obj.metadata, ctx.obj.environment)
     _execute_cmd(
         batch.list_jobs, ctx.obj.flow.name, run_id, user, my_runs, ctx.obj.echo
@@ -100,6 +101,7 @@ def kill(ctx, run_id, user, my_runs):
     "Metaflow."
 )
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option("--executable", help="Executable requirement for AWS Batch.")
@@ -146,12 +148,21 @@ def kill(ctx, run_id, user, my_runs):
     help="Activate designated number of elastic fabric adapter devices. "
     "EFA driver must be installed and instance type compatible with EFA",
 )
+@click.option(
+    "--aws-batch-tag",
+    "aws_batch_tags",
+    multiple=True,
+    default=None,
+    help="AWS tags. Format: key=value, multiple allowed",
+)
 @click.option("--use-tmpfs", is_flag=True, help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-tempdir", is_flag=True, help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-size", help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-path", help="tmpfs requirement for AWS Batch.")
-# TODO: Maybe remove it altogether since it's not used here
-@click.option("--ubf-context", default=None, type=click.Choice([None, "ubf_control"]))
+# NOTE: ubf-context is not explicitly used, but @parallel decorator tries to pass this so keep it for now
+@click.option(
+    "--ubf-context", default=None, type=click.Choice(["none", UBF_CONTROL, UBF_TASK])
+)
 @click.option("--host-volumes", multiple=True)
 @click.option("--efs-volumes", multiple=True)
 @click.option(
@@ -183,6 +194,7 @@ def kill(ctx, run_id, user, my_runs):
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
@@ -199,6 +211,7 @@ def step(
     swappiness=None,
     inferentia=None,
     efa=None,
+    aws_batch_tags=None,
     use_tmpfs=None,
     tmpfs_tempdir=None,
     tmpfs_size=None,
@@ -271,11 +284,19 @@ def step(
         "metaflow_version"
     ]
 
+    env = {"METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0])}
+
+    if aws_batch_tags is not None:
+        # We do not need to validate these again,
+        # as they come supplied by the batch decorator which already performed validation.
+        batch_tags = {}
+        for item in list(aws_batch_tags):
+            key, value = item.split("=")
+            batch_tags[key] = value
+
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]
     if env_deco:
-        env = env_deco[0].attributes["vars"]
-    else:
-        env = {}
+        env.update(env_deco[0].attributes["vars"])
 
     # Add the environment variables related to the input-paths argument
     if split_vars:
@@ -315,6 +336,7 @@ def step(
                 step_name,
                 step_cli,
                 task_spec,
+                code_package_metadata,
                 code_package_sha,
                 code_package_url,
                 ctx.obj.flow_datastore.TYPE,
@@ -336,6 +358,7 @@ def step(
                 host_volumes=host_volumes,
                 efs_volumes=efs_volumes,
                 use_tmpfs=use_tmpfs,
+                aws_batch_tags=batch_tags,
                 tmpfs_tempdir=tmpfs_tempdir,
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
@@ -344,7 +367,7 @@ def step(
                 log_options=log_options,
                 num_parallel=num_parallel,
             )
-    except Exception as e:
+    except Exception:
         traceback.print_exc()
         _sync_metadata()
         sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)

metaflow/plugins/aws/batch/batch_client.py

@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
-from collections import defaultdict, deque
+from collections import defaultdict
 import copy
 import random
-import select
-import sys
 import time
 import hashlib
 
@@ -89,6 +87,9 @@ class BatchJob(object):
         # Multinode
         if getattr(self, "num_parallel", 0) >= 1:
             num_nodes = self.num_parallel
+            # We need this task-id set so that all the nodes are aware of the control
+            # task's task-id. These "MF_" variables populate the `current.parallel` namedtuple
+            self.environment_variable("MF_PARALLEL_CONTROL_TASK_ID", self._task_id)
             main_task_override = copy.deepcopy(self.payload["containerOverrides"])
 
             # main

metaflow/plugins/aws/batch/batch_decorator.py

@@ -1,34 +1,33 @@
 import os
-import sys
 import platform
-import requests
+import sys
 import time
 
-from metaflow import util
 from metaflow import R, current
-
 from metaflow.decorators import StepDecorator
-from metaflow.plugins.resources_decorator import ResourcesDecorator
-from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
-from metaflow.metadata import MetaDatum
-from metaflow.metadata.util import sync_local_metadata_to_datastore
+from metaflow.metadata_provider import MetaDatum
+from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
-    ECS_S3_ACCESS_IAM_ROLE,
-    BATCH_JOB_QUEUE,
     BATCH_CONTAINER_IMAGE,
     BATCH_CONTAINER_REGISTRY,
-    ECS_FARGATE_EXECUTION_ROLE,
+    BATCH_DEFAULT_TAGS,
+    BATCH_JOB_QUEUE,
     DATASTORE_LOCAL_DIR,
+    ECS_FARGATE_EXECUTION_ROLE,
+    ECS_S3_ACCESS_IAM_ROLE,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
 )
+from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.sidecar import Sidecar
 from metaflow.unbounded_foreach import UBF_CONTROL
 
-from .batch import BatchException
 from ..aws_utils import (
     compute_resource_attributes,
     get_docker_registry,
     get_ec2_instance_metadata,
+    validate_aws_tag,
 )
+from .batch import BatchException
 
 
 class BatchDecorator(StepDecorator):
@@ -71,6 +70,9 @@ class BatchDecorator(StepDecorator):
         A swappiness value of 0 causes swapping not to happen unless absolutely
         necessary. A swappiness value of 100 causes pages to be swapped very
         aggressively. Accepted values are whole numbers between 0 and 100.
+    aws_batch_tags: Dict[str, str], optional, default None
+        Sets arbitrary AWS tags on the AWS Batch compute environment.
+        Set as string key-value pairs.
     use_tmpfs : bool, default False
         This enables an explicit tmpfs mount for this step. Note that tmpfs is
         not available on Fargate compute environments
@@ -88,15 +90,15 @@ class BatchDecorator(StepDecorator):
         Alias for inferentia. Use only one of the two.
     efa : int, default 0
         Number of elastic fabric adapter network devices to attach to container
-    ephemeral_storage: int, default None
-        The total amount, in GiB, of ephemeral storage to set for the task (21-200)
+    ephemeral_storage : int, default None
+        The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
         This is only relevant for Fargate compute environments
     log_driver: str, optional, default None
         The log driver to use for the Amazon ECS container.
     log_options: List[str], optional, default None
         List of strings containing options for the chosen log driver. The configurable values
         depend on the `log driver` chosen. Validation of these options is not supported yet.
-        Example usage: ["awslogs-group:aws/batch/job"]
+        Example: [`awslogs-group:aws/batch/job`]
     """
 
     name = "batch"
@@ -117,25 +119,30 @@ class BatchDecorator(StepDecorator):
         "host_volumes": None,
         "efs_volumes": None,
         "use_tmpfs": False,
+        "aws_batch_tags": None,
         "tmpfs_tempdir": True,
         "tmpfs_size": None,
         "tmpfs_path": "/metaflow_temp",
         "ephemeral_storage": None,
         "log_driver": None,
         "log_options": None,
+        "executable": None,
     }
     resource_defaults = {
         "cpu": "1",
         "gpu": "0",
         "memory": "4096",
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(BatchDecorator, self).__init__(attributes, statically_defined)
+    # Conda environment support
+    supports_conda_environment = True
+    target_platform = "linux-64"
 
+    def init(self):
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
             # If metaflow-config specifies a docker image, just use that.
@@ -174,6 +181,29 @@ class BatchDecorator(StepDecorator):
         if self.attributes["trainium"] is not None:
             self.attributes["inferentia"] = self.attributes["trainium"]
 
+        if not isinstance(BATCH_DEFAULT_TAGS, dict) and not all(
+            isinstance(k, str) and isinstance(v, str)
+            for k, v in BATCH_DEFAULT_TAGS.items()
+        ):
+            raise BatchException(
+                "BATCH_DEFAULT_TAGS environment variable must be Dict[str, str]"
+            )
+
+        if self.attributes["aws_batch_tags"] is not None:
+            if not isinstance(self.attributes["aws_batch_tags"], dict) and not all(
+                isinstance(k, str) and isinstance(v, str)
+                for k, v in self.attributes["aws_batch_tags"].items()
+            ):
+                raise BatchException("aws_batch_tags must be Dict[str, str]")
+        else:
+            self.attributes["aws_batch_tags"] = {}
+
+        if BATCH_DEFAULT_TAGS:
+            self.attributes["aws_batch_tags"] = {
+                **BATCH_DEFAULT_TAGS,
+                **self.attributes["aws_batch_tags"],
+            }
+
         # clean up the alias attribute so it is not passed on.
         self.attributes.pop("trainium", None)
 
@@ -206,6 +236,11 @@ class BatchDecorator(StepDecorator):
         if self.attributes["tmpfs_path"] and self.attributes["tmpfs_path"][0] != "/":
             raise BatchException("'tmpfs_path' needs to be an absolute path")
 
+        # Validate Batch tags
+        if self.attributes["aws_batch_tags"]:
+            for key, val in self.attributes["aws_batch_tags"].items():
+                validate_aws_tag(key, val)
+
     def runtime_init(self, flow, graph, package, run_id):
         # Set some more internal state.
         self.flow = flow
@@ -227,10 +262,20 @@ class BatchDecorator(StepDecorator):
             # to execute on AWS Batch anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["batch", "step"]
+            cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
-            cli_args.command_options.update(self.attributes)
+            # skip certain keys as CLI arguments
+            _skip_keys = ["aws_batch_tags"]
+            cli_args.command_options.update(
+                {k: v for k, v in self.attributes.items() if k not in _skip_keys}
+            )
             cli_args.command_options["run-time-limit"] = self.run_time_limit
+
+            # Pass the supplied AWS batch tags to the step CLI cmd
+            cli_args.command_options["aws-batch-tag"] = [
+                "%s=%s" % (k, v) for k, v in self.attributes["aws_batch_tags"].items()
+            ]
             if not R.use_r():
                 cli_args.entrypoint[0] = sys.executable
 
@@ -261,8 +306,8 @@ class BatchDecorator(StepDecorator):
         # metadata. A rudimentary way to detect non-local execution is to
         # check for the existence of AWS_BATCH_JOB_ID environment variable.
 
+        meta = {}
         if "AWS_BATCH_JOB_ID" in os.environ:
-            meta = {}
             meta["aws-batch-job-id"] = os.environ["AWS_BATCH_JOB_ID"]
             meta["aws-batch-job-attempt"] = os.environ["AWS_BATCH_JOB_ATTEMPT"]
             meta["aws-batch-ce-name"] = os.environ["AWS_BATCH_CE_NAME"]
@@ -275,6 +320,10 @@ class BatchDecorator(StepDecorator):
             # Metaflow would be running the container agent compatible with
             # version V4.
             # https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint.html
+
+            # TODO: Remove dependency on requests
+            import requests
+
             try:
                 logs_meta = (
                     requests.get(url=os.environ["ECS_CONTAINER_METADATA_URI_V4"])
@@ -290,21 +339,16 @@ class BatchDecorator(StepDecorator):
             instance_meta = get_ec2_instance_metadata()
             meta.update(instance_meta)
 
-            entries = [
-                MetaDatum(
-                    field=k,
-                    value=v,
-                    type=k,
-                    tags=["attempt_id:{0}".format(retry_count)],
-                )
-                for k, v in meta.items()
-            ]
-            # Register book-keeping metadata for debugging.
-            metadata.register_metadata(run_id, step_name, task_id, entries)
-
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+            # Start spot termination monitor sidecar.
+            current._update_env(
+                {"spot_termination_notice": "/tmp/spot_termination_notice"}
+            )
+            self._spot_monitor_sidecar = Sidecar("spot_termination_monitor")
+            self._spot_monitor_sidecar.start()
+
         num_parallel = int(os.environ.get("AWS_BATCH_JOB_NUM_NODES", 0))
         if num_parallel >= 1 and ubf_context == UBF_CONTROL:
             # UBF handling for multinode case
@@ -322,6 +366,21 @@ class BatchDecorator(StepDecorator):
 
         if num_parallel >= 1:
             _setup_multinode_environment()
+            # current.parallel.node_index will be correctly available over here.
+            meta.update({"parallel-node-index": current.parallel.node_index})
+
+        if len(meta) > 0:
+            entries = [
+                MetaDatum(
+                    field=k,
+                    value=v,
+                    type=k,
+                    tags=["attempt_id:{0}".format(retry_count)],
+                )
+                for k, v in meta.items()
+            ]
+            # Register book-keeping metadata for debugging.
+            metadata.register_metadata(run_id, step_name, task_id, entries)
 
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_retries
@@ -342,6 +401,7 @@ class BatchDecorator(StepDecorator):
 
         try:
             self._save_logs_sidecar.terminate()
+            self._spot_monitor_sidecar.terminate()
         except:
             # Best effort kill
             pass
@@ -378,7 +438,7 @@ class BatchDecorator(StepDecorator):
                             len(flow._control_mapper_tasks),
                         )
                     )
-            except Exception as e:
+            except Exception:
                 pass
         raise Exception(
             "Batch secondary workers did not finish in %s seconds" % TIMEOUT
@@ -387,9 +447,16 @@ class BatchDecorator(StepDecorator):
     @classmethod
     def _save_package_once(cls, flow_datastore, package):
         if cls.package_url is None:
-            cls.package_url, cls.package_sha = flow_datastore.save_data(
-                [package.blob], len_hint=1
-            )[0]
+            if not FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+                cls.package_url, cls.package_sha = flow_datastore.save_data(
+                    [package.blob], len_hint=1
+                )[0]
+                cls.package_metadata = package.package_metadata
+            else:
+                # Blocks until the package is uploaded
+                cls.package_url = package.package_url()
+                cls.package_sha = package.package_sha()
+                cls.package_metadata = package.package_metadata
 
 
 def _setup_multinode_environment():

metaflow/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.py

@@ -50,24 +50,27 @@ class AwsSecretsManagerSecretsProvider(SecretsProvider):
         The secret payload from AWS is EITHER a string OR a binary blob.
 
         If the secret contains a string payload ("SecretString"):
-        - if the `parse_secret_string_as_json` option is True (default):
+        - if the `json` option is True (default):
             {SecretString} will be parsed as a JSON. If successfully parsed, AND the JSON contains a
             top-level object, each entry K/V in the object will also be converted to an entry in the result. V will
             always be casted to a string (if not already a string).
-        - If `parse_secret_string_as_json` option is False:
-            {SecretString} will be returned as a single entry in the result, with the key being the secret_id.
+        - If `json` option is False:
+            {SecretString} will be returned as a single entry in the result, where the key is either:
+                - the `secret_id`, OR
+                - the value set by `options={"env_var_name": custom_env_var_name}`.
 
-        Otherwise, the secret contains a binary blob payload ("SecretBinary"). In this case
-        - The result dic contains '{SecretName}': '{SecretBinary}', where {SecretBinary} is a base64-encoded string
+        Otherwise, if the secret contains a binary blob payload ("SecretBinary"):
+        - The result dict contains '{SecretName}': '{SecretBinary}', where {SecretBinary} is a base64-encoded string.
 
-        All keys in the result are sanitized to be more valid environment variable names. This is done on a best effort
+        All keys in the result are sanitized to be more valid environment variable names. This is done on a best-effort
         basis. Further validation is expected to be done by the invoking @secrets decorator itself.
 
-        :param secret_id: ARN or friendly name of the secret
-        :param options: unused
-        :param role: AWS IAM Role ARN to assume before reading the secret
-        :return: dict of environment variables. All keys and values are strings.
+        :param secret_id: ARN or friendly name of the secret.
+        :param options: Dictionary of additional options. E.g., `options={"env_var_name": custom_env_var_name}`.
+        :param role: AWS IAM Role ARN to assume before reading the secret.
+        :return: Dictionary of environment variables. All keys and values are strings.
         """
+
         import botocore
         from metaflow.plugins.aws.aws_client import get_aws_client
 

metaflow/plugins/aws/step_functions/dynamo_db_client.py

@@ -1,8 +1,5 @@
-import os
 import time
 
-import requests
-
 from metaflow.metaflow_config import SFN_DYNAMO_DB_TABLE
 
 

metaflow/plugins/aws/step_functions/production_token.py

@@ -3,7 +3,7 @@ import os
 import random
 import string
 import zlib
-from itertools import dropwhile, islice
+from itertools import dropwhile
 
 from metaflow.util import to_bytes