torch.futures¶
此包提供 Future
类型,它封装了异步执行和一组实用函数,用于简化对 Future
对象的操作。目前,Future
类型主要由 分布式 RPC 框架 使用。
- class torch.futures.Future(*, devices=None)¶
围绕
torch._C.Future
的包装器,它封装了可调用对象的异步执行,例如rpc_async()
。它还公开了用于添加回调函数和设置结果的一组 API。警告
GPU 支持是测试版功能,可能会发生变化。
- add_done_callback(callback)[源代码]¶
将给定的回调函数追加到此
Future
,该函数将在Future
完成时运行。可以将多个回调函数添加到同一个Future
,但无法保证它们的执行顺序。回调函数必须接受一个参数,该参数是对此Future
的引用。回调函数可以使用value()
方法获取值。请注意,如果此Future
已经完成,则给定的回调函数将在线运行。我们建议您使用
then()
方法,因为它提供了一种在回调函数完成后的同步方式。add_done_callback
的成本可能更低,如果您的回调函数不返回值。但是,then()
和add_done_callback
在幕后使用相同的回调函数注册 API。关于 GPU 张量,此方法的行为与
then()
相同。- 参数
callback (
Future
) – 接受一个参数的Callable
,该参数是对此Future
的引用。
注意
请注意,如果回调函数抛出异常,无论是通过原始未来完成时出现异常并调用
fut.wait()
,还是通过回调函数中的其他代码,都必须仔细处理错误处理。例如,如果此回调函数稍后完成其他未来,则这些未来不会标记为已完成,并出现错误,并且用户负责独立处理这些未来的完成/等待。- 示例:
>>> def callback(fut): ... print("This will run after the future has finished.") ... print(fut.wait()) >>> fut = torch.futures.Future() >>> fut.add_done_callback(callback) >>> fut.set_result(5) This will run after the future has finished. 5
- done()[源代码]¶
如果此
Future
已完成,则返回True
。如果Future
具有结果或异常,则它已完成。如果值包含驻留在 GPU 上的张量,即使填充这些张量的异步内核尚未在设备上完成运行,
Future.done()
也会返回True
,因为在此阶段结果已经可以使用,前提是执行了适当的同步(参见wait()
)。- 返回类型
- set_exception(result)[source]¶
为该
Future
设置异常,这将标记该Future
为已完成并出现错误,并触发所有附加的回调。请注意,在调用该Future
上的 wait()/value() 时,此处设置的异常将被内联引发。- 参数
result (BaseException) – 该
Future
的异常。
- 示例:
>>> fut = torch.futures.Future() >>> fut.set_exception(ValueError("foo")) >>> fut.wait() Traceback (most recent call last): ... ValueError: foo
- set_result(result)[source]¶
设置该
Future
的结果,这将标记该Future
为已完成并触发所有附加的回调。请注意,一个Future
无法被标记为已完成两次。如果结果包含驻留在 GPU 上的张量,则即使填充这些张量的异步内核尚未在设备上完成运行,也可以调用此方法,前提是在调用此方法时将这些内核排队的流设置为当前流。简而言之,在启动这些内核后立即调用此方法是安全的,无需任何其他同步,只要您在这两者之间不更改流即可。此方法将记录所有相关当前流上的事件,并使用它们来确保对该
Future
的所有使用者进行适当的调度。- 参数
result (object) – 该
Future
的结果对象。
- 示例:
>>> import threading >>> import time >>> def slow_set_future(fut, value): ... time.sleep(0.5) ... fut.set_result(value) >>> fut = torch.futures.Future() >>> t = threading.Thread( ... target=slow_set_future, ... args=(fut, torch.ones(2) * 3) ... ) >>> t.start() >>> print(fut.wait()) tensor([3., 3.]) >>> t.join()
- then(callback)[source]¶
将给定的回调函数追加到该
Future
,该函数将在Future
完成时运行。多个回调可以添加到同一个Future
,但无法保证它们执行的顺序(要强制执行特定顺序,请考虑链式操作:fut.then(cb1).then(cb2)
)。回调必须接受一个参数,即对该Future
的引用。回调函数可以使用value()
方法获取值。请注意,如果该Future
已经完成,则给定的回调将立即内联运行。如果
Future
的值包含驻留在 GPU 上的张量,则回调可能会在填充这些张量的异步内核尚未在设备上完成执行时被调用。但是,回调将使用一些专用流作为当前流(从全局池中获取)来调用,这些流将与这些内核同步。因此,回调对这些张量执行的任何操作都将在内核完成之后在设备上调度。换句话说,只要回调不切换流,它就可以安全地操作结果,无需任何其他同步。这类似于wait()
的非阻塞行为。类似地,如果回调返回一个包含驻留在 GPU 上的张量的值,那么它即使在生成这些张量的内核仍在设备上运行时也可以这样做,只要回调在执行期间没有更改流即可。如果要更改流,则必须小心地将它们与原始流重新同步,即在调用回调时为当前的流。
- 参数
callback (
Callable
) – 一个Callable
,它将该Future
作为唯一参数。- 返回
一个新的
Future
对象,它保存了callback
的返回值,并在给定的callback
完成时被标记为已完成。- 返回类型
Future[S]
注意
请注意,如果回调函数抛出异常,无论是由于原始 future 完成时出现异常并调用了
fut.wait()
,还是由于回调中的其他代码,由then
返回的 future 将被相应的错误标记。但是,如果此回调随后完成其他 future,则这些 future 不会被错误标记为已完成,用户有责任独立处理这些 future 的完成/等待。- 示例:
>>> def callback(fut): ... print(f"RPC return value is {fut.wait()}.") >>> fut = torch.futures.Future() >>> # The inserted callback will print the return value when >>> # receiving the response from "worker1" >>> cb_fut = fut.then(callback) >>> chain_cb_fut = cb_fut.then( ... lambda x : print(f"Chained cb done. {x.wait()}") ... ) >>> fut.set_result(5) RPC return value is 5. Chained cb done. None
- torch.futures.collect_all(futures)[source]¶
将提供的
Future
对象收集到一个单独的组合Future
中,该Future
在所有子 future 完成时完成。- 示例:
>>> fut0 = torch.futures.Future() >>> fut1 = torch.futures.Future() >>> fut = torch.futures.collect_all([fut0, fut1]) >>> fut0.set_result(0) >>> fut1.set_result(1) >>> fut_list = fut.wait() >>> print(f"fut0 result = {fut_list[0].wait()}") fut0 result = 0 >>> print(f"fut1 result = {fut_list[1].wait()}") fut1 result = 1