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 и распределяет задачи в этом порядке, пока на узле не будет задано указанное количество задач. Графические процессоры на узлах распределяются между задачами согласно спецификации путем установки переменной окружения CUDA_VISIBLE_DEVICES.

get_task_info

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

get_task_info()

Возвращает имя задания и task_id для вызывающего процесс.

Возвращает имя задания и индекс задачи для процесса, вызывающего эту функцию, в соответствии с его рангом и спецификацией кластера. Имя задания и индекс задачи задаются после создания кластера с помощью 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
)

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

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

При необходимости вызывающие стороны могут указать task_type и task_id, чтобы запросить количество ускорителей для конкретной задачи TensorFlow. Это поддерживает неоднородные среды, в которых количество ядер ускорителей на узел различается.