def dtensor_initialize_multi_client(
enable_coordination_service: Optional[bool] = False) -> None:
"""Initializes Multi Client DTensor.
The following environment variables controls the behavior of this function.
If the variables are unset, DTensor will be configured to run in single-client
mode.
- DTENSOR_CLIENT_ID: integer, between 0 to num_clients - 1, to identify the
client id of the current process.
- DTENSOR_NUM_CLIENTS: integer, the number of clients.
- DTENSOR_JOB_NAME: string, a hostname like string for the name of the dtensor
job. The job name is used by TensorFlow in the job name section of
the DeviceSpec.
- DTENSOR_JOBS: string, a comma separated list. Each item in the list is
of format `{hostname}:{port}` and the items must be sorted in alphabet
order. The implication is the RPC port numbers of the clients from
the same host must be ordered by the client ID.
Examples of valid DTENSOR_JOBS values:
- 4 clients on localhost:
`localhost:10000,localhost:10001,localhost:10002,localhost:10003`
- 2 clients on host1, 2 clients on host2
`host1:10000,host1:10001,host2:10000,host2:10003`
Args:
enable_coordination_service: If true, enable distributed coordination
service to make sure that workers know the devices on each other, a
prerequisite for data transfer through cross-worker rendezvous.
"""
global _in_multi_client_mode
assert context.executing_eagerly()
_in_multi_client_mode = api.job_name() != 'localhost'
if not _in_multi_client_mode and api.num_clients() != 1:
raise ValueError(
'DTENSOR_NUM_CLIENTS is set and not 1, while DTENSOR_JOB_NAME is '
'set to localhost for single client mode.')
# Collective GRPC servers are only necessary in multi-client setup.
# Single clients can use local mode of collectives.
if _in_multi_client_mode:
if api.jobs() is None:
raise ValueError(
'DTENSOR_JOBS environment variable is required when'
'using multi-client to properly set up communications between servers'
)
multi_client_util.initialize_multi_client_cluster(
job_name=api.job_name(),
dtensor_jobs=api.jobs(),
client_id=api.client_id(),
collective_leader=api.full_job_name(task_id=0),
enable_coordination_service=enable_coordination_service,
protocol='grpc')
# Make sure the server change is fully propagated before returning.
context.ensure_initialized()
context.async_wait()
context.context()._clear_caches() # pylint: disable=protected-access
def dtensor_initialize_multi_client(
enable_coordination_service: Optional[bool] = False) -> None:
"""Initializes Multi Client DTensor.
The following environment variables controls the behavior of this function.
If the variables are unset, DTensor will be configured to run in single-client
mode.
- DTENSOR_CLIENT_ID: integer, between 0 to num_clients - 1, to identify the
client id of the current process. The default value is 0.
- DTENSOR_NUM_CLIENTS: integer, the number of clients. The default value is 1.
- DTENSOR_JOB_NAME: string, a hostname like string for the name of the dtensor
job. The default is `localhost` when number of clients is 1, and `worker`
when the number of clients is greater than 1.
The job name controls the job name section of the TensorFlow DeviceSpecs,
e.g., `job:worker` in `/job:worker/replica:0/task:0/device:TPU:0` when
the job name is `worker`.
- DTENSOR_JOBS: string, a comma separated list. Each item in the list is
of format `{hostname}:{port}` and the items must be sorted in alphabet
order. The implication is the RPC port numbers of the clients from
the same host must be ordered by the client ID.
Examples of valid DTENSOR_JOBS values:
- 4 clients on localhost:
`localhost:10000,localhost:10001,localhost:10002,localhost:10003`
- 2 clients on host1, 2 clients on host2
`host1:10000,host1:10001,host2:10000,host2:10003`
Args:
enable_coordination_service: If true, enable distributed coordination
service to make sure that workers know the devices on each other, a
prerequisite for data transfer through cross-worker rendezvous.
"""
assert context.executing_eagerly()
# Collective GRPC servers are only necessary in multi-client setup.
# Single clients can use local mode of collectives.
if api.num_clients() > 1:
multi_client_util.initialize_multi_client_cluster(
job_name=api.job_name(),
dtensor_jobs=api.jobs(),
client_id=api.client_id(),
collective_leader=api.full_job_name(task_id=0),
enable_coordination_service=enable_coordination_service)
# Make sure the server change is fully propagated before returning.
context.ensure_initialized()
context.async_wait()
context.context()._clear_caches() # pylint: disable=protected-access
def create_distributed_mesh(mesh_dims: List[Tuple[str, int]],
mesh_name: str = '',
num_global_devices: Optional[int] = None,
num_clients: Optional[int] = None,
client_id: Optional[int] = None,
device_type: str = 'CPU') -> layout.Mesh:
"""Creates a single- or multi-client mesh.
For CPU and GPU meshes, users can choose to use fewer local devices than what
is available. If any argument is missing, it will be extracted from
environment variables. The default values for these environment variables
create a mesh using all devices (common for unit tests).
For TPU meshes, users should not specify any of the nullable arguments. The
DTensor runtime will set these arguments automatically, using all TPU cores
available in the entire cluster.
Args:
mesh_dims: A list of (dim_name, dim_size) tuples.
mesh_name: Name of the created mesh. Defaults to ''.
num_global_devices: Number of devices in the DTensor cluster. Defaults to
the corresponding environment variable.
num_clients: Number of clients in the DTensor cluster. Defaults to the
corresponding environment variable, DTENSOR_NUM_CLIENTS.
client_id: This client's ID. Defaults to the corresponding environment
variable, DTENSOR_CLIENT_ID.
device_type: Type of device to build the mesh for. Defaults to 'CPU'.
Returns:
A mesh created from specified or default arguments.
"""
dim_names, shape = zip(*mesh_dims)
if device_type.upper() in ['CPU', 'GPU']:
# For CPU and GPU meshes, user-specified args take precedence over env vars.
# This is particularly useful on single clients when users want to create
# meshes that use fewer logical devices than what's available.
if num_global_devices is None:
num_global_devices = api.num_global_devices(device_type)
if num_global_devices <= 0:
raise ValueError(
f'num_global_devices ({num_global_devices}) must be > 0')
if num_global_devices != np.prod(shape):
raise ValueError(
f'num_global_devices ({num_global_devices}) must be '
f'equal to total size of the mesh of shape {shape}')
if num_clients is None:
num_clients = api.num_clients()
if num_clients <= 0:
raise ValueError(f'num_clients ({num_clients}) must be > 0')
if _in_multi_client_mode is None and num_clients > 1:
raise ValueError(
'Invalid multi-client topology, run dtensor.initialize_multi_client() first'
)
if client_id is None:
client_id = api.client_id()
if client_id < 0:
raise ValueError(f'client_id ({client_id}) must be >= 0')
if client_id >= num_clients:
raise ValueError(
f'client_id ({client_id}) must be < {num_clients}')
if num_global_devices % num_clients != 0:
raise ValueError(
f'num_global_devices ({num_global_devices}) must be '
f'divisible by num_clients ({num_clients})')
num_local_devices = num_global_devices // num_clients
# It's allowed to create a CPU or GPU mesh using fewer logical devices than
# what's available. If so, just use the first N logical devices.
num_available_devices = api.num_local_devices(device_type)
if num_local_devices > num_available_devices:
raise ValueError(
f'Not enough devices; {num_local_devices} needed, '
f'only {num_available_devices} available')
local_devices = api.local_devices(device_type,
client_id)[:num_local_devices]
global_device_ids = np.arange(num_global_devices).reshape(shape)
flattened = np.ravel(global_device_ids).tolist()
start_idx = num_local_devices * client_id
local_device_ids = flattened[start_idx:start_idx + num_local_devices]
mesh = layout.Mesh(dim_names=dim_names,
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=local_devices,
mesh_name=mesh_name)
_print_context(num_global_devices, num_clients, client_id, device_type,
mesh)
return mesh
if device_type.upper() == 'TPU':
# TPU meshes can only be configured through environment variables that
# reflect the actual TPU topology. Do not let users specify custom args.
if num_global_devices is not None:
raise ValueError(
f'Do not specify num_global_devices for {device_type.upper()} meshes. '
'It will be filled in automatically from environmental variables.'
'See api.py for the list of environmental variables for DTensor.'
)
if num_clients is not None:
raise ValueError(
f'Do not specify num_clients for {device_type.upper()} meshes. '
'It will be filled in automatically from environmental variables.'
'See api.py for the list of environmental variables for DTensor.'
)
if client_id is not None:
raise ValueError(
f'Do not specify client_id for {device_type.upper()} meshes. '
'It will be filled in automatically from environmental variables.'
'See api.py for the list of environmental variables for DTensor.'
)
mesh = tpu_util.create_tpu_mesh(dim_names, shape, mesh_name)
_print_context(api.num_global_devices(device_type), api.num_clients(),
api.client_id(), device_type, mesh)
return mesh
raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')
def dtensor_initialize_tpu_system(enable_coordination_service=False):
"""Initialize the TPU devices.
Args:
enable_coordination_service: If true, enable distributed coordination
service to make sure that workers know the devices on each other, a
prerequisite for data transfer through cross-worker rendezvous.
Raises:
RuntimeError: If running inside a tf.function.
NotFoundError: If no TPU devices found in eager mode.
"""
assert context.executing_eagerly()
in_multi_client_mode = api.job_name() != "localhost"
# Collective GRPC servers are only necessary in mutli-client setup.
# Single clients (e.g. Forge) can use local mode of collectives.
if in_multi_client_mode:
if api.jobs() is None:
raise ValueError(
"DTENSOR_JOBS environment variable is required when"
"using multi-client to properly set up communications between servers"
)
multi_client_util.initialize_multi_client_cluster(
job_name=api.job_name(),
dtensor_jobs=api.jobs(),
client_id=api.client_id(),
collective_leader=api.full_job_name(task_id=0),
enable_coordination_service=enable_coordination_service)
# Make sure the server change is fully propagated before attempting to run
# the core ID merging logic below.
context.ensure_initialized()
context.async_wait()
context.context()._clear_caches() # pylint: disable=protected-access
@function.defun
def _tpu_init_fn():
return gen_dtensor_ops.configure_and_initialize_global_tpu()
try:
with ops.device("/job:" + api.full_job_name() + "/device:TPU_SYSTEM:0"): # pylint: disable=protected-access
my_core_ids = _tpu_init_fn()
logging.info("TPU core IDs: %s", my_core_ids)
context.initialize_logical_devices()
# Configure virtual CPUs that is 1:1 mapped to TPU cores.
context.context().set_logical_cpu_devices(
len(api.local_devices(_TPU_DEVICE_TYPE)),
tf_device.DeviceSpec(
job=api.job_name(), replica=0, task=api.client_id()).to_string())
# `my_core_ids` contains the IDs of TPU cores attached to this host.
#
# To generate correct and efficient XLA AllReduce group assignment, we must
# merge these arrays from all hosts and broadcast the result back to all
# hosts, so all hosts can use these mappings in their MLIR passes.
#
# This is essentially doing what WaitForDistributedTpuOp and
# SetGlobalTPUArrayOp do, in our multi-client environment.
task_id = api.client_id()
num_tasks = api.num_clients()
num_devices = api.num_global_devices(_TPU_DEVICE_TYPE)
num_devices_per_task = int(num_devices / num_tasks)
# Create a one-time use mesh and layout just for merging core IDs.
mesh = layout_lib.Mesh([_MESH_DIM_X],
*_create_device_array((num_devices,),
_TPU_DEVICE_TYPE,
api.client_id()))
layout = layout_lib.Layout([_MESH_DIM_X, layout_lib.UNSHARDED], mesh)
device = dtensor_device.DTensorDevice(meshes=[mesh])
logging.info("TPU core locations: %s",
device.tpu_core_ids_to_locations(my_core_ids))
# At this point, we don't know which cores are attached to other hosts.
# The core ID mappings in the runtime haven't been set yet.
#
# The core ID merging AllReduce below is carefully written so it works
# without needing correct core mappings to be set in the runtime. We will
# use this AllReduce's result to set the core ID mappings, and all future
# user-initiated AllReduces will use the mappings.
#
# The runtime is hard-coded to ignore core ID mappings on this AllReduce.
all_core_ids = np.zeros([num_devices], dtype=np.int32)
for i in range(len(my_core_ids)):
all_core_ids[task_id * num_devices_per_task + i] = my_core_ids[i]
# Only one local device gets valid input: 8 local core IDs among
# (num_tasks - 1) * 8 zeros. The 8 core IDs are set using task ID as offset.
# The other 7 local devices get zero inputs. All devices on all host
# participate in one AllReduce, whose result will be core IDs arranged by
# task-device ordinals.
all_core_ids = constant_op.constant([all_core_ids])
zeros = array_ops.zeros_like(all_core_ids)
all_core_ids = [all_core_ids] + [zeros] * (num_devices_per_task - 1)
with ops.device(device.name):
all_core_ids = device.pack(all_core_ids, layout)
all_core_ids = math_ops.reduce_sum(all_core_ids, axis=[0])
unpacked_all_tpu_ids = device.unpack(all_core_ids)
all_core_ids = list(unpacked_all_tpu_ids[0].numpy())
logging.info("All TPU core IDs: %s", all_core_ids)
# Set the default core ID mappings in the runtime for legacy code and tests.
#
# Legacy code and tests create TPU meshes directly without using the
# `create_tpu_mesh` function below. Those meshes have global device IDs
# equal to TF task-device ordinals. The `all_core_ids` array happens to
# arrange core IDs by TF task-device ordinals. Using this array on those
# meshes guarantee correct although inefficient results.
device.set_tpu_core_ids("", all_core_ids)
# Remember enough global, immutable information to be able to build any ring
# we want prescribed by `create_tpu_mesh` in the future.
global _all_core_ids
_all_core_ids = all_core_ids
all_core_locations = device.tpu_core_ids_to_locations(all_core_ids)
all_core_locations = [
_CoreLocation(l[0], l[1], l[2], l[3]) for l in all_core_locations
]
global _all_core_locations
_all_core_locations = all_core_locations
logging.info("All TPU core locations: %s", all_core_locations)
tpu_topology = _create_tpu_topology(all_core_locations, num_tasks,
num_devices_per_task)
global _tpu_topology
_tpu_topology = tpu_topology
logging.vlog(1, "TPU Topology: %s, %s", tpu_topology.mesh_shape,
tpu_topology.device_coordinates)
global _dtensor_device
_dtensor_device = device
context.async_wait()
except errors.InvalidArgumentError as e:
raise errors.NotFoundError(
None, None, "Initialization failed, no valid TPUs found. " + str(e))
except errors.InternalError as e:
logging.error("Hit internal error during TPU system initialization. "
+ "It is likely hareware failure. \nPlease check the error "
+ "messages above to see whether that's the case. \nIf so, "
+ "consider to restart the job or try another machine.")
raise e
# Optionally exchange heartbeats between workers every minute.
if in_multi_client_mode and api.heartbeat_enabled():
logging.info(
"Starting DTensor heartbeat service exchanging signals every 10 minutes"
)
heartbeat.start(period=180)
# Clear out the eager context caches since the memory is invalid now.
logging.info("Clearing out eager caches")
context.context()._clear_caches() # pylint: disable=protected-access
def initialize_multi_client_cluster(job_name: str,
dtensor_jobs: List[str],
client_id: int,
collective_leader: str,
port: Optional[int] = None,
enable_coordination_service: bool = False):
"""Initialize GRPC servers and collectives for multi-client DTensor setup.
While single clients (e.g. Forge) can use local mode of collectives, GRPC
servers are necessary in mutli-client setup. This function can be used to
initialize a cluster and enable collective ops.
NOTE: this function must be called in an eager context.
Args:
job_name: The job name used by all clients in the DTensor cluster.
dtensor_jobs: A list of the DTensor client jobs participating in the
cluster. Must be strings of the form "hostname:port".
client_id: The ID of the DTensor client this function is being called in.
collective_leader: The job/task that will be used to run collectives.
port: The port this client's GRPC server will run on.
enable_coordination_service: If true, enable distributed coordination
service to make sure that workers know the devices on each other, a
prerequisite for data transfer through cross-worker rendezvous.
Raises:
RuntimeError: If running inside a tf.function.
"""
global _is_multi_client_initialized
assert context.executing_eagerly()
if _is_multi_client_initialized:
raise ValueError("Multi-client mode has already been initialized.")
if api.num_clients() <= 1:
raise ValueError(
"DTENSOR_NUM_CLIENTS must be set greater than 1 for multi-client mode."
)
if not api.jobs() or len(api.jobs()) <= 1:
raise ValueError(
"DTENSOR_JOBS environment variable is required when using multi-client "
"mode to properly set up communications between servers.")
if len(api.jobs()) != api.num_clients():
raise ValueError(
"DTENSOR_JOBS environment variable must be configured with the same "
"number of items as DTENSOR_NUM_CLIENTS.")
if not collective_leader.startswith("/job:"):
collective_leader = "/job:" + collective_leader
context.context().configure_collective_ops(
collective_leader=collective_leader)
if enable_coordination_service:
context.context().configure_coordination_service(
service_type="standalone", service_leader=collective_leader)
config_proto = context.get_config()
config_proto.experimental.collective_group_leader = collective_leader
# Construct server def from the host directly instead of relying on
# TF_CONFIG.
cluster_def = cluster_pb2.ClusterDef()
# Note that we will currently rely on the sorted string of job name as the
# order of assigning task ids. This might be brittle once we have jobs
# across multiple cells.
cluster_def.job.add(name=job_name, tasks=dict(enumerate(dtensor_jobs)))
server_def = tensorflow_server_pb2.ServerDef(
cluster=cluster_def,
default_session_config=config_proto,
job_name=job_name,
task_index=client_id,
protocol=remote_utils.get_default_communication_protocol(),
port=port)
server_def.default_session_config.rpc_options.num_channels_per_target = 4
server_def.default_session_config.experimental.recv_buf_max_chunk = -1
logging.info("Enabling collectives with server_def: %s", server_def)
context.context().enable_collective_ops(server_def)
context.ensure_initialized()
_is_multi_client_initialized = True
def start(period: int) -> threading.Event:
"""Starts a persistent thread exchanging heartbeats between workers.
Args:
period: Heartbeat interval in seconds. Heartbeat timeout is set to the
larger of `period` - 10 and 2s.
Returns:
A threading.Event object. Users can choose to call its set() method to shut
down the heartbeat service gracefully. This isn't necessary in most cases,
because the heartbeat service automatically shuts down at successful program
exit through atexit handlers. But in situations when atexit handlers are not
invoked, such as when multiprocessing processes exit in tests, users can
manually request a shutdown.
"""
global _heartbeat_timer
if _heartbeat_timer is not None:
logging.warning(
'A heartbeat thread is already running, skipping this one.')
return _heartbeat_timer
task_id = api.client_id()
num_tasks = api.num_clients()
# Worker 0 generates a random token. All other workers receive that token.
if task_id == 0:
token = np.random.randint(0,
pow(2, 16) - 1) # reserve the other 16 bits
signal = np.full([num_tasks], token, dtype=np.int32)
else:
signal = np.zeros([num_tasks], dtype=np.int32)
logging.info('Initial heartbeat signal: %s', signal)
device = tf_device.DeviceSpec(job=api.job_name(),
replica=0,
task=task_id,
device_type='CPU',
device_index=0)
# Always use 0 for group and instance keys to reduce unnecessary
# collective hangs and simplify failure analysis. This also avoid
# collision with normal collectives.
with ops.device(device):
signal = all_reduce(constant_op.constant(signal),
group_size=num_tasks,
group_key=0,
instance_key=0,
timeout=max(period - 10, 2)).numpy()
logging.info('Merged heartbeat signal %s', signal)
# The merged signal should have equal elements. If not, some worker(s) may be
# out of sync, and we should terminate all workers.
if task_id == 0:
if not np.all(signal == token):
logging.fatal('Merged heartbeat signal has value != %d', token)
else:
if len(set(signal)) != 1:
logging.fatal('Merged heartbeat signal has unequal elements')
token = signal[0]
# On normal main process exit, set the timer to stop the heartbeat thread.
_heartbeat_timer = threading.Event()
def stop_heartbeat():
logging.info('Stopping the heartbeat thread')
_heartbeat_timer.set()
# Give the threads some time to clean up.
time.sleep(max(period // 10, 2))
atexit.register(stop_heartbeat)
# Start the persistent heartbeat thread.
thread = threading.Thread(
target=_heartbeat,
args=[period, _heartbeat_timer, token, num_tasks, task_id, device],
daemon=True)
thread.start()
return _heartbeat_timer
def create_distributed_mesh(mesh_dims: List[Tuple[str, int]],
mesh_name: str = '',
local_devices: Optional[List[str]] = None,
device_type: Optional[str] = None) -> layout.Mesh:
"""Creates a distributed mesh.
This is similar to `create_mesh`, but with a different set of arguments to
create a mesh that spans evenly across a multi-client DTensor cluster.
For CPU and GPU meshes, users can choose to use fewer local devices than what
is available `local_devices`.
For TPU, only meshes that uses all TPU cores is supported by the DTensor
runtime.
Args:
mesh_dims: A list of (dim_name, dim_size) tuples.
mesh_name: Name of the created mesh. Defaults to ''.
local_devices: String representations of devices to use. This is the device
part of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available local
logical devices.
device_type: Type of device to build the mesh for. Defaults to 'CPU'.
Supported values are 'CPU', 'GPU', 'TPU'.
Returns:
A mesh that spans evenly across all DTensor clients in the cluster.
"""
dim_names, shape = zip(*mesh_dims)
if device_type and device_type.upper() == 'TPU':
# TODO(b/185940495): Allow multi-mesh and partial on TPU.
# TPU meshes can only be configured through environment variables that
# reflect the actual TPU topology. Do not let users specify custom args.
if local_devices is not None:
raise ValueError(
f'Do not specify devices for {device_type.upper()} meshes. '
f'Using a partial list of devices for {device_type.upper()} '
f'is not supported.')
device_specs, device_type = _make_device_specs(local_devices, device_type)
if device_type.upper() in ['CPU', 'GPU']:
# For CPU and GPU meshes, user-specified args take precedence over env vars.
# This is particularly useful on single clients when users want to create
# meshes that use fewer logical devices than what's available.
if api.num_clients() > 1 and not multi_client_util.is_initialized():
raise ValueError('Invalid multi-client topology, please run '
'dtensor.initialize_multi_client() first.')
local_spec = tf_device.DeviceSpec(job=api.job_name(),
replica=0,
task=api.client_id())
device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
# Assumes identical number of local devices per client.
num_global_devices = len(device_specs) * api.num_clients()
if np.prod(shape) != num_global_devices:
raise ValueError(
f'Global number of devices '
f'({len(device_specs)} per client * {api.num_clients()} clients '
f'= {num_global_devices}) must be '
f'equal to total size of the mesh of shape {shape}')
global_device_ids = np.arange(num_global_devices).reshape(shape)
flattened = np.ravel(global_device_ids).tolist()
start_idx = len(device_specs) * api.client_id()
local_device_ids = flattened[start_idx:start_idx + len(device_specs)]
mesh = layout.Mesh(dim_names=dim_names,
global_device_ids=global_device_ids,
local_device_ids=local_device_ids,
local_devices=device_specs,
mesh_name=mesh_name)
_print_context(num_global_devices, api.num_clients(), api.client_id(),
device_type, mesh)
return mesh
if device_type.upper() == 'TPU':
mesh = tpu_util.create_tpu_mesh(dim_names, shape, mesh_name)
_print_context(api.num_global_devices(device_type), api.num_clients(),
api.client_id(), device_type, mesh)
return mesh
raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')
