SRC/API/PYTHON: add pylint checks in CI

.gitlab/test_python.sh

@@ -55,6 +55,9 @@ pip3 install --break-system-packages .
 pip3 install --break-system-packages pytest
 pip3 install --break-system-packages pytest-timeout
 pip3 install --break-system-packages zmq
+pip3 install --break-system-packages pylint
+
+pylint src/ examples/ test/python
 
 echo "==== Running ETCD server ===="
 etcd_port=$(get_next_tcp_port)
@@ -82,9 +85,11 @@ python3 query_mem_example.py
 export NIXL_TELEMETRY_ENABLE=1
 blocking_send_recv_port=$(get_next_tcp_port)
 
-python3 blocking_send_recv_example.py --mode="target" --ip=127.0.0.1 --port="$blocking_send_recv_port"&
+python3 examples/python/blocking_send_recv_example.py --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 examples/python/blocking_send_recv_example.py --mode="initiator" --ip=127.0.0.1 --port="$blocking_send_recv_port"
+
+python3 examples/python/query_mem_example.py
 
 python3 telemetry_reader.py --telemetry_path /tmp/initiator &
 telePID=$!

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,6 +16,8 @@
 # limitations under the License.
 
 import argparse
+import enum
+import sys
 
 import torch
 
@@ -25,6 +27,12 @@
 logger = get_logger(__name__)
 
 
+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)
@@ -66,7 +74,7 @@ def parse_args():
     reg_descs = agent.register_memory(tensors)
     if not reg_descs:  # Same as reg_descs if successful
         logger.error("Memory registration failed.")
-        exit()
+        sys.exit(BlockingSendRecvErrCodes.MEM_REG_FAILED.value)
 
     # Target code
     if args.mode == "target":
@@ -116,25 +124,25 @@ def parse_args():
 
         if not xfer_handle:
             logger.error("Creating transfer failed.")
-            exit()
+            sys.exit(BlockingSendRecvErrCodes.TRANSFER_FAILED.value)
 
         state = agent.transfer(xfer_handle)
         if state == "ERR":
             logger.error("Posting transfer failed.")
-            exit()
+            sys.exit(BlockingSendRecvErrCodes.TRANSFER_FAILED.value)
         while True:
             state = agent.check_xfer_state(xfer_handle)
             if state == "ERR":
                 logger.error("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)):
                 logger.error("Data verification failed for tensor %d.", i)
-                exit()
+                sys.exit(BlockingSendRecvErrCodes.DATA_VERIFICATION_FAILED.value)
         logger.info("%s Data verification passed", args.mode)
 
     if args.mode != "target":

examples/python/nixl_api_example.py

@@ -15,19 +15,30 @@
 # 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
 
+# pylint: disable=import-error,no-name-in-module
 import nixl._utils as nixl_utils
+from examples.python import util
 from nixl._api import nixl_agent, nixl_agent_config
 from nixl.logging import get_logger
 
 # Configure logging
 logger = get_logger(__name__)
 
 
+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
@@ -113,30 +124,15 @@
     )
     if not xfer_handle_1:
         logger.error("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":
-                    logger.error("Transfer got to Error state.")
-                    exit()
-                elif state == "DONE":
-                    init_done = True
-                    logger.info("Initiator done")
-
-            if not target_done:
-                if nixl_agent1.check_remote_xfer_done("initiator", b"UUID1"):
-                    target_done = True
-                    logger.info("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(
@@ -173,11 +169,11 @@
     )
     if not local_prep_handle or not remote_prep_handle:
         logger.error("Preparing transfer side handles failed.")
-        exit()
+        sys.exit(NixlApiExampleErrCodes.PREP_TRANSFER_SIDE_HANDLES_FAILED.value)
 
     if not xfer_handle_2:
         logger.error("Make prepped transfer failed.")
-        exit()
+        sys.exit(NixlApiExampleErrCodes.MAKE_PREPPED_TRANSFER_FAILED.value)
 
     state = nixl_agent2.transfer(xfer_handle_2)
     assert state != "ERR"
@@ -192,7 +188,7 @@
             state = nixl_agent2.check_xfer_state(xfer_handle_2)
             if state == "ERR":
                 logger.error("Transfer got to Error state.")
-                exit()
+                sys.exit(NixlApiExampleErrCodes.TRANSFER_FAILED.value)
             elif state == "DONE":
                 init_done = True
                 logger.info("Initiator done")

examples/python/nixl_gds_example.py

@@ -15,6 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 import os
 import sys
 
@@ -26,13 +27,21 @@
 logger = get_logger(__name__)
 
 
+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:
         logger.error("Please specify file path in argv")
-        exit(0)
+        sys.exit(NixlGdsExampleErrCodes.MISSING_FILE_PATH.value)
 
     logger.info("Using NIXL Plugins from:\n%s", os.environ["NIXL_PLUGIN_DIR"])
 
@@ -86,7 +95,7 @@
     )
     if not xfer_handle_1:
         logger.error("Creating transfer failed.")
-        exit()
+        sys.exit(NixlGdsExampleErrCodes.CREATE_TRANSFER_FAILED.value)
 
     state = nixl_agent1.transfer(xfer_handle_1)
     assert state != "ERR"
@@ -97,7 +106,7 @@
         state = nixl_agent1.check_xfer_state(xfer_handle_1)
         if state == "ERR":
             logger.error("Transfer got to Error state.")
-            exit()
+            sys.exit(NixlGdsExampleErrCodes.TRANSFER_FAILED.value)
         elif state == "DONE":
             done = True
             logger.info("Initiator done")
@@ -108,7 +117,7 @@
     )
     if not xfer_handle_2:
         logger.error("Creating transfer failed.")
-        exit()
+        sys.exit(NixlGdsExampleErrCodes.INIT_XFER_FAILED.value)
 
     state = nixl_agent1.transfer(xfer_handle_2)
     assert state != "ERR"
@@ -119,7 +128,7 @@
         state = nixl_agent1.check_xfer_state(xfer_handle_2)
         if state == "ERR":
             logger.error("Transfer got to Error state.")
-            exit()
+            sys.exit(NixlGdsExampleErrCodes.TRANSFER_FAILED.value)
         elif state == "DONE":
             done = True
             logger.info("Initiator done")