fixup branch for upstream

Author: tstamler

examples/python/remote_storage_example/README.md

@@ -108,3 +108,21 @@ The system automatically selects the best available storage backend:
 1. Initiator sends memory descriptors to target
 2. Target performs storage-to-memory or memory-to-storage operations
 3. Data is transferred between initiator and target memory
+
+Remote reads are implemented as a read from storage followed by a network write.
+
+Remote writes are implemented as a read from network following by a storage write.
+
+### Pipelining 
+
+To improve performance of the remote storage server, we can pipeline operations to network and storage. This pipelining allows multiple threads to handle each request. However, in order to maintain correctness, the order of network and storage must happen in order for each individual remote storage operation. To do this, we implemented a simple pipelining scheme. 
+
+![Remote Operation Pipelines](storage_pipelines.png)
+
+### Performance Tips
+
+For high-speed storage and network hardware, you may need to tweak performance with a couple of environment variables.
+
+First, for optimal GDS performance, ensure you are using the GDS_MT backend with default concurrency. Additionally, you can use the cufile options described in the [GDS README](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/cuda_gds/README.md). 
+
+On the network side, remote reads from VRAM to DRAM can be limited by UCX rail selection. This can be tweaked by setting UCX_MAX_RMA_RAILS=1. However, with larger batch or message sizes, this might limit bandwidth and a higher number of rails might be needed.  

examples/python/remote_storage_example/agent_file.in

@@ -36,7 +36,7 @@ def execute_transfer(my_agent, local_descs, remote_descs, remote_name, operation
     my_agent.release_xfer_handle(handle)
 
 
-def remote_storage_transfer(my_agent, my_mem_descs, operation, remote_agent_name):
+def remote_storage_transfer(my_agent, my_mem_descs, operation, remote_agent_name, iterations):
     """Initiate remote memory transfer."""
     if operation != "READ" and operation != "WRITE":
         logger.error("Invalid operation, exiting")
@@ -47,14 +47,14 @@ def remote_storage_transfer(my_agent, my_mem_descs, operation, remote_agent_name
     else:
         operation = b"READ"
 
-    iterations = b"0100"
+    iterations_str = bytes(f"{iterations:04d}", "utf-8")
     # Send the descriptors that you want to read into or write from
     logger.info(f"Sending {operation} request to {remote_agent_name}")
     test_descs_str = my_agent.get_serialized_descs(my_mem_descs)
 
     start_time = time.time()
 
-    my_agent.send_notif(remote_agent_name, operation + iterations + test_descs_str)
+    my_agent.send_notif(remote_agent_name, operation + iterations_str + test_descs_str)
 
     while not my_agent.check_remote_xfer_done(remote_agent_name, b"COMPLETE"):
         continue
@@ -113,8 +113,7 @@ def pipeline_reads(my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs,
             s+=1
             n+=1
 
-
-        done, not_done = concurrent.futures.wait(futures, return_when=concurrent.futures.ALL_COMPLETED)
+        _, not_done = concurrent.futures.wait(futures, return_when=concurrent.futures.ALL_COMPLETED)
         assert not not_done
 
 def pipeline_writes(my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations):
@@ -142,17 +141,14 @@ def pipeline_writes(my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs
             s+=1
             n+=1
 
-
-        done, not_done = concurrent.futures.wait(futures, return_when=concurrent.futures.ALL_COMPLETED)
+        _, not_done = concurrent.futures.wait(futures, return_when=concurrent.futures.ALL_COMPLETED)
         assert not not_done
 
 def handle_remote_transfer_request(my_agent, my_mem_descs, my_file_descs):
     """Handle remote memory and storage transfers as target."""
     # Wait for initiator to send list of memory descriptors
     notifs = my_agent.get_new_notifs()
 
-#    logger.info("Waiting for a remote transfer request...")
-
     while len(notifs) == 0:
         notifs = my_agent.get_new_notifs()
 
@@ -182,52 +178,47 @@ def handle_remote_transfer_request(my_agent, my_mem_descs, my_file_descs):
         # Send completion notification to initiator
         my_agent.send_notif(req_agent, b"COMPLETE")
 
-#    logger.info("One transfer test complete.")
-
-
-def run_client(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, agents_file):
+def run_client(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, agents_file, iterations):
     logger.info("Client initialized, ready for local transfer test...")
 
     # For sample purposes, write to and then read from local storage
     logger.info("Starting local transfer test...")
 
     start_time = time.time()
 
-    for i in range (1, 100):
+    for i in range (1, iterations):
         execute_transfer(
             my_agent,
             nixl_mem_reg_descs.trim(),
             nixl_file_reg_descs.trim(),
             my_agent.name,
             "WRITE",
-#            ["POSIX"]
             ["GDS_MT"]
         )
 
     end_time = time.time()
 
     elapsed = end_time - start_time
 
-    logger.info(f"Time for 100 WRITE iterations: {elapsed} seconds")
+    logger.info(f"Time for {iterations} WRITE iterations: {elapsed} seconds")
 
     start_time = time.time()
 
-    for i in range (1, 100):
+    for i in range (1, iterations):
         execute_transfer(
             my_agent,
             nixl_mem_reg_descs.trim(),
             nixl_file_reg_descs.trim(),
             my_agent.name,
             "READ",
-#            ["POSIX"]
             ["GDS_MT"]
         )
 
     end_time = time.time()
 
     elapsed = end_time - start_time
 
-    logger.info(f"Time for 100 READ iterations: {elapsed} seconds")
+    logger.info(f"Time for {iterations} READ iterations: {elapsed} seconds")
 
     logger.info("Local transfer test complete")
 
@@ -238,10 +229,10 @@ def run_client(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, agents_file):
     # For sample purposes, write to and then read from each target agent
     for target_agent in target_agents:
         remote_storage_transfer(
-            my_agent, nixl_mem_reg_descs.trim(), "WRITE", target_agent
+            my_agent, nixl_mem_reg_descs.trim(), "WRITE", target_agent, iterations
         )
         remote_storage_transfer(
-            my_agent, nixl_mem_reg_descs.trim(), "READ", target_agent
+            my_agent, nixl_mem_reg_descs.trim(), "READ", target_agent, iterations
         )
 
     logger.info("Remote transfer test complete")
@@ -254,7 +245,6 @@ def run_storage_server(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs):
             my_agent, nixl_mem_reg_descs.trim(), nixl_file_reg_descs.trim()
         )
 
-
 if __name__ == "__main__":
     parser = nsu.get_base_parser()
     parser.add_argument(
@@ -276,6 +266,7 @@ def run_storage_server(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs):
         type=str,
         help="File containing list of target agents (only needed for client)",
     )
+    parser.add_argument("--iterations", type=int, default=100, help="Number of iterations for each transfer")
     args = parser.parse_args()
 
     mem = "DRAM"
@@ -299,7 +290,7 @@ def run_storage_server(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs):
             parser.error("--agents_file is required when role is client")
         try:
             run_client(
-                my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, args.agents_file
+                my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, args.agents_file, args.iterations
             )
         finally:
             nsu.cleanup_resources(