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, этот метод будет ничем не делать.