[源碼解析] 深度學習分佈式訓練框架 horovod (16) --- 彈性訓練之Worker生命週期

[源碼解析] 深度學習分佈式訓練框架 horovod (16) --- 彈性訓練之Worker生命週期

0x00 摘要

Horovod 是Uber於2017年發佈的一個易於使用的高性能的分佈式訓練框架,在業界得到了廣泛應用。

本系列將通過源碼分析來帶領大家瞭解 Horovod。本文是第十六篇,看看 horovod 彈性訓練中 worker 的生命週期。

我們先給出一個邏輯圖,大家先有一個粗略的瞭解,本圖左側是 Driver 部分,右側是一個 Worker。

本系列其他文章鏈接如下:

[源碼解析] 深度學習分佈式訓練框架 Horovod (1) --- 基礎知識

[源碼解析] 深度學習分佈式訓練框架 horovod (2) --- 從使用者角度切入

[源碼解析] 深度學習分佈式訓練框架 horovod (3) --- Horovodrun背後做了什麼

[源碼解析] 深度學習分佈式訓練框架 horovod (4) --- 網絡基礎 & Driver

[源碼解析] 深度學習分佈式訓練框架 horovod (5) --- 融合框架

[源碼解析] 深度學習分佈式訓練框架 horovod (6) --- 後臺線程架構

[源碼解析] 深度學習分佈式訓練框架 horovod (7) --- DistributedOptimizer

[源碼解析] 深度學習分佈式訓練框架 horovod (8) --- on spark

[源碼解析] 深度學習分佈式訓練框架 horovod (9) --- 啓動 on spark

[源碼解析] 深度學習分佈式訓練框架 horovod (10) --- run on spark

[源碼解析] 深度學習分佈式訓練框架 horovod (11) --- on spark --- GLOO 方案

[源碼解析] 深度學習分佈式訓練框架 horovod (12) --- 彈性訓練總體架構

[源碼解析] 深度學習分佈式訓練框架 horovod (13) --- 彈性訓練之 Driver

[源碼解析] 深度學習分佈式訓練框架 horovod (14) --- 彈性訓練發現節點 & State

[源碼解析] 深度學習分佈式訓練框架 horovod (15) --- 廣播 & 通知

0x01 Worker 是什麼

首先,我們要看看 worker 是什麼。爲了可以單獨成文,本章節回憶了很多之前的知識,看過之前文章的同學可以跳過。

1.1 角色

”訓練“ 是通過計算梯度下降的方式利用數據來迭代地優化神經網絡參數,最終輸出網絡模型的過程。

我們首先要看看彈性訓練中的角色設定。

Horovod 的彈性訓練包含兩個角色,driver 進程和 worker 進程。driver 進程運行在 CPU 節點上,worker 進程可運行在 CPU 或者 GPU 節點上。在 Horovod 中,訓練進程是平等的參與者,每個 worker 進程既負責梯度的分發,也負責具體的梯度計算。

這兩個角色和 Spark 的 Driver -- Executor 依然很類似。Driver 進程就可以認爲是 Spark 的 Driver,或者說是 master 節點。Worker 就類似於 Spark 的 Executor。

具體如圖:

                 +------------------------------+
                 |                              |
                 |            Driver            |
                 |                              |
                 |                              |
                 +-----+-------+--------+-------+
                       ^       ^        ^
                       |       |        |
                       |       |        |
         +-------------+       |        +--------------+
         |                     |                       |
         |                     |                       |
         |                     |                       |
         v                     v                       v
+--------+----+        +-------+------+           +----+--------+
|  Worker     |        |  Worker      |           |  Worker     |
|             |        |              |           |             |
|      host1  |        |      host2   |           |     host3   |
+-------------+        +--------------+           +-------------+

1.2 職責

角色的職責如下:

master(主節點)職責:

  • 負責實時探活 worker(工作節點)是否有變化,掉線情況;
  • 負責實時監控 host 是否有變化;
  • 負責分配任務到存活的worker(工作節點);
  • 在有進程失敗導致 AllReduce 調用失敗的 情況下,master 通過 blacklist 機制 組織剩下的活着的進程構造一個新的環。
  • 如果有新 host 加入,則生成新的 worker,新 worker 和 舊 worker 一起構造成一個新的環。

worker(工作節點)職責:

  • 負責彙報(其實是被動的,沒有主動機制)當前worker(工作節點)的狀態(就是訓練完成情況);

  • 負責在該worker(工作節點)負責的數據上執行訓練。

1.3 組網機制

Horovod 在單機的多個 GPU 上採用 NCCL 來通信,在多機之間通過 ring-based AllReduce 算法進行通信。

Horovod 的彈性訓練是指多機的彈性訓練。在多機的 ring-based 通信中的每個 worker 節點有一個左鄰和一個右鄰,每個 worker 只會向它的右鄰居發送數據,並從左鄰居接受數據。

1.3.1 通信環

Driver 進程用於幫助 worker 調用 gloo 構造 AllReduce 通信環。

當 Horovod 在調用 Gloo 來構造通信域時,Horovod 需要給 Gloo 創建一個帶有 KVStore 的 RendezvousServer,其中 KVStore 用於存儲 通信域內每個節點的 host 地址給其在邏輯通信環分配的序號 rank 等信息。

構建過程如下:

  • Driver 進程創建帶有 KVStore 的 RendezvousServer,即這個 RendezvousServer 運行在 Horovod 的 driver 進程裏

  • Driver 進程拿到所有 worker 進程節點的 IP 地址和 GPU 卡數信息後,會將其寫入RendezvousServer 的 KVStore 中。

  • 每個 worker 節點會通過調用 gloo 從而 請求 RendezvousServer 獲取自己的鄰居節點信息(ip,port...),從而構造通信域。

1.3.2 彈性構建

當有 worker 失敗或者新的 worker 加入訓練時,每個 worker 會停止當前的訓練,記錄當前模型迭代的步數,並嘗試重新初始化 AllReduce 的通信域

1.3.2.1 Driver 監控

因爲 driver 進程一直在監控 worker 的狀態 和 host 節點情況,所以

  • 當 host 變化時候,當驅動進程通過節點發現腳本發現一個節點被標記爲新增或者移除時,它將發送一個通知到 所有workers,在下一次 state.commit() 或者更輕量的 state.check_host_updates() 被調用時,會拋出一個 HostsUpdateInterrupt 異常。
  • 當有 worker 失敗時,driver 會重新獲取存活的 worker 的host,
1.3.2.2 Driver 重新構建

爲了不讓其他 worker 進程退出,Horovod 會捕獲 gloo 拋出的異常,並將異常傳遞給封裝的 Python API。進而 driver 會重新配置 RendezvousServer,從而讓 worker 節點能重新構造通信域。所以 Horovod 是可容錯的。

如果有 host 變化,worker 進程可以通過這個 rendezvous 來構造新的通信域。當新的通信域構造成功後,rank=0 的 worker 會將自身的模型廣播給其他 worker,然後接着上次停止的迭代步數開始訓練。

組網機制如下:

                         +-------------------------------+
                         | Driver                        |
                         |                               |
                         |   +------------------------+  |
                         |   | RendezvousServer       |  |
                         |   |                        |  |
                         |   |                        |  |
                         |   |   host1, host2, host3  |  |
                         |   +------------------------+  |
                         +-------------------------------+
                                ^       ^        ^
                                |       |        |
                                |       |        |
                  +-------------+       |        +--------------+
                  |                     |                       |
                  |                     |                       |
                  |                     |                       |
                  v                     v                       v
         +--------+----+        +-------+------+           +----+--------+
         |  Worker     |        |  Worker      |           |  Worker     |
+------> |             +------> |              +---------> |             | +------+
|        |      host1  |        |      host2   |           |     host3   |        |
|        +-------------+        +--------------+           +-------------+        |
|                                                                                 |
|                                                                                 |
|                                                                                 v
<--------------------------------------------------------------------------------+

所以,本文就看看 Worker 的一個整體生命流程。

0x02 總體生命流程

在 Driver 的 launch_gloo_elastic 之中,如下代碼負責啓動 worker。

  • command 就是傳入的可執行命令,比如 python train.py。經過 get_run_command 之後就得到了 env python train.py 之類的樣子,就是加上環境變量,可以運行了。
  • exec_command 類似如下: exec_command = _exec_command_fn(settings),就是基於各種配置來生成可以執行命令環境
run_command = get_run_command(command, server_ip, nics, global_rendezv_port, elastic=True)

create_worker = _create_elastic_worker_fn(exec_command, run_command, env, event)

driver.start(settings.num_proc, create_worker)

res = driver.get_results()

可以清晰看出來三個詳細的過程(因爲 訓練過程是在 worker 內部,所以 driver 分析沒有深入此部分):

  • _create_elastic_worker_fn 是配置過程
  • driver.start 是啓動過程
  • driver.get_results 就是得到 & 註冊 運行結果

我們接下來就按照這三個過程詳細分析一下。

0x03 配置過程

配置過程 是由 _create_elastic_worker_fn 完成,就是提供一個在某個環境下運行某個命令的能力

_create_elastic_worker_fn 分爲兩部分:

  • _slot_info_to_command_fn 會建立 slot_info_to_command,套路和之前文章中類似,就是把各種 horovod 環境變量和運行命令 run_command 糅合起來,得到一個可以在“某個 host and slot” 之上運行的命令文本
  • 返回 create_worker。
    • create_worker 是利用 exec_command 和 命令文本 構建的函數。
    • exec_command 我們在之前介紹過,就是提供了一種運行命令的能力,或者說是運行環境;
    • 所以 create_worker 就是提供一個在某個環境下運行某個命令的能力;
# 得到一個 可以在“某個 host and slot” 之上運行的命令文本
def _slot_info_to_command_fn(run_command, env):
     def slot_info_to_command(slot_info):
        """
        Given a slot_info, creates a command used by gloo to launch a single job.
        :param slot_info: host and slot to execute the run command on
        """
        env_vars = create_slot_env_vars(slot_info)
        horovod_rendez_env = " ".join(
            [f"{k}={str(v)}" for k, v in env_vars.items()])

        return '{horovod_env} {env} {run_command}' .format(
            horovod_env=horovod_rendez_env,
            env=' '.join(['%s=%s' % (key, quote(value)) for key, value in env.items()
                          if env_util.is_exportable(key)]),
            run_command=run_command)

    return slot_info_to_command

def _create_elastic_worker_fn(exec_command, run_command, env, event):
    get_command_with_env = _slot_info_to_command_fn(run_command, env)

    # 提供一個在某個環境下運行某個命令的能力
    def create_worker(slot_info, events):
        command = get_command_with_env(slot_info)
        events = [event] + (events or [])
        return exec_command(command, slot_info, events)
    return create_worker

所以,最終得到的 create_worker 是:

               command (python train.py)
                          +
                          |
                          |
          get_run_command |
                          |
                          v

               run_command(env python train.py) #得到命令類似 env python train.py
                          +
                          |
                          |
 _slot_info_to_command_fn |
                          |
                          v

             {horovod_env} {env} {run_command} #得到命令類似 horovod_env env python train.py
                          +
                          |
                          |
             exec_command |
                          |
                          |
                          v

create_worker = exec_command({horovod_env} {env} {run_command})#得到在某個環境下運行某個命令的能力

這樣,create_worker 就直接可以運行訓練代碼了。

0x04 啓動過程

create_worker = _create_elastic_worker_fn 提供了一個在某個環境下運行某個命令的能力,因爲 create_worker 方法內部已經包括了執行命令和執行環境,就是說,只要運行create_worker,就可以自動訓練。下面我們就利用這個能力來啓動 worker。

啓動過程基本都是在 ElasticDriver 類 的 start 方法中完成。

4.1 總體邏輯

以下邏輯都運行在 ElasticDriver 之中。

  • 首先,會把 上面生成的 create_worker 賦值給 self._create_worker_fn。
  • 其次,會調用 _activate_workers 啓動多個 worker,其中包括:
    • 先使用 wait_for_available_slots 等待 min_np 數目的可用的 hosts。之前分析過此函數,就是 無限循環等待,如果 avail_slots >= min_np and avail_hosts >= min_hosts 纔會返回。
    • 使用 _update_host_assignments 來得到 slots;
    • 使用 _start_worker_processes 來啓動多個 worker;
def _activate_workers(self, min_np):
    current_hosts = self.wait_for_available_slots(min_np)
    pending_slots = self._update_host_assignments(current_hosts)
    self._worker_registry.reset(self.world_size())
    self._start_worker_processes(pending_slots)

def start(self, np, create_worker_fn):
    self._create_worker_fn = create_worker_fn
    self._activate_workers(np)

下面逐一看看。

4.2 賦值

第一步就是上面生成的 create_worker 賦值給 self._create_worker_fn。

                 command (python train.py)
                            +
                            |
                            |
            get_run_command |
                            |
                            v

                 run_command(env python train.py)
                            +
                            |
                            |
   _slot_info_to_command_fn |
                            |                                          +-----------------------------+
                            v                                          |                             |
                                                                       |                             |
               {horovod_env} {env} {run_command}                       |     ElasticDriver           |
                            +                                          |                             |
                            |                                          |                             |
                            |                                          |                             |
               exec_command |                                          |                             |
                            |                                          |                             |
                            |                                          |                             |
                            v                                       1  |                             |
+---------------------------+------------------------------------+     |                             |
| create_worker = exec_command({horovod_env} {env} {run_command})| <------------+  _create_worker_fn |
+----------------------------------------------------------------+     |                             |
                                                                       |                             |
                                                                       +-----------------------------+

手機如下:

4.3 獲取 host 信息

接下來要使用 _update_host_assignments 來得到 slots,具體分爲兩步:

首先構建 host 和 rank 之間的分配狀況。

其次

4.3.1 更新 host 和 rank

_update_host_assignments 函數中會根據 最新的 host 信息,重新構建 rendezvous,比如:

self._rendezvous.init(host_assignments_list)

具體邏輯是:

  • 獲取 活躍的slot :active_slots
  • 獲取 host 分配情況;
  • 確保每個 worker 有前驅者,就是可以傳遞狀態,構成環;
  • 調用 self._rendezvous.init 重新構造 rendezvous;
  • 分配 rank 和 slot 的關係;
  • 返回 pending_slots,就是分配的slot之中,不在 活躍slot列表 active_slots 中的。不活躍的就是接下來可以啓動 新 worker 的
def _update_host_assignments(self, current_hosts):
    # Determine the slots that are already filled so we do not respawn these processes
    # 獲取 活躍的slot
    active_slots = set([(host, slot_info.local_rank)
                        for host, slots in self._host_assignments.items()
                        for slot_info in slots])

    # Adjust the host assignments to account for added / removed hosts
    host_assignments, host_assignments_list = self._get_host_assignments(current_hosts)

    if len(self._host_assignments) > 0:
        # Ensure that at least one previously active host is still assigned, otherwise there is no
        # way to sync the state to the new workers
        prev_hosts = self._host_assignments.keys()
        next_hosts = host_assignments.keys()
        if not prev_hosts & next_hosts:
            raise RuntimeError('No hosts from previous set remaining, unable to broadcast state.')

    self._host_assignments = host_assignments
    self._world_size = len(host_assignments_list)
    
    self._rendezvous.init(host_assignments_list) # 重新構造 rendezvous

    # Rank assignments map from world rank to slot info
    rank_assignments = {}
    for slot_info in host_assignments_list:
        rank_assignments[slot_info.rank] = slot_info
    self._rank_assignments = rank_assignments

    # Get the newly assigned slots that need to be started
    pending_slots = [slot_info
                     for host, slots in self._host_assignments.items()
                     for slot_info in slots
                     if (host, slot_info.local_rank) not in active_slots]
    return pending_slots

4.3.2 獲取 host 和 rank

其中,host信息是通過 _get_host_assignments 來完成

def _get_host_assignments(self, current_hosts):
    # Adjust the host assignments to account for added / removed hosts
    host_list = [hosts.HostInfo(host, current_hosts.get_slots(host))
                 for host in current_hosts.host_assignment_order]
    host_assignments_list = hosts.get_host_assignments(host_list, self._min_np, self._max_np)
    host_assignments = defaultdict(list)
    for slot_info in host_assignments_list:
        host_assignments[slot_info.hostname].append(slot_info)
    return host_assignments, host_assignments_list

_get_host_assignments 調用get_host_assignments具體完成業務。

get_host_assignments 會依據 host 和 process capacities (slots) 來給 Horovod 之中的進程分配,即給出一個 horovod rank 和 slot 的對應關係。設置了幾個 np,就有幾個 slot。

給出的分配方案類似如下,這樣就知道了哪個rank對應於哪個host上的哪個slot

[
  SlotInfo(hostname='h1', rank=0, local_rank=0, cross_rank=0, size=2, local_size=2, coress_size=1),
	SlotInfo(hostname='h2', rank=1, local_rank=0, cross_rank=0, size=2, local_size=2, coress_size=1),
]

代碼如下:

def get_host_assignments(hosts, min_np, max_np=None):
    """Assign hosts with process capacities (slots) to ranks in the Horovod process.
    This function will try to allocate as many as possible processes on the same host to leverage local network.

    :param hosts: list of HostInfo objects describing host and slot capacity
    :type hosts: list[HostInfo]
    :param min_np: minimum number of processes to be allocated
    :param max_np: (optional) maximum number of processes to be allocated
    :return: a list of the allocation of process on hosts in a `SlotInfo` object.
    :rtype: list[SlotInfo]
    """
    host_ranks = []
    cross_ranks = collections.defaultdict(dict)
    rank = 0
    # 依據 hosts 信息構建 rank, local rank, cross rank(hierarchical allreduce所需要)
    for host_info in hosts:
        ranks = []
        for local_rank in range(host_info.slots):
            if rank == max_np:
                break

            ranks.append(rank)
            rank += 1

            cross_ranks_at_local = cross_ranks[local_rank]
            cross_ranks_at_local[host_info.hostname] = len(cross_ranks_at_local)

        host_ranks.append((host_info, ranks))

    world_size = rank

    # 給出一個 horovod rank 和 slot 的對應關係。返回一個alloc_list,每個SlotInfo包括各種rank信息
    alloc_list = []
    for host_info, ranks in host_ranks:
        local_size = len(ranks)
        for local_rank, rank in enumerate(ranks):
            cross_ranks_at_local = cross_ranks[local_rank]
            cross_rank = cross_ranks_at_local[host_info.hostname]
            cross_size = len(cross_ranks_at_local)

            alloc_list.append(
                SlotInfo(
                    hostname=host_info.hostname,
                    rank=rank,
                    local_rank=local_rank,
                    cross_rank=cross_rank,
                    size=world_size,
                    local_size=local_size,
                    cross_size=cross_size))

    return alloc_list

4.3.3 拓展邏輯

目前拓展邏輯如下,經過 4.2,4.3 兩步之後,ElasticDriver 中的一些變量被賦值了,我們簡化了上圖中的 exec_command 如下(第一步在上圖中):

                                                           From Host discovery
                                                                   +
+------------------------------+                                   |
| ElasticDriver                |                                   |
|                              |                                   |  2
|                              |       wait_for_available_slots    |
|    wait_for_available_slots  |                                   |
|                              |                                   v
|   _update_host_assignments   |                              current_hosts
|                              |                                   +
|          _rendezvous         |                                   |
|                              |       _update_host_assignments    |
|         _host_assignments    |                                   |  3
|                              |          self._rendezvous.init    |
|         _world_size          |                                   |
|                              |                                   |
|         _rank_assignments    |                                   |
|                              |                                   v
|         _create_worker_fn+---------+                        pending_slots
|                              |     |
|                              |     |
+------------------------------+     |
                                     |
                                     |
                                     v
 +-----------------------------------+------------+
 | exec_command({horovod_env} {env} {run_command})|
 +------------------------------------------------+

4.4 啓動

_start_worker_processes 完成了啓動過程,邏輯遞增如下。

  • create_worker_fn 就是使用 之前賦值的 create_worker = exec_command({horovod_env} {env} {run_command}) ;
  • run_worker() 中 執行 create_worker_fn(slot_info, [shutdown_event, host_event]) 就是 運行訓練代碼;
  • threading.Thread(target=run_worker) 就是在一個 thread 之中運行訓練代碼;
  • _start_worker_processes 就是在多個 thread 之中運行多份訓練代碼;
def _start_worker_processes(self, pending_slots):
    for slot_info in pending_slots:
        self._start_worker_process(slot_info)

def _start_worker_process(self, slot_info):
    create_worker_fn = self._create_worker_fn
    shutdown_event = self._shutdown
    host_event = self._host_manager.get_host_event(slot_info.hostname)

    def run_worker():
        res = create_worker_fn(slot_info, [shutdown_event, host_event])
        exit_code, timestamp = res
        self._handle_worker_exit(slot_info, exit_code, timestamp)

    thread = threading.Thread(target=run_worker)
    thread.daemon = True
    thread.start()
    self._results.expect(thread)

啓動之後,邏輯如下,可以看到,經過 4 步之後,啓動了 count(slot_info)這麼多的 Thread,每個 Thread 之中,有一個 _create_worker_fn 運行在一個 Slot 之上:

+------------------------------+
| ElasticDriver                |                           From Host discovery
|                              |                                   +
|    _start_worker_processes   |                                   |
|                              |                                   |
|    wait_for_available_slots  |                                   |  2
|                              |       wait_for_available_slots    |
|   _update_host_assignments   |                                   |
|                              |                                   v
|          _rendezvous         |                              current_hosts
|                              |                                   +
|         _host_assignments    |                                   |
|                              |       _update_host_assignments    |
|         _world_size          |                                   |  3
|                              |          self._rendezvous.init    |
|         _rank_assignments    |                                   |
|                              |                                   |
|         _create_worker_fn+---------+                             |
|                              |     |                             v
|         _worker_registry     |     |                        pending_slots
|                              |     |                             +
|                              |     |                             |
+------------------------------+     |    _start_worker_processes  | 4
                                     |                             |
                                     |                             |
                                     v                             |
 +-----------------------------------+------------+                |
 | exec_command({horovod_env} {env} {run_command})|                |
 +---+--------------------------------------------+                |
     ^                                                             v
     |                               +-----------------+-----------+------------------+
     |                               |                 |                              |
     |                               |                 |                              |
     |                               |                 v                              v
 +-----------------------------------+---+     +-------+---------+         +----------+-+
 |   |       Thread 1                    |     | Thread 2        |         | Thread n   |
 | +----------------------------------+  |     |                 |         |            |
 | | |               run_worker       |  |     |                 |         |            |
 | | |                                |  |     |   run_worker    | ......  | run_worker |
 | | |                                |  |     |                 |         |            |
 | | +----+ create_worker_fn( slot 1 )|  |     |     slot 2      |         |  slot n    |
 | |                                  |  |     |                 |         |            |
 | +----------------------------------+  |     |                 |         |            |
 +---------------------------------------+     +-----------------+         +------------+

手機如下:

0x05 運行過程

運行過程其實在上面已經提到了,就是在 Thread 之中運行 exec_command({horovod_env} {env} {run_command})這就是調用用戶代碼進行訓練

thread = threading.Thread(target=run_worker)
thread.daemon = True
thread.start()
self._results.expect(thread)

這裏要說明的是 self._results = ResultsRecorder() 。具體在其中記錄每一個運行的 Thread。

class ResultsRecorder(object):
    def __init__(self):
        self._error_message = None
        self._worker_results = {}
        self._worker_threads = queue.Queue()

    def expect(self, worker_thread):
        self._worker_threads.put(worker_thread)

於是我們的邏輯變成如下,self._results 裏面記錄了所有的 Threads:

   +------------------------------+
   | ElasticDriver                |                           From Host discovery
   |                              |                                   +
   |                              |                                   |
   |                              |                                   |
   |    _start_worker_processes   |                                   |  2
   |                              |       wait_for_available_slots    |
   |    wait_for_available_slots  |                                   |
   |                              |                                   v
   |   _update_host_assignments   |                              current_hosts
   |                              |                                   +
   |          _rendezvous         |                                   |
   |                              |       _update_host_assignments    |
   |         _host_assignments    |                                   |  3
   |                              |          self._rendezvous.init    |
   |         _world_size          |                                   |
   |                              |                                   |
   |         _rank_assignments    |                                   |
   |                              |                                   v
   |         _create_worker_fn+---------+                        pending_slots
   |                              |     |                             +
+---------+  _results             |     |                             |
|  |                              |     |    _start_worker_processes  | 4
|  +------------------------------+     |                             |
|                                       |                             |
|                                       v                             |
|   +-----------------------------------+------------+                |
|   | exec_command({horovod_env} {env} {run_command})|                |
|   +---+--------------------------------------------+                |
|       ^                                                             v
|       |                               +-----------------+-----------+------------------+
|       |                               |                 |                              |
|       |                               |                 |                              |
|       |                               |                 v                              v
|   +-----------------------------------+---+     +-------+---------+         +----------+-+
|   |   |       Thread 1                    |     | Thread 2        |         | Thread n   |
|   | +----------------------------------+  |     |                 |         |            |
|   | | |               run_worker       |  |     |                 |         |            |
|   | | |                                |  |     |   run_worker    | ......  | run_worker |
|   | | |                                |  |     |                 |         |            |
|   | | +----+ create_worker_fn( slot 1 )|  |     |     slot 2      |         |  slot n    |
|   | |                                  |  |     |                 |         |            |
|   | +----------------------------------+  |     |                 |         |            |
|   +---------------------------------------+     +--------+--------+         +-------+----+
|                  ^                                       ^                          ^
|                  |                                       |                          |
+------------------+---------------------------------------+--------------------------+

手機如下:

0x06 註冊,結果 & 協調

本節主要對應架構圖中的下面三部分。

此部分的邏輯層次在後續介紹的 容錯機制之上。容錯機制是在 Worker 內部,此部分是在 Worker 之上

6.1 Worker 的邏輯層次

worker 是在 訓練腳本之上的階段,即 Worker 來運行使用 "python train.py" 來運行訓練腳本。

所以,上圖我們簡化如下:

   +------------------------------+
   | ElasticDriver                |                           From Host discovery
   |                              |                                   +
   |     _start_worker_processes  |                                   |
   |                              |                                   |
   |     wait_for_available_slots |                                   |  2
   |                              |       wait_for_available_slots    |
   |    _update_host_assignments  |                                   |
   |                              |                                   v
   |           _rendezvous        |                              current_hosts
   |                              |                                   +
   |          _host_assignments   |                                   |
   |                              |       _update_host_assignments    |
   |          _world_size         |                                   |  3
   |                              |          self._rendezvous.init    |
   |          _rank_assignments   |                                   |
   |                              |                                   |
   |          _create_worker_fn   |                                   |
   |                              |                                   v
   |          _worker_registry    |                              pending_slots
   |                              |                                   +
+---------+  _results             |                                   |
|  |                              |          _start_worker_processes  | 4
|  +------------------------------+                                   |
|                                                                     v
|                                       +-----------------+-----------+------------------+
|                                       |                 |                              |
|                                       |                 |                              |
|                                       |                 v                              v
|   +-----------------------------------+---+     +-------+---------+       +------------+-+
|   | Thread 1                       slot 1 |     | Thread 2        |       | Thread n     |
|   |                                       |     |                 |       |              |
|   |  +---------------------------------+  |     | slot 2          |       | slot n       |
|   |  | worker 1                        |  |     |  +-----------+  | ......|  +---------+ |
|   |  |                                 |  |     |  |worker 2   |  |       |  |worker n | |
|   |  |          Python train.py        |  |     |  |           |  |       |  |         | |
|   |  |                                 |  |     |  |           |  |       |  |         | |
|   |  +---------------------------------+  |     |  +-----------+  |       |  +---------+ |
|   +---------------------------------------+     +-----------------+       +--------------+
|                  ^                                       ^                          ^
|                  |                                       |                          |
+------------------+---------------------------------------+--------------------------+

手機如下:

6.2 worker 運行階段

於是我們提出了新問題如下:

  • worker 的運行,怎麼纔算一個階段?一共有幾個階段(狀態)?
  • Driver 根據什麼特徵來記錄 Worker 的運行結果?

從源碼中我們可以看到,Worker 有三個狀態如下:

READY = 'READY'
SUCCESS = 'SUCCESS'
FAILURE = 'FAILURE'

所以,Worker 可以分爲四個階段,RUNNING 是我自己加上去的,就是運行訓練腳本這個過程,官方沒有這個狀態,但是我覺得這樣應該更清晰。而 SUCCESS 和 FAILURE 就是腳本運行成功與否。

          Worker
            +
            |
            |
            |
            v
          READY
            +
            |
            |
            v
   +--------+---------+
   | RUNNINZG         |
   |                  |
   |  Python train.py |
   |                  |
   +--------+---------+
            |
            |
            |
            v
   +--------+--------+
   |                 |
   |                 |
   v                 v

SUCCESS           FAILURE

我們接下來看看運行階段。

6.2.1 進入 C++ 世界

當 Driver 初始化 / resume(比如收到了HostsUpdateInterrupt ) 時候,就會調用到 hvd.init。

進入 hvd.init 有幾個調用途徑(按照下面1,2,3 順序邏輯進行):

  1. 依靠 WorkerStateRegistry . _barrier : 作用是當所有 worker 完成之後,會進一步處理。有三個途徑會觸發這個barrier:
    • start 一個 worker,worker會 hvd.init,進而調用了 gloo in c++,進而 聯繫 rendezvous,rendezvous 通知 driver,進而在 WorkerStateRegistry 設置自己的狀態是 READY,如果達到了 min_np,則會觸發了 _barrier途徑 1);
    • 新發現一個host,從而導致觸發一個 HostsUpdateInterrupt,worker 捕獲這個異常之後,進而會 reset,reset 時候會調用 hvd.init,進而和上述一樣,最終觸發_barrier途徑 2);
    • worker失敗,會調用會 _handle_worker_exit,進而在 WorkerStateRegistry 設置自己的狀態是 FAILURE,會觸發了 _barrier途徑 3);
  2. _barrier 繼續執行時候,會調用構建時候設置的 handler,即 _action 函數,進而調用到了 _on_workers_recorded,最終調用到了 self._driver.resume()
  3. resume 函數 會self._activate_workers(self._min_np),最終就是重新生成(也許是部分,根據 pending_slots決定)worker

6.2.2 構建 Gloo

前文我們提到過,在 python 世界調用 hvd.init 的時候,會進入到 C++世界,這時候如果編譯了 GLOO,就建立了一個 GlooContext。

namespace {

// All the Horovod state that must be stored globally per-process.
HorovodGlobalState horovod_global;

#if HAVE_GLOO
GlooContext gloo_context;
#endif

GlooContext 就得到了一個與 RendezvousServer 通訊的接口。

6.2.2.1 去 Rendezvous 獲取信息

從 GlooContext::Initialize 可以知道,需要獲取大量配置信息,其中最重要的就是 rank 等信息。

這些信息是在 RendezvousServer 設置存儲的。所以 GlooContext 會去 RendezvousServer 進行 http 交互,下面代碼還是比較清晰易懂的。

#define HOROVOD_GLOO_GET_RANK_AND_SIZE "rank_and_size"

void GlooContext::Initialize(const std::string& gloo_iface) {
  // Create a device for communication
  // TODO(sihan): Add support for multiple interfaces:
  //  https://github.com/facebookincubator/gloo/issues/190
  attr device_attr;
  device_attr.iface = gloo_iface;

  device_attr.ai_family = AF_UNSPEC;
  auto dev = CreateDevice(device_attr);
  auto timeout = GetTimeoutFromEnv();

  auto host_env = std::getenv(HOROVOD_HOSTNAME);
  std::string hostname = host_env != nullptr ? std::string(host_env) : std::string("localhost");

  int rank = GetIntEnvOrDefault(HOROVOD_RANK, 0);
  int size = GetIntEnvOrDefault(HOROVOD_SIZE, 1);
  int local_rank = GetIntEnvOrDefault(HOROVOD_LOCAL_RANK, 0);
  int local_size = GetIntEnvOrDefault(HOROVOD_LOCAL_SIZE, 1);
  int cross_rank = GetIntEnvOrDefault(HOROVOD_CROSS_RANK, 0);
  int cross_size = GetIntEnvOrDefault(HOROVOD_CROSS_SIZE, 1);

  auto rendezvous_addr_env = std::getenv(HOROVOD_GLOO_RENDEZVOUS_ADDR);
  auto rendezvous_port = GetIntEnvOrDefault(HOROVOD_GLOO_RENDEZVOUS_PORT, -1);

  bool elastic = GetBoolEnvOrDefault(HOROVOD_ELASTIC, false);
  if (elastic && reset_) {
    std::string server_addr = rendezvous_addr_env;
    std::string scope = HOROVOD_GLOO_GET_RANK_AND_SIZE;
    HTTPStore init_store(server_addr, rendezvous_port, scope, rank);

    auto key = hostname + ":" + std::to_string(local_rank);
    std::vector<char> result = init_store.get(key);
    std::string s(result.begin(), result.end());
    std::stringstream ss(s);

    int last_rank = rank;
    int last_size = size;
    int last_local_rank = local_rank;
    int last_local_size = local_size;
    int last_cross_rank = cross_rank;
    int last_cross_size = cross_size;

    rank = ParseNextInt(ss);

    size = ParseNextInt(ss);
    local_rank = ParseNextInt(ss);
    local_size = ParseNextInt(ss);
    cross_rank = ParseNextInt(ss);
    cross_size = ParseNextInt(ss);

    SetEnv(HOROVOD_RANK, std::to_string(rank).c_str());
    SetEnv(HOROVOD_SIZE, std::to_string(size).c_str());
    SetEnv(HOROVOD_LOCAL_RANK, std::to_string(local_rank).c_str());
    SetEnv(HOROVOD_LOCAL_SIZE, std::to_string(local_size).c_str());
    SetEnv(HOROVOD_CROSS_RANK, std::to_string(cross_rank).c_str());
    SetEnv(HOROVOD_CROSS_SIZE, std::to_string(cross_size).c_str());
  }

  ctx = Rendezvous(HOROVOD_GLOO_GLOBAL_PREFIX,
                   rendezvous_addr_env, rendezvous_port,
                   rank, size, dev, timeout);
  local_ctx = Rendezvous(HOROVOD_GLOO_LOCAL_PREFIX + hostname,
                         rendezvous_addr_env, rendezvous_port,
                         local_rank, local_size, dev, timeout);
  cross_ctx = Rendezvous(HOROVOD_GLOO_CROSS_PREFIX + std::to_string(local_rank),
                         rendezvous_addr_env, rendezvous_port,
                         cross_rank, cross_size, dev, timeout);
}
6.2.2.2 RendezvousServer

RendezvousServer 對外提供了 GET 方法。

# GET methods
GET_RANK_AND_SIZE = 'rank_and_size'

def _get_value(self, scope, key):
    if scope == GET_RANK_AND_SIZE:
        host, local_rank = key.split(':')
        return self._get_rank_and_size(host, int(local_rank))

    return super(RendezvousHandler, self)._get_value(scope, key)

ElasticRendezvousHandler 是 RendezvousServer 的響應handler,其中 ElasticRendezvousHandler._get_rank_and_size 函數是:

def _get_rank_and_size(self, host, local_rank):
    driver.record_ready(host, local_rank)
    slot_info = driver.get_slot_info(host, local_rank)
    return slot_info.to_response_string().encode('ascii')

這裏會調用到 driver.record_ready,就是通知 Driver,現在有一個 worker 是 READY 狀態。

def record_ready(self, host, slot):
    self._worker_registry.record_ready(host, slot)

6.2.3 進入READY 狀態

當調用 hvd.init -----> GlooContext 建立 之後,會與RendezvousServer 通訊,這之後 Worker 就進入到 READY 狀態。

我們需要繼續深化下,看看一個 worker 從開始運行到 READY 狀態 之間都發生了什麼。

  1. Worker 開始調用 python train.py;

  2. 在 train.py 之中,調用 hvd.init(),此方法會深入到 C++ 世界,從而生成了 GlooContext;

  3. GlooContext 之中,會從環境變量之中得到 Rendezvous Server 的ip, port,進而調用 init_store 生成一個 HTTPStore;

  4. 調用 init_store.get(hostname + ":" + std::to_string(local_rank)) 向 Rendezvous Server 發送請求,要求獲得本worker 的 rank 對應的 各種配置(local_rank, cross_rank...,因爲 Rendezvous Server 可能會重新初始化從而重新分配);

  5. ElasticRendezvousHandler 是 響應函數,其中會 調用 driver.record_ready(host, local_rank) 從而在 WorkerStateRegistry 的 READY 字典中記錄下來,worker 2 已經是 READY 了。

  6. 會調用 driver.get_slot_info(host, local_rank) 從 driver 獲得 slot info;

  7. 此時,Worker 的狀態就是 READY(其實 Worker 本身沒有這個狀態,只是 WorkerStateRegistry 有這個狀態);

  8. ElasticRendezvousHandler 會返回 slot info 到 worker 的 C++ 世界;

  9. 在 worker 的 C++ 世界 之中繼續執行,把 slot info 返回給 GlooContext,進行各種設置;

具體邏輯圖如下:

Python                                                                        +                                                C++
                                                                              |
                +--------------------+      +----------------------+          |
                | ElasticDriver      |      | RendezvousServer     |          |
+------------>  |                    |      |                      |          |
|               |   _rendezvous +---------> |       handler_cls    |          |
|               |                    |      |             +        |          |
|               |                    |      +----------------------+          |
|         +-------+ _worker_registry |                    |                   |
|         |     |                    |                    v                   |
|         |     +--------+-----------+    +---------------+---------------+   |
|         |              ^                | ElasticRendezvousHandler      |
|         |              |                |  +---------------------+      |   |     4     get(rank)              +---------------+
|         |              |                |  |  _get_rank_and_size +<------------------------------------------+ | HTTPStore     |
|         |            6 | get_slot_info  |  |                     |      |   |                                  |               |
|         |              +-------------------+  driver  +------------------------------------------------------> |               |
|         |                               |  |                     |      |   | 8 slot_info (rank,local_rank...) +--+------------+
|         |        +-------------------------+                     |      |   |                                     |
|         |        |  5  record_ready     |  +---------------------+      |   |                                     |     ^
|         |        |                      +-------------------------------+   |                                     |     |
|         v        v                    ----------------------------------+   |                                     |     |
|   +-----+--------+---------------+    |               Worker            |   |                                     |     |
|   | WorkerStateRegistry          |    |                 +               |   |                                     |     |
|   |                              |    |                 |               |   |                                   9 |     | 3
|   |                              |    |                 |  1            |   |                                     |     |
|   |      _host_manager           |    |                 v               |   |                                     |     |
|   |                              |    |           Python train.py       |   |                                     |     |
+--------+ _driver                 |    |                 +               |   |                       +----------------------------+
    |                              |    |                 |               |   |                       | GlooContext |     |        |
    |      _barrier                |    |                 |               |   |                       |             |     |        |
    |                              |    |                 v               |   |  create an instance   |             |     +        |
    |   +---------------------+    |    |             hvd.init() +----------------------------------> |             | init_store   |
    |   |  _workers           |    |    |                                 |   |          2            |             v              |
    |   |                     |    | 7  |                                 |   |                       |  gloo::Context(rank)       |
    |   |   READY[]+----------------------------------> READY             |   |                       |                            |
    |   |   SUCCESS[]         |    |    |                 +               |   |                       |  gloo::Context(local_rank) |
    |   |   FAILURE[]         |    |    |                 |               |   |                       |                            |
    |   |                     |    |              run_fn(train_function)  |   |                       |  gloo::Context(cross_rank) |
    |   +---------------------+    |    |                 |               |   |                       |                            |
    +------------------------------+    |                 |               |   |                       +----------------------------+
                                        |                 v               |   |
                                        |        +--------+--------+      |   |
                                        |        |                 |      |   |
                                        |        |                 |      |   |
                                        |        v                 v      |   |
                                        |                                 |   |
                                        |     SUCCESS           FAILURE   |   |
                                        |                                 |   |
                                        +---------------------------------+   +

手機如下:

至此,Worker 就可以開始運行了。

6.3 WorkerStateRegistry

WorkerStateRegistry 的作用是 註冊運行結果,然後進行對應協調

其主要成員變量是:

  • _driver :用來聯繫 Driver,因爲會調用 driver 來做處理;

  • _host_manager :用來發現 host;

  • _workers : 紀錄每種狀態的worker,狀態包括:'READY','SUCCESS','FAILURE';

    • def count(self, state):
          return len(self._workers[state])
      
  • _states : 紀錄 worker 的狀態;

  • _barrier : 作用是當所有 worker 完成之後,會進一步處理;

具體定義如下:

class WorkerStateRegistry(object):
    def __init__(self, driver, host_manager, reset_limit=None, verbose=False):
        self._driver = driver
        self._host_manager = host_manager
        self._reset_limit = reset_limit
        self._reset_count = 0
        self._lock = threading.Lock()
        self._states = {}
        self._workers = defaultdict(set)
        self._barrier = None
        self._rendezvous_id = 0
        self._verbose = verbose
        self._size = 0

6.3.1 初始化

WorkerStateRegistry 在 Driver 之中 進行初始化,並且把自己設置爲 Driver 的一個成員變量,這樣 Driver 就可以方便調用:

self._worker_registry = WorkerStateRegistry(self, self._host_manager, reset_limit=reset_limit)

6.3.2 啓動

在 master 啓動所有 worker 之前,會調用 reset。

def _activate_workers(self, min_np):
    current_hosts = self.wait_for_available_slots(min_np)
    pending_slots = self._update_host_assignments(current_hosts)
    self._worker_registry.reset(self.world_size()) # 這裏reset
    self._start_worker_processes(pending_slots)

reset 函數中有複雜邏輯。

有兩個問題:

  • 爲什麼要有 _barrier?原因是:大部分機器學習算法機制是需要當所有 worker(或者若干worker) 完成之後,纔會進一步處理,所以需要等待
    • 這裏 barrier 的參數 parties 具體數值是 self.world_size(),就是說,只有等到barrier 內部計數達到 self.world_size() 時候,就會激發 self._action 函數。
    • 每個worker 結束時候,都會調用到 _handle_worker_exit,最終會 self._barrier.wait()
    • 這樣,當所有 worker 都結束時候,barrier 會激發 self._action 函數。
  • 設置的 _action 起到什麼作用?其作用是:根據本次訓練結果,進一步控制,決定下一步動作

代碼如下:

def reset(self, size):
    with self._lock:
        self._states.clear()
        self._workers.clear()
        self._barrier = threading.Barrier(parties=size, action=self._action)
        self._rendezvous_id += 1
        self._size = size

6.3.3 worker 結束

當 worker 結束時候,會回到 Driver 設置的 _handle_worker_exit。根據 exit_code 來決定是調用 success 函數還是 failure 函數。

def  _handle_worker_exit(self, slot_info, exit_code, timestamp):
    if not self.has_rank_assignment(slot_info.hostname, slot_info.local_rank):
        return

    if exit_code == 0:
        rendezvous_id = self._worker_registry.record_success(slot_info.hostname, slot_info.local_rank)
    else:
        rendezvous_id = self._worker_registry.record_failure(slot_info.hostname, slot_info.local_rank)

    if self.finished() and self._worker_registry.last_rendezvous() == rendezvous_id:
        name = '{}[{}]'.format(slot_info.hostname, slot_info.local_rank)
        self._results.add_result(name, (exit_code, timestamp))

從而調用到 WorkerStateRegistry 之中。

def record_ready(self, host, slot):
    return self._record_state(host, slot, READY)

def record_success(self, host, slot):
    return self._record_state(host, slot, SUCCESS)

def record_failure(self, host, slot):
    return self._record_state(host, slot, FAILURE)

_record_state 函數會使用 self._workers[state].add(key) 來紀錄狀態,並且調用 _wait。

def _record_state(self, host, slot, state):
    if self._driver.finished():
        return self._rendezvous_id
    if self._host_manager.is_blacklisted(host):
        return self._rendezvous_id

    key = (host, slot)
    with self._lock:
        if key in self._states:
            if state == FAILURE:
                self._barrier.reset()

        if key not in self._states or state == FAILURE:
            self._states[key] = state
            self._workers[state].add(key)

        rendezvous_id = self._rendezvous_id

    rendezvous_id = self._wait(key, state, rendezvous_id)
    return rendezvous_id

_wait 會並且調用 self._barrier.wait() 來等待,這是爲了等待其他 worker 的信息,最後一起處理

def _wait(self, key, state, rendezvous_id):
    while True:
        try:
            self._barrier.wait()
            return rendezvous_id
        except threading.BrokenBarrierError:
            if self._barrier.broken:
                # Timeout or other non-recoverable error, so exit
                raise

            # Barrier has been reset
            with self._lock:
                # Check to make sure the reset was not caused by a change of state for this key
                rendezvous_id = self._rendezvous_id
                saved_state = self._states.get(key, state)
                if saved_state != state:
                    # This worker changed its state, so do not attempt to wait again to avoid double-counting
                    raise RuntimeError('State {} overridden by {}'.format(state, saved_state))

6.3.4 進一步控制

_action 函數會在所有worker 結束之後,進行判斷,控制。

def _action(self):
    self._on_workers_recorded()

_on_workers_recorded 函數會完成控制邏輯。

  • 判斷是否有一個 worker 成功,如果有一個 worker 成功了,就關閉其他process,結束訓練;因爲此時所有 worker 都已經運行結束,所以只要有一個 worker 成功,就可以跳出循環;
  • 如果所有的 worker 都失敗了,就結束訓練;
  • 把失敗的 worker 紀錄到 黑名單;
  • 如果所有的 host 都在黑名單,則結束訓練;
  • 如果已經到了最大重試數目,則結束訓練;
  • 否則調用 _driver.resume() 重啓訓練,因爲已經 commit 了,所以會自動恢復訓練;

具體代碼如下:

def _on_workers_recorded(self):
    # Check for success state, if any process succeeded, shutdown all other processes
    if self.count(SUCCESS) > 0:
        self._driver.stop()
        return

    # Check that all processes failed, indicating that processing should stop
    if self.count(FAILURE) == self._size:
        self._driver.stop()
        return

    # Check for failures, and add them to the blacklisted hosts list
    failures = self.get(FAILURE)
    for host, slot in failures:
        self._host_manager.blacklist(host)

    # If every active host is blacklisted, then treat this as job failure
    if all([self._host_manager.is_blacklisted(host) for host, slot in self.get_recorded_slots()]):
        self._driver.stop()
        return

    # Check that we have already reset the maximum number of allowed times
    if self._reset_limit is not None and self._reset_count >= self._reset_limit:
      
self._driver.stop(error_message=constants.RESET_LIMIT_EXCEEDED_MESSAGE.format(self._reset_limit))
        return

    try:
        self._reset_count += 1
        self._driver.resume()
    except Exception:
        self._driver.stop()

6.4 Driver.resume 場景

resume 的作用 就是 所有都重新來過。

def resume(self):
    self._activate_workers(self._min_np)

我們之前分析的場景是:一個 worker 從開始運行到 READY 狀態 之間都發生了什麼。

現在,我們加上一個情形:就是當 Driver 在 resume 的時候,發現居然有新節點,隨即啓動了一個新 worker 3

  1. Worker 2 開始調用 python train.py;

  2. 在 train.py 之中,調用 hvd.init(),此方法會深入到 C++ 世界,從而生成了 GlooContext;

  3. GlooContext 之中,會從環境變量之中得到 Rendezvous Server 的ip, port,進而調用 init_store 生成一個 HTTPStore;

  4. 調用 init_store.get(hostname + ":" + std::to_string(local_rank)) 向 Rendezvous Server 發送請求,要求獲得本worker 的 rank 對應的 各種配置(local_rank, cross_rank...,因爲 Rendezvous Server 可能會重新初始化從而重新分配);

  5. ElasticRendezvousHandler 是 響應函數,其中會 調用 driver.record_ready(host, local_rank) 從而在 WorkerStateRegistry 的 READY 字典中記錄下來,worker 2 已經是 READY 了。

  6. 會調用 driver.get_slot_info(host, local_rank) 從 driver 獲得 slot info;

  7. 把 slot info 返回給 worker 中的 Http_store;

  8. 在 worker 2 之中繼續執行,把 slot info 返回給 GlooContext,進行各種設置;

  9. 我們接着 第 5 項繼續進行;record_ready 之中 會調用 rendezvous_id = self._wait(key, state, rendezvous_id) 來在 WorkerStateRegistry . _barrier 之上等待; _barrier 的類型是threading.Barrier(parties=size, action=self._action)

  10. 如果 READY 的 worker 數目達到了 Horovod 設置的 min-np,就是可以啓動的最小 worker 數目, _barrier 就結束使命,就 broken,繼續執行;

  11. _barrier 繼續執行時候,會調用構建時候設置的 handler,即 _action 函數,進而調用到了 _on_workers_recorded,最終調用到了 self._driver.resume()

  12. ElasticDriver 的 resume 函數調用到了 _activate_workers,其定義如下,可以看到,如果此時 discovery 腳本已經發現了新節點,進而返回了 pending_slotspending_slots 就是可以在這些 slot 之上啓動新 worker 的,於是 就會 調用 _start_worker_processes

    1. def _activate_workers(self, min_np):
          current_hosts = self.wait_for_available_slots(min_np)
          pending_slots = self._update_host_assignments(current_hosts)
          self._worker_registry.reset(self.world_size())
          self._start_worker_processes(pending_slots)
      
  13. _start_worker_processes 會開啓一個新的 worker : Worker 3;

  14. Worker 3 也執行 python train.py,至此,新worker啓動完畢;

  15. 回到worker 2,如果訓練結束,則會依據訓練結果,返回 SUCCESS 或者 FAILURE 到 Driver ;

  16. Driver 會調用 _handle_worker_exit 對訓練結果進行處理;

至此,新的邏輯完成。

       +--------------------------------------------+
       |      +------------------------------+      |        +----------------------+
       |      | WorkerStateRegistry          |      | 15     |  ElasticDriver       |          +----------------------+
       |      |                              |      +----->  |     _rendezvous +-------------> | RendezvousServer     |
       |      |     _driver +-----------+    |               |                      |          |                      |
       |      |        ^                |    +<----------------+ _worker_registry   | 12 init  |       handler_cls    |
       |      |        |                |    |               |                      +--------> |             +        |
       |      |        | 10 broken      |    |   11  resume  |                      |          +----------------------+
       |      |        |                +----------------------> _activate_workers  |                        |
       |      |        |                     |      +----------+_handle_worker_exit |                        v
       |      |        +          9 wait     |      |        |                      |        +---------------+---------------+
       |      |     _barrier <---------------------------+   +--+-----------+-------+        | ElasticRendezvousHandler      |
       |      |                              |      |    |      |           ^                |  +---------------------+      |         4     get(rank)
       |      | +--------------------------+ |      |    |      |           |                |  |  _get_rank_and_size +<-----------------------------------------------------+
       |      | | _workers                 | +<-----+    |      |         6 | get_slot_info  |  |                     |      |                                               |
       |      | |                          | |  16       |      |           +-------------------+  driver  +-----------------------------------------------------------+     |
       |      | |SUCCESS[]                 | |           |      |                            |  |                     |      |     7 slot_info (rank,local_rank...)    |     |
       |      | |READY[(host 1, slot1)] <----------------+--------------------------------------+                     |      |                                         |     |
       |      | |FAILURE[]                 | |                  |     5 record_ready         |  +---------------------+      |                                         |     |
       |      | +--------------------------+ |                  |                            +-------------------------------+                                         |     |
       +      +------------------------------+                  |                                                                                                      |     |
_handle_worker_exit                                             +          Host 1                                                                                      +     +
 +-----+-----------------------------------------------------+ rsh +---------+----------------------------------------------------------------------------------------+ Socket +------+
       |                                                        +   Host 3   |    Host 2                                                              Host 2           +     +
       |                                                        |            |             +-----------------+---------------+  Python +   C++                         |     |
       |                                                        +            |             |  Worker 2       | 1     slot 2  |         |                               v     |
       |                          13  _start_worker_processes(pending_slots) |             |                 v       host 2  |         |                                     |
       |                                                        +            |             |           Python train.py       |         |                            +--------+---+
       |                                                        |            |             |                 +               |         |                            | HTTPStore  |
       |                                                        v            |             |                 |               |         |                            +---+-----+--+
       |                                                                     |             |                 v               |         |                                |     ^
       |                                 +---------------------------+       |             |             hvd.init()          |         |   2                          8 |     | 3
       |                                 | Worker 3                  |       |             |                 +      +------------------+-------+                        |     |
       |                                 |           +       slot 3  |       |             |                 |               |                 |          +----------------------------+
       |                                 |           |       host 3  |       |             |                 +               |         +       |          | GlooContext |     |        |
       |                                 |       14  |               |       |             |               READY             |         |       |          |             |     |        |
       |                                 |           v               |       |             |                 +               |         |       |          |             |     +        |
       |                                 |                           |       |             |                 |               |         |       |          |             | init_store   |
       |                                 |      Python train.py      |       |             |                 v               |         |       |          |             v              |
       |                                 |                           |       |             |        hvd.elastic.run(state)   |         |       |          |  gloo::Context(rank)       |
       |                                 |                           |       |             |                 +               |         |       |          |                            |
       |                                 +---------------------------+       |             |                 +               |         |       +--------> |  gloo::Context(local_rank) |
       |                                                                     |             |         run_fn(train_function)  |         |                  |                            |
       |                                                                     |             |                 +               |         |                  |  gloo::Context(cross_rank) |
       |                                                                     |             |                 v               |         |                  |                            |
       |                                                                     |             |        +--------+--------+      |         |                  +----------------------------+
       |                                                                     |             |        |                 |      |         |
       |                                                                     |             |        v                 v      |         +
       |                                                                     |             |     SUCCESS           FAILURE   |
       |                                                                     +             |         +                 +     |
       |                                                                                   |         |                 |     |
       |                                                                                   +---------------------------------+
       |                                                                                             |                 |
       |                                                                                             |                 |
       +--------------------------------------------------------------------------<------------------+-----------------+

手機如下:

我們可以擴大部分細節看看,就是Driver + 一個 Worker,這樣看得更清楚:

                                                                                                    +
+------------------------------+           +----------------------+                                 |
| WorkerStateRegistry          |           |  ElasticDriver       |                        Host 1   |  Host 2
|                              |10 resume  |                      |                                 |                  1
|      _driver +-------------------------->+      start/resume    +-----------------------------------------------------------+
|                              |           |                      |               _activate_workers |                         |
|        ^                     |           |                      | <-------------+                 |                         |
|        |                     |           +---+---+--------------+ get_slot_info |                 |       +---------------------------------+
|      9 | broken              |               |   |                              |                 |       |  Worker 2       |       slot 2  |
|        |                     |               |   |      ^                       |                 |       |                 v       host 2  |
|        |                     |               |   |      |                       |                 |       |           Python train.py       |
|        +                     |               |   |      |            +----------+-----------+     |       |                 +               |
|     _barrier                 |               |   |      |            | RendezvousServer     |     |       |                 |               |
|        ^                     |               |   |      |            |                      |     |   2   |                 v               |
|      8 | wait                |               |   |      |            |   _get_rank_and_size | <---------------------+   hvd.init()          |
|        |                     |               |   |      |            |                      |     |       |                 +               |
|        |                     |               |   |      |            +-----------+----------+     |       |                 |               |
| +------+-------------------+ |               |   |      |                        |                |       |                 +               |
| |_workers                  | |               |   |      |                        |                |       |               READY             |
| |                          | |               |   |      |                        |                |       |                 +               |
| |SUCCESS[]    <------------------------------+   |      |                        |                |       |                 |               |
| |                          | | 6 record_success  |      |                        |                |       |                 v               |
| |                          | |                   |      |                        |                |       |        hvd.elastic.run(state)   |
| |READY[(host 1, slot1)] <--------------------------------------------------------+                |       |                 +               |
| |                          | |                   |      |          3 record_ready                 |       |                 +               |
| |FAILURE[]   <-----------------------------------+      |                                         |       |         run_fn(train_function)  |
| |                          | | 7 record_failure         |                                         |       |                 +               |
| +--------------------------+ |                          |                                         |       |                 v               |
+------------------------------+                          |                                         |       |        +--------+--------+      |
                                                          |                                         |       |        |                 |      |
                                                          |                                         |       |        v                 v      |
                                                          |                                         |       |     SUCCESS           FAILURE   |
                                                          |                                         |       |         +                 +     |
                                                          |                                         |       |         |                 |     |
                                                          |                                         |       +---------------------------------+
                                                          |                                         |                 |                 |
                                                          | _handle_worker_exit                     |               4 |                5|
                                                          |                                         |                 |                 |
                                                          +-----------------------------------------------------<-----+-----------------+
                                                                                                    |
                                                                                                    +

手機如下:

至此,worker部分分析完畢,下一篇我們看看如何處理錯誤。

0xEE 個人信息

★★★★★★關於生活和技術的思考★★★★★★

微信公衆賬號:羅西的思考

如果您想及時得到個人撰寫文章的消息推送,或者想看看個人推薦的技術資料,敬請關注。

在這裏插入圖片描述

0xFF 參考

PaddlePaddle Fluid:彈性深度學習在Kubernetes中的實踐

Horovod 彈性訓練

ElasticDL調用 Horovod 在Kubernetes上實現彈性 AllReduce(一)

Kubernetes-native 彈性分佈式深度學習系統

雲原生的彈性 AI 訓練系列之一:基於 AllReduce 的彈性分佈式訓練實踐

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章