Ray 中的类型提示#

从 Ray 2.48 开始，Ray 为远程函数和 Actor 提供了全面的 Python 类型提示支持。这使得在分布式 Ray 应用程序中能够获得更好的 IDE 支持、静态类型检查和更高的代码可维护性。

概述#

在大多数情况下，Ray 应用程序可以在不修改现有代码的情况下使用类型提示。Ray 会自动处理标准远程函数和基本 Actor 用法的类型推断。例如，远程函数支持标准的 Python 类型注解，无需额外配置。@ray.remote 装饰器会保留原始函数签名和类型信息。

import ray

@ray.remote
def add_numbers(x: int, y: int) -> int:
    return x + y

# Type hints work seamlessly with remote function calls
a = add_numbers.remote(5, 3)
print(ray.get(a))

然而，某些模式，尤其是在使用 Actor 时，需要特定的方法来确保正确的类型注解。

模式 1：使用 ray.remote 作为函数来构建 Actor#

直接使用 ray.remote 函数来创建 Actor 类，而不是使用 @ray.remote 装饰器。这将保留原始类类型，并允许类型推断正常工作。例如，在这种情况下，原始类类型是 DemoRay，Actor 类类型是 ActorClass[DemoRay]。

import ray
from ray.actor import ActorClass

class DemoRay:
    def __init__(self, init: int):
        self.init = init

    @ray.method
    def calculate(self, v1: int, v2: int) -> int:
        return self.init + v1 + v2

ActorDemoRay: ActorClass[DemoRay] = ray.remote(DemoRay)
# DemoRay is the original class type, ActorDemoRay is the ActorClass[DemoRay] type

创建 ActorClass[DemoRay] 类型后，我们可以通过调用 ActorDemoRay.remote(1) 来使用它来实例化一个 Actor。它会返回一个 ActorProxy[DemoRay] 类型，它代表一个 Actor 句柄。

此句柄将为 Actor 方法提供类型提示，包括它们的参数和返回类型。


actor: ActorProxy[DemoRay] = ActorDemoRay.remote(1)

def func(actor: ActorProxy[DemoRay]) -> int:
    b: ObjectRef[int] = actor.calculate.remote(1, 2)
    return ray.get(b)

a = func.remote()
print(ray.get(a))

为什么我们需要这样做？

在 Ray 中，@ray.remote 装饰器表示类 T 的实例是 Actor，每个 Actor 都在自己的 Python 进程中运行。然而，@ray.remote 装饰器会将类 T 转换为 ActorClass[T] 类型，而不是原始类类型。

不幸的是，IDE 和静态类型检查器无法推断 ActorClass[T] 的原始类型 T。为了解决这个问题，使用 ray.remote(T) 会显式返回一个新的泛型类 ActorClass[T] 类型，同时保留原始类类型。

模式 2：使用 @ray.method 装饰器来处理远程方法#

将 @ray.method 装饰器添加到 Actor 的方法中，以便通过 ActorProxy[T] 类型获取 Actor 远程方法的类型提示，包括它们的参数和返回类型。

from ray.actor import ActorClass, ActorProxy

class DemoRay:
    def __init__(self, init: int):
        self.init = init

    @ray.method
    def calculate(self, v1: int, v2: int) -> int:
        return self.init + v1 + v2

ActorDemoRay: ActorClass[DemoRay] = ray.remote(DemoRay)
actor: ActorProxy[DemoRay] = ActorDemoRay.remote(1)
# IDEs will be able to correctly list the remote methods of the actor
# and provide type hints for the arguments and return values of the remote methods
a: ObjectRef[int] = actor.calculate.remote(1, 2)
print(ray.get(a))

注意

我们希望在不使用 @ray.method 装饰器的情况下实现远程方法的类型提示。如果社区成员有任何想法，我们欢迎提交 PR。