SRC/API/PYTHON: fix pylint warnings in src, examples and tests

Signed-off-by: Roie Danino <[email protected]>

Author: roiedanino

examples/python/__init__.py

@@ -16,12 +16,20 @@
 # limitations under the License.
 
 import argparse
+import enum
+import sys
 
 import torch
 
 from nixl._api import nixl_agent, nixl_agent_config
 
 
+class BlockingSendRecvErrCodes(enum.Enum):
+    MEM_REG_FAILED = 1
+    TRANSFER_FAILED = 2
+    DATA_VERIFICATION_FAILED = 3
+
+
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--ip", type=str, required=True)
@@ -63,7 +71,7 @@ def parse_args():
     reg_descs = agent.register_memory(tensors)
     if not reg_descs:  # Same as reg_descs if successful
         print("Memory registration failed.")
-        exit()
+        sys.exit(BlockingSendRecvErrCodes.MEM_REG_FAILED.value)
 
     # Target code
     if args.mode == "target":
@@ -113,25 +121,25 @@ def parse_args():
 
         if not xfer_handle:
             print("Creating transfer failed.")
-            exit()
+            sys.exit(BlockingSendRecvErrCodes.TRANSFER_FAILED.value)
 
         state = agent.transfer(xfer_handle)
         if state == "ERR":
             print("Posting transfer failed.")
-            exit()
+            sys.exit(BlockingSendRecvErrCodes.TRANSFER_FAILED.value)
         while True:
             state = agent.check_xfer_state(xfer_handle)
             if state == "ERR":
                 print("Transfer got to Error state.")
-                exit()
+                sys.exit(BlockingSendRecvErrCodes.TRANSFER_FAILED.value)
             elif state == "DONE":
                 break
 
         # Verify data after read
         for i, tensor in enumerate(tensors):
             if not torch.allclose(tensor, torch.ones(10)):
                 print(f"Data verification failed for tensor {i}.")
-                exit()
+                sys.exit(BlockingSendRecvErrCodes.DATA_VERIFICATION_FAILED.value)
         print(f"{args.mode} Data verification passed - {tensors}")
 
     if args.mode != "target":

examples/python/blocking_send_recv_example.py

@@ -16,13 +16,25 @@
 # limitations under the License.
 
 import os
+import sys
+import enum
 
 import numpy as np
 import torch
 
 import nixl._utils as nixl_utils
 from nixl._api import nixl_agent, nixl_agent_config
 
+from . import util
+
+
+class NixlApiExampleErrCodes(enum.Enum):
+    CREATE_TRANSFER_FAILED = 1
+    PREP_TRANSFER_SIDE_HANDLES_FAILED = 2
+    MAKE_PREPPED_TRANSFER_FAILED = 3
+    TRANSFER_FAILED = 4
+
+
 if __name__ == "__main__":
     buf_size = 256
     # Allocate memory and register with NIXL
@@ -104,30 +116,13 @@
     )
     if not xfer_handle_1:
         print("Creating transfer failed.")
-        exit()
+        sys.exit(NixlApiExampleErrCodes.CREATE_TRANSFER_FAILED.value)
 
     # test multiple postings
     for _ in range(2):
         state = nixl_agent2.transfer(xfer_handle_1)
         assert state != "ERR"
-
-        target_done = False
-        init_done = False
-
-        while (not init_done) or (not target_done):
-            if not init_done:
-                state = nixl_agent2.check_xfer_state(xfer_handle_1)
-                if state == "ERR":
-                    print("Transfer got to Error state.")
-                    exit()
-                elif state == "DONE":
-                    init_done = True
-                    print("Initiator done")
-
-            if not target_done:
-                if nixl_agent1.check_remote_xfer_done("initiator", b"UUID1"):
-                    target_done = True
-                    print("Target done")
+        util.wait_for_transfer_completion(nixl_agent2, nixl_agent1, xfer_handle_1, b"UUID1")
 
     # prep transfer mode
     local_prep_handle = nixl_agent2.prep_xfer_dlist(
@@ -165,11 +160,11 @@
     )
     if not local_prep_handle or not remote_prep_handle:
         print("Preparing transfer side handles failed.")
-        exit()
+        sys.exit(NixlApiExampleErrCodes.PREP_TRANSFER_SIDE_HANDLES_FAILED.value)
 
     if not xfer_handle_2:
         print("Make prepped transfer failed.")
-        exit()
+        sys.exit(NixlApiExampleErrCodes.MAKE_PREPPED_TRANSFER_FAILED.value)
 
     state = nixl_agent2.transfer(xfer_handle_2)
     assert state != "ERR"
@@ -184,7 +179,7 @@
             state = nixl_agent2.check_xfer_state(xfer_handle_2)
             if state == "ERR":
                 print("Transfer got to Error state.")
-                exit()
+                sys.exit(NixlApiExampleErrCodes.TRANSFER_FAILED.value)
             elif state == "DONE":
                 init_done = True
                 print("Initiator done")

examples/python/nixl_api_example.py

@@ -17,17 +17,27 @@
 
 import os
 import sys
+import enum
 
 import nixl._utils as nixl_utils
 from nixl._api import nixl_agent, nixl_agent_config
 
+
+class NixlGdsExampleErrCodes(enum.Enum):
+    MISSING_FILE_PATH = 1
+    CREATE_TRANSFER_FAILED = 2
+    TRANSFER_FAILED = 3
+    INIT_XFER_FAILED = 4
+    DATA_VERIFICATION_FAILED = 5
+
+
 if __name__ == "__main__":
     buf_size = 16 * 4096
     # Allocate memory and register with NIXL
 
     if len(sys.argv) < 2:
         print("Please specify file path in argv")
-        exit(0)
+        sys.exit(NixlGdsExampleErrCodes.MISSING_FILE_PATH.value)
 
     print("Using NIXL Plugins from:")
     print(os.environ["NIXL_PLUGIN_DIR"])
@@ -79,7 +89,7 @@
     )
     if not xfer_handle_1:
         print("Creating transfer failed.")
-        exit()
+        sys.exit(NixlGdsExampleErrCodes.CREATE_TRANSFER_FAILED.value)
 
     state = nixl_agent1.transfer(xfer_handle_1)
     assert state != "ERR"
@@ -90,7 +100,7 @@
         state = nixl_agent1.check_xfer_state(xfer_handle_1)
         if state == "ERR":
             print("Transfer got to Error state.")
-            exit()
+            sys.exit(NixlGdsExampleErrCodes.TRANSFER_FAILED.value)
         elif state == "DONE":
             done = True
             print("Initiator done")
@@ -101,7 +111,7 @@
     )
     if not xfer_handle_2:
         print("Creating transfer failed.")
-        exit()
+        sys.exit(NixlGdsExampleErrCodes.INIT_XFER_FAILED.value)
 
     state = nixl_agent1.transfer(xfer_handle_2)
     assert state != "ERR"
@@ -112,7 +122,7 @@
         state = nixl_agent1.check_xfer_state(xfer_handle_2)
         if state == "ERR":
             print("Transfer got to Error state.")
-            exit()
+            sys.exit(NixlGdsExampleErrCodes.TRANSFER_FAILED.value)
         elif state == "DONE":
             done = True
             print("Initiator done")

examples/python/nixl_gds_example.py

@@ -16,13 +16,18 @@
 # limitations under the License.
 
 import argparse
+import enum
 import os
 
 import nixl._utils as nixl_utils
 from nixl._api import nixl_agent, nixl_agent_config
 from nixl._bindings import nixlNotFoundError
+from . import util
 
 
+class PartialMdExampleErrCodes(enum.Enum):
+    TRANSFER_FAILED = 1
+
 def exchange_target_metadata(
     target_agent,
     init_agent,
@@ -62,8 +67,21 @@ def invalidate_target_metadata(
     else:
         target_agent.invalidate_local_metadata(ip_addr, init_port)
 
+def allocate_strings(malloc_addrs: list, buf_size: int, num_strings: int):
+    target_strs = []
+    for _ in range(num_strings):
+        addr1 = nixl_utils.malloc_passthru(buf_size)
+        target_strs.append((addr1, buf_size, 0, "test"))
+        malloc_addrs.append(addr1)
 
-if __name__ == "__main__":
+    return target_strs
+
+def xfer(init_agent, xfer_handle, target_agent):
+    state = init_agent.transfer(xfer_handle)
+    assert state != "ERR"
+    util.wait_for_transfer_completion(init_agent, target_agent, xfer_handle, b"UUID1")
+
+def main():
     buf_size = 256
     # Allocate memory and register with NIXL
 
@@ -102,17 +120,8 @@ def invalidate_target_metadata(
 
     malloc_addrs = []
 
-    target_strs1 = []
-    for _ in range(10):
-        addr1 = nixl_utils.malloc_passthru(buf_size)
-        target_strs1.append((addr1, buf_size, 0, "test"))
-        malloc_addrs.append(addr1)
-
-    target_strs2 = []
-    for _ in range(10):
-        addr1 = nixl_utils.malloc_passthru(buf_size)
-        target_strs2.append((addr1, buf_size, 0, "test"))
-        malloc_addrs.append(addr1)
+    target_strs1 = allocate_strings(malloc_addrs, buf_size, 10)
+    target_strs2 = allocate_strings(malloc_addrs, buf_size, 10)
 
     target_reg_descs1 = target_agent.get_reg_descs(target_strs1, "DRAM", is_sorted=True)
     target_reg_descs2 = target_agent.get_reg_descs(target_strs2, "DRAM", is_sorted=True)
@@ -126,11 +135,7 @@ def invalidate_target_metadata(
     agent_config2 = nixl_agent_config(True, True, init_port)
     init_agent = nixl_agent("initiator", agent_config2)
 
-    init_strs = []
-    for _ in range(10):
-        addr1 = nixl_utils.malloc_passthru(buf_size)
-        init_strs.append((addr1, buf_size, 0, "test"))
-        malloc_addrs.append(addr1)
+    init_strs = allocate_strings(malloc_addrs, buf_size, 10)
 
     init_reg_descs = init_agent.get_reg_descs(init_strs, "DRAM", is_sorted=True)
     init_xfer_descs = init_reg_descs.trim()
@@ -157,27 +162,7 @@ def invalidate_target_metadata(
         "READ", init_xfer_descs, target_xfer_descs1, "target", b"UUID1"
     )
 
-    state = init_agent.transfer(xfer_handle_1)
-    assert state != "ERR"
-
-    target_done = False
-    init_done = False
-
-    while (not init_done) or (not target_done):
-        if not init_done:
-            state = init_agent.check_xfer_state(xfer_handle_1)
-            if state == "ERR":
-                print("Transfer got to Error state.")
-                exit()
-            elif state == "DONE":
-                init_done = True
-                print("Initiator done")
-
-        if not target_done:
-            if target_agent.check_remote_xfer_done("initiator", b"UUID1"):
-                target_done = True
-                print("Target done")
-
+    xfer(init_agent, xfer_handle_1, target_agent)
     # Second set of descs was not sent, should fail
     try:
         xfer_handle_2 = init_agent.initialize_xfer(
@@ -214,26 +199,7 @@ def invalidate_target_metadata(
         else:
             ready = True
 
-    state = init_agent.transfer(xfer_handle_2)
-    assert state != "ERR"
-
-    target_done = False
-    init_done = False
-
-    while (not init_done) or (not target_done):
-        if not init_done:
-            state = init_agent.check_xfer_state(xfer_handle_2)
-            if state == "ERR":
-                print("Transfer got to Error state.")
-                exit()
-            elif state == "DONE":
-                init_done = True
-                print("Initiator done")
-
-        if not target_done:
-            if target_agent.check_remote_xfer_done("initiator", b"UUID1"):
-                target_done = True
-                print("Target done")
+    xfer(init_agent, xfer_handle_2, target_agent)
 
     init_agent.release_xfer_handle(xfer_handle_1)
     init_agent.release_xfer_handle(xfer_handle_2)
@@ -250,3 +216,7 @@ def invalidate_target_metadata(
     del target_agent
 
     print("Test Complete.")
+
+
+if __name__ == "__main__":
+    main()

examples/python/partial_md_example.py

@@ -15,9 +15,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pylint: disable=broad-exception-caught
+
 import os
 import sys
 import tempfile
+import traceback
 
 try:
     from nixl._api import nixl_agent, nixl_agent_config
@@ -27,7 +30,8 @@
     print("NIXL API missing install NIXL.")
     NIXL_AVAILABLE = False
 
-if __name__ == "__main__":
+
+def main():
     print("NIXL queryMem Python API Example")
     print("=" * 40)
 
@@ -111,8 +115,6 @@
 
     except Exception as e:
         print(f"Error in example: {e}")
-        import traceback
-
         traceback.print_exc()
 
     finally:
@@ -122,3 +124,7 @@
             if os.path.exists(temp_file_path):
                 os.unlink(temp_file_path)
                 print(f"Removed: {temp_file_path}")
+
+
+if __name__ == "__main__":
+    main()