有没有一种优雅的方法可以用pytest测试射线远程功能？

Question

假设你有这样一个函数：

@ray.remote
def remote_function():
    return 1

你也许可以这样测试它：

def test_remote_function():
    ray.init()
    x=ray.get(remote_function.remote())
    assert x==1

但这意味着初始化 ray 而实际上并不需要它（我只是想测试函数本身）。在我的例子中，它甚至不必是异步的或线程化的。

我在 atm 上所做的是访问其受保护的包装函数：

def test_remote_function():
    assert remote_function._function() == 1

但这感觉很老套，而且 linter 不服从我:)

我想要的是像这样的 pytest fixture:

def test_remote_function(ray):
    x=ray.get(remote_function.remote())
    assert x==1

或

def test_remote_function(ray_sync):
    x=ray.get(remote_function.remote())
    assert x==1

明确地使其同步（实际上不使用射线）。

我已经看到 ray 的内部测试有某种 pytest 固定装置，但我猜没有公开。

有没有人知道更好的方法？

Answer 1

基于Robert Nishihara's 的解决方案：

@pytest.fixture(scope="module")
def ray_fix():
    ray.init(num_cpus=1)
    yield None
    ray.shutdown()

像这样使用它：

def test_remote_function(ray_fix):
    x=ray.get(remote_function.remote())
    assert x==1

这样射线仍然是 used/initialized，但每个模块只有一次。

在我的例子中，这增加了大约 2 秒的总测试时间，可以忽略不计。

如果 ray 带有一组默认的 pytest fixtures，那就太好了。

感谢提示！

Is there an elegant way to test ray remote functions with pytest?