metaflow 2.8.1__tar.gz → 2.8.3__tar.gz

{metaflow-2.8.1/metaflow.egg-info → metaflow-2.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.8.1
+Version: 2.8.3
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -73,8 +73,3 @@ There are several ways to get in touch with us:
 ## Contributing
 
 We welcome contributions to Metaflow. Please see our [contribution guide](https://docs.metaflow.org/introduction/contributing-to-metaflow) for more details.
-
-### Code style
-
-We use [black](https://black.readthedocs.io/en/stable/) as a code formatter. The easiest way to ensure your commits are always formatted with the correct version of `black` it is to use [pre-commit](https://pre-commit.com/): install it and then run `pre-commit install` once in your local copy of the repo.
-

{metaflow-2.8.1 → metaflow-2.8.3}/README.md

@@ -48,8 +48,3 @@ There are several ways to get in touch with us:
 ## Contributing
 
 We welcome contributions to Metaflow. Please see our [contribution guide](https://docs.metaflow.org/introduction/contributing-to-metaflow) for more details.
-
-### Code style
-
-We use [black](https://black.readthedocs.io/en/stable/) as a code formatter. The easiest way to ensure your commits are always formatted with the correct version of `black` it is to use [pre-commit](https://pre-commit.com/): install it and then run `pre-commit install` once in your local copy of the repo.
-

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/client/core.py

@@ -265,6 +265,7 @@ class MetaflowObject(object):
         self._parent = _parent
         self._path_components = None
         self._attempt = attempt
+        self._namespace_check = _namespace_check
 
         if self._attempt is not None:
             if self._NAME not in ["task", "artifact"]:
@@ -315,7 +316,7 @@ class MetaflowObject(object):
         self._user_tags = frozenset(self._object.get("tags") or [])
         self._system_tags = frozenset(self._object.get("system_tags") or [])
 
-        if _namespace_check and not self.is_in_namespace():
+        if self._namespace_check and not self.is_in_namespace():
             raise MetaflowNamespaceMismatch(current_namespace)
 
     def _get_object(self, *path_components):
@@ -330,7 +331,8 @@ class MetaflowObject(object):
         """
         Iterate over all child objects of this object if any.
 
-        Note that only children present in the current namespace are returned.
+        Note that only children present in the current namespace are returned iff
+        _namespace_check is set.
 
         Returns
         -------
@@ -338,7 +340,8 @@ class MetaflowObject(object):
             Iterator over all children
         """
         query_filter = {}
-        if current_namespace:
+        # skip namespace filtering if _namespace_check is False
+        if self._namespace_check and current_namespace:
             query_filter = {"any_tags": current_namespace}
 
         unfiltered_children = self._metaflow.metadata.get_object(
@@ -381,6 +384,10 @@ class MetaflowObject(object):
             if all(tag in child.tags for tag in tags):
                 yield child
 
+    def _ipython_key_completions_(self):
+        """Returns available options for ipython auto-complete."""
+        return [child.id for child in self._filtered_children()]
+
     @classmethod
     def _url_token(cls):
         return "%ss" % cls._NAME
@@ -444,7 +451,10 @@ class MetaflowObject(object):
         obj = self._get_child(id)
         if obj:
             return _CLASSES[self._CHILD_CLASS](
-                attempt=self._attempt, _object=obj, _parent=self
+                attempt=self._attempt,
+                _object=obj,
+                _parent=self,
+                _namespace_check=self._namespace_check,
             )
         else:
             raise KeyError(id)

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/cmd/configure_cmd.py

@@ -249,13 +249,13 @@ def configure_s3_datastore(existing_env):
         show_default=True,
     )
     # Set Amazon S3 folder for datatools.
-    env["METAFLOW_DATATOOLS_SYSROOT_S3"] = click.prompt(
-        cyan("[METAFLOW_DATATOOLS_SYSROOT_S3]")
+    env["METAFLOW_DATATOOLS_S3ROOT"] = click.prompt(
+        cyan("[METAFLOW_DATATOOLS_S3ROOT]")
         + yellow(" (optional)")
         + " Amazon S3 folder for Metaflow datatools "
         + "(s3://<bucket>/<prefix>).",
         default=existing_env.get(
-            "METAFLOW_DATATOOLS_SYSROOT_S3",
+            "METAFLOW_DATATOOLS_S3ROOT",
             os.path.join(env["METAFLOW_DATASTORE_SYSROOT_S3"], "data"),
         ),
         show_default=True,

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/cmd/main_cli.py

@@ -4,7 +4,8 @@ from metaflow._vendor import click
 
 from metaflow.extension_support.cmd import process_cmds, resolve_cmds
 from metaflow.plugins.datastores.local_storage import LocalStorage
-from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
+from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, CONTACT_INFO
+from metaflow.metaflow_version import get_version
 
 from .util import echo_always
 
@@ -80,24 +81,18 @@ def start(ctx):
     echo("Metaflow ", fg="magenta", bold=True, nl=False)
 
     if ctx.invoked_subcommand is None:
-        echo("(%s): " % metaflow.__version__, fg="magenta", bold=False, nl=False)
+        echo("(%s): " % get_version(), fg="magenta", bold=False, nl=False)
     else:
-        echo("(%s)\n" % metaflow.__version__, fg="magenta", bold=False)
+        echo("(%s)\n" % get_version(), fg="magenta", bold=False)
 
     if ctx.invoked_subcommand is None:
         echo("More data science, less engineering\n", fg="magenta")
 
-        # metaflow URL
-        echo("http://docs.metaflow.org", fg="cyan", nl=False)
-        echo(" - Read the documentation")
-
-        # metaflow chat
-        echo("http://chat.metaflow.org", fg="cyan", nl=False)
-        echo(" - Chat with us")
-
-        # metaflow help email
-        echo("help@metaflow.org", fg="cyan", nl=False)
-        echo("        - Get help by email\n")
+        lnk_sz = max(len(lnk) for lnk in CONTACT_INFO.values()) + 1
+        for what, lnk in CONTACT_INFO.items():
+            echo("%s%s" % (lnk, " " * (lnk_sz - len(lnk))), fg="cyan", nl=False)
+            echo("- %s" % what)
+        echo("")
 
         print(ctx.get_help())
 

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/current.py

@@ -2,6 +2,8 @@ from collections import namedtuple
 import os
 from typing import Any, Optional
 
+from metaflow.metaflow_config import TEMPDIR
+
 Parallel = namedtuple("Parallel", ["main_ip", "num_nodes", "node_index"])
 
 
@@ -17,6 +19,7 @@ class Current(object):
         self._username = None
         self._metadata_str = None
         self._is_running = False
+        self._tempdir = TEMPDIR
 
         def _raise(ex):
             raise ex
@@ -230,6 +233,18 @@ class Current(object):
         """
         return self._tags
 
+    @property
+    def tempdir(self) -> str:
+        """
+        Currently configured temp dir.
+
+        Returns
+        -------
+        str
+            temp dir.
+        """
+        return self._tempdir
+
 
 # instantiate the Current singleton. This will be populated
 # by task.MetaflowTask before a task is executed.

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/datastore/datastore_set.py

@@ -22,15 +22,14 @@ class TaskDataStoreSet(object):
         prefetch_data_artifacts=None,
         allow_not_done=False,
     ):
-
-        task_datastores = flow_datastore.get_latest_task_datastores(
+        self.task_datastores = flow_datastore.get_latest_task_datastores(
             run_id, steps=steps, pathspecs=pathspecs, allow_not_done=allow_not_done
         )
 
         if prefetch_data_artifacts:
             # produce a set of SHA keys to prefetch based on artifact names
             prefetch = set()
-            for ds in task_datastores:
+            for ds in self.task_datastores:
                 prefetch.update(ds.keys_for_artifacts(prefetch_data_artifacts))
             # ignore missing keys
             prefetch.discard(None)
@@ -43,9 +42,10 @@ class TaskDataStoreSet(object):
 
         self.pathspec_index_cache = {}
         self.pathspec_cache = {}
-        for ds in task_datastores:
-            self.pathspec_index_cache[ds.pathspec_index] = ds
-            self.pathspec_cache[ds.pathspec] = ds
+        if not allow_not_done:
+            for ds in self.task_datastores:
+                self.pathspec_index_cache[ds.pathspec_index] = ds
+                self.pathspec_cache[ds.pathspec] = ds
 
     def get_with_pathspec(self, pathspec):
         return self.pathspec_cache.get(pathspec, None)
@@ -54,7 +54,7 @@ class TaskDataStoreSet(object):
         return self.pathspec_index_cache.get(pathspec_index, None)
 
     def __iter__(self):
-        for v in self.pathspec_cache.values():
+        for v in self.task_datastores:
             yield v
 
 

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/datastore/flow_datastore.py

@@ -172,7 +172,7 @@ class FlowDataStore(object):
         else:
             latest_to_fetch = latest_started_attempts & done_attempts
         latest_to_fetch = [
-            (v[0], v[1], v[2], v[3], data_objs[v], "r", allow_not_done)
+            (v[0], v[1], v[2], v[3], data_objs.get(v), "r", allow_not_done)
             for v in latest_to_fetch
         ]
         return list(itertools.starmap(self.get_task_datastore, latest_to_fetch))
@@ -187,7 +187,6 @@ class FlowDataStore(object):
         mode="r",
         allow_not_done=False,
     ):
-
         return TaskDataStore(
             self,
             run_id,

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/extension_support/__init__.py

@@ -283,6 +283,7 @@ _extension_points = [
     "exceptions",
     "toplevel",
     "cmd",
+    "alias",
 ]
 
 

metaflow-2.8.3/metaflow/extension_support/integrations.py

@@ -0,0 +1,141 @@
+import importlib
+import traceback
+
+from metaflow.metaflow_config_funcs import from_conf
+
+from . import _ext_debug, get_modules
+
+# This file is similar in functionality to the cmd.py file. Please refer to that
+# one for more information on what the functions do.
+
+
+def process_integration_aliases(module_globals):
+    _resolve_relative_paths(module_globals)
+
+    all_aliases = _get_ext_aliases(module_globals)
+    all_aliases_dict = {}
+
+    toggle_alias = []
+    list_of_aliases = from_conf("ENABLED_INTEGRATION_ALIAS")
+
+    try:
+        modules_to_import = get_modules("alias")
+        # This is like multiload_all but we load globals independently since we just care
+        # about the TOGGLE and ENABLED values
+        for m in modules_to_import:
+            for n, o in m.module.__dict__.items():
+                if n == "TOGGLE_INTEGRATION_ALIAS":
+                    toggle_alias.extend(o)
+                elif n == "ENABLED_INTEGRATION_ALIAS":
+                    list_of_aliases = o
+            _resolve_relative_paths(m.module.__dict__)
+            all_aliases.extend(_get_ext_aliases(m.module.__dict__))
+    except Exception as e:
+        _ext_debug(
+            "\tWARNING: ignoring all integration aliases due to error during import: %s"
+            % e
+        )
+        print(
+            "WARNING: Integration aliases did not load -- ignoring all of them which "
+            "may not be what you want: %s" % e
+        )
+        traceback.print_exc()
+
+    # At this point, we have _all_aliases populated with all the tuples
+    # (name, module_class) from all the aliases in all the extensions (if any)
+    # We build a dictionary taking the latest presence for each name (so plugins
+    # override metaflow core)
+    for name, obj_path in all_aliases:
+        _ext_debug("    Adding integration alias '%s' from '%s'" % (name, obj_path))
+        all_aliases_dict[name] = obj_path
+
+    # Resolve the ENABLED_INTEGRATION_ALIAS variable. The rules are the following:
+    #  - if ENABLED_INTEGRATION_ALIAS is non None, it means it was either set directly
+    #    by the user in a configuration file, on the command line or by an extension.
+    #    In that case we honor those wishes and completely ignore the extensions' toggles.
+    #  - if ENABLED_INTEGRATION_ALIAS is None, we populate it with everything included
+    #    here and in all the extensions and use the TOGGLE_ list to produce the final list.
+    # The rationale behind this is to support both a configuration option where the
+    # aliases enabled are explicitly listed (typical in a lot of software) but also to
+    # support a "configuration-less" version where the installation of the extensions
+    # determines what is activated.
+    if list_of_aliases is None:
+        list_of_aliases = list(all_aliases_dict) + toggle_alias
+
+    _ext_debug("    Resolving metaflow integration aliases")
+    _ext_debug("        Raw list is: %s" % str(list_of_aliases))
+
+    set_of_aliases = set()
+    for p in list_of_aliases:
+        if p.startswith("-"):
+            set_of_aliases.discard(p[1:])
+        elif p.startswith("+"):
+            set_of_aliases.add(p[1:])
+        else:
+            set_of_aliases.add(p)
+    _ext_debug("        Resolved list is: %s" % str(set_of_aliases))
+
+    for name in set_of_aliases:
+        obj_path = all_aliases_dict.get(name, None)
+        if obj_path is None:
+            raise ValueError(
+                "Configuration requested integration alias '%s' but no such alias "
+                "is available" % name
+            )
+        path, obj_name = obj_path.rsplit(".", 1)
+        try:
+            alias_module = importlib.import_module(path)
+        except ImportError:
+            raise ValueError(
+                "Cannot locate integration alias '%s' at '%s'" % (name, path)
+            )
+
+        obj = getattr(alias_module, obj_name, None)
+        if obj is None:
+            raise ValueError(
+                "Cannot locate '%s' object for integration alias at '%s'"
+                % (obj_name, path)
+            )
+        _ext_debug("        Added integration alias '%s' from '%s'" % (name, obj_path))
+        module_globals[name] = obj
+
+
+def _get_ext_aliases(module_globals):
+    return module_globals.get("ALIASES_DESC", [])
+
+
+def _set_ext_aliases(module_globals, value):
+    module_globals["ALIASES_DESC"] = value
+
+
+def _resolve_relative_paths(module_globals):
+    # We want to modify all the relevant lists so that the relative paths
+    # are made fully qualified paths for the modules
+    pkg_path = module_globals["__package__"]
+    pkg_components = pkg_path.split(".")
+
+    def resolve_path(class_path):
+        # Converts a relative class_path to an absolute one considering that the
+        # relative class_path is present in a package pkg_path
+        if class_path[0] == ".":
+            i = 1
+            # Check for multiple "." at the start of the class_path
+            while class_path[i] == ".":
+                i += 1
+            if i > len(pkg_components):
+                raise ValueError(
+                    "Path '%s' exits out of Metaflow module at %s"
+                    % (class_path, pkg_path)
+                )
+            return (
+                ".".join(pkg_components[: -i + 1] if i > 1 else pkg_components)
+                + class_path[i - 1 :]
+            )
+        return class_path
+
+    _set_ext_aliases(
+        module_globals,
+        list(
+            map(lambda p: (p[0], resolve_path(p[1])), _get_ext_aliases(module_globals))
+        ),
+    )

metaflow-2.8.3/metaflow/integrations.py

@@ -0,0 +1,29 @@
+# This file can contain "shortcuts" to other parts of Metaflow (integrations)
+
+# This is an alternative to providing an extension package where you would define
+# these aliases in the toplevel file.
+
+# It follows a similar pattern to plugins so that the these integration aliases can be
+# turned on and off and avoid exposing things that are not necessarily needed/wanted.
+
+from metaflow.extension_support.integrations import process_integration_aliases
+
+# To enable an alias `metaflow.alias.get_s3_client` to
+# `metaflow.plugins.aws.aws_client.get_aws_client`, use the following:
+#
+# ALIASES_DESC = [("get_s3_client", ".plugins.aws.aws_client.get_aws_client")]
+#
+# ALIASES_DESC is a list of tuples:
+#  - name: name of the integration alias
+#  - obj: object it points to
+#
+# Aliases can be enabled or disabled through configuration or extensions:
+#  - ENABLED_INTEGRATION_ALIAS: list of alias names to enable.
+#  - TOGGLE_INTEGRATION_ALIAS: if ENABLED_INTEGRATION_ALIAS is not set anywhere
+#    (environment variable, configuration or extensions), list of integration aliases
+#    to toggle (+<name> or <name> enables, -<name> disables) to build
+#    ENABLED_INTEGRATION_ALIAS from the concatenation of the names in
+#    ALIASES_DESC (concatenation of the names here as well as in the extensions).
+
+# Keep this line and make sure ALIASES_DESC is above this line.
+process_integration_aliases(globals())

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/metaflow_config.py

@@ -96,6 +96,8 @@ DATATOOLS_S3ROOT = from_conf(
     else None,
 )
 
+TEMPDIR = from_conf("TEMPDIR", ".")
+
 DATATOOLS_CLIENT_PARAMS = from_conf("DATATOOLS_CLIENT_PARAMS", {})
 if S3_ENDPOINT_URL:
     DATATOOLS_CLIENT_PARAMS["endpoint_url"] = S3_ENDPOINT_URL
@@ -193,8 +195,25 @@ DEFAULT_CONTAINER_IMAGE = from_conf("DEFAULT_CONTAINER_IMAGE")
 # Default container registry
 DEFAULT_CONTAINER_REGISTRY = from_conf("DEFAULT_CONTAINER_REGISTRY")
 
+###
+# Organization customizations
+###
 UI_URL = from_conf("UI_URL")
 
+# Contact information displayed when running the `metaflow` command.
+# Value should be a dictionary where:
+#  - key is a string describing contact method
+#  - value is a string describing contact itself (email, web address, etc.)
+# The default value shows an example of this
+CONTACT_INFO = from_conf(
+    "CONTACT_INFO",
+    {
+        "Read the documentation": "http://docs.metaflow.org",
+        "Chat with us": "http://chat.metaflow.org",
+        "Get help by email": "help@metaflow.org",
+    },
+)
+
 ###
 # AWS Batch configuration
 ###
@@ -260,6 +279,8 @@ KUBERNETES_CONTAINER_IMAGE = from_conf(
 KUBERNETES_CONTAINER_REGISTRY = from_conf(
     "KUBERNETES_CONTAINER_REGISTRY", DEFAULT_CONTAINER_REGISTRY
 )
+# Toggle for trying to fetch EC2 instance metadata
+KUBERNETES_FETCH_EC2_METADATA = from_conf("KUBERNETES_FETCH_EC2_METADATA", False)
 
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/metaflow_environment.py

@@ -162,7 +162,7 @@ class MetaflowEnvironment(object):
         ]
         return cmds
 
-    def get_environment_info(self):
+    def get_environment_info(self, include_ext_info=False):
         global version_cache
         if version_cache is None:
             version_cache = metaflow_version.get_version()
@@ -187,9 +187,10 @@ class MetaflowEnvironment(object):
             env["metaflow_r_version"] = R.metaflow_r_version()
             env["r_version"] = R.r_version()
             env["r_version_code"] = R.r_version_code()
-        # Information about extension modules (to load them in the proper order)
-        ext_key, ext_val = dump_module_info()
-        env[ext_key] = ext_val
+        if include_ext_info:
+            # Information about extension modules (to load them in the proper order)
+            ext_key, ext_val = dump_module_info()
+            env[ext_key] = ext_val
         return env
 
     def executable(self, step_name):

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/package.py

@@ -151,7 +151,7 @@ class MetaflowPackage(object):
 
     def _add_info(self, tar):
         info = tarfile.TarInfo(os.path.basename(INFO_FILE))
-        env = self.environment.get_environment_info()
+        env = self.environment.get_environment_info(include_ext_info=True)
         buf = BytesIO()
         buf.write(json.dumps(env).encode("utf-8"))
         buf.seek(0)

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/plugins/airflow/airflow.py

@@ -54,7 +54,6 @@ AIRFLOW_DEPLOY_TEMPLATE_FILE = os.path.join(os.path.dirname(__file__), "dag.py")
 
 
 class Airflow(object):
-
     TOKEN_STORAGE_ROOT = "mf.airflow"
 
     def __init__(

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/plugins/argo/argo_workflows.py

@@ -18,6 +18,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_NODE_SELECTOR,
     KUBERNETES_SANDBOX_INIT_SCRIPT,
     KUBERNETES_SECRETS,
+    KUBERNETES_FETCH_EC2_METADATA,
     S3_ENDPOINT_URL,
     AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
     DATASTORE_SYSROOT_AZURE,
@@ -785,6 +786,7 @@ class ArgoWorkflows(object):
                         "METAFLOW_DEFAULT_METADATA": DEFAULT_METADATA,
                         "METAFLOW_CARD_S3ROOT": CARD_S3ROOT,
                         "METAFLOW_KUBERNETES_WORKLOAD": 1,
+                        "METAFLOW_KUBERNETES_FETCH_EC2_METADATA": KUBERNETES_FETCH_EC2_METADATA,
                         "METAFLOW_RUNTIME_ENVIRONMENT": "kubernetes",
                         "METAFLOW_OWNER": self.username,
                     },

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/plugins/argo/argo_workflows_cli.py

@@ -8,7 +8,7 @@ from hashlib import sha1
 
 from metaflow import JSONType, current, decorators, parameters
 from metaflow._vendor import click
-from metaflow.metaflow_config import SERVICE_VERSION_CHECK
+from metaflow.metaflow_config import SERVICE_VERSION_CHECK, UI_URL
 from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.package import MetaflowPackage
 from metaflow.plugins.environment_decorator import EnvironmentDecorator
@@ -511,3 +511,13 @@ def trigger(obj, run_id_file=None, **kwargs):
         "(run-id *{run_id}*).".format(name=obj.workflow_name, run_id=run_id),
         bold=True,
     )
+
+    run_url = (
+        "%s/%s/%s" % (UI_URL.rstrip("/"), obj.flow.name, run_id) if UI_URL else None
+    )
+
+    if run_url:
+        obj.echo(
+            "See the run in the UI at %s" % run_url,
+            bold=True,
+        )

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/plugins/aws/aws_utils.py

@@ -20,8 +20,13 @@ def get_ec2_instance_metadata():
     # for non-AWS deployments.
     # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html
     try:
+        # Set a very aggressive timeout, as the communication is happening in the same subnet,
+        # there should not be any significant delay in the response.
+        # Having a long default timeout here introduces unnecessary delay in launching tasks when the
+        # instance is unreachable.
         instance_meta = requests.get(
-            url="http://169.254.169.254/latest/dynamic/instance-identity/document"
+            url="http://169.254.169.254/latest/dynamic/instance-identity/document",
+            timeout=(1, 10),
         ).json()
         meta["ec2-instance-id"] = instance_meta.get("instanceId")
         meta["ec2-instance-type"] = instance_meta.get("instanceType")

{metaflow-2.8.1 → metaflow-2.8.3}/metaflow/plugins/aws/batch/batch.py

@@ -179,6 +179,10 @@ class Batch(object):
         env={},
         attrs={},
         host_volumes=None,
+        use_tmpfs=None,
+        tmpfs_tempdir=None,
+        tmpfs_size=None,
+        tmpfs_path=None,
         num_parallel=0,
     ):
         job_name = self._job_name(
@@ -201,6 +205,14 @@ class Batch(object):
             .image(image)
             .iam_role(iam_role)
             .execution_role(execution_role)
+            .cpu(cpu)
+            .gpu(gpu)
+            .memory(memory)
+            .shared_memory(shared_memory)
+            .max_swap(max_swap)
+            .swappiness(swappiness)
+            .inferentia(inferentia)
+            .timeout_in_secs(run_time_limit)
             .job_def(
                 image,
                 iam_role,
@@ -210,17 +222,14 @@ class Batch(object):
                 max_swap,
                 swappiness,
                 inferentia,
+                memory=memory,
                 host_volumes=host_volumes,
+                use_tmpfs=use_tmpfs,
+                tmpfs_tempdir=tmpfs_tempdir,
+                tmpfs_size=tmpfs_size,
+                tmpfs_path=tmpfs_path,
                 num_parallel=num_parallel,
             )
-            .cpu(cpu)
-            .gpu(gpu)
-            .memory(memory)
-            .shared_memory(shared_memory)
-            .max_swap(max_swap)
-            .swappiness(swappiness)
-            .inferentia(inferentia)
-            .timeout_in_secs(run_time_limit)
             .task_id(attrs.get("metaflow.task_id"))
             .environment_variable("AWS_DEFAULT_REGION", self._client.region())
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
@@ -248,6 +257,11 @@ class Batch(object):
                 AWS_SECRETS_MANAGER_DEFAULT_REGION,
             )
 
+        tmpfs_enabled = use_tmpfs or (tmpfs_size and use_tmpfs is None)
+
+        if tmpfs_enabled and tmpfs_tempdir:
+            job.environment_variable("METAFLOW_TEMPDIR", tmpfs_path)
+
         # Skip setting METAFLOW_DATASTORE_SYSROOT_LOCAL because metadata sync between the local user
         # instance and the remote AWS Batch instance assumes metadata is stored in DATASTORE_LOCAL_DIR
         # on the remote AWS Batch instance; this happens when METAFLOW_DATASTORE_SYSROOT_LOCAL
@@ -300,6 +314,10 @@ class Batch(object):
         swappiness=None,
         inferentia=None,
         host_volumes=None,
+        use_tmpfs=None,
+        tmpfs_tempdir=None,
+        tmpfs_size=None,
+        tmpfs_path=None,
         num_parallel=0,
         env={},
         attrs={},
@@ -333,6 +351,10 @@ class Batch(object):
             env=env,
             attrs=attrs,
             host_volumes=host_volumes,
+            use_tmpfs=use_tmpfs,
+            tmpfs_tempdir=tmpfs_tempdir,
+            tmpfs_size=tmpfs_size,
+            tmpfs_path=tmpfs_path,
             num_parallel=num_parallel,
         )
         self.num_parallel = num_parallel