Настройка

Средство запуска

Программа запуска, поставляемая с TorchElastic, подходит для большинства сценариев использования (см. torchrun (эластичный запуск)). Вы можете реализовать собственное средство запуска, программно создав агент и передав ему спецификации для ваших рабочих процессов, как показано ниже.

# my_launcher.py

if __name__ == "__main__":
    args = parse_args(sys.argv[1:])
    rdzv_handler = RendezvousHandler(...)
    spec = WorkerSpec(
        local_world_size=args.nproc_per_node,
        fn=trainer_entrypoint_fn,
        args=(trainer_entrypoint_fn args.fn_args,...),
        rdzv_handler=rdzv_handler,
        max_restarts=args.max_restarts,
        monitor_interval=args.monitor_interval,
    )

    agent = LocalElasticAgent(spec, start_method="spawn")
    try:
        run_result = agent.run()
        if run_result.is_failed():
            print(f"worker 0 failed with: run_result.failures[0]")
        else:
            print(f"worker 0 return value is: run_result.return_values[0]")
    except Exception ex:
        # handle exception

Обработчик rendezvous

Чтобы реализовать собственный rendezvous, расширьте torch.distributed.elastic.rendezvous.RendezvousHandler и реализуйте его методы.

После реализации вы можете передать собственный обработчик rendezvous спецификации рабочих процессов при создании агента.

spec = WorkerSpec(
    rdzv_handler=MyRendezvousHandler(params),
    ...
)
elastic_agent = LocalElasticAgent(spec, start_method=start_method)
elastic_agent.run(spec.role)

Обработчик метрик

TorchElastic генерирует метрики уровня платформы (см. раздел Метрики). По умолчанию метрики отправляются в /dev/null, поэтому вы их не увидите. Чтобы передавать метрики в службу обработки метрик вашей инфраструктуры, реализуйте torch.distributed.elastic.metrics.MetricHandler и configure в собственном средстве запуска.

# my_launcher.py

import torch.distributed.elastic.metrics as metrics

class MyMetricHandler(metrics.MetricHandler):
    def emit(self, metric_data: metrics.MetricData):
        # push metric_data to your metric sink

def main():
    metrics.configure(MyMetricHandler())

    spec = WorkerSpec(...)
    agent = LocalElasticAgent(spec)
    agent.run()

Обработчик событий

TorchElastic поддерживает запись событий (см. раздел События). Модуль событий определяет API, позволяющий записывать события и реализовывать пользовательский EventHandler. EventHandler используется для публикации событий, возникающих во время выполнения torchelastic, в различные источники, например AWS CloudWatch. По умолчанию используется torch.distributed.elastic.events.NullEventHandler, который игнорирует события. Чтобы настроить пользовательский обработчик событий, необходимо реализовать интерфейс torch.distributed.elastic.events.EventHandler и configure в собственном средстве запуска.

# my_launcher.py

import torch.distributed.elastic.events as events

class MyEventHandler(events.EventHandler):
    def record(self, event: events.Event):
        # process event

def main():
    events.configure(MyEventHandler())

    spec = WorkerSpec(...)
    agent = LocalElasticAgent(spec)
    agent.run()