ob-metaflow 2.10.7.4__py2.py3-none-any.whl → 2.10.9.2__py2.py3-none-any.whl

metaflow/plugins/cards/exception.py

@@ -116,7 +116,7 @@ class UnresolvableDatastoreException(MetaflowException):
         super(UnresolvableDatastoreException, self).__init__(msg)
 
 
-class IncorrectArguementException(MetaflowException):
+class IncorrectArgumentException(MetaflowException):
     headline = (
         "`get_cards` function requires a `Task` object or pathspec as an argument"
     )
@@ -138,3 +138,22 @@ class IncorrectPathspecException(MetaflowException):
             % pthspec
         )
         super().__init__(msg=msg, lineno=None)
+
+
+class ComponentOverwriteNotSupportedException(MetaflowException):
+    headline = "Component overwrite is not supported"
+
+    def __init__(self, component_id, card_id, card_type):
+        id_str = ""
+        if card_id is not None:
+            id_str = "id='%s'" % card_id
+        msg = (
+            "Card component overwrite is not supported. "
+            "Component with id %s already exists in the @card(type='%s', %s). \n"
+            "Instead of calling `current.card.components[ID] = MyComponent`. "
+            "You can overwrite the entire component Array by calling "
+            "`current.card.components = [MyComponent]`"
+        ) % (component_id, card_type, id_str)
+        super().__init__(
+            msg=msg,
+        )

metaflow/plugins/datastores/azure_storage.py

@@ -250,9 +250,8 @@ class _AzureRootClient(object):
 class AzureStorage(DataStoreStorage):
     TYPE = "azure"
 
+    @check_azure_deps
     def __init__(self, root=None):
-        # cannot decorate __init__... invoke it with dummy decoratee
-        check_azure_deps(lambda: 0)
         super(AzureStorage, self).__init__(root)
         self._tmproot = ARTIFACT_LOCALROOT
         self._default_scope_token = None

metaflow/plugins/datastores/gs_storage.py

@@ -145,9 +145,8 @@ class _GSRootClient(object):
 class GSStorage(DataStoreStorage):
     TYPE = "gs"
 
+    @check_gs_deps
     def __init__(self, root=None):
-        # cannot decorate __init__... invoke it with dummy decoratee
-        check_gs_deps(lambda: 0)
         super(GSStorage, self).__init__(root)
         self._tmproot = ARTIFACT_LOCALROOT
         self._root_client = None

metaflow/plugins/datastores/s3_storage.py

@@ -2,7 +2,7 @@ import os
 
 from itertools import starmap
 
-from metaflow.plugins.datatools.s3.s3 import S3, S3Client, S3PutObject
+from metaflow.plugins.datatools.s3.s3 import S3, S3Client, S3PutObject, check_s3_deps
 from metaflow.metaflow_config import DATASTORE_SYSROOT_S3, ARTIFACT_LOCALROOT
 from metaflow.datastore.datastore_storage import CloseAfterUse, DataStoreStorage
 
@@ -18,6 +18,7 @@ except:
 class S3Storage(DataStoreStorage):
     TYPE = "s3"
 
+    @check_s3_deps
     def __init__(self, root=None):
         super(S3Storage, self).__init__(root)
         self.s3_client = S3Client()

metaflow/plugins/datatools/s3/s3.py

@@ -51,15 +51,25 @@ from .s3util import (
 if TYPE_CHECKING:
     from metaflow.client import Run
 
-try:
-    import boto3
-    from boto3.s3.transfer import TransferConfig
 
-    DOWNLOAD_FILE_THRESHOLD = 2 * TransferConfig().multipart_threshold
-    DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1
-    boto_found = True
-except:
-    boto_found = False
+def _check_and_init_s3_deps():
+    try:
+        import boto3
+        from boto3.s3.transfer import TransferConfig
+    except (ImportError, ModuleNotFoundError):
+        raise MetaflowException("You need to install 'boto3' in order to use S3.")
+
+
+def check_s3_deps(func):
+    """The decorated function checks S3 dependencies (as needed for AWS S3 storage backend).
+    This includes boto3.
+    """
+
+    def _inner_func(*args, **kwargs):
+        _check_and_init_s3_deps()
+        return func(*args, **kwargs)
+
+    return _inner_func
 
 
 TEST_INJECT_RETRYABLE_FAILURES = int(
@@ -498,6 +508,7 @@ class S3(object):
     def get_root_from_config(cls, echo, create_on_absent=True):
         return DATATOOLS_S3ROOT
 
+    @check_s3_deps
     def __init__(
         self,
         tmproot: str = TEMPDIR,
@@ -508,9 +519,6 @@ class S3(object):
         encryption: Optional[str] = S3_SERVER_SIDE_ENCRYPTION,
         **kwargs
     ):
-        if not boto_found:
-            raise MetaflowException("You need to install 'boto3' in order to use S3.")
-
         if run:
             # 1. use a (current) run ID with optional customizations
             if DATATOOLS_S3ROOT is None:
@@ -875,6 +883,11 @@ class S3(object):
         `S3Object`
             An S3Object corresponding to the object requested.
         """
+        from boto3.s3.transfer import TransferConfig
+
+        DOWNLOAD_FILE_THRESHOLD = 2 * TransferConfig().multipart_threshold
+        DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1
+
         url, r = self._url_and_range(key)
         src = urlparse(url)
 

metaflow/plugins/env_escape/client.py

@@ -37,6 +37,7 @@ from .data_transferer import DataTransferer, ObjReference
 from .exception_transferer import load_exception
 from .override_decorators import LocalAttrOverride, LocalException, LocalOverride
 from .stub import create_class
+from .utils import get_canonical_name
 
 BIND_TIMEOUT = 0.1
 BIND_RETRY = 0
@@ -336,7 +337,7 @@ class Client(object):
     def get_local_class(self, name, obj_id=None):
         # Gets (and creates if needed), the class mapping to the remote
         # class of name 'name'.
-        name = self._get_canonical_name(name)
+        name = get_canonical_name(name, self._aliases)
         if name == "function":
             # Special handling of pickled functions. We create a new class that
             # simply has a __call__ method that will forward things back to
@@ -398,17 +399,6 @@ class Client(object):
             local_instance = local_class(self, remote_class_name, obj_id)
         return local_instance
 
-    def _get_canonical_name(self, name):
-        # We look at the aliases looking for the most specific match first
-        base_name = self._aliases.get(name)
-        if base_name is not None:
-            return base_name
-        for idx in reversed([pos for pos, char in enumerate(name) if char == "."]):
-            base_name = self._aliases.get(name[:idx])
-            if base_name is not None:
-                return ".".join([base_name, name[idx + 1 :]])
-        return name
-
     def _communicate(self, msg):
         if os.getpid() != self._active_pid:
             raise RuntimeError(

metaflow/plugins/env_escape/client_modules.py

@@ -8,6 +8,7 @@ import sys
 from .consts import OP_CALLFUNC, OP_GETVAL, OP_SETVAL
 from .client import Client
 from .override_decorators import LocalException
+from .utils import get_canonical_name
 
 
 def _clean_client(client):
@@ -23,6 +24,7 @@ class _WrappedModule(object):
             r"^%s\.([a-zA-Z_][a-zA-Z0-9_]*)$" % prefix.replace(".", r"\.")  # noqa W605
         )
         self._exports = {}
+        self._aliases = exports["aliases"]
         for k in ("classes", "functions", "values"):
             result = []
             for item in exports[k]:
@@ -43,6 +45,11 @@ class _WrappedModule(object):
             return self._prefix
         if name in ("__file__", "__path__"):
             return self._client.name
+
+        # Make the name canonical because the prefix is also canonical.
+        name = get_canonical_name(self._prefix + "." + name, self._aliases)[
+            len(self._prefix) + 1 :
+        ]
         if name in self._exports["classes"]:
             # We load classes lazily
             return self._client.get_local_class("%s.%s" % (self._prefix, name))
@@ -87,6 +94,7 @@ class _WrappedModule(object):
             "_client",
             "_exports",
             "_exception_classes",
+            "_aliases",
         ):
             object.__setattr__(self, name, value)
             return
@@ -95,6 +103,11 @@ class _WrappedModule(object):
             # module when loading
             object.__setattr__(self, name, value)
             return
+
+        # Make the name canonical because the prefix is also canonical.
+        name = get_canonical_name(self._prefix + "." + name, self._aliases)[
+            len(self._prefix) + 1 :
+        ]
         if name in self._exports["values"]:
             self._client.stub_request(
                 None, OP_SETVAL, "%s.%s" % (self._prefix, name), value
@@ -126,7 +139,7 @@ class ModuleImporter(object):
 
     def find_module(self, fullname, path=None):
         if self._handled_modules is not None:
-            if fullname in self._handled_modules:
+            if get_canonical_name(fullname, self._aliases) in self._handled_modules:
                 return self
             return None
         if any([fullname.startswith(prefix) for prefix in self._module_prefixes]):
@@ -224,24 +237,15 @@ class ModuleImporter(object):
                 self._handled_modules[prefix] = _WrappedModule(
                     self, prefix, exports, formed_exception_classes, self._client
                 )
-        fullname = self._get_canonical_name(fullname)
-        module = self._handled_modules.get(fullname)
+        canonical_fullname = get_canonical_name(fullname, self._aliases)
+        # Modules are created canonically but we need to return something for any
+        # of the aliases.
+        module = self._handled_modules.get(canonical_fullname)
         if module is None:
             raise ImportError
         sys.modules[fullname] = module
         return module
 
-    def _get_canonical_name(self, name):
-        # We look at the aliases looking for the most specific match first
-        base_name = self._aliases.get(name)
-        if base_name is not None:
-            return base_name
-        for idx in reversed([pos for pos, char in enumerate(name) if char == "."]):
-            base_name = self._aliases.get(name[:idx])
-            if base_name is not None:
-                return ".".join([base_name, name[idx + 1 :]])
-        return name
-
 
 def create_modules(python_executable, pythonpath, max_pickle_version, path, prefixes):
     # This is an extra verification to make sure we are not trying to use the

metaflow/plugins/env_escape/server.py

@@ -53,7 +53,7 @@ from .override_decorators import (
     RemoteExceptionSerializer,
 )
 from .exception_transferer import dump_exception
-from .utils import get_methods
+from .utils import get_methods, get_canonical_name
 
 BIND_TIMEOUT = 0.1
 BIND_RETRY = 1
@@ -61,7 +61,6 @@ BIND_RETRY = 1
 
 class Server(object):
     def __init__(self, config_dir, max_pickle_version):
-
         self._max_pickle_version = data_transferer.defaultProtocol = max_pickle_version
         try:
             mappings = importlib.import_module(".server_mappings", package=config_dir)
@@ -108,6 +107,11 @@ class Server(object):
             for alias in aliases:
                 a = self._aliases.setdefault(alias, base_name)
                 if a != base_name:
+                    # Technically we could have a that aliases b and b that aliases c
+                    # and then a that aliases c. This would error out in that case
+                    # even though it is valid. It is easy for the user to get around
+                    # this by listing aliases in the same order so we don't support
+                    # it for now.
                     raise ValueError(
                         "%s is an alias to both %s and %s" % (alias, base_name, a)
                     )
@@ -155,12 +159,13 @@ class Server(object):
         parent_to_child = {}
 
         for ex_name, ex_cls in self._known_exceptions.items():
+            ex_name_canonical = get_canonical_name(ex_name, self._aliases)
             parents = []
             for base in ex_cls.__mro__[1:]:
                 if base is object:
                     raise ValueError(
-                        "Exported exceptions not rooted in a builtin exception are not supported: %s"
-                        % ex_name
+                        "Exported exceptions not rooted in a builtin exception "
+                        "are not supported: %s." % ex_name
                     )
                 if base.__module__ == "builtins":
                     # We found our base exception
@@ -168,17 +173,19 @@ class Server(object):
                     break
                 else:
                     fqn = ".".join([base.__module__, base.__name__])
-                    if fqn in self._known_exceptions:
-                        parents.append(fqn)
-                        children = parent_to_child.setdefault(fqn, [])
-                        children.append(ex_name)
+                    canonical_fqn = get_canonical_name(fqn, self._aliases)
+                    if canonical_fqn in self._known_exceptions:
+                        parents.append(canonical_fqn)
+                        children = parent_to_child.setdefault(canonical_fqn, [])
+                        children.append(ex_name_canonical)
                     else:
                         raise ValueError(
                             "Exported exception %s has non exported and non builtin parent "
-                            "exception: %s" % (ex_name, fqn)
+                            "exception: %s. Known exceptions: %s"
+                            % (ex_name, fqn, str(self._known_exceptions))
                         )
-            name_to_parent_count[ex_name] = len(parents) - 1
-            name_to_parents[ex_name] = parents
+            name_to_parent_count[ex_name_canonical] = len(parents) - 1
+            name_to_parents[ex_name_canonical] = parents
 
         # We now form the exceptions and put them in self._known_exceptions in
         # the proper order (topologically)

metaflow/plugins/env_escape/utils.py

@@ -20,3 +20,15 @@ def get_methods(class_object):
         elif isinstance(attribute, classmethod):
             all_methods["___c___%s" % name] = inspect.getdoc(attribute)
     return all_methods
+
+
+def get_canonical_name(name, aliases):
+    # We look at the aliases looking for the most specific match first
+    base_name = aliases.get(name)
+    if base_name is not None:
+        return base_name
+    for idx in reversed([pos for pos, char in enumerate(name) if char == "."]):
+        base_name = aliases.get(name[:idx])
+        if base_name is not None:
+            return ".".join([base_name, name[idx + 1 :]])
+    return name

metaflow/plugins/gcp/gs_exceptions.py

@@ -2,4 +2,4 @@ from metaflow.exception import MetaflowException
 
 
 class MetaflowGSPackageError(MetaflowException):
-    headline = "Missing required packages google-cloud-storage google-auth"
+    headline = "Missing required packages 'google-cloud-storage' and 'google-auth'"

metaflow/plugins/gcp/gs_utils.py

@@ -34,7 +34,7 @@ def _check_and_init_gs_deps():
 
 
 def check_gs_deps(func):
-    """The decorated function checks GS dependencies (as needed for Azure storage backend). This includes
+    """The decorated function checks GS dependencies (as needed for Google Cloud storage backend). This includes
     various GCP SDK packages, as well as a Python version of >=3.7
     """
 

metaflow/plugins/kubernetes/kubernetes.py

@@ -4,6 +4,7 @@ import os
 import re
 import shlex
 import time
+import copy
 from typing import Dict, List, Optional
 import uuid
 from uuid import uuid4
@@ -174,6 +175,10 @@ class Kubernetes(object):
         persistent_volume_claims=None,
         tolerations=None,
         labels=None,
+        annotations=None,
+        num_parallel=0,
+        attrs={},
+        port=None,
     ):
         if env is None:
             env = {}
@@ -213,6 +218,9 @@ class Kubernetes(object):
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
                 persistent_volume_claims=persistent_volume_claims,
+                num_parallel=num_parallel,
+                attrs=attrs,
+                port=port,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -266,6 +274,7 @@ class Kubernetes(object):
             # see get_datastore_root_from_config in datastore/local.py).
         )
 
+        self.num_parallel = num_parallel
         # Temporary passing of *some* environment variables. Do not rely on this
         # mechanism as it will be removed in the near future
         for k, v in config_values():
@@ -341,7 +350,7 @@ class Kubernetes(object):
             sigmoid = 1.0 / (1.0 + math.exp(-0.01 * secs_since_start + 9.0))
             return 0.5 + sigmoid * 30.0
 
-        def wait_for_launch(job):
+        def wait_for_launch(job, child_jobs):
             status = job.status
             echo(
                 "Task is starting (%s)..." % status,
@@ -351,11 +360,38 @@ class Kubernetes(object):
             t = time.time()
             start_time = time.time()
             while job.is_waiting:
-                new_status = job.status
-                if status != new_status or (time.time() - t) > 30:
-                    status = new_status
+                # new_status = job.status
+                if status != job.status or (time.time() - t) > 30:
+                    if not child_jobs:
+                        child_statuses = ""
+                    else:
+                        status_keys = set(
+                            [child_job.status for child_job in child_jobs]
+                        )
+                        status_counts = [
+                            (
+                                status,
+                                len(
+                                    [
+                                        child_job.status == status
+                                        for child_job in child_jobs
+                                    ]
+                                ),
+                            )
+                            for status in status_keys
+                        ]
+                        child_statuses = " (parallel node status: [{}])".format(
+                            ", ".join(
+                                [
+                                    "{}:{}".format(status, num)
+                                    for (status, num) in sorted(status_counts)
+                                ]
+                            )
+                        )
+
+                    status = job.status
                     echo(
-                        "Task is starting (%s)..." % status,
+                        "Task is starting (status %s)... %s" % (status, child_statuses),
                         "stderr",
                         job_id=job.id,
                     )
@@ -367,8 +403,9 @@ class Kubernetes(object):
         stdout_tail = get_log_tailer(stdout_location, self._datastore.TYPE)
         stderr_tail = get_log_tailer(stderr_location, self._datastore.TYPE)
 
+        child_jobs = []
         # 1) Loop until the job has started
-        wait_for_launch(self._job)
+        wait_for_launch(self._job, child_jobs)
 
         # 2) Tail logs until the job has finished
         tail_logs(

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -107,6 +107,26 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@click.option(
+    "--labels",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
+@click.option(
+    "--annotations",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
+@click.option("--ubf-context", default=None, type=click.Choice([None, "ubf_control"]))
+@click.option(
+    "--num-parallel",
+    default=0,
+    type=int,
+    help="Number of parallel nodes to run as a multi-node job.",
+)
+@click.option("--port", default=None, help="port number")
 @click.pass_context
 def step(
     ctx,
@@ -132,6 +152,10 @@ def step(
     run_time_limit=None,
     persistent_volume_claims=None,
     tolerations=None,
+    labels=None,
+    annotations=None,
+    num_parallel=None,
+    port=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -177,11 +201,17 @@ def step(
         )
         time.sleep(minutes_between_retries * 60)
 
+    step_args = " ".join(util.dict_to_cli_options(kwargs))
+    num_parallel = num_parallel or 0
+    if num_parallel and num_parallel > 1:
+        # For multinode, we need to add a placeholder that can be mutated by the caller
+        step_args += " [multinode-args]"
+
     step_cli = "{entrypoint} {top_args} step {step} {step_args}".format(
         entrypoint="%s -u %s" % (executable, os.path.basename(sys.argv[0])),
         top_args=" ".join(util.dict_to_cli_options(ctx.parent.parent.params)),
         step=step_name,
-        step_args=" ".join(util.dict_to_cli_options(kwargs)),
+        step_args=step_args,
     )
 
     # Set log tailing.
@@ -207,6 +237,10 @@ def step(
                 ),
             )
 
+    attrs = {
+        "metaflow.task_id": kwargs["task_id"],
+        "requires_passwordless_ssh": any([getattr(deco, "requires_passwordless_ssh", False) for deco in node.decorators]),
+    }
     try:
         kubernetes = Kubernetes(
             datastore=ctx.obj.flow_datastore,
@@ -245,6 +279,11 @@ def step(
                 env=env,
                 persistent_volume_claims=persistent_volume_claims,
                 tolerations=tolerations,
+                labels=labels,
+                annotations=annotations,
+                num_parallel=num_parallel,
+                port=port,
+                attrs=attrs,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -2,6 +2,7 @@ import json
 import os
 import platform
 import sys
+import time
 
 from metaflow import current
 from metaflow.decorators import StepDecorator
@@ -20,10 +21,12 @@ from metaflow.metaflow_config import (
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
     KUBERNETES_TOLERATIONS,
     KUBERNETES_SERVICE_ACCOUNT,
+    KUBERNETES_PORT,
 )
 from metaflow.plugins.resources_decorator import ResourcesDecorator
 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 ..aws.aws_utils import get_docker_registry, get_ec2_instance_metadata
 from .kubernetes import KubernetesException, parse_kube_keyvalue_list
@@ -88,6 +91,8 @@ class KubernetesDecorator(StepDecorator):
     persistent_volume_claims: Dict[str, str], optional
         A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
         volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
+    port: int, optional
+        Number of the port to specify in the Kubernetes job object
     """
 
     name = "kubernetes"
@@ -110,6 +115,7 @@ class KubernetesDecorator(StepDecorator):
         "tmpfs_size": None,
         "tmpfs_path": "/metaflow_temp",
         "persistent_volume_claims": None,  # e.g., {"pvc-name": "/mnt/vol", "another-pvc": "/mnt/vol2"}
+        "port": None,
     }
     package_url = None
     package_sha = None
@@ -195,6 +201,8 @@ class KubernetesDecorator(StepDecorator):
             if not self.attributes["tmpfs_size"]:
                 # default tmpfs behavior - https://man7.org/linux/man-pages/man5/tmpfs.5.html
                 self.attributes["tmpfs_size"] = int(self.attributes["memory"]) // 2
+        if not self.attributes["port"]:
+            self.attributes["port"] = KUBERNETES_PORT
 
     # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
@@ -216,12 +224,6 @@ class KubernetesDecorator(StepDecorator):
                 "Kubernetes. Please use one or the other.".format(step=step)
             )
 
-        for deco in decos:
-            if getattr(deco, "IS_PARALLEL", False):
-                raise KubernetesException(
-                    "@kubernetes does not support parallel execution currently."
-                )
-
         # Set run time limit for the Kubernetes job.
         self.run_time_limit = get_run_time_limit_for_task(decos)
         if self.run_time_limit < 60:
@@ -432,6 +434,27 @@ class KubernetesDecorator(StepDecorator):
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+        num_parallel = int(os.environ.get("WORLD_SIZE", 0))
+        if num_parallel >= 1:
+            if ubf_context == UBF_CONTROL:
+                control_task_id = current.task_id
+                top_task_id = control_task_id.replace("control-", "")
+                mapper_task_ids = [control_task_id] + [
+                    "%s-node-%d" % (top_task_id, node_idx)
+                    for node_idx in range(1, num_parallel)
+                ]
+                flow._control_mapper_tasks = [
+                    "%s/%s/%s" % (run_id, step_name, mapper_task_id)
+                    for mapper_task_id in mapper_task_ids
+                ]
+                flow._control_task_is_mapper_zero = True
+            else:
+                worker_job_rank = int(os.environ["RANK"])
+                os.environ["RANK"] = str(worker_job_rank + 1)
+
+        if num_parallel >= 1:
+            _setup_multinode_environment()
+
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_retries
     ):
@@ -459,9 +482,53 @@ class KubernetesDecorator(StepDecorator):
             # Best effort kill
             pass
 
+        if is_task_ok and len(getattr(flow, "_control_mapper_tasks", [])) > 1:
+            self._wait_for_mapper_tasks(flow, step_name)
+
+    def _wait_for_mapper_tasks(self, flow, step_name):
+        """
+        When launching multinode task with UBF, need to wait for the secondary
+        tasks to finish cleanly and produce their output before exiting the
+        main task. Otherwise, the main task finishing will cause secondary nodes
+        to terminate immediately, and possibly prematurely.
+        """
+        from metaflow import Step  # avoid circular dependency
+
+        TIMEOUT = 600
+        last_completion_timeout = time.time() + TIMEOUT
+        print("Waiting for batch secondary tasks to finish")
+        while last_completion_timeout > time.time():
+            time.sleep(2)
+            try:
+                step_path = "%s/%s/%s" % (flow.name, current.run_id, step_name)
+                tasks = [task for task in Step(step_path)]
+                if len(tasks) == len(flow._control_mapper_tasks):
+                    if all(
+                        task.finished_at is not None for task in tasks
+                    ):  # for some reason task.finished fails
+                        return True
+                else:
+                    print(
+                        "Waiting for all parallel tasks to finish. Finished: {}/{}".format(
+                            len(tasks),
+                            len(flow._control_mapper_tasks),
+                        )
+                    )
+            except Exception as e:
+                pass
+        raise Exception(
+            "Batch secondary workers did not finish in %s seconds" % TIMEOUT
+        )
+
     @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]
+
+def _setup_multinode_environment():
+    import socket
+    os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(os.environ["MASTER_ADDR"])
+    os.environ["MF_PARALLEL_NUM_NODES"] = os.environ["WORLD_SIZE"]
+    os.environ["MF_PARALLEL_NODE_INDEX"] = os.environ["RANK"]