reference/generated/tensordict._ucxx.TensorDictServer

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictServer

class tensordict._ucxx.TensorDictServer(port: int, *, consolidated: bool = True)

Bases: object

A one-to-many listener that accepts connections and yields TensorDictPipe instances.

Warning

TensorDictServer is experimental: its API can change in a minor release without a deprecation cycle.

Wraps ucxx.create_listener to provide a server that can handle multiple clients, each getting their own TensorDictPipe. It needs the optional ucxx package. Import it from tensordict._ucxx.

Examples

Each client connects with TensorDictPipe.connect(). This example is not tested: it needs ucxx and client processes.

import asyncio
from tensordict._ucxx import TensorDictServer

async def handler(pipe):
    async for td in pipe:
        print(td["obs"].shape)

async def main():
    server = TensorDictServer(port=13337)
    await server.serve(handler)  # runs until server.close()

asyncio.run(main())

async aclose() → None

Close the server (async).

close() → None

Close the server and stop accepting connections.

async serve(on_connect: Callable[[TensorDictPipe], Awaitable[None]]) → None

Accept connections and call on_connect for each new client.

This runs indefinitely until close() is called.

Parameters:

on_connect – an async callable that receives a TensorDictPipe for each new client connection.