Try bootstrapping tpu env from env vars (#4499)

* Try bootstrapping tpu env from env vars

* Try bootstrapping tpu env from env vars

* Formatting

* Mock env vars

* Use TPU_SKIP_MDS_QUERY

---------

Co-authored-by: root <root@t1v-n-804806aa-w-0.us-central2-b.c.tpu-pytorch.internal>

Author: steventk-g

test/pjrt/test_experimental_tpu.py

@@ -46,14 +46,15 @@ def test_task_id(self, task_id, expected):
 
     self.assertEqual(i, expected)
 
-  def test_tpu_env(self):
+  def test_tpu_env_from_gce_metadata(self):
     tpu_env_yaml = textwrap.dedent("""
       ACCELERATOR_TYPE: 'v4-16'
       CHIPS_PER_HOST_BOUNDS: '2,2,1'
       HOST_BOUNDS: '1,1,2'
       TPU_CHIPS_PER_PROCESS_BOUNDS: '2,2,1'
       TPU_PROCESS_BOUNDS: '1,1,2'
       ZONE: 'us-central2-b'
+      WORKER_ID: '0'
     """)
 
     with mock.patch.object(tpu, '_get_metadata', return_value=tpu_env_yaml):
@@ -67,8 +68,43 @@ def test_tpu_env(self):
             'TPU_CHIPS_PER_PROCESS_BOUNDS': '2,2,1',
             'TPU_PROCESS_BOUNDS': '1,1,2',
             'ZONE': 'us-central2-b',
+            'WORKER_ID': '0'
         })
 
+  @parameterized.named_parameters(
+      ('all-vars-set', {
+          xenv.TPU_SKIP_MDS_QUERY: '1',
+          xenv.TPU_ACCELERATOR_TYPE: 'v4-16',
+          xenv.TPU_PROCESS_BOUNDS: '1,2,2',
+          xenv.TPU_HOST_BOUNDS: '1,1,2',
+          xenv.TPU_CHIPS_PER_PROCESS_BOUNDS: '2,2,1',
+          xenv.TPU_CHIPS_PER_HOST_BOUNDS: '2,1,1',
+          xenv.CLOUD_TPU_TASK_ID: '1',
+          xenv.TPU_WORKER_ID: '0'
+      }, {
+          xenv.ACCELERATOR_TYPE: 'v4-16',
+          xenv.TPU_CHIPS_PER_PROCESS_BOUNDS: '2,2,1',
+          xenv.TPU_PROCESS_BOUNDS: '1,2,2',
+          xenv.WORKER_ID: '1'
+      }),
+      ('defaults-only', {
+          xenv.TPU_SKIP_MDS_QUERY: '1',
+          xenv.TPU_ACCELERATOR_TYPE: 'v4-16',
+          xenv.TPU_HOST_BOUNDS: '1,1,2',
+          xenv.TPU_CHIPS_PER_HOST_BOUNDS: '2,1,1',
+          xenv.TPU_WORKER_ID: '0'
+      }, {
+          xenv.ACCELERATOR_TYPE: 'v4-16',
+          xenv.TPU_CHIPS_PER_PROCESS_BOUNDS: '2,1,1',
+          xenv.TPU_PROCESS_BOUNDS: '1,1,2',
+          xenv.WORKER_ID: '0'
+      }),
+  )
+  def test_tpu_env_from_env_vars(self, envs, expected):
+    with mock.patch.dict(os.environ, envs, clear=True):
+      tpu_env = tpu.get_tpu_env()
+    self.assertDictEqual(tpu_env, expected)
+
   @parameterized.named_parameters(
       ('one_host', 't1v-n-ea9d3291-w-0:12345:10.130.0.31', ['localhost']),
       (

torch_xla/core/xla_env_vars.py

@@ -16,6 +16,12 @@
 GPU_NUM_DEVICES = 'GPU_NUM_DEVICES'
 CPU_NUM_DEVICES = 'CPU_NUM_DEVICES'
 CLOUD_TPU_TASK_ID = 'CLOUD_TPU_TASK_ID'
+ACCELERATOR_TYPE = 'ACCELERATOR_TYPE'
+WORKER_ID = 'WORKER_ID'
+TPU_SKIP_MDS_QUERY = 'TPU_SKIP_MDS_QUERY'
+TPU_ACCELERATOR_TYPE = 'TPU_ACCELERATOR_TYPE'
+TPU_WORKER_ID = 'TPU_WORKER_ID'
+TPU_WORKER_HOSTNAMES = 'TPU_WORKER_HOSTNAMES'
 TPU_HOST_BOUNDS = 'TPU_HOST_BOUNDS'
 TPU_CHIPS_PER_HOST_BOUNDS = 'TPU_CHIPS_PER_HOST_BOUNDS'
 TPU_MESH_CTLER_ADDR = 'TPU_MESH_CONTROLLER_ADDRESS'

torch_xla/experimental/tpu.py

@@ -3,6 +3,7 @@
 import os
 import re
 from typing import Dict, NamedTuple, Optional, List, Tuple
+from typing_extensions import TypedDict
 import requests
 import yaml
 
@@ -32,6 +33,13 @@
 }
 
 
+class TpuEnv(TypedDict):
+  accelerator_type: str
+  tpu_process_bounds: str
+  tpu_chips_per_process_bound: str
+  worker_id: int
+
+
 class MeshShape(NamedTuple):
   """Represents a TPU mesh shape (e.g. '2,2,1' or '1,1,1')"""
   x: int
@@ -65,7 +73,6 @@ def _get_metadata(key: str) -> str:
 def process_bounds_size(default: int = 1) -> int:
   """Returns number of processes across all TPU hosts."""
   process_bounds = xu.getenv_as(xenv.TPU_PROCESS_BOUNDS, str)
-
   return MeshShape.from_string(
       process_bounds).size if process_bounds else default
 
@@ -81,10 +88,28 @@ def task_id() -> Optional[int]:
   return xu.getenv_as(xenv.CLOUD_TPU_TASK_ID, int)
 
 
-def get_tpu_env() -> Dict[str, str]:
+def _using_env_vars() -> bool:
+  return xu.getenv_as(xenv.TPU_SKIP_MDS_QUERY, str, False)
+
+
+def build_tpu_env_from_vars() -> TpuEnv:
+  metadata = dict()
+  metadata[xenv.ACCELERATOR_TYPE] = xu.getenv_as(xenv.TPU_ACCELERATOR_TYPE, str)
+  metadata[xenv.TPU_PROCESS_BOUNDS] = xu.getenv_as(
+      xenv.TPU_PROCESS_BOUNDS, str, xu.getenv_as(xenv.TPU_HOST_BOUNDS, str))
+  metadata[xenv.TPU_CHIPS_PER_PROCESS_BOUNDS] = xu.getenv_as(
+      xenv.TPU_CHIPS_PER_PROCESS_BOUNDS, str,
+      xu.getenv_as(xenv.TPU_CHIPS_PER_HOST_BOUNDS, str))
+  metadata[xenv.WORKER_ID] = xu.getenv_as(xenv.CLOUD_TPU_TASK_ID, str,
+                                          xu.getenv_as(xenv.TPU_WORKER_ID, str))
+  return metadata
+
+
+def get_tpu_env() -> TpuEnv:
   """Fetches and parses `tpu-env` metadata field."""
+  if _using_env_vars():
+    return build_tpu_env_from_vars()
   metadata = _get_metadata('tpu-env')
-
   return yaml.load(metadata, yaml.Loader)
 
 
@@ -94,19 +119,22 @@ def version() -> int:
   except requests.HTTPError as e:
     raise EnvironmentError('Failed to get TPU metadata') from e
 
-  match = re.match(r'^v(\d)-(\d+)$', env['ACCELERATOR_TYPE'])
+  match = re.match(r'^v(\d)-(\d+)$', env[xenv.ACCELERATOR_TYPE])
   return int(match.groups()[0])
 
 
 def get_worker_ips() -> List[str]:
   """Returns ordered list of TPU worker IPs from TPU metadata."""
-  metadata = _get_metadata('worker-network-endpoints')
-
-  # Workers have format 'hostname:uid:ip,hostname:uid:ip,...'
-  workers = metadata.split(',')
-  ips = [worker.split(':')[2] for worker in workers]
-
-  return ips if len(ips) > 1 else ['localhost']
+  if _using_env_vars():
+    hostnames_string = xu.getenv_as(xenv.TPU_WORKER_HOSTNAMES, str, '')
+    # String has the format 'host-name-1,host-name-2,...,host-name-n'
+    hostnames = hostnames_string.split(',')
+  else:
+    hostnames_string = _get_metadata('worker-network-endpoints')
+    # Workers have format 'hostname:uid:ip,hostname:uid:ip,...'
+    workers = hostnames_string.split(',')
+    hostnames = [worker.split(':')[2] for worker in workers]
+  return hostnames if len(hostnames) > 1 else ['localhost']
 
 
 def configure_one_chip_topology() -> None:
@@ -135,8 +163,8 @@ def configure_topology(local_rank: int,
   """
   tpu_env = get_tpu_env()
 
-  accelerator_type = tpu_env['ACCELERATOR_TYPE']
-  if tpu_env['ACCELERATOR_TYPE'].startswith('v4'):
+  accelerator_type = tpu_env[xenv.ACCELERATOR_TYPE]
+  if version() == 4:
     # Process bounds with 4 chips per process
     default_process_bounds = MeshShape.from_string(
         tpu_env[xenv.TPU_PROCESS_BOUNDS])
@@ -156,7 +184,7 @@ def configure_topology(local_rank: int,
                         ','.join(str(dim) for dim in process_bounds))
 
   # Assume each TPU has the same number of local processes with the same ports
-  worker_id = int(tpu_env['WORKER_ID'])
+  worker_id = int(tpu_env[xenv.WORKER_ID])
   os.environ.setdefault(xenv.CLOUD_TPU_TASK_ID,
                         str(worker_id * local_world_size + local_rank))
 
@@ -186,7 +214,7 @@ def discover_master_worker_ip(use_localhost: bool = True) -> str:
     return 'localhost'
 
   tpu_env = get_tpu_env()
-  current_worker_id = int(tpu_env['WORKER_ID'])
+  current_worker_id = int(tpu_env[xenv.WORKER_ID])
   t = torch.tensor([current_worker_id], device=xm.xla_device())
   xm.collective_broadcast([t])
   xm.mark_step()