no-op

PiperOrigin-RevId: 548556494

Author: anuartb

tfx/components/infra_validator/model_server_runners/local_docker_runner.py

@@ -14,8 +14,9 @@
 """Module for LocalDockerModelServerRunner."""
 
 import os
+import subprocess
 import time
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 from absl import logging
 import docker
@@ -39,7 +40,7 @@ def _make_docker_client(config: infra_validator_pb2.LocalDockerConfig):
   return docker.DockerClient(**params)
 
 
-def _find_host_port(ports: Dict[str, Any], container_port: int) -> str:
+def _find_host(ports: Dict[str, Any], container_port: int) -> Tuple[str, str]:
   """Find host port from container port mappings.
 
   `ports` is a nested dictionary of the following structure:
@@ -65,10 +66,15 @@ def _find_host_port(ports: Dict[str, Any], container_port: int) -> str:
   Raises:
     ValueError: No corresponding host port was found.
   """
+  command = '/sbin/ip route|awk \'/default/ { print $3 }\''
+  ip_address = subprocess.check_output(
+      command, shell=True).decode('utf-8').strip()
+  logging.info('Container IP address: %s', ip_address)
+
   mappings = ports.get('{}/tcp'.format(container_port), [])
   for mapping in mappings:
     if mapping['HostIp'] == '0.0.0.0':
-      return mapping['HostPort']
+      return ip_address, mapping['HostPort']
   else:
     raise ValueError(
         'No HostPort found for ContainerPort={} (all port mappings: {})'
@@ -146,9 +152,11 @@ def WaitUntilRunning(self, deadline):
         continue
       # The container is running :)
       if status == 'running':
-        host_port = _find_host_port(self._container.ports,
-                                    self._serving_binary.container_port)
-        self._endpoint = 'localhost:{}'.format(host_port)
+        ip_address, host_port = _find_host(
+            self._container.ports,
+            self._serving_binary.container_port)
+
+        self._endpoint = '{}:{}'.format(ip_address, host_port)
         return
       # Docker status is one of {'created', 'restarting', 'running', 'removing',
       # 'paused', 'exited', or 'dead'}. Status other than 'created' and

tfx/examples/chicago_taxi_pipeline/taxi_pipeline_simple_airflow_e2e_test.py

@@ -100,6 +100,7 @@ def _CheckPendingTasks(self, pending_task_names: Sequence[str]) -> Set[str]:
 
   def setUp(self):
     super().setUp()
+    #
     # setup airflow_home in a temp directory, config and init db.
     self._airflow_home = self.tmp_dir
     self.enter_context(

tfx/examples/penguin/penguin_pipeline_local_infraval_e2e_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """E2E Tests for tfx.examples.penguin.penguin_pipeline_local_infraval."""
 
+#
 import os
 
 from absl.testing import parameterized