SRC/API/PYTHON: add pylint checks in CI

.gitlab/test_python.sh

@@ -64,9 +64,9 @@ etcd --listen-client-urls ${NIXL_ETCD_ENDPOINTS} --advertise-client-urls ${NIXL_
 sleep 5
 
 echo "==== Running python tests ===="
-python3 examples/python/nixl_api_example.py
-python3 examples/python/partial_md_example.py
-python3 examples/python/partial_md_example.py --etcd
+python3 -m examples.python.nixl_api_example
+python3 -m examples.python.partial_md_example
+python3 -m examples.python.partial_md_example --etcd
 pytest test/python
 
 python3 test/python/prep_xfer_perf.py list
@@ -75,11 +75,10 @@ python3 test/python/prep_xfer_perf.py array
 echo "==== Running python examples ===="
 blocking_send_recv_port=$(get_next_tcp_port)
 
-cd examples/python
-python3 blocking_send_recv_example.py --mode="target" --ip=127.0.0.1 --port="$blocking_send_recv_port"&
+python3 -m examples.python.blocking_send_recv_example --mode="target" --ip=127.0.0.1 --port="$blocking_send_recv_port"&
 sleep 5
-python3 blocking_send_recv_example.py --mode="initiator" --ip=127.0.0.1 --port="$blocking_send_recv_port"
+python3 -m examples.python.blocking_send_recv_example --mode="initiator" --ip=127.0.0.1 --port="$blocking_send_recv_port"
 
-python3 query_mem_example.py
+python3 -m examples.python.query_mem_example
 
 pkill etcd

.pre-commit-config.yaml

@@ -40,6 +40,14 @@ repos:
         args: ["--max-line-length=88", "--select=C,E,F,W,B,B950", "--extend-ignore=E203,E501,W503"]
         types_or: [python, cython]
 
+  - repo: https://github.com/pylint-dev/pylint
+    rev: v3.3.7
+    hooks:
+      - id: pylint
+        args: []
+        additional_dependencies: [torch, numpy, pytest]
+        types_or: [python, cython]
+
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.4
     hooks:

examples/python/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

examples/python/blocking_send_recv_example.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/nixl_api_example.py

@@ -15,14 +15,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 import os
+import sys
 
 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,15 @@
     )
     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 +162,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 +181,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_gds_example.py

@@ -15,19 +15,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 import os
 import sys
 
 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")