feat: nixl_connect: Improve Concurrency Support (#4433)

Signed-off-by: J Wyman <[email protected]>

Author: whoisj

components/src/dynamo/sglang/request_handlers/multimodal/encode_worker_handler.py

@@ -159,7 +159,7 @@ async def generate(
             # Create descriptor for the multimodal data
             descriptor = connect.Descriptor(precomputed_embeddings)
 
-            with self._connector.create_readable(descriptor) as readable:
+            with await self._connector.create_readable(descriptor) as readable:
                 request.serialized_request = readable.metadata()
 
                 logger.debug(f"Request: {request.model_dump_json()}")
@@ -184,6 +184,5 @@ async def async_init(self, runtime: DistributedRuntime):
         # Create and initialize a dynamo connector for this worker.
         # We'll needs this to move data between this worker and remote workers efficiently.
         self._connector = connect.Connector()
-        await self._connector.initialize()
 
         logger.info("Startup completed.")

components/src/dynamo/sglang/request_handlers/multimodal/worker_handler.py

@@ -77,7 +77,6 @@ def __init__(self):
     async def initialize(self):
         """Initialize the connector for embeddings processing"""
         self._connector = connect.Connector()
-        await self._connector.initialize()
 
     async def process_embeddings(self, request: SglangMultimodalRequest):
         """Process embeddings from serialized request"""
@@ -103,7 +102,6 @@ async def process_embeddings(self, request: SglangMultimodalRequest):
                 "Connector is None - this should not happen after initialization"
             )
             self._connector = connect.Connector()
-            await self._connector.initialize()
 
         read_op = await self._connector.begin_read(
             request.serialized_request, descriptor

components/src/dynamo/trtllm/encode_helper.py

@@ -241,7 +241,7 @@ async def process_embedding_request(
 
         # Create readable operation with main embeddings tensor (works for both formats)
         descriptor = nixl_connect.Descriptor(encodings)
-        with connector.create_readable(descriptor) as readable_op:
+        with await connector.create_readable(descriptor) as readable_op:
             # Get the metadata for the readable operation
             op_metadata = readable_op.metadata()
 

components/src/dynamo/trtllm/main.py

@@ -332,7 +332,6 @@ async def init(runtime: DistributedRuntime, config: Config):
     connector = None
     logging.info("Initializing NIXL Connect.")
     connector = nixl_connect.Connector()
-    await connector.initialize()
 
     dump_config(
         config.dump_config_to, {"engine_args": engine_args, "dynamo_args": config}

components/src/dynamo/vllm/multimodal_handlers/encode_worker_handler.py

@@ -69,7 +69,6 @@ async def async_init(self, runtime: DistributedRuntime):
         # Create and initialize a dynamo connector for this worker.
         # We'll needs this to move data between this worker and remote workers efficiently.
         self._connector = connect.Connector()
-        await self._connector.initialize()
         logger.info("Encode worker startup completed.")
 
     async def generate(
@@ -130,7 +129,7 @@ async def generate(
             request.embeddings_shape = tuple(embeddings.shape)
             descriptor = connect.Descriptor(embeddings_cpu)
 
-            with self._connector.create_readable(descriptor) as readable:
+            with await self._connector.create_readable(descriptor) as readable:
                 request.serialized_request = readable.metadata()
                 # Clear the image URL as hint that the image is passed as embeddings.
                 request.multimodal_input.image_url = None

components/src/dynamo/vllm/multimodal_handlers/worker_handler.py

@@ -52,7 +52,6 @@ def __init__(
     async def async_init(self, runtime: DistributedRuntime):
         """Async initialization - connector needs async setup"""
         self._connector = connect.Connector()
-        await self._connector.initialize()
         logger.info("Multimodal Decode Worker async initialization completed.")
 
     async def generate(self, request: vLLMMultimodalRequest, context):
@@ -138,7 +137,6 @@ async def async_init(self, runtime: DistributedRuntime):
         """Async initialization for connector that requires async setup"""
         # Initialize the connector asynchronously
         self._connector = connect.Connector()
-        await self._connector.initialize()
         logger.info("Multimodal PD Worker async initialization completed.")
 
     async def generate(self, request: vLLMMultimodalRequest, context):

docs/api/nixl_connect/connector.md

@@ -47,7 +47,6 @@ The metadata contains required information (identifiers, keys, etc.) which enabl
     @async_on_start
     async def async_init(self):
       self.connector = dynamo.nixl_connect.Connector()
-      await self.connector.initialize()
 ```
 
 > [!Tip]
@@ -109,7 +108,7 @@ Use [`.wait_for_completion()`](write_operation.md#wait_for_completion) to block
 ### `create_readable`
 
 ```python
-def create_readable(
+async def create_readable(
     self,
     local_descriptors: Descriptor | list[Descriptor],
 ) -> ReadableOperation:
@@ -130,7 +129,7 @@ Use [`.wait_for_completion()`](readable_operation.md#wait_for_completion) to blo
 ### `create_writable`
 
 ```python
-def create_writable(
+async def create_writable(
     self,
     local_descriptors: Descriptor | list[Descriptor],
 ) -> WritableOperation:
@@ -151,6 +150,15 @@ Use [`.wait_for_completion()`](writable_operation.md#wait_for_completion) to blo
 
 ## Properties
 
+### `hostname`
+
+```python
+@property
+def hostname(self) -> str:
+```
+
+Gets the name of the current worker's host.
+
 ### `is_cuda_available`
 
 ```python
@@ -169,22 +177,6 @@ def name(self) -> str | None:
 
 Gets the Dynamo component name used by the connector.
 
-### `namespace`
-
-```python
-@property
-def namespace(self) -> str:
-```
-
-Gets the Dynamo namespace used by the connector.
-
-### `runtime`
-
-```python
-def runtime(self) -> dynamo.runtime.DistributedRuntime:
-```
-
-Gets the Dynamo distributed runtime instance associated with the connector.
 
 ## Related Classes
 

docs/api/nixl_connect/read_operation.md

@@ -38,7 +38,7 @@ therefore the operation should be awaited until completed unless cancellation is
     ) -> None:
       descriptor = dynamo.nixl_connect.Descriptor(local_tensor)
 
-      with self.connector.begin_read(descriptor, remote_metadata) as read_op:
+      with await self.connector.begin_read(remote_metadata, descriptor) as read_op:
         # Wait for the operation to complete writing data from the remote worker to local_tensor.
         await read_op.wait_for_completion()
 ```

docs/api/nixl_connect/readable_operation.md

@@ -37,7 +37,7 @@ therefore the operation should be awaited until completed unless cancellation is
     ) -> None:
       descriptor = dynamo.nixl_connect.Descriptor(local_tensor)
 
-      with self.connector.create_readable(descriptor) as read_op:
+      with await self.connector.create_readable(descriptor) as read_op:
         op_metadata = read_op.metadata()
 
         # Send the metadata to the remote worker via sideband communication.

docs/api/nixl_connect/writable_operation.md

@@ -38,7 +38,7 @@ Cancellation is handled asynchronously.
     ) -> None:
       descriptor = dynamo.nixl_connect.Descriptor(local_tensor)
 
-      with self.connector.create_writable(descriptor) as write_op:
+      with await self.connector.create_writable(descriptor) as write_op:
         op_metadata = write_op.metadata()
 
         # Send the metadata to the remote worker via sideband communication.