Merge pull request #2362 from dhruvesh09/r0.23.0

Wait indefinitely for pod to complete in Kubernetes launcher.

Author: dhruvesh09

tfx/orchestration/launcher/kubernetes_component_launcher.py

@@ -138,13 +138,15 @@ def _run_executor(self, execution_id: int,
             'Failed to created container executor pod!\nReason: %s\nBody: %s' %
             (e.reason, e.body))
 
+    # Wait up to 300 seconds for the pod to move from pending to another status.
     logging.info('Waiting for pod "%s:%s" to start.', namespace, pod_name)
     self._wait_pod(
         core_api,
         pod_name,
         namespace,
         exit_condition_lambda=_pod_is_not_pending,
-        condition_description='non-pending status')
+        condition_description='non-pending status',
+        timeout_sec=300)
 
     logging.info('Start log streaming for pod "%s:%s".', namespace, pod_name)
     try:
@@ -162,12 +164,14 @@ def _run_executor(self, execution_id: int,
     for log in logs:
       logging.info(log.decode().rstrip('\n'))
 
+    # Wait indefinitely for the pod to complete.
     resp = self._wait_pod(
         core_api,
         pod_name,
         namespace,
         exit_condition_lambda=_pod_is_done,
-        condition_description='done state')
+        condition_description='done state',
+        timeout_sec=0)
 
     if resp.status.phase == kube_utils.PodPhase.FAILED.value:
       raise RuntimeError('Pod "%s:%s" failed with status "%s".' %
@@ -254,7 +258,7 @@ def _wait_pod(self,
                 namespace: Text,
                 exit_condition_lambda: Callable[[client.V1Pod], bool],
                 condition_description: Text,
-                timeout_sec: int = 300) -> client.V1Pod:
+                timeout_sec: int) -> client.V1Pod:
     """Wait for a POD to meet an exit condition.
 
     Args:
@@ -265,7 +269,8 @@ def _wait_pod(self,
         for a POD to exit. The function returns True to exit.
       condition_description: The description of the exit condition which will be
         set in the error message if the wait times out.
-      timeout_sec: The seconds for the function to wait. Defaults to 300s.
+      timeout_sec: The seconds for the function to wait. Waits indefinitely if
+        value is 0.
 
     Returns:
       The POD object which meets the exit condition.
@@ -279,8 +284,8 @@ def _wait_pod(self,
       logging.info(resp.status.phase)
       if exit_condition_lambda(resp):
         return resp
-      elapse_time = datetime.datetime.utcnow() - start_time
-      if elapse_time.seconds >= timeout_sec:
+      elapsed_time = datetime.datetime.utcnow() - start_time
+      if timeout_sec != 0 and elapsed_time.seconds >= timeout_sec:
         raise RuntimeError(
             'Pod "%s:%s" does not reach "%s" within %s seconds.' %
             (namespace, pod_name, condition_description, timeout_sec))