tf.distribute.cluster_resolver.SlurmClusterResolver

ClusterResolver для систем с диспетчером рабочих нагрузок Slurm.

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

tf.distribute.cluster_resolver.SlurmClusterResolver(
    jobs=None,
    port_base=8888,
    gpus_per_node=None,
    gpus_per_task=None,
    tasks_per_node=None,
    auto_set_gpu=True,
    rpc_layer='grpc'
)

Это реализация ClusterResolver для кластеров Slurm. Это позволяет указать задачи и количество задач, количество задач на узел, количество графических процессоров на каждом узле и количество графических процессоров для каждой задачи. Он извлекает системные атрибуты из переменных окружения Slurm, определяет имена выделенных вычислительных узлов, строит кластер и возвращает объект ClusterResolver, который можно использовать для распределенного TensorFlow.

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 на основе последних данных о группе экземпляров.

Это возвращает объект ClusterSpec для использования на основе информации из указанных параметров инициализации и переменных окружения Slurm. Спецификация кластера решается каждый раз при вызове этой функции. Разрешитель извлекает имена хостов узлов с помощью scontrol и упаковывает задачи в этом порядке, пока узел a не будет иметь количество задач, равное спецификации. Графические процессоры на узлах распределяются по задачам в соответствии со спецификацией путем установки переменной среды CUDA_VISIBLE_DEVICES.

get_task_info

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

get_task_info()

Возвращает имя задачи и идентификатор задачи для процесса, вызвавшего это.

Это возвращает имя задачи и индекс задачи для процесса, вызвавшего эту функцию в соответствии с его рангом и спецификацией кластера. Имя задачи и индекс задачи устанавливаются после создания кластера с помощью cluster_spec, в противном случае по умолчанию возвращается None.

master

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

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

Возвращает строку master для подключения к мастеру TensorFlow.

num_accelerators

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

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

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

Это возвращает количество ядер ускорителей (например, графических процессоров и TPU), доступных на задачу.

Мы также позволяем вызывающим сторонам указать тип_задачи и идентификатор_задачи, если они хотят нацелиться на определенную задачу TensorFlow для запроса количества ускорителей. Это поддерживает гетерогенные среды, где количество ядер ускорителей на хосте отличается.