Update Dependant Code

This change updates all of the Dynamo usages of nixl_connect to adopt the changes in the previous commit.

- Removal of `Connector.initialize()` calls.
- Addition of the `await` keyword to `Connector.create_readable()` and `.create_writable()` calls.

Signed-off-by: J Wyman <jwyman@nvidia.com>

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

@@ -279,7 +279,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

@@ -46,7 +46,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):
@@ -126,7 +125,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):

examples/multimodal/components/encode_worker.py

@@ -125,7 +125,7 @@ async def generate(
             request.embeddings_shape = tuple(embeddings.shape)
             descriptor = connect.Descriptor(embeddings)
 
-            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
@@ -158,7 +158,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("Startup completed.")
 

examples/multimodal/components/video_encode_worker.py

@@ -153,7 +153,7 @@ async def generate(
             request.embeddings_shape = tuple(tensor_for_descriptor.shape)
             descriptor = connect.Descriptor(tensor_for_descriptor)
 
-            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.video_url = None
@@ -199,7 +199,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("Startup completed.")