隐语分布式底座: 基于Ray的分布式联邦执行引擎

作者:王清,周爱辉

1. Ray & SecretFlow

1.1 Ray

Ray是由UC Berkeley RISELab发起,蚂蚁集团参与共同建设的新一代高性能分布式计算系统。历经Ray社区和蚂蚁集团不断的共建与打磨,Ray已经发展成为蚂蚁集团的计算基础设施底盘(作为高性能计算的分布式研发框架)。Ray的大量的基础能力比如分布式调用、服务发现、故障恢复、自动扩缩容、无感迁移、云原生等能够被低成本复用,强大的内核能力使得蚂蚁对新的计算引擎的构建变得非常简单,上游计算引擎开发者只需要关注计算pattern内的事情,专注做核心的价值。

目前Ray在蚂蚁集团内部衍生出了多种计算平台,并服务了多种业务。具体见Ray不可挡,扬帆远航

1.2 SecretFlow

SecretFlow(隐语)是一款具有隐私保护能力的数据分析和机器学习的统一框架。SecretFlow支持包括MPC、FL、TEE、HE、DP在内的多种主流隐私计算技术,以期通过良好的分层设计及开箱即用的隐私保护数据分析及机器学习等功能,降低开发者使用门槛。

SecretFlow的架构如下:

1.3 隐语为什么选择Ray

隐语旨在提供隐私保护下的机器学习和数据分析的统一框架,面临的场景天然具备跨机构分布式的特性。作为众多隐私保护技术的提供者,隐语希望提供明密文混合编程的统一视角,这给隐语带来了如下挑战:

  1. 同构和异构逻辑交织

在水平联邦场景下,多个客户端往往是同构计算,而服务端与客户端则属于异构逻辑。在垂直场景下,比如拆分学习,有label方和无label方的逻辑往往是异构的,无label的多个参与方则往往是同构的。在多方安全计算模型下,参与方的逻辑主要是同构的;而同态加密模型下,密钥持有方和计算方逻辑则是异构的。以上场景不一而足,隐语需要有一种简洁高效的表达方式来支撑各种场景下的分布式逻辑表达。

  1. 复杂的计算模型

既有有状态的计算,也有无状态的计算。比如在数据处理中,我们希望数据是无状态、不可更改的;在模型训练中,模型是有状态、可更改的,通过迭代更新内部状态。

为了支持大规模数据和模型,隐语还需要支持多种并行计算模型,包括数据并行、模型并行、混合并行等。

  1. 同步和异步计算

机构内的并行计算比如Parameter Server模式可能是同步计算,而机构之间由于不同机构的计算资源、数据量、模型都可能存在差异,异步计算可以带来效率的提升。为了最大化计算效率,隐语既需要支持同步计算,也需要支持灵活的异步计算。

  1. 如何保护数据安全

除了算法本身的安全设计外,隐语底座也需要加强安全保护,包括但不限于机构间的身份认证、跨机构数据流动的安全性等。

在分布式机器学习领域,TensorFlow、PyTorch、MXNet等框架已经取得了巨大的成功。面临这些挑战,隐语是否可以直接基于这些AI框架进行修改从而站在巨人的肩膀上?答案是不太可行。原因有多方面:

  1. 以上多个框架本身是割裂的,这意味着比如如果选择了TensorFlow,则其他框架无法使用(除非每个框架都改一遍,但是代价太高)。

  2. 这些AI框架在设计之初,面临的问题和隐私保护计算并不大相同,隐私保护计算场景的一个基本设定是原始数据是固定在机构内的不能任意流动,而TensorFlow、PyTorch等设计之初就不考虑数据流动受限的场景,而侵入式去修改这些框架,其工作量不亚于从头实现,而且也难以享受这些框架本身升级迭代的好处。

  3. 最后,面临复杂的同构异构逻辑交织、复杂计算模型、同步异步计算等,这些AI框架本身也不是很胜任。

当Ray进入隐语的视角之后,隐语发现Ray能够比较好的帮助隐语克服上述挑战。Ray提供了一套非常简洁但是强有力的分布式原语(remote、get),同时支持无状态(function)和有状态(actor)计算模型,用户可以在Ray提供的driver视角下,非常方便而自然地进行分布式编程。而这些特性恰好和隐语的需求很贴合。

因此,针对隐私计算安全领域的特殊需求,基于隐语的业务场景和Ray的分布式基建能力,蚂蚁Ray团队和隐语团队共同构建了新一代基于Ray的隐私分布式计算底座。

2. 隐语分布式底座模式一: 共享Ray集群

2.1 基于Ray构建隐语的分布式底座

隐语提供了多种明文密文设备,比如PYU(明文计算设备)、SPU(多方安全计算设备)、HEU(同态加密设备)等,这些设备通常由一个或者多个参与方的进程构成。用户在设备的基础上进行编程,而这些设备的底座则是Ray。

PYU本质上是Ray的一个function或者actor,提供了python runtime,从而用户可以在PYU的基础上执行任意python代码。SPU/HEU则由多个Ray的actor构成。设备之间的数据流动则是基于Ray的remote object。

隐语使用单一Ray集群作为分布式底座,每个机构都有若干个节点,所有机构的节点整体构成一个Ray集群,我们称之为第一分布式底座。由于Ray既支持多机器集群,也支持单进程的模拟集群,用户可以快速在单机上进行分布式的模拟,而且代码从单机迁移到多机集群是零成本,这给隐语带来了很好的易用性。

2.2 增强Ray

由于,虽然Ray的通用性抽象能给隐语屏蔽掉繁杂的分布式细节,但是Ray自身的隐私和安全性问题,仍然是面临的挑战。因此,为了满足隐私计算领域的需求,蚂蚁Ray团队对Ray进行大幅度的安全能力增强,提供了一个具有隐私安全的Ray版本, Secure-Ray。

类型 问题 解决方案
执行漏洞 代码动态分发和执行 限制代码分发,only load code from local
恶意resource攻击 resource鉴权
pickle序列化漏洞 支持白名单限制反序列化类型
通信安全 单点Redis和Redis conntion 去除Redis
ray通信仅支持局域网 网络改造: hostname,grpc跨域
跨机构通信安全问题 grpc tls + proxy

如上表所示,Ray上的安全问题大体可分为两类:执行漏洞和通信安全问题。在Ray里的执行漏洞有代码动态分发和执行,恶意的resource攻击以及pickle反序列化攻击;而通信安全问题主要是跨机构间各种复杂的协议交互所引入的安全问题。对于通信安全问题,一般都有较为通用的解决方案,在此不多赘述,这里仅介绍Ray对执行漏洞类型的问题的解决方案。

  1. 限制代码分发

Ray默认会将待执行的代码自动分发到执行的机器上执行,并且这种执行是被动的。即发起者发起一次Ray的远程调用,执行者将无条件执行该段代码,这在跨机构的隐私计算中是不可接受的。我们通过对Ray进行安全改造,以加入限制代码的分发机制来保证任一待执行的远程执行代码,均是从本地机器上加载而来,以达到安全性保证。

  1. resource鉴权

Ray的按资源调度策略,会将需要该资源的任务调度到拥有该资源的机器上执行;同时,各个节点上报自己拥有的资源给全局调度器。例如,A机构的一台机器上有X资源,如果B机构发起一次任务调用,该任务需要X资源,则调度器会将该任务调度到A机构中。那么恶意机构就有可能通过虚假地声明自己所拥有某个资源,来欺骗调度器,以获取到任务并查看。我们通过给resource上报过程增加鉴权验证信息,以防止机构恶意声明不属于自己的资源的安全问题。

  1. 白名单限制反序列化类型

因为Ray需要进行动态代码执行,因此使用了cloudpickle作为基础序列化工具。众所周知的是,pickle反序列化均能被注入恶意代码进行执行。为了解决这个问题,我们通过添加自定义的restrict_loads,限制仅能从用户所定义的白名单中反序列化类型,而不属于白名单中的类型,是会被严格拒绝执行反序列化操作。

3. 隐语分布式底座模式二: 多控制器和多Ray集群

3.1 共享单一Ray集群的不足

单Ray集群架构中,每个机构属于Ray集群中的若干节点,这会带来了更多的安全性挑战。虽然我们对安全性做了增强,但是仍然存在一些不易解决的安全问题。举例如下:

  1. Ray的driver作为代码驱动器,负责给机构节点下发计算任务。如何确保driver不作恶是一个难题。

  2. 除了driver外,如何保护机构节点的被动任务调度接口不被恶意攻击。

  3. ray.get可以不受限地从其他机构获取数据,如何防范恶意的数据拉取。

以上等等问题一直困扰着隐语。为了继续提升隐语的安全性,我们从基础安全架构出发,进而提出了RayFed。

3.2 解决方案:RayFed

RayFed作为Ray的生态组件之一,面向隐私和联邦领域,致力于提供可安全使用的,跨机构的分布式联邦执行引擎。相比于在Ray里打上各式各样安全补丁的方案,RayFed从多个维度考虑出发,以跨机构安全性为核心要求,重新设计。

  1. 作为Ray的生态组件

Ray作为通用性分布式计算引擎基座,可基于Ray开发各种计算模式的引擎,因此,基于Ray开发的各种引擎一起组成了Ray丰富的生态组件。近些年,Ray社区在AI领域积累了很多一流的生态组件,如:

  • 数据处理:RayData, Mars, Dask, Modin, etc.

  • 分布式训练: RayTrain, XGBoost, LightGBM, Horovod, PyTorch Lighting, etc.

  • 模型推理和服务: RayServe, Seldon.

  • 其他:如RayTune, RLlib等

而RayFed是作为Ray的生态组件之一,主要提供跨机构的分布式联邦执行能力。

  1. 多Ray集群架构,拒绝安全性漏洞

由于Ray引擎复杂的分布式协议和通信的存在,想从根本上解决是安全性问题是不切实际的,因此,RayFed基于多Ray集群架构,先从物理上隔离Ray集群,避免Ray复杂的分布式协议和通信引入的任何安全性漏洞问题。我们再提供非常有限的跨机构间通信的协议,以进行跨机构任务协调和数据传输。

  1. 能力继承,有增无减
  • 提供一个统一的,上帝编程视角的,并且single-machine-like programming的编程范式,以便于能够自然地编写跨机构的分布式程序;

  • 可以帮助更好地管理和使用联邦Ray集群;

  • 提供了具备高性能的安全的以太网通信能力;并且能在弱网络环境中稳定运行;

  • 面向可信设备友好的调度和使用能力;

  • 基础底座能力:高性能的RPC,分布式故障处理和恢复,任务调度和异构资源管理,以及高性能的数据传输等。(powed by Ray)

  1. 一个FL的例子

上述代码将在2个机构内,以多控制器方式执行。每个机构将执行指定在其机构内的ray task。跨机构传递的ray objects将会被加密,通过动态构图时插入的send/recv barriers完成数据传输。所有fed.get()的语句将会触发一个boradcast barrier。由于是多控制器设计,不同机构看到的DAG全图应该是完全一致的。

3.3 新的隐语分布式底座

如上图所示,基于Ray的分布式引擎和RayFed的跨机构执行能力,构成了隐语的分布式底座。RayFed有着灵活的设备资源的管理与调度能力,它可以对更底层的不同的隐私设备进行灵活的管理和调度。隐语以RayFed为底座,构建了一套核心的Device Flow Framework编程抽象,以供上层不同的AI和BI的框架使用。

单个Ray集群的架构给隐语提供了优秀的模拟体验,用户可以低成本的进行单机模型,并且无缝从单机模拟切换至多机执行。但与此同时,也带了较多的安全挑战。RayFed提出了多控制器多ray集群的模式,带来了安全性的提升,同时集成了Ray的各种原生优点。

在RayFed的基础上,隐语同时支持两种模式供用户使用。第一种是单Ray集群的架构,该架构适合用户进行本地单机或者多机进行模拟仿真,用户可以快速验证算法正确性,易上手。第二种是多控制器多ray集群架构,该架构适合生产环境使用,提供了跨机构下安全可靠的计算调度和执行能力。与此同时,隐语提供了非常便捷的操作,用户可以几乎无成本的从模拟模式切换至生产模式,仅需要在初始化的时候打开开关即可,底座架构的切换对用户是透明的。

4. Ray面向隐私领域的展望

作为通用分布式计算引擎,Ray社区在专注于AI领域的同时,也密切关注行业发展和新的趋势。Ray社区接下来也会在隐私和联邦领域上进行更多的探索,以帮助用户屏蔽分布式,跨机构执行等繁杂细节,以及以更简单的方式构建其隐私计算或联邦学习的框架。在接下来的半年到一年内,我们有如下计划:

Ray AI Pipeline Integration

Ray在AI生态提供了非常丰富的组件,而隐语的主流计算类似也是AI型。因此,我们需要让机构内部的AI计算场景可以直接利用Ray的AI生态,而不需要再做重复的繁杂工作,以构建机构内全链路的AI pipeline。隐语将在机构内的数据处理,分布式训练等阶段使用Ray Datasets, Mars和Ray Train等;如有需求,我们还可以在机构内部使用Ray Tune进行超参搜索等。

Ray AI Pipeline FL化

除了在机构内使用Ray AI pipeline生态组件之外,跨机构联邦的机器学习和深度学习也是需要解决的问题。我们期望通过将Ray AI pipeline的各种生态组件联邦化来解决这个问题。如,Ray Datasets不仅可以在机构内部进行分布式表示,同时还可以跨机构表示shard数据集;Ray Train不仅能在机构内部进行多机多卡的分布式训练,同时还能跨机构联邦化训练等等。

跨机构通信和网关增强

在大规模联邦训练和MPC的生产实践中,对跨机构通信和网关提出了巨大的挑战,不仅需要提供统一分布式代理和网关以屏蔽跨机构通信细节,非常高性能的以太网传输能力外,还需要应对以太网的恶劣环境来保证生产的稳定性和可用性等。Ray团队将和隐语团队一起,开发RayFed的统一分布式代理和网关,并使用更加native的方式,实现跨机构通信能力,包括跨机构协议交互和大数据的传输。

全面集成可信执行环境

由于可信执行环境是隐私计算必不可少的一个环节,开箱即用地使用可信执行环境和秘密计算设备是RayFed始终追求的目标。目前Ray已经完成了单机Occlum环境的测试工作。在接下来的一个季度中,蚂蚁集团的Ray团队,隐语团队和Occlum团队将联合进行Ray与可信执行环境的分布式集成。

最后:

也欢迎大家访问Ray中文问答论坛:https://ray.osanswer.net/ 针对技术问题进行深入交流!

3 Likes

如何将ray替换成具有隐私安全的Ray版本,Secure-Ray呢

当前架构下,使用开源的Ray就是安全的。