DGL Operator:基于AWS DGL和Kubernetes的图神经网络训练控制器

简介:DGL Operator是由360智能工程部 AI平台团队开发并开源的一款基于Kubernetes云原生技术栈和AWS DGL图神经网络算法框架的训练控制器。目前已开源在Github:https://github.com/Qihoo360/dgl-operator

PS:,"360技术"是360团队的线下技术聚合平台,致力于提供有价值的技术干货,点关注哦!

名词解释

  • 工作负载(Overload):该名词为逻辑概念,等效应用于Kubernetes中一个Pod的物理概念。工作负载意指实际进行工作产出(即切图与训练)的单位。

  • 训练任务(Job):逻辑上由一组工作负载(或物理上由一组Pod)所承担的,包含切图的宏观DGL训练生命周期。

  • Pod:Kubernetes术语,最小的调度单位,可包含多个容器。

  • initContainer:Kubernetes术语,一个Pod可以存在多个有限生命周期的容器,这些容器必须在主容器之前成功执行完毕,且必须按顺序执行。一般情况下,initContainer多用于执行主容器的前置依赖工作。

  • Worker Pod:在Kubernetes中进行分布式训练时,实际承担训练工作负载的Pod。

  • Partitioner Pod:在Kubernetes中进行切图时,实际承担切图工作负载的Pod。

  • Launcher Pod:在Kubernetes中控制整个DGL训练生命周期的Pod。其中该Pod包含多个initContainer容器,主容器负责对多个Worker Pod下发训练命令。相比于Worker Pod,该Pod为轻量级资源,不需求GPU、高CPU或高内存资源。dglrun:在容器内,控制DGL切图、训练等环节的工作流脚本文件。

  • ipconfig:DGL分布式训练时,供给Launcher Pod容器内dglrun的配置文件,以记录各个Worker的相关参数,按“Pod名字IP地址 端口 GPU需求个数”格式排列。

  • kubexec.sh:在Kubernetes中远程执行命令时的帮助脚本。

  • 单机切图:使用DGL API在一个Partitioner Pod内进行全图切分。

  • 分布式切图:使用实验性的ParMETIS在多个Worker Pod内进行分布式图切分。

1

背景

近两年来,图神经网络(GNN)在搜索、推荐、知识图谱等领域取得了卓有成效的进展。但随之而来的问题是如何快速构建GNN模型。TensorFlow、PyTorch、MXNet等深度学习框架在CNN和RNN领域提供了大量开箱即用的API,但在构建GNN网络方面却捉襟见肘。在此背景下,纽约大学和亚马逊联合开发了DGL框架。DGL在TensorFlow、PyTorch、MXNet框架之上,定制并封装了一系列针对图的API,并做了大量的性能优化,使得开箱即用的GNN构建成为可能。

 

然而,在工业场景下,算法工程师或数据科学家们基于DGL开发和构建GNN模型时,常常需要处理数千万甚至数十亿个节点或边的大图。所以2020年亚马逊推出了DGL分布式训练模式。

 

DGL原生提供的分布式策略为,在多台物理机或VMs的环境下,采用完全分布式的方法,将数据和计算同时分布在一组计算资源中。DGL会将一张图切分为多张子图,集群中的每台机器各自负责一张子图,执行并行化计算。开发者需要在集群所有机器上运行相同的DGL训练脚本,并在同样的机器上运行服务器以将分区数据提供给训练器。

 

DGL分布式训练的挑战

DGL的分布式训练模式在实际应用中,仍面临以下挑战:

1. 需要提前启动好多台物理机或VMs环境;手动收集IP、端口等信息,记录于ipconfig内。

2. 多台物理机或VMs的环境中,需要事先设置好ssh免密互信。

3. 大图划分为多张子图的过程中,需要在主节点手动触发脚本命令;并且如果想尝试分布式切分子图,则需要自行安装ParMETIS的解决方案组件。

4. 分布式训练过程中,也需要在主节点手动触发脚本命令,无法自动化流程。

5. 训练完成后,小部分资源需要手动释放;多台物理机或VMs需要手动划归其他业务使用,无法做到自动化释放资源。

 

DGL Operator的解决方案

为了解决DGL原生分布式模式的挑战,我们基于Kubernetes生态实现了DGL Operator。它依托Kubernetes控制器设计理念,托管整个DGL训练的生命周期及相关配置文件。

 

当用户提出DGL分布式训练请求时,DGLOperator自动调度符合要求的宿主机环境,启动包含镜像的Pod;然后根据预先定义好的ipconfig文件结构,自动生成训练时需要的配置文件;之后按照DGL工作流的逻辑定义(即dglrun)进行大图切分,再自动触发无主的、完全分布式的训练过程;最后,持久化相应的模型产出等,自动释放计算、存储资源。

2

DGL Operator技术栈

什么是Kubernetes?

Kubernetes是一个对容器化应用程序进行自动化部署、扩缩容及管理的开源系统。它的出现不仅主导了容器化应用程序编排的市场,更改变了以往的开发与运维的方式。在Kubernetes中,每一个工程师都可以定义容器之间或Pod之间的拓扑关系,Pod节点的个数及资源的使用量;并且能够快速实现水平扩容、金丝雀发布等复杂的运维或部署操作。Kubernetes从顶层设计上,充分考量了容器化应用程序的需求,将物理机资源打散、池化、统筹规划,从逻辑上形成了一个“大池子”,动态地供不同的工作负载去执行不同生命周期的任务。

 

什么是Operator?

Kubernetes在默认调度单位下(即Pod),具有完善的无状态应用支持;对于有状态的容器化应用程序,我们可以使用Kubernetes自带的StatefulSet,它能够管理应用的拓扑顺序及存储状态。不过在Kubernetes的生态中,还有一种更加灵活,且编程友好的解决方案 —— Operator。Operator机制能让开发者根据有状态的特点,去定义预期应用规范,以及协调实际应用状态的控制器。

 

控制器核心作用和本质在于纠偏。在整个有状态应用的运行过程中,Kubernetes会不断将该应用的实时状态汇报给控制器,构成反馈回路。当应用的实时状态不满足我们的预期时,控制器将根据用户定义的协调机制进行调整,直到维持预期的状态,并持续至应用结束。

 

DGL Operator是如何使用的?

DGL Operator就是依靠上述Kubernetes及Operator机制,将DGL的图神经网络训练过程抽象成一个有状态的应用程序,然后在训练任务内部按步骤进行切图、分布式训练触发、资源释放等。

 

用户只需要在现有的Kubernetes集群中,定义并提交DGL Operator的部署YAML,即可在Kubernetes集群中新增DGLJob自定义资源,拉取镜像并部署Operator组件。用户只负责书写分布式训练的业务代码,根据DGL Operator提供的Dockerfile逻辑,将DGL分布式工作流脚本(即dglrun)打入镜像;用户通过提交带有参数及镜像路径的YAML文件即可自动启动分布式训练。

3

API示例

容器镜像

用户的镜像需要必须包含DGL及相关依赖,并且需要包含我们提供的dglrun脚本。SSH相关内容是不需要的。

 

提交训练任务的YAML定义

apiVersion: qihoo.net/v1alpha1
kind: DGLJob
metadata:
  name: dgl-graphsage
  namespace: dgl-operator
spec:
  cleanPodPolicy: Running
  partitionMode: DGL-API
  dglReplicaSpecs:
    Launcher:
      replicas: 1
      template:
        spec:
          containers:
          - image: dgloperator/graphsage:v0.1.0
            name: dgl-graphsage
            command:
            - dglrun
            args:
            - --graph-name
            - graphsage
            # partition arguments
            - --partition-entry-point
            - code/load_and_partition_graph.py
            - --num-partitions
            - "2"
            # training arguments
            - --train-entry-point
            - code/train_dist.py
            - --num-epochs
            - "1"
            - --batch-size
            - "1000"
    Worker:
      replicas: 2
      template:
        spec:
          containers:
          - image: dgloperator/graphsage:v0.1.0
            name: dgl-graphsage
  • “cleanPodPolicy”可以选择性地配置为“Running”、“None”、“All”,表明当Pod的生命周期结束时,删除该Pod的策略。

  • “Launcher”和“Worker”的内容遵循“PodTemplateSpec”定义,所以用户可以使用原生的“spec”键值对定义。

  • “partitionMode”可以被配置为“ParMETIS”、“DGL-API”,来表明如何切分图。“DGL-API”使用原生的DGL API“dgl.distributed.partition_graph”;“ParMETIS”使用DGL官方推荐的ParMETIS解决方案。“DGL-API”由于精度更高,被定义为默认切图方式。

提交后产出的Launcher定义

kind: Pod
apiVersion: v1
metadata:
  name: dgl-graphsage-launcher
spec:
  volumes:
    - name: kube-volume
      emptyDir: {}
    - name: dataset-volume
      emptyDir: {}
    - name: config-volume
      configMap:
        name: dgl-graphsage-config
        items:
          - key: kubexec.sh
            path: kubexec.sh
            mode: 365
          - key: hostfile
            path: hostfile
            mode: 292
          - key: partfile
            path: partfile
            mode: 292
  initContainers:
    - name: kubectl-download
      image: 'dgloperator/kubectl-download:v0.1.0'
      volumeMounts:
        - name: kube-volume
          mountPath: /opt/kube
      imagePullPolicy: Always
    - name: watcher-loop-partitioner
      image: 'dgloperator/watcher-loop:v0.1.0'
      env:
        - name: WATCHERFILE
          value: /etc/dgl/partfile
        - name: WATCHERMODE
          value: finished
      volumeMounts:
        - name: config-volume
          mountPath: /etc/dgl
        - name: dataset-volume
          mountPath: /dgl_workspace/dataset
      imagePullPolicy: Always
    - name: watcher-loop-worker
      image: 'dgloperator/watcher-loop:v0.1.0'
      env:
        - name: WATCHERFILE
          value: /etc/dgl/hostfile
        - name: WATCHERMODE
          value: ready
      volumeMounts:
        - name: config-volume
          mountPath: /etc/dgl
      imagePullPolicy: Always
  containers:
    - name: dgl-graphsage
      image: 'dgloperator/graphsage:v0.1.0'
      command:
        - dglrun
      args:
        - '--graph-name'
        - graphsage
        - '--partition-entry-point'
        - code/load_and_partition_graph.py
        - '--num-partitions'
        - '2'
        - '--balance-train'
        - '--balance-edges'
        - '--train-entry-point'
        - code/train_dist.py
        - '--num-epochs'
        - '1'
        - '--batch-size'
        - '1000'
        - '--num-trainers'
        - '1'
        - '--num-samplers'
        - '4'
        - '--num-servers'
        - '1'
      volumeMounts:
        - name: kube-volume
          mountPath: /opt/kube
        - name: config-volume
          mountPath: /etc/dgl
        - name: dataset-volume
          mountPath: /dgl_workspace/dataset
      imagePullPolicy: Always
  restartPolicy: Never


提交后产出的Partitioner定义

kind: Pod
apiVersion: v1
metadata:
  name: dgl-graphsage-partitioner
spec:
  volumes:
    - name: config-volume
      configMap:
        name: dgl-graphsage-config
        items:
          - key: kubexec.sh
            path: kubexec.sh
            mode: 365
          - key: hostfile
            path: hostfile
            mode: 292
          - key: partfile
            path: partfile
            mode: 292
          - key: leadfile
            path: leadfile
            mode: 292
    - name: kube-volume
      emptyDir: {}
  initContainers:
    - name: kubectl-download
      image: 'dgloperator/kubectl-download:v0.1.0'
      volumeMounts:
        - name: kube-volume
          mountPath: /opt/kube
      imagePullPolicy: Always
  containers:
    - name: dgl-graphsage
      image: 'dgloperator/graphsage:v0.1.0'
      env:
        - name: DGL_OPERATOR_PHASE_ENV
          value: Partitioner
      volumeMounts:
        - name: config-volume
          mountPath: /etc/dgl
        - name: kube-volume
          mountPath: /opt/kube
      imagePullPolicy: Always
  restartPolicy: Never

提交后产出的Worker定义

kind: Pod
apiVersion: v1
metadata:
  name: dgl-graphsage-worker-0
spec:
  volumes:
    - name: shm-volume
      emptyDir:
        medium: Memory
        sizeLimit: 10G # sizeLimit will auto set as the half of limited memory amount
    - name: config-volume
      configMap:
        name: dgl-graphsage-config
        items:
          - key: kubexec.sh
            path: kubexec.sh
            mode: 365
          - key: hostfile
            path: hostfile
            mode: 292
          - key: partfile
            path: partfile
            mode: 292
          - key: leadfile
            path: leadfile
            mode: 292
  containers:
    - name: dgl-graphsage
      image: 'dgloperator/graphsage:v0.1.0'
      command:
        - sleep
      args:
        - 365d
      ports:
        - name: dglserver
          containerPort: 30050
          protocol: TCP
      volumeMounts:
        - name: shm-volume
          mountPath: /dev/shm
        - name: config-volume
          mountPath: /etc/dgl
      imagePullPolicy: Always

4

架构及流程设计

DGL Operator每一次接收到YAML请求,并实施的DGL训练任务的生命周期中,实际具有两层工作流,Operator侧工作流和dglrun侧工作流。这两层工作流在不同的层级,控制着内、外两个串行流程,独立工作以完成整个DGL训练的生命周期。本部分将从不同类型的训练任务角度出发,梳理完整的生命周期。

 

准备工作:切图

DGL内置单机切图的API,特点是更加精准、训练效果更好,但负责切图的工作负载需要将全图载入内存,容易受到硬件条件的限制。该模式适合中小规模的图或数据集。

 

DGL官方也推荐使用ParMETIS,进行分布式切图,无需担心单机工作负载的内存瓶颈,将工作量分担给多个worker。但相对于统一的单机切图模式,分布式切图会丢失一定的精准度。该模式适合大或超大规模的图或数据集。

 

准备工作:DGLJob

其中,跟踪整个DGL分布式训练的实时状态,多个工作负载的实时状态,以及记录预期状态的数据结构,均存储于自定义资源DGLJob中。用户提交一个YAML,便会在Kubernetes后台生成一条DGLJob数据。

 

生命周期:单机切图与分布式训练

  • Operator侧工作流

  1. 创建ConfigMap,其包含有“kubexec.sh”与“ipconfig”。

  2. 创建RBAC资源,包括“Role”、“ServiceAccount”及“RoleBinding”,从而具备远程执行的条件。

  3. 创建隶属于Launcher Pod,且名为“kubectl-download”的initContainer容器。该容器负责下载“kubectl”二进制文件,并将其移动到“emptyDir”卷内。

  4. 创建一个Partitioner Pod,进行单机切图。

  5. 创建隶属于Launcher Pod,且名为“watcher-loop-partitioner”的initContainer容器。该容器负责等待Partitioner Pod的成功结束。

  6. 按需创建多个Worker Pod。

  7. 创建隶属于Launcher Pod,且名为“watcher-loop-worker”的initContainer容器。该容器负责,等待所有Worker Pod创建成功,并准备就绪。

  8. 启动Launcher Pod的主容器,运行dglrun命令及加载相应参数。

  9. 等待该DGLJob成功结束,自动清理全部Worker Pod,释放资源。

  • dglrun侧工作流

  1. Partitioner Pod主容器启动后,

    1. dglrun触发“kubectl cp”命令,将各个子图数据从Partitioner Pod统一转存到Launcher Pod。

    2. dglrun运行用户的切图代码,存储各个子图数据。

  2. Partitioner Pod将会变成完成状态。

  3. Launcher Pod主容器启动后,

    1. 子图数据分发完毕后,dglrun触发分布式训练命令,随即各个Worker Pod开始协同工作,进行训练。

    2. dglrun触发“kubectl cp”命令,将各个子图数据从Launcher Pod分发到相应的Worker Pod内。

  4. 训练完成后,Launcher Pod将会变成完成状态;各个Worker Pod将会被删除。

 

生命周期:分布式切图与分布式训练

  • Operator侧工作流

  1. 创建ConfigMap,其包含有“kubexec.sh”与“ipconfig”。

  2. 创建RBAC资源,包括“Role”、“ServiceAccount”及“RoleBinding”,从而具备远程执行的条件。

  3. 创建隶属于Launcher Pod,且名为“kubectl-download”的initContainer容器。该容器负责下载“kubectl”二进制文件,并将其移动到“emptyDir”卷内。

  4. 按需创建多个Worker Pod。这些Pod既负责切图的工作负载,也负责训练的工作负载。

  5. 创建隶属于Launcher Pod,且名为“watcher-loop-worker”的initContainer容器。该容器负责,等待所有Worker Pod创建成功,并准备就绪。

  6. 启动Launcher Pod的主容器,运行dglrun命令及加载相应参数。其中dglrun既负责下发“ParMETIS”分布式切图命令,也负责下发分布式训练命令。

  7. 等待该DGLJob成功结束,自动清理全部Worker Pod,释放资源。

  • dglrun侧工作流

  1. Launcher Pod主容器启动后,

    1. 分布式切图完毕后,dglrun触发分布式训练命令,随即各个Worker Pod开始协同工作,进行训练。

    2. dglrun触发“ParMETIS”分布式切图命令,在各个Worker Pod内切图。

  2. 训练完成后,Launcher Pod将会变成完成状态;各个Worker Pod将会被删除。

5

总结

自Kubernetes诞生之后,它凭借着基于镜像的复用、隔离的完备支持,系统化的资源定义,以及多样化、易拓展的调度策略,接过了大规模集群编排的大旗;它也使得各大公司利用这项技术,梳理并盘活自身物理级资源与资产,为在其之上研究新一代大规模机器学习系统提供了坚实的基础。DGL Operator在Kubernetes生态之上,通过自动ipconfig生成,到自动切图、自动触发分布式命令到最后的释放相应的计算、存储资源,解决了DGL分布式训练自动化的易用性和复杂性挑战,逐步践行MLOps的理念。

 

在开源社区,Kubeflow是推动“MLon Kubernetes”的先行者和领先者,该社区维护着使用广泛的Operator组件,如TensorFlow Operator(即TF Operator)、PyTorch Operator及MPI Operator等。这些Operator从框架和分布式策略支持的角度,完善了机器学习工作负载的自动化运行,使得用户只需要提交一个包含训练代码镜像地址的YAML文件,即可将自己的分布式模型运行在Kubernetes之中。不仅推广了基于各个机器学习、深度学习框架的Operator理念,也沉淀了许多公共的实践思路。这些产物在一定程度上继续影响着各个框架,挖掘AI软件及硬件们,实现更快的速度、更高的利用率、更强的容错能力。

 

DGL Operator也是在这样的背景下,基于公司内使用大规模图神经网络训练的实际需求,设计开发并正式在公司内上线了DGL Operator。近期,我们开源了这部分工作并把DGL Operator的设计提交到了Kubeflow社区,引起了广泛的讨论。希望能为云原生环境下分布式图神经网络学习提供一个简单易用的工具。

 

我们同样希望广大的开源爱好者尤其是DGL的使用者,可以尝试使用DGL Operator或参与到本项目开发中,为图神经网络的发展做出一些贡献。

最后再宣传一下我们开源的代码仓库,欢迎大家试用:

https://github.com/Qihoo360/dgl-operator

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值