ob-metaflow 2.11.4.2__tar.gz → 2.11.10.2__tar.gz

{ob-metaflow-2.11.4.2/ob_metaflow.egg-info → ob-metaflow-2.11.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow
-Version: 2.11.4.2
+Version: 2.11.10.2
 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: ob-metaflow-stubs==2.11.4.2; extra == "stubs"
+Requires-Dist: ob-metaflow-stubs==2.11.10.2; 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.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/cli.py

@@ -473,13 +473,6 @@ def logs(obj, input_path, stdout=None, stderr=None, both=None, timestamps=False)
     help="Pathspec of the origin task for this task to clone. Do "
     "not execute anything.",
 )
-@click.option(
-    "--clone-wait-only/--no-clone-wait-only",
-    default=False,
-    show_default=True,
-    help="If specified, waits for an external process to clone the task",
-    hidden=True,
-)
 @click.option(
     "--clone-run-id",
     default=None,
@@ -519,7 +512,6 @@ def step(
     retry_count=None,
     max_user_code_retries=None,
     clone_only=None,
-    clone_wait_only=False,
     clone_run_id=None,
     decospecs=None,
     ubf_context="none",
@@ -575,7 +567,6 @@ def step(
             task_id,
             clone_only,
             retry_count,
-            wait_only=clone_wait_only,
         )
     else:
         task.run_step(
@@ -719,6 +710,13 @@ def common_run_options(func):
     hidden=True,
     help="If specified, allows this call to be called in parallel",
 )
+@click.option(
+    "--resume-identifier",
+    default=None,
+    show_default=True,
+    hidden=True,
+    help="If specified, it identifies the task that started this resume call. It is in the form of {step_name}-{task_id}",
+)
 @click.argument("step-to-rerun", required=False)
 @cli.command(help="Resume execution of a previous run of this flow.")
 @common_run_options
@@ -736,6 +734,7 @@ def resume(
     max_log_size=None,
     decospecs=None,
     run_id_file=None,
+    resume_identifier=None,
 ):
     before_run(obj, tags, decospecs + obj.environment.decospecs())
 
@@ -789,10 +788,15 @@ def resume(
         max_workers=max_workers,
         max_num_splits=max_num_splits,
         max_log_size=max_log_size * 1024 * 1024,
+        resume_identifier=resume_identifier,
     )
     write_run_id(run_id_file, runtime.run_id)
+    runtime.print_workflow_info()
     runtime.persist_constants()
-    runtime.execute()
+    if clone_only:
+        runtime.clone_original_run()
+    else:
+        runtime.execute()
 
 
 @tracing.cli_entrypoint("cli/run")
@@ -845,6 +849,7 @@ def run(
     write_run_id(run_id_file, runtime.run_id)
 
     obj.flow._set_constants(obj.graph, kwargs)
+    runtime.print_workflow_info()
     runtime.persist_constants()
     runtime.execute()
 

ob-metaflow-2.11.10.2/metaflow/clone_util.py

@@ -0,0 +1,71 @@
+import time
+from .metadata import MetaDatum
+
+
+def clone_task_helper(
+    flow_name,
+    clone_run_id,
+    run_id,
+    step_name,
+    clone_task_id,
+    task_id,
+    flow_datastore,
+    metadata_service,
+    origin_ds_set=None,
+    attempt_id=0,
+):
+    # 1. initialize output datastore
+    output = flow_datastore.get_task_datastore(
+        run_id, step_name, task_id, attempt=attempt_id, mode="w"
+    )
+    output.init_task()
+
+    origin_run_id, origin_step_name, origin_task_id = (
+        clone_run_id,
+        step_name,
+        clone_task_id,
+    )
+    # 2. initialize origin datastore
+    origin = None
+    if origin_ds_set:
+        origin = origin_ds_set.get_with_pathspec(
+            "{}/{}/{}".format(origin_run_id, origin_step_name, origin_task_id)
+        )
+    else:
+        origin = flow_datastore.get_task_datastore(
+            origin_run_id, origin_step_name, origin_task_id
+        )
+    metadata_tags = ["attempt_id:{0}".format(attempt_id)]
+    output.clone(origin)
+    _ = metadata_service.register_task_id(
+        run_id,
+        step_name,
+        task_id,
+        attempt_id,
+    )
+    metadata_service.register_metadata(
+        run_id,
+        step_name,
+        task_id,
+        [
+            MetaDatum(
+                field="origin-task-id",
+                value=str(origin_task_id),
+                type="origin-task-id",
+                tags=metadata_tags,
+            ),
+            MetaDatum(
+                field="origin-run-id",
+                value=str(origin_run_id),
+                type="origin-run-id",
+                tags=metadata_tags,
+            ),
+            MetaDatum(
+                field="attempt",
+                value=str(attempt_id),
+                type="attempt",
+                tags=metadata_tags,
+            ),
+        ],
+    )
+    output.done()

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/cmd/develop/stub_generator.py

@@ -30,6 +30,7 @@ from typing import (
 from metaflow import FlowSpec, step
 from metaflow.debug import debug
 from metaflow.decorators import Decorator, FlowDecorator
+from metaflow.extension_support import get_aliased_modules
 from metaflow.graph import deindent_docstring
 from metaflow.metaflow_version import get_version
 
@@ -116,6 +117,7 @@ class StubGenerator:
 
         self._write_generated_for = include_generated_for
         self._pending_modules = ["metaflow"]  # type: List[str]
+        self._pending_modules.extend(get_aliased_modules())
         self._root_module = "metaflow."
         self._safe_modules = ["metaflow.", "metaflow_extensions."]
 

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/cmd/develop/stubs.py

@@ -323,17 +323,26 @@ def get_packages_for_stubs() -> Tuple[List[Tuple[str, str]], List[str]]:
         return [], []
 
     dist_list = []
-    for dist in _metadata_package.distributions():
+
+    # We check the type because if the user has multiple importlib metadata, for
+    # some reason it shows up multiple times.
+    interesting_dists = [
+        d
+        for d in _metadata_package.distributions()
         if any(
             [
-                pkg == "metaflow-stubs"
-                for pkg in (dist.read_text("top_level.txt") or "").split()
+                p == "metaflow-stubs"
+                for p in (d.read_text("top_level.txt") or "").split()
             ]
-        ):
-            # This is a package we care about
-            root_path = dist.locate_file("metaflow-stubs").as_posix()
-            dist_list.append((dist.metadata["Name"], root_path))
-            all_paths.discard(root_path)
+        )
+        and isinstance(d, _metadata_package.PathDistribution)
+    ]
+
+    for dist in interesting_dists:
+        # This is a package we care about
+        root_path = dist.locate_file("metaflow-stubs").as_posix()
+        dist_list.append((dist.metadata["Name"], root_path))
+        all_paths.discard(root_path)
     return dist_list, list(all_paths)
 
 

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/decorators.py

@@ -518,7 +518,8 @@ def _init_flow_decorators(
         else:
             # Each "non-multiple" flow decorator is only allowed to have one set of options
             deco_flow_init_options = {
-                option: deco_options[option] for option in deco.options
+                option: deco_options[option.replace("-", "_")]
+                for option in deco.options
             }
         for deco in decorators:
             deco.flow_init(

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/metaflow_config.py

@@ -26,7 +26,7 @@ DEFAULT_METADATA = from_conf("DEFAULT_METADATA", "local")
 DEFAULT_MONITOR = from_conf("DEFAULT_MONITOR", "nullSidecarMonitor")
 DEFAULT_PACKAGE_SUFFIXES = from_conf("DEFAULT_PACKAGE_SUFFIXES", ".py,.R,.RDS")
 DEFAULT_AWS_CLIENT_PROVIDER = from_conf("DEFAULT_AWS_CLIENT_PROVIDER", "boto3")
-DEFAULT_GCP_CLIENT_PROVIDER = from_conf("DEFAULT_ACP_CLIENT_PROVIDER", "gcp-default")
+DEFAULT_GCP_CLIENT_PROVIDER = from_conf("DEFAULT_GCP_CLIENT_PROVIDER", "gcp-default")
 DEFAULT_SECRETS_BACKEND_TYPE = from_conf("DEFAULT_SECRETS_BACKEND_TYPE")
 DEFAULT_SECRETS_ROLE = from_conf("DEFAULT_SECRETS_ROLE")
 
@@ -82,6 +82,9 @@ S3_SERVER_SIDE_ENCRYPTION = from_conf("S3_SERVER_SIDE_ENCRYPTION")
 # so setting it to 0 means each operation will be tried once.
 S3_RETRY_COUNT = from_conf("S3_RETRY_COUNT", 7)
 
+# Number of concurrent S3 processes for parallel operations.
+S3_WORKER_COUNT = from_conf("S3_WORKER_COUNT", 64)
+
 # Number of retries on *transient* failures (such as SlowDown errors). Note
 # that if after S3_TRANSIENT_RETRY_COUNT times, all operations haven't been done,
 # it will try up to S3_RETRY_COUNT again so the total number of tries can be up to
@@ -142,6 +145,15 @@ DATATOOLS_LOCALROOT = from_conf(
 # Secrets Backend - AWS Secrets Manager configuration
 AWS_SECRETS_MANAGER_DEFAULT_REGION = from_conf("AWS_SECRETS_MANAGER_DEFAULT_REGION")
 
+# Secrets Backend - GCP Secrets name prefix. With this, users don't have
+# to specify the full secret name in the @secret decorator.
+#
+# Note that it makes a difference whether the prefix ends with a slash or not
+# E.g. if secret name passed to @secret decorator is mysecret:
+# - "projects/1234567890/secrets/" -> "projects/1234567890/secrets/mysecret"
+# - "projects/1234567890/secrets/foo-" -> "projects/1234567890/secrets/foo-mysecret"
+GCP_SECRET_MANAGER_PREFIX = from_conf("GCP_SECRET_MANAGER_PREFIX")
+
 # The root directory to save artifact pulls in, when using S3 or Azure
 ARTIFACT_LOCALROOT = from_conf("ARTIFACT_LOCALROOT", os.getcwd())
 
@@ -312,6 +324,8 @@ KUBERNETES_FETCH_EC2_METADATA = from_conf("KUBERNETES_FETCH_EC2_METADATA", False
 KUBERNETES_PORT = from_conf("KUBERNETES_PORT", None)
 # Shared memory in MB to use for this step
 KUBERNETES_SHARED_MEMORY = from_conf("KUBERNETES_SHARED_MEMORY", None)
+# Default port number to open on the pods
+KUBERNETES_PORT = from_conf("KUBERNETES_PORT", None)
 
 
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
@@ -450,6 +464,7 @@ def get_pinned_conda_libs(python_version, datastore_type):
     elif datastore_type == "gs":
         pins["google-cloud-storage"] = ">=2.5.0"
         pins["google-auth"] = ">=2.11.0"
+        pins["google-cloud-secret-manager"] = ">=2.10.0"
     elif datastore_type == "local":
         pass
     else:

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/metaflow_environment.py

@@ -129,7 +129,7 @@ class MetaflowEnvironment(object):
             )
         elif datastore_type == "gs":
             cmds.append(
-                "%s -m pip install google-cloud-storage google-auth simple-gcp-object-downloader -qqq"
+                "%s -m pip install google-cloud-storage google-auth simple-gcp-object-downloader google-cloud-secret-manager -qqq"
                 % self._python()
             )
         else:

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/package.py

@@ -88,9 +88,10 @@ class MetaflowPackage(object):
             # if path and (path[0] == '.' or './' in path):
             #    continue
             for fname in files:
-                if fname[0] == ".":
-                    continue
-                if any(fname.endswith(suffix) for suffix in suffixes):
+                if (fname[0] == "." and fname in suffixes) or (
+                    fname[0] != "."
+                    and any(fname.endswith(suffix) for suffix in suffixes)
+                ):
                     p = os.path.join(path, fname)
                     yield p, p[prefixlen:]
 

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/parameters.py

@@ -342,10 +342,10 @@ class Parameter(object):
             )
 
     def __repr__(self):
-        return "metaflow.Parameter(name=%s, kwargs=%s)" % (name, kwargs)
+        return "metaflow.Parameter(name=%s, kwargs=%s)" % (self.name, self.kwargs)
 
     def __str__(self):
-        return "metaflow.Parameter(name=%s, kwargs=%s)" % (name, kwargs)
+        return "metaflow.Parameter(name=%s, kwargs=%s)" % (self.name, self.kwargs)
 
     def option_kwargs(self, deploy_mode):
         kwargs = self.kwargs

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/__init__.py

@@ -120,6 +120,10 @@ SECRETS_PROVIDERS_DESC = [
         "aws-secrets-manager",
         ".aws.secrets_manager.aws_secrets_manager_secrets_provider.AwsSecretsManagerSecretsProvider",
     ),
+    (
+        "gcp-secret-manager",
+        ".gcp.gcp_secret_manager_secrets_provider.GcpSecretManagerSecretsProvider",
+    ),
 ]
 
 AZURE_CLIENT_PROVIDERS_DESC = [

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/airflow/airflow.py

@@ -17,6 +17,7 @@ from metaflow.metaflow_config import (
     AIRFLOW_KUBERNETES_KUBECONFIG_FILE,
     AIRFLOW_KUBERNETES_STARTUP_TIMEOUT_SECONDS,
     AWS_SECRETS_MANAGER_DEFAULT_REGION,
+    GCP_SECRET_MANAGER_PREFIX,
     AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
     CARD_AZUREROOT,
     CARD_GSROOT,
@@ -408,6 +409,8 @@ class Airflow(object):
             env[
                 "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION"
             ] = AWS_SECRETS_MANAGER_DEFAULT_REGION
+        if GCP_SECRET_MANAGER_PREFIX:
+            env["METAFLOW_GCP_SECRET_MANAGER_PREFIX"] = GCP_SECRET_MANAGER_PREFIX
 
         env.update(additional_mf_variables)
 

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/argo/argo_workflows.py

@@ -32,6 +32,7 @@ from metaflow.metaflow_config import (
     DATATOOLS_S3ROOT,
     DEFAULT_METADATA,
     DEFAULT_SECRETS_BACKEND_TYPE,
+    GCP_SECRET_MANAGER_PREFIX,
     KUBERNETES_FETCH_EC2_METADATA,
     KUBERNETES_LABELS,
     KUBERNETES_NAMESPACE,
@@ -1313,6 +1314,7 @@ class ArgoWorkflows(object):
             env[
                 "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION"
             ] = AWS_SECRETS_MANAGER_DEFAULT_REGION
+            env["METAFLOW_GCP_SECRET_MANAGER_PREFIX"] = GCP_SECRET_MANAGER_PREFIX
 
             # support for Azure
             env[
@@ -1384,12 +1386,14 @@ class ArgoWorkflows(object):
             # Set shared_memory to 0 if it isn't specified. This results
             # in Kubernetes using it's default value when the pod is created.
             shared_memory = resources.get("shared_memory", 0)
+            port = resources.get("port", None)
+            if port:
+                port = int(port)
 
             tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
 
             if tmpfs_enabled and tmpfs_tempdir:
                 env["METAFLOW_TEMPDIR"] = tmpfs_path
-
             # Create a ContainerTemplate for this node. Ideally, we would have
             # liked to inline this ContainerTemplate and avoid scanning the workflow
             # twice, but due to issues with variable substitution, we will have to
@@ -1448,6 +1452,9 @@ class ArgoWorkflows(object):
                         kubernetes_sdk.V1Container(
                             name=self._sanitize(node.name),
                             command=cmds,
+                            ports=[kubernetes_sdk.V1ContainerPort(container_port=port)]
+                            if port
+                            else None,
                             env=[
                                 kubernetes_sdk.V1EnvVar(name=k, value=str(v))
                                 for k, v in env.items()
@@ -1996,8 +2003,10 @@ class ArgoWorkflows(object):
                                                 # Technically, we don't need to create
                                                 # a payload carry-on and can stuff
                                                 # everything within the body.
-                                                data_template="{{ .Input.body.payload.%s | toJson }}"
-                                                % v,
+                                                # NOTE: We need the conditional logic in order to successfully fall back to the default value
+                                                # when the event payload does not contain a key for a parameter.
+                                                data_template='{{ if (hasKey $.Input.body.payload "%s") }}{{- (.Input.body.payload.%s | toJson) -}}{{- else -}}{{ (fail "use-default-instead") }}{{- end -}}'
+                                                % (v, v),
                                                 # Unfortunately the sensor needs to
                                                 # record the default values for
                                                 # the parameters - there doesn't seem
@@ -2926,4 +2935,4 @@ class LifecycleHook(object):
         return self.payload
 
     def __str__(self):
-        return json.dumps(self.payload, indent=4)
+        return json.dumps(self.payload, indent=4)

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/argo/argo_workflows_decorator.py

@@ -72,6 +72,7 @@ class ArgoWorkflowsInternalDecorator(StepDecorator):
         meta["argo-workflow-name"] = os.environ["ARGO_WORKFLOW_NAME"]
         meta["argo-workflow-namespace"] = os.environ["ARGO_WORKFLOW_NAMESPACE"]
         meta["auto-emit-argo-events"] = self.attributes["auto-emit-argo-events"]
+        meta["argo-workflow-template-owner"] = os.environ["METAFLOW_OWNER"]
         entries = [
             MetaDatum(
                 field=k, value=v, type=k, tags=["attempt_id:{0}".format(retry_count)]

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/aws/batch/batch.py

@@ -192,6 +192,9 @@ class Batch(object):
         tmpfs_size=None,
         tmpfs_path=None,
         num_parallel=0,
+        ephemeral_storage=None,
+        log_driver=None,
+        log_options=None,
     ):
         job_name = self._job_name(
             attrs.get("metaflow.user"),
@@ -240,6 +243,9 @@ class Batch(object):
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
                 num_parallel=num_parallel,
+                ephemeral_storage=ephemeral_storage,
+                log_driver=log_driver,
+                log_options=log_options,
             )
             .task_id(attrs.get("metaflow.task_id"))
             .environment_variable("AWS_DEFAULT_REGION", self._client.region())
@@ -353,6 +359,9 @@ class Batch(object):
         num_parallel=0,
         env={},
         attrs={},
+        ephemeral_storage=None,
+        log_driver=None,
+        log_options=None,
     ):
         if queue is None:
             queue = next(self._client.active_job_queues(), None)
@@ -390,6 +399,9 @@ class Batch(object):
             tmpfs_size=tmpfs_size,
             tmpfs_path=tmpfs_path,
             num_parallel=num_parallel,
+            ephemeral_storage=ephemeral_storage,
+            log_driver=log_driver,
+            log_options=log_options,
         )
         self.num_parallel = num_parallel
         self.job = job.execute()

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/aws/batch/batch_cli.py

@@ -154,6 +154,25 @@ def kill(ctx, run_id, user, my_runs):
 @click.option("--ubf-context", default=None, type=click.Choice([None, "ubf_control"]))
 @click.option("--host-volumes", multiple=True)
 @click.option("--efs-volumes", multiple=True)
+@click.option(
+    "--ephemeral-storage",
+    default=None,
+    type=int,
+    help="Ephemeral storage (for AWS Batch only)",
+)
+@click.option(
+    "--log-driver",
+    default=None,
+    type=str,
+    help="Log driver for AWS ECS container",
+)
+@click.option(
+    "--log-options",
+    default=None,
+    type=str,
+    multiple=True,
+    help="Log options for the chosen log driver",
+)
 @click.option(
     "--num-parallel",
     default=0,
@@ -186,6 +205,9 @@ def step(
     tmpfs_path=None,
     host_volumes=None,
     efs_volumes=None,
+    ephemeral_storage=None,
+    log_driver=None,
+    log_options=None,
     num_parallel=None,
     **kwargs
 ):
@@ -317,6 +339,9 @@ def step(
                 tmpfs_tempdir=tmpfs_tempdir,
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
+                ephemeral_storage=ephemeral_storage,
+                log_driver=log_driver,
+                log_options=log_options,
                 num_parallel=num_parallel,
             )
     except Exception as e:

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/aws/batch/batch_client.py

@@ -158,6 +158,9 @@ class BatchJob(object):
         tmpfs_size,
         tmpfs_path,
         num_parallel,
+        ephemeral_storage,
+        log_driver,
+        log_options,
     ):
         # identify platform from any compute environment associated with the
         # queue
@@ -195,6 +198,25 @@ class BatchJob(object):
             "propagateTags": True,
         }
 
+        log_options_dict = {}
+        if log_options:
+            if isinstance(log_options, str):
+                log_options = [log_options]
+            for each_log_option in log_options:
+                k, v = each_log_option.split(":", 1)
+                log_options_dict[k] = v
+
+        if log_driver or log_options:
+            job_definition["containerProperties"]["logConfiguration"] = {}
+        if log_driver:
+            job_definition["containerProperties"]["logConfiguration"][
+                "logDriver"
+            ] = log_driver
+        if log_options:
+            job_definition["containerProperties"]["logConfiguration"][
+                "options"
+            ] = log_options_dict
+
         if platform == "FARGATE" or platform == "FARGATE_SPOT":
             if num_parallel > 1:
                 raise BatchJobException("Fargate does not support multinode jobs.")
@@ -210,6 +232,10 @@ class BatchJob(object):
             job_definition["containerProperties"]["networkConfiguration"] = {
                 "assignPublicIp": "ENABLED"
             }
+            if ephemeral_storage:
+                job_definition["containerProperties"]["ephemeralStorage"] = {
+                    "sizeInGiB": ephemeral_storage
+                }
 
         if platform == "EC2" or platform == "SPOT":
             if "linuxParameters" not in job_definition["containerProperties"]:
@@ -254,6 +280,10 @@ class BatchJob(object):
                     job_definition["containerProperties"]["linuxParameters"][
                         "maxSwap"
                     ] = int(max_swap)
+            if ephemeral_storage:
+                raise BatchJobException(
+                    "The ephemeral_storage parameter is only available for FARGATE compute environments"
+                )
 
         if inferentia:
             if not (isinstance(inferentia, (int, unicode, basestring))):
@@ -315,6 +345,10 @@ class BatchJob(object):
                         {"sourceVolume": name, "containerPath": container_path}
                     )
 
+        if use_tmpfs and (platform == "FARGATE" or platform == "FARGATE_SPOT"):
+            raise BatchJobException(
+                "tmpfs is not available for Fargate compute resources"
+            )
         if use_tmpfs or (tmpfs_size and not use_tmpfs):
             if tmpfs_size:
                 if not (isinstance(tmpfs_size, (int, unicode, basestring))):
@@ -442,6 +476,9 @@ class BatchJob(object):
         tmpfs_size,
         tmpfs_path,
         num_parallel,
+        ephemeral_storage,
+        log_driver,
+        log_options,
     ):
         self.payload["jobDefinition"] = self._register_job_definition(
             image,
@@ -461,6 +498,9 @@ class BatchJob(object):
             tmpfs_size,
             tmpfs_path,
             num_parallel,
+            ephemeral_storage,
+            log_driver,
+            log_options,
         )
         return self
 

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/aws/batch/batch_decorator.py

@@ -72,7 +72,8 @@ class BatchDecorator(StepDecorator):
         necessary. A swappiness value of 100 causes pages to be swapped very
         aggressively. Accepted values are whole numbers between 0 and 100.
     use_tmpfs : bool, default False
-        This enables an explicit tmpfs mount for this step.
+        This enables an explicit tmpfs mount for this step. Note that tmpfs is
+        not available on Fargate compute environments
     tmpfs_tempdir : bool, default True
         sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
     tmpfs_size : int, optional, default None
@@ -83,8 +84,19 @@ class BatchDecorator(StepDecorator):
         Path to tmpfs mount for this step. Defaults to /metaflow_temp.
     inferentia : int, default 0
         Number of Inferentia chips required for this step.
+    trainium : int, default None
+        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)
+        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"]
     """
 
     name = "batch"
@@ -100,6 +112,7 @@ class BatchDecorator(StepDecorator):
         "max_swap": None,
         "swappiness": None,
         "inferentia": None,
+        "trainium": None,  # alias for inferentia
         "efa": None,
         "host_volumes": None,
         "efs_volumes": None,
@@ -107,6 +120,9 @@ class BatchDecorator(StepDecorator):
         "tmpfs_tempdir": True,
         "tmpfs_size": None,
         "tmpfs_path": "/metaflow_temp",
+        "ephemeral_storage": None,
+        "log_driver": None,
+        "log_options": None,
     }
     resource_defaults = {
         "cpu": "1",
@@ -146,6 +162,21 @@ class BatchDecorator(StepDecorator):
                     self.attributes["image"],
                 )
 
+        # Alias trainium to inferentia and check that both are not in use.
+        if (
+            self.attributes["inferentia"] is not None
+            and self.attributes["trainium"] is not None
+        ):
+            raise BatchException(
+                "only specify a value for 'inferentia' or 'trainium', not both."
+            )
+
+        if self.attributes["trainium"] is not None:
+            self.attributes["inferentia"] = self.attributes["trainium"]
+
+        # clean up the alias attribute so it is not passed on.
+        self.attributes.pop("trainium", None)
+
     # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
     # to understand where these functions are invoked in the lifecycle of a
     # Metaflow flow.

{ob-metaflow-2.11.4.2 → ob-metaflow-2.11.10.2}/metaflow/plugins/aws/step_functions/step_functions.py

@@ -839,6 +839,9 @@ class StepFunctions(object):
                 attrs=attrs,
                 host_volumes=resources["host_volumes"],
                 efs_volumes=resources["efs_volumes"],
+                ephemeral_storage=resources["ephemeral_storage"],
+                log_driver=resources["log_driver"],
+                log_options=resources["log_options"],
             )
             .attempts(total_retries + 1)
         )