tf.distribute.experimental.PreemptionWatcher

Отслеживает сигнал о вытеснении и сохраняет его.

tf.distribute.experimental.PreemptionWatcher()

Примечание: в настоящее время поддерживается только среда Borg TPU с TPUClusterResolver.

Этот класс позволяет отслеживать сигнал о вытеснении во время обучения на TPU. Он запускает фоновый поток для наблюдения за процессом обучения и пытается получить сообщение о вытеснении от службы координации. При вытеснении рабочий процесс, подвергшийся вытеснению, записывает сообщение о вытеснении в службу координации. Поэтому непустое сообщение о вытеснении означает, что произошло вытеснение.

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

keep_running = True
preemption_watcher = None
while keep_running:
  try:
    # Initialize TPU cluster and stratygy.
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(resolver)
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.TPUStrategy(resolver)

    # PreemptionWatcher must be created after connected to cluster.
    preemption_watcher = tf.distribute.experimental.PreemptionWatcher()
    train_model(strategy)
    keep_running = False
  except Exception as e:
    if preemption_watcher and preemption_watcher.preemption_message:
      preemption_watcher.block_until_worker_exit()
      keep_running = True
    else:
      raise e

Методы

block_until_worker_exit

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

block_until_worker_exit()

Блокирует координатор до завершения работы процессов.

В некоторых редких случаях в течение льготного периода перед вытеснением может возникнуть другая ошибка. Из-за этого координатор повторно подключится к тем же рабочим процессам TPU, которые позднее будут остановлены. Это не позволит координатору подключиться к новым рабочим процессам TPU и приведет к полному перезапуску. Чтобы избежать такой ситуации, этот метод блокирует повторное подключение координатора до завершения работы процессов. Для платформ, отличных от TPU, этот метод ничего не делает.