metaflow 2.12.25__tar.gz → 2.12.27__tar.gz

{metaflow-2.12.25/metaflow.egg-info → metaflow-2.12.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.25
+Version: 2.12.27
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -13,7 +13,6 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: MacOS :: MacOS X
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.5
 Classifier: Programming Language :: Python :: 3.6
 Classifier: Programming Language :: Python :: 3.7
 Classifier: Programming Language :: Python :: 3.8
@@ -21,12 +20,13 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.12.25; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.12.27; extra == "stubs"
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)
 

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/cli.py

@@ -701,12 +701,12 @@ def resume(
     runtime.persist_constants()
 
     if runner_attribute_file:
-        with open(runner_attribute_file, "w") as f:
+        with open(runner_attribute_file, "w", encoding="utf-8") as f:
             json.dump(
                 {
                     "run_id": runtime.run_id,
                     "flow_name": obj.flow.name,
-                    "metadata": get_metadata(),
+                    "metadata": obj.metadata.metadata_str(),
                 },
                 f,
             )
@@ -779,12 +779,12 @@ def run(
     runtime.persist_constants()
 
     if runner_attribute_file:
-        with open(runner_attribute_file, "w") as f:
+        with open(runner_attribute_file, "w", encoding="utf-8") as f:
             json.dump(
                 {
                     "run_id": runtime.run_id,
                     "flow_name": obj.flow.name,
-                    "metadata": get_metadata(),
+                    "metadata": obj.metadata.metadata_str(),
                 },
                 f,
             )

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/client/core.py

@@ -16,6 +16,7 @@ from typing import (
     List,
     NamedTuple,
     Optional,
+    TYPE_CHECKING,
     Tuple,
 )
 
@@ -37,6 +38,9 @@ from metaflow.util import cached_property, is_stringish, resolve_identity, to_un
 from ..info_file import INFO_FILE
 from .filecache import FileCache
 
+if TYPE_CHECKING:
+    from metaflow.metadata import MetadataProvider
+
 try:
     # python2
     import cPickle as pickle
@@ -82,28 +86,16 @@ def metadata(ms: str) -> str:
         get_metadata()).
     """
     global current_metadata
-    infos = ms.split("@", 1)
-    types = [m.TYPE for m in METADATA_PROVIDERS]
-    if infos[0] in types:
-        current_metadata = [m for m in METADATA_PROVIDERS if m.TYPE == infos[0]][0]
-        if len(infos) > 1:
-            current_metadata.INFO = infos[1]
-    else:
-        # Deduce from ms; if starts with http, use service or else use local
-        if ms.startswith("http"):
-            metadata_type = "service"
-        else:
-            metadata_type = "local"
-        res = [m for m in METADATA_PROVIDERS if m.TYPE == metadata_type]
-        if not res:
-            print(
-                "Cannot find a '%s' metadata provider -- "
-                "try specifying one explicitly using <type>@<info>",
-                metadata_type,
-            )
-            return get_metadata()
-        current_metadata = res[0]
-        current_metadata.INFO = ms
+    provider, info = _metadata(ms)
+    if provider is None:
+        print(
+            "Cannot find a metadata provider -- "
+            "try specifying one explicitly using <type>@<info>",
+        )
+        return get_metadata()
+    current_metadata = provider
+    if info:
+        current_metadata.INFO = info
     return get_metadata()
 
 
@@ -127,7 +119,7 @@ def get_metadata() -> str:
     """
     if current_metadata is False:
         default_metadata()
-    return "%s@%s" % (current_metadata.TYPE, current_metadata.INFO)
+    return current_metadata.metadata_str()
 
 
 def default_metadata() -> str:
@@ -268,9 +260,16 @@ class MetaflowObject(object):
         _object: Optional["MetaflowObject"] = None,
         _parent: Optional["MetaflowObject"] = None,
         _namespace_check: bool = True,
+        _metaflow: Optional["Metaflow"] = None,
         _current_namespace: Optional[str] = None,
+        _current_metadata: Optional[str] = None,
     ):
-        self._metaflow = Metaflow()
+        # the default namespace is activated lazily at the first
+        # get_namespace(). The other option of activating
+        # the namespace at the import time is problematic, since there
+        # may be other modules that alter environment variables etc.
+        # which may affect the namespace setting.
+        self._metaflow = Metaflow(_current_metadata) or _metaflow
         self._parent = _parent
         self._path_components = None
         self._attempt = attempt
@@ -390,6 +389,7 @@ class MetaflowObject(object):
                     attempt=self._attempt,
                     _object=obj,
                     _parent=self,
+                    _metaflow=self._metaflow,
                     _namespace_check=self._namespace_check,
                     _current_namespace=(
                         self._current_namespace if self._namespace_check else None
@@ -505,6 +505,7 @@ class MetaflowObject(object):
                 attempt=self._attempt,
                 _object=obj,
                 _parent=self,
+                _metaflow=self._metaflow,
                 _namespace_check=self._namespace_check,
                 _current_namespace=(
                     self._current_namespace if self._namespace_check else None
@@ -552,7 +553,25 @@ class MetaflowObject(object):
             _current_namespace=ns,
         )
 
-    _UNPICKLE_FUNC = {"2.8.4": _unpickle_284, "2.12.4": _unpickle_2124}
+    def _unpickle_21227(self, data):
+        if len(data) != 5:
+            raise MetaflowInternalError(
+                "Unexpected size of array: {}".format(len(data))
+            )
+        pathspec, attempt, md, ns, namespace_check = data
+        self.__init__(
+            pathspec=pathspec,
+            attempt=attempt,
+            _namespace_check=namespace_check,
+            _current_metadata=md,
+            _current_namespace=ns,
+        )
+
+    _UNPICKLE_FUNC = {
+        "2.8.4": _unpickle_284,
+        "2.12.4": _unpickle_2124,
+        "2.12.27": _unpickle_21227,
+    }
 
     def __setstate__(self, state):
         """
@@ -595,10 +614,11 @@ class MetaflowObject(object):
         # checking for the namespace even after unpickling since we will know which
         # namespace to check.
         return {
-            "version": "2.12.4",
+            "version": "2.12.27",
             "data": [
                 self.pathspec,
                 self._attempt,
+                self._metaflow.metadata.metadata_str(),
                 self._current_namespace,
                 self._namespace_check,
             ],
@@ -670,7 +690,7 @@ class MetaflowObject(object):
         origin_pathspec = None
         if self._NAME == "run":
             latest_step = next(self.steps())
-            if latest_step:
+            if latest_step and latest_step.task:
                 # If we had a step
                 task = latest_step.task
                 origin_run_id = [
@@ -1888,9 +1908,10 @@ class Run(MetaflowObject):
         # TODO: A more optimized way of figuring out if a run has remote steps (and thus a codepackage) available.
         # This might require changes to the metadata-service as well.
         for step in self:
-            code = step.task.code
-            if code:
-                return code
+            if step.task:
+                code = step.task.code
+                if code:
+                    return code
 
     @property
     def data(self) -> Optional[MetaflowData]:
@@ -2152,7 +2173,7 @@ class Run(MetaflowObject):
         Trigger, optional
             Container of triggering events
         """
-        if "start" in self:
+        if "start" in self and self["start"].task:
             meta = self["start"].task.metadata_dict.get("execution-triggers")
             if meta:
                 return Trigger(json.loads(meta))
@@ -2287,17 +2308,16 @@ class Metaflow(object):
         if it has at least one run in the namespace.
     """
 
-    def __init__(self):
-        # the default namespace is activated lazily at the first object
-        # invocation or get_namespace(). The other option of activating
-        # the namespace at the import time is problematic, since there
-        # may be other modules that alter environment variables etc.
-        # which may affect the namescape setting.
-        if current_namespace is False:
-            default_namespace()
-        if current_metadata is False:
-            default_metadata()
-        self.metadata = current_metadata
+    def __init__(self, _current_metadata: Optional[str] = None):
+        if _current_metadata:
+            provider, info = _metadata(_current_metadata)
+            self.metadata = provider
+            if info:
+                self.metadata.INFO = info
+        else:
+            if current_metadata is False:
+                default_metadata()
+            self.metadata = current_metadata
 
     @property
     def flows(self) -> List[Flow]:
@@ -2334,7 +2354,7 @@ class Metaflow(object):
         all_flows = all_flows if all_flows else []
         for flow in all_flows:
             try:
-                v = Flow(_object=flow)
+                v = Flow(_object=flow, _metaflow=self)
                 yield v
             except MetaflowNamespaceMismatch:
                 continue
@@ -2358,7 +2378,26 @@ class Metaflow(object):
         Flow
             Flow with the given name.
         """
-        return Flow(name)
+        return Flow(name, _metaflow=self)
+
+
+def _metadata(ms: str) -> Tuple[Optional["MetadataProvider"], Optional[str]]:
+    infos = ms.split("@", 1)
+    types = [m.TYPE for m in METADATA_PROVIDERS]
+    if infos[0] in types:
+        provider = [m for m in METADATA_PROVIDERS if m.TYPE == infos[0]][0]
+        if len(infos) > 1:
+            return provider, infos[1]
+        return provider, None
+    # Deduce from ms; if starts with http, use service or else use local
+    if ms.startswith("http"):
+        metadata_type = "service"
+    else:
+        metadata_type = "local"
+    res = [m for m in METADATA_PROVIDERS if m.TYPE == metadata_type]
+    if not res:
+        return None, None
+    return res[0], ms
 
 
 _CLASSES["flow"] = Flow

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/metadata/metadata.py

@@ -76,6 +76,12 @@ class ObjectOrder:
 
 @with_metaclass(MetadataProviderMeta)
 class MetadataProvider(object):
+    TYPE = None
+
+    @classmethod
+    def metadata_str(cls):
+        return "%s@%s" % (cls.TYPE, cls.INFO)
+
     @classmethod
     def compute_info(cls, val):
         """

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/argo/argo_client.py

@@ -10,6 +10,14 @@ class ArgoClientException(MetaflowException):
     headline = "Argo Client error"
 
 
+class ArgoResourceNotFound(MetaflowException):
+    headline = "Resource not found"
+
+
+class ArgoNotPermitted(MetaflowException):
+    headline = "Operation not permitted"
+
+
 class ArgoClient(object):
     def __init__(self, namespace=None):
         self._client = KubernetesClient()
@@ -140,9 +148,7 @@ class ArgoClient(object):
             if e.status == 404:
                 return None
             else:
-                raise ArgoClientException(
-                    json.loads(e.body)["message"] if e.body is not None else e.reason
-                )
+                raise wrap_api_error(e)
 
     def delete_workflow_template(self, name):
         """
@@ -164,9 +170,7 @@ class ArgoClient(object):
             if e.status == 404:
                 return None
             else:
-                raise ArgoClientException(
-                    json.loads(e.body)["message"] if e.body is not None else e.reason
-                )
+                raise wrap_api_error(e)
 
     def terminate_workflow(self, name):
         client = self._client.get()
@@ -428,6 +432,18 @@ class ArgoClient(object):
         except client.rest.ApiException as e:
             if e.status == 404:
                 return None
-            raise ArgoClientException(
-                json.loads(e.body)["message"] if e.body is not None else e.reason
-            )
+            raise wrap_api_error(e)
+
+
+def wrap_api_error(error):
+    message = (
+        json.loads(error.body)["message"] if error.body is not None else error.reason
+    )
+    # catch all
+    ex = ArgoClientException(message)
+    if error.status == 404:
+        # usually handled outside this function as most cases want to return None instead.
+        ex = ArgoResourceNotFound(message)
+    if error.status == 403:
+        ex = ArgoNotPermitted(message)
+    return ex

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/argo/argo_workflows.py

@@ -2288,7 +2288,9 @@ class ArgoWorkflows(object):
             and k not in set(ARGO_WORKFLOWS_ENV_VARS_TO_SKIP.split(","))
         }
         return [
-            Template("error-msg-capture-hook").container(
+            Template("error-msg-capture-hook")
+            .service_account_name(resources["service_account"])
+            .container(
                 to_camelcase(
                     kubernetes_sdk.V1Container(
                         name="main",
@@ -2431,7 +2433,7 @@ class ArgoWorkflows(object):
         Use Slack's Block Kit to add general information about the environment and
         execution metadata, including a link to the UI and an optional message.
         """
-        ui_link = "%s%s/argo-{{workflow.name}}" % (UI_URL, self.flow.name)
+        ui_link = "%s/%s/argo-{{workflow.name}}" % (UI_URL.rstrip("/"), self.flow.name)
         # fmt: off
         if getattr(current, "project_name", None):
             # Add @project metadata when available.

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/argo/argo_workflows_cli.py

@@ -226,12 +226,12 @@ def create(
     validate_tags(tags)
 
     if deployer_attribute_file:
-        with open(deployer_attribute_file, "w") as f:
+        with open(deployer_attribute_file, "w", encoding="utf-8") as f:
             json.dump(
                 {
                     "name": obj.workflow_name,
                     "flow_name": obj.flow.name,
-                    "metadata": get_metadata(),
+                    "metadata": obj.metadata.metadata_str(),
                 },
                 f,
             )
@@ -657,7 +657,7 @@ def trigger(obj, run_id_file=None, deployer_attribute_file=None, **kwargs):
             json.dump(
                 {
                     "name": obj.workflow_name,
-                    "metadata": get_metadata(),
+                    "metadata": obj.metadata.metadata_str(),
                     "pathspec": "/".join((obj.flow.name, run_id)),
                 },
                 f,
@@ -1016,13 +1016,7 @@ def validate_run_id(
 
     if project_name is not None:
         # Verify we are operating on the correct project.
-        # Perform match with separators to avoid substrings matching
-        # e.g. 'test_proj' and 'test_project' should count as a mismatch.
-        project_part = "%s." % sanitize_for_argo(project_name)
-        if (
-            current.get("project_name") != project_name
-            and project_part not in workflow_name
-        ):
+        if current.get("project_name") != project_name:
             raise RunIdMismatch(
                 "The workflow belongs to the project *%s*. "
                 "Please use the project decorator or --name to target the correct project"
@@ -1030,13 +1024,7 @@ def validate_run_id(
             )
 
         # Verify we are operating on the correct branch.
-        # Perform match with separators to avoid substrings matching.
-        # e.g. 'user.tes' and 'user.test' should count as a mismatch.
-        branch_part = ".%s." % sanitize_for_argo(branch_name)
-        if (
-            current.get("branch_name") != branch_name
-            and branch_part not in workflow_name
-        ):
+        if current.get("branch_name") != branch_name:
             raise RunIdMismatch(
                 "The workflow belongs to the branch *%s*. "
                 "Please use --branch, --production or --name to target the correct branch"

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/aws/batch/batch_decorator.py

@@ -134,6 +134,10 @@ class BatchDecorator(StepDecorator):
     package_sha = None
     run_time_limit = None
 
+    # Conda environment support
+    supports_conda_environment = True
+    target_platform = "linux-64"
+
     def __init__(self, attributes=None, statically_defined=False):
         super(BatchDecorator, self).__init__(attributes, statically_defined)
 

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -162,7 +162,7 @@ def create(
                 {
                     "name": obj.state_machine_name,
                     "flow_name": obj.flow.name,
-                    "metadata": get_metadata(),
+                    "metadata": obj.metadata.metadata_str(),
                 },
                 f,
             )
@@ -502,7 +502,7 @@ def trigger(obj, run_id_file=None, deployer_attribute_file=None, **kwargs):
             json.dump(
                 {
                     "name": obj.state_machine_name,
-                    "metadata": get_metadata(),
+                    "metadata": obj.metadata.metadata_str(),
                     "pathspec": "/".join((obj.flow.name, run_id)),
                 },
                 f,

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/azure/azure_tail.py

@@ -70,7 +70,7 @@ class AzureTail(object):
         if data is None:
             return None
         if data:
-            buf = BytesIO(data)
+            buf = BytesIO(self._tail + data)
             self._pos += len(data)
             self._tail = b""
             return buf

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/cards/card_decorator.py

@@ -117,6 +117,7 @@ class CardDecorator(StepDecorator):
         # `'%s-%s'%(evt_name,step_name)` ensures that we capture this once per @card per @step.
         # Since there can be many steps checking if event is registered for `evt_name` will only make it check it once for all steps.
         # Hence, we have `_is_event_registered('%s-%s'%(evt_name,step_name))`
+        self._is_runtime_card = False
         evt = "%s-%s" % (evt_name, step_name)
         if not self._is_event_registered(evt):
             # We set the total count of decorators so that we can use it for

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/gcp/gs_tail.py

@@ -11,15 +11,15 @@ from metaflow.plugins.gcp.gs_utils import parse_gs_full_path
 class GSTail(object):
     def __init__(self, blob_full_uri):
         """Location should be something like gs://<bucket_name>/blob"""
-        bucket_name, blob_name = parse_gs_full_path(blob_full_uri)
-        if not blob_name:
+        self.bucket_name, self.blob_name = parse_gs_full_path(blob_full_uri)
+        if not self.blob_name:
             raise MetaflowException(
                 msg="Failed to parse blob_full_uri into gs://<bucket_name>/<blob_name> (got %s)"
                 % blob_full_uri
             )
         client = get_gs_storage_client()
-        bucket = client.bucket(bucket_name)
-        self._blob_client = bucket.blob(blob_name)
+        self.bucket = client.bucket(self.bucket_name)
+        self._blob_client = self.bucket.blob(self.blob_name)
         self._pos = 0
         self._tail = b""
 
@@ -46,7 +46,11 @@ class GSTail(object):
     def _make_range_request(self):
         try:
             # Yes we read to the end... memory blow up is possible. We can improve by specifying length param
-            return self._blob_client.download_as_bytes(start=self._pos)
+            # NOTE: We must re-instantiate the whole client here due to a behavior with the GS library,
+            # otherwise download_as_bytes will simply return the same content for consecutive requests with the same attributes,
+            # even if the blob has grown in size.
+            blob_client = self.bucket.blob(self.blob_name)
+            return blob_client.download_as_bytes(start=self._pos)
         except NotFound:
             return None
         except ClientError as e:
@@ -63,7 +67,7 @@ class GSTail(object):
         if data is None:
             return None
         if data:
-            buf = BytesIO(data)
+            buf = BytesIO(self._tail + data)
             self._pos += len(data)
             self._tail = b""
             return buf

{metaflow-2.12.25 → metaflow-2.12.27}/metaflow/plugins/kubernetes/kubernetes.py

@@ -710,12 +710,23 @@ class Kubernetes(object):
         wait_for_launch(self._job)
 
         # 2) Tail logs until the job has finished
+        self._output_final_logs = False
+
+        def _has_updates():
+            if self._job.is_running:
+                return True
+            # Make sure to output final tail for a job that has finished.
+            if not self._output_final_logs:
+                self._output_final_logs = True
+                return True
+            return False
+
         tail_logs(
             prefix=prefix(),
             stdout_tail=stdout_tail,
             stderr_tail=stderr_tail,
             echo=echo,
-            has_log_updates=lambda: self._job.is_running,
+            has_log_updates=_has_updates,
         )
         # 3) Fetch remaining logs
         #

{metaflow-2.12.25 → metaflow-2.12.27}/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
@@ -104,6 +105,9 @@ class KubernetesDecorator(StepDecorator):
     compute_pool : str, optional, default None
         Compute pool to be used for for this step.
         If not specified, any accessible compute pool within the perimeter is used.
+    hostname_resolution_timeout: int, default 10 * 60
+        Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
+        Only applicable when @parallel is used.
     """
 
     name = "kubernetes"
@@ -130,11 +134,16 @@ class KubernetesDecorator(StepDecorator):
         "port": None,
         "compute_pool": None,
         "executable": None,
+        "hostname_resolution_timeout": 10 * 60,
     }
     package_url = None
     package_sha = None
     run_time_limit = None
 
+    # Conda environment support
+    supports_conda_environment = True
+    target_platform = "linux-64"
+
     def __init__(self, attributes=None, statically_defined=False):
         super(KubernetesDecorator, self).__init__(attributes, statically_defined)
 
@@ -386,7 +395,7 @@ class KubernetesDecorator(StepDecorator):
             cli_args.command_args.append(self.package_url)
 
             # skip certain keys as CLI arguments
-            _skip_keys = ["compute_pool"]
+            _skip_keys = ["compute_pool", "hostname_resolution_timeout"]
             # --namespace is used to specify Metaflow namespace (a different
             # concept from k8s namespace).
             for k, v in self.attributes.items():
@@ -478,7 +487,9 @@ class KubernetesDecorator(StepDecorator):
             num_parallel = flow._parallel_ubf_iter.num_parallel
 
         if num_parallel and num_parallel > 1:
-            _setup_multinode_environment()
+            _setup_multinode_environment(
+                ubf_context, self.attributes["hostname_resolution_timeout"]
+            )
             # current.parallel.node_index will be correctly available over here.
             meta.update({"parallel-node-index": current.parallel.node_index})
             if ubf_context == UBF_CONTROL:
@@ -542,18 +553,44 @@ class KubernetesDecorator(StepDecorator):
 
 
 # TODO: Unify this method with the multi-node setup in @batch
-def _setup_multinode_environment():
-    # TODO [FIXME SOON]
-    # Even if Kubernetes may deploy control pods before worker pods, there is always a
-    # possibility that the worker pods may start before the control. In the case that this happens,
-    # the worker pods will not be able to resolve the control pod's IP address and this will cause
-    # the worker pods to fail. This function should account for this in the near future.
+def _setup_multinode_environment(ubf_context, hostname_resolution_timeout):
     import socket
 
+    def _wait_for_hostname_resolution(max_wait_timeout=10 * 60):
+        """
+        keep trying to resolve the hostname of the control task until the hostname is resolved
+        or the max_wait_timeout is reached. This is a workaround for the issue where the control
+        task is not scheduled before the worker task and the worker task fails because it cannot
+        resolve the hostname of the control task.
+        """
+        start_time = time.time()
+        while True:
+            try:
+                return socket.gethostbyname(os.environ["MF_MASTER_ADDR"])
+            except socket.gaierror:
+                if time.time() - start_time > max_wait_timeout:
+                    raise MetaflowException(
+                        "Failed to get host by name for MF_MASTER_ADDR after waiting for {} seconds.".format(
+                            max_wait_timeout
+                        )
+                    )
+                time.sleep(1)
+
     try:
-        os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(
-            os.environ["MF_MASTER_ADDR"]
-        )
+        # Even if Kubernetes may deploy control pods before worker pods, there is always a
+        # possibility that the worker pods may start before the control. In the case that this happens,
+        # the worker pods will not be able to resolve the control pod's IP address and this will cause
+        # the worker pods to fail. So if the worker pods are requesting a hostname resolution, we will
+        # make it wait for the name to be resolved within a reasonable timeout period.
+        if ubf_context != UBF_CONTROL:
+            os.environ["MF_PARALLEL_MAIN_IP"] = _wait_for_hostname_resolution(
+                hostname_resolution_timeout
+            )
+        else:
+            os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(
+                os.environ["MF_MASTER_ADDR"]
+            )
+
         os.environ["MF_PARALLEL_NUM_NODES"] = os.environ["MF_WORLD_SIZE"]
         os.environ["MF_PARALLEL_NODE_INDEX"] = (
             str(0)