precommit

Author: tstamler

examples/python/remote_storage_example/README.md

@@ -113,16 +113,16 @@ 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 
+### 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. 
+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). 
+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.  
+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/nixl_p2p_storage_example.py

@@ -28,15 +28,21 @@
 logger = get_logger(__name__)
 
 
-def execute_transfer(my_agent, local_descs, remote_descs, remote_name, operation, use_backends = []):
+def execute_transfer(
+    my_agent, local_descs, remote_descs, remote_name, operation, use_backends=[]
+):
 
-    handle = my_agent.initialize_xfer(operation, local_descs, remote_descs, remote_name, backends=use_backends)
+    handle = my_agent.initialize_xfer(
+        operation, local_descs, remote_descs, remote_name, backends=use_backends
+    )
     my_agent.transfer(handle)
     nsu.wait_for_transfer(my_agent, handle)
     my_agent.release_xfer_handle(handle)
 
 
-def remote_storage_transfer(my_agent, my_mem_descs, operation, remote_agent_name, iterations):
+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")
@@ -88,62 +94,145 @@ def connect_to_agents(my_agent, agents_file):
 
     return target_agents
 
-def pipeline_reads(my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations):
+
+def pipeline_reads(
+    my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations
+):
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
         n = 0
         s = 0
         futures = []
 
         while n < iterations and s < iterations:
-            
+
             if s == 0:
-                futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, my_file_descs, my_agent.name, "READ"))
-                s+=1
+                futures.append(
+                    executor.submit(
+                        execute_transfer,
+                        my_agent,
+                        my_mem_descs,
+                        my_file_descs,
+                        my_agent.name,
+                        "READ",
+                    )
+                )
+                s += 1
                 continue
 
             if s == iterations:
-                futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, sent_descs, req_agent, "WRITE"))
-                n+=1
+                futures.append(
+                    executor.submit(
+                        execute_transfer,
+                        my_agent,
+                        my_mem_descs,
+                        sent_descs,
+                        req_agent,
+                        "WRITE",
+                    )
+                )
+                n += 1
                 continue
 
             # Do two storage and network in parallel
-            futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, my_file_descs, my_agent.name, "READ"))
-            futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, sent_descs, req_agent, "WRITE"))
-            s+=1
-            n+=1
+            futures.append(
+                executor.submit(
+                    execute_transfer,
+                    my_agent,
+                    my_mem_descs,
+                    my_file_descs,
+                    my_agent.name,
+                    "READ",
+                )
+            )
+            futures.append(
+                executor.submit(
+                    execute_transfer,
+                    my_agent,
+                    my_mem_descs,
+                    sent_descs,
+                    req_agent,
+                    "WRITE",
+                )
+            )
+            s += 1
+            n += 1
 
-        _, 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):
+
+def pipeline_writes(
+    my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations
+):
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
         n = 0
         s = 0
         futures = []
 
         while n < iterations and s < iterations:
-            
+
             if s == 0:
-                futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, sent_descs, req_agent, "READ"))
-                s+=1
+                futures.append(
+                    executor.submit(
+                        execute_transfer,
+                        my_agent,
+                        my_mem_descs,
+                        sent_descs,
+                        req_agent,
+                        "READ",
+                    )
+                )
+                s += 1
                 continue
 
             if s == iterations:
-                futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, my_file_descs, my_agent.name, "WRITE"))
-                n+=1
+                futures.append(
+                    executor.submit(
+                        execute_transfer,
+                        my_agent,
+                        my_mem_descs,
+                        my_file_descs,
+                        my_agent.name,
+                        "WRITE",
+                    )
+                )
+                n += 1
                 continue
 
             # Do two storage and network in parallel
-            futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, sent_descs, req_agent, "READ"))
-            futures.append(executor.submit(execute_transfer, my_agent, my_mem_descs, my_file_descs, my_agent.name, "WRITE"))
-            s+=1
-            n+=1
+            futures.append(
+                executor.submit(
+                    execute_transfer,
+                    my_agent,
+                    my_mem_descs,
+                    sent_descs,
+                    req_agent,
+                    "READ",
+                )
+            )
+            futures.append(
+                executor.submit(
+                    execute_transfer,
+                    my_agent,
+                    my_mem_descs,
+                    my_file_descs,
+                    my_agent.name,
+                    "WRITE",
+                )
+            )
+            s += 1
+            n += 1
 
-        _, 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
@@ -171,29 +260,36 @@ def handle_remote_transfer_request(my_agent, my_mem_descs, my_file_descs):
         sent_descs = my_agent.deserialize_descs(recv_msg[8:])
 
         if operation == "READ":
-            pipeline_reads(my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations)
+            pipeline_reads(
+                my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations
+            )
         elif operation == "WRITE":
-            pipeline_writes(my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations)
-    
+            pipeline_writes(
+                my_agent, req_agent, my_mem_descs, my_file_descs, sent_descs, iterations
+            )
+
         # Send completion notification to initiator
         my_agent.send_notif(req_agent, b"COMPLETE")
 
-def run_client(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, agents_file, iterations):
+
+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, iterations):
+    for i in range(1, iterations):
         execute_transfer(
             my_agent,
             nixl_mem_reg_descs.trim(),
             nixl_file_reg_descs.trim(),
             my_agent.name,
             "WRITE",
-            ["GDS_MT"]
+            ["GDS_MT"],
         )
 
     end_time = time.time()
@@ -204,14 +300,14 @@ def run_client(my_agent, nixl_mem_reg_descs, nixl_file_reg_descs, agents_file, i
 
     start_time = time.time()
 
-    for i in range (1, iterations):
+    for i in range(1, iterations):
         execute_transfer(
             my_agent,
             nixl_mem_reg_descs.trim(),
             nixl_file_reg_descs.trim(),
             my_agent.name,
             "READ",
-            ["GDS_MT"]
+            ["GDS_MT"],
         )
 
     end_time = time.time()
@@ -245,6 +341,7 @@ 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(
@@ -266,7 +363,12 @@ 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")
+    parser.add_argument(
+        "--iterations",
+        type=int,
+        default=100,
+        help="Number of iterations for each transfer",
+    )
     args = parser.parse_args()
 
     mem = "DRAM"
@@ -290,7 +392,11 @@ 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, args.iterations
+                my_agent,
+                nixl_mem_reg_descs,
+                nixl_file_reg_descs,
+                args.agents_file,
+                args.iterations,
             )
         finally:
             nsu.cleanup_resources(

examples/python/remote_storage_example/nixl_storage_utils/common.py

@@ -20,6 +20,7 @@
 
 import argparse
 import os
+
 import torch
 
 import nixl._utils as nixl_utils
@@ -76,7 +77,9 @@ def setup_memory_and_files(agent, batch_size, buf_size, fileprefix, mem="DRAM"):
             my_mem_list.append(nixl_utils.malloc_passthru(buf_size))
             nixl_mem_reg_list.append((my_mem_list[-1], buf_size, 0, str(i)))
 
-        my_file_list.append(os.open(f"{fileprefix}_{i}", os.O_RDWR | os.O_CREAT | os.O_DIRECT))
+        my_file_list.append(
+            os.open(f"{fileprefix}_{i}", os.O_RDWR | os.O_CREAT | os.O_DIRECT)
+        )
         nixl_file_reg_list.append((0, buf_size, my_file_list[-1], str(i)))
 
     nixl_mem_reg_descs = agent.register_memory(nixl_mem_reg_list, mem)