tf.distribute.cluster_resolver.TPUClusterResolver

Резольвер кластера для Google Cloud TPUs.

Наследуется от: ClusterResolver

tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=None,
    zone=None,
    project=None,
    job_name='worker',
    coordinator_name=None,
    coordinator_address=None,
    credentials='default',
    service=None,
    discovery_url=None
)

Используется в ноутбуках

Это реализация резольверов кластеров для сервиса Google Cloud TPU.

TPUClusterResolver поддерживает следующие различные среды: Google Compute Engine, Google Kubernetes Engine, внутренние среды Google.

Он может быть передан в tf.distribute.TPUStrategy для поддержки обучения TF2 на Cloud TPUs.

cluster_spec = tf.train.ClusterSpec({
    "ps": ["localhost:2222", "localhost:2223"],
    "worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
})

# SimpleClusterResolver is used here for illustration; other cluster
# resolvers may be used for other source of task type/id.
simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
                                        task_id=0)

...

if cluster_resolver.task_type == 'worker' and cluster_resolver.task_id == 0:
  # Perform something that's only applicable on 'worker' type, id 0. This
  # block will run on this particular instance since we've specified this
  # task to be a 'worker', id 0 in above cluster resolver.
else:
  # Perform something that's only applicable on other ids. This block will
  # not run on this particular instance.

cluster_spec = tf.train.ClusterSpec({
    "ps": ["localhost:2222", "localhost:2223"],
    "worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
})

# SimpleClusterResolver is used here for illustration; other cluster
# resolvers may be used for other source of task type/id.
simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
                                        task_id=1)

...

if cluster_resolver.task_type == 'worker':
  # Perform something that's only applicable on workers. This block
  # will run on this particular instance since we've specified this task to
  # be a worker in above cluster resolver.
elif cluster_resolver.task_type == 'ps':
  # Perform something that's only applicable on parameter servers. This
  # block will not run on this particular instance.

Методы

cluster_spec

Просмотреть исходный код

cluster_spec()

Возвращает объект ClusterSpec на основе последних данных о TPU.

Мы получаем информацию из API GCE каждый раз, когда вызывается этот метод.

connect

Просмотреть исходный код

@staticmethod
connect(
    tpu=None, zone=None, project=None
)

Инициализирует TPU и возвращает TPUClusterResolver.

Этот API подключится к удаленному кластеру TPU и инициализирует аппаратное обеспечение TPU. Пример использования:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
    tpu='')

Он может рассматриваться как удобная обёртка следующего кода:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

get_coordination_service_leader

Просмотреть исходный код

get_coordination_service_leader()

Возвращает расположение сервиса координации.

Сервис координации должен находиться на TPU worker0.

get_job_name

Просмотреть исходный код

get_job_name()

get_master

Просмотреть исходный код

get_master()

get_tpu_system_metadata

Просмотреть исходный код

get_tpu_system_metadata()

Возвращает метаданные системы TPU.

Пользователи могут вызвать этот метод, чтобы получить некоторые сведения о системе TPU, такие как общее количество ядер, количество TPU-рабочих узлов и устройства. Например:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tpu_system_metadata = resolver.get_tpu_system_metadata()
num_hosts = tpu_system_metadata.num_hosts

master

Просмотреть исходный код

master(
    task_type=None, task_id=None, rpc_layer=None
)

Получает строку Master, используемую для сеанса.

В нормальном случае это возвращает путь grpc (grpc://1.2.3.4:8470) первого экземпляра в ClusterSpec, возвращённого функцией cluster_spec.

Если при построении TPUClusterResolver используется имя, отличное от TPU, оно будет возвращено вместо этого (например, если значение аргумента tpus при построении этого TPUClusterResolver было 'grpc://10.240.1.2:8470', будет возвращено 'grpc://10.240.1.2:8470').

num_accelerators

Просмотреть исходный код

num_accelerators(
    task_type=None, task_id=None, config_proto=None
)

Возвращает количество TPU-ядер на рабочий узел.

Подключается к мастеру, перечисляет все имеющиеся устройства на мастер-узле и подсчитывает их. Также проверяет, что количество устройств на хост в кластере одинаковое, прежде чем возвращать количество TPU-ядер на хост.

set_tpu_topology

Просмотреть исходный код

set_tpu_topology(
    serialized_tpu_topology
)

Устанавливает информацию о топологии TPU, хранящуюся в этом резолвере.

__enter__

Просмотреть исходный код

__enter__()

__exit__

Просмотреть исходный код

__exit__(
    type, value, traceback
)