Merge remote-tracking branch 'leofang/cluster' into HEAD

Author: ksimpson-work

.github/actions/test/action.yml

@@ -14,6 +14,13 @@ runs:
       shell: bash --noprofile --norc -xeuo pipefail {0}
       run: nvidia-smi
 
+    # The cache action needs this
+    - name: Install zstd
+      shell: bash --noprofile --norc -xeuo pipefail {0}
+      run: |
+        apt update
+        apt install zstd
+
     - name: Download bindings build artifacts
       uses: actions/download-artifact@v4
       with:

.github/workflows/ci-gh.yml

@@ -26,6 +26,7 @@ jobs:
         upload-enabled:
           - false
         python-version:
+          - "3.13"
           - "3.12"
           - "3.11"
           - "3.10"

.github/workflows/gh-build-and-test.yml

@@ -76,17 +76,19 @@ jobs:
   test:
     # TODO: improve the name once a separate test matrix is defined
     name: Test (CUDA ${{ inputs.cuda-version }})
-    # TODO: enable testing once linux-aarch64 & win-64 GPU runners are up
+    # TODO: enable testing once win-64 GPU runners are up
     if: ${{ (github.repository_owner == 'nvidia') &&
-             startsWith(inputs.host-platform, 'linux-x64') }}
+             startsWith(inputs.host-platform, 'linux') }}
     permissions:
       id-token: write # This is required for configure-aws-credentials
       contents: read  # This is required for actions/checkout
-    runs-on: ${{ (inputs.host-platform == 'linux-x64' && 'linux-amd64-gpu-v100-latest-1') }}
-    # TODO: use a different (nvidia?) container, or just run on bare image
+    runs-on: ${{ (inputs.host-platform == 'linux-x64' && 'linux-amd64-gpu-v100-latest-1') ||
+                 (inputs.host-platform == 'linux-aarch64' && 'linux-arm64-gpu-a100-latest-1') }}
+    # Our self-hosted runners require a container
+    # TODO: use a different (nvidia?) container
     container:
       options: -u root --security-opt seccomp=unconfined --privileged --shm-size 16g
-      image: condaforge/miniforge3:latest
+      image: ubuntu:22.04
       env:
         NVIDIA_VISIBLE_DEVICES: ${{ env.NVIDIA_VISIBLE_DEVICES }}
     needs:

cuda_core/cuda/core/_version.py

@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"

cuda_core/cuda/core/experimental/__init__.py

@@ -9,3 +9,8 @@
 from cuda.core.experimental._linker import Linker, LinkerOptions
 from cuda.core.experimental._program import Program
 from cuda.core.experimental._stream import Stream, StreamOptions
+from cuda.core.experimental._system import System
+
+system = System()
+__import__("sys").modules[__spec__.name + ".system"] = system
+del System

cuda_core/cuda/core/experimental/_linker.py

@@ -443,7 +443,7 @@ def link(self, target_type) -> ObjectCode:
         return ObjectCode(bytes(code), target_type)
 
     def get_error_log(self) -> str:
-        """ Get the error log generated by the linker.
+        """Get the error log generated by the linker.
 
         Returns
         -------

cuda_core/cuda/core/experimental/_module.py

@@ -5,7 +5,7 @@
 import importlib.metadata
 
 from cuda import cuda
-from cuda.core.experimental._utils import handle_return
+from cuda.core.experimental._utils import handle_return, precondition
 
 _backend = {
     "old": {
@@ -106,30 +106,43 @@ class ObjectCode:
 
     """
 
-    __slots__ = ("_handle", "_code_type", "_module", "_loader", "_sym_map")
+    __slots__ = ("_handle", "_backend_version", "_jit_options", "_code_type", "_module", "_loader", "_sym_map")
     _supported_code_type = ("cubin", "ptx", "ltoir", "fatbin")
 
     def __init__(self, module, code_type, jit_options=None, *, symbol_mapping=None):
         if code_type not in self._supported_code_type:
             raise ValueError
         _lazy_init()
+
+        # handle is assigned during _lazy_load
         self._handle = None
+        self._jit_options = jit_options
+
+        self._backend_version = "new" if (_py_major_ver >= 12 and _driver_ver >= 12000) else "old"
+        self._loader = _backend[self._backend_version]
+
+        self._code_type = code_type
+        self._module = module
+        self._sym_map = {} if symbol_mapping is None else symbol_mapping
 
-        backend = "new" if (_py_major_ver >= 12 and _driver_ver >= 12000) else "old"
-        self._loader = _backend[backend]
+    # TODO: do we want to unload in a finalizer? Probably not..
 
+    def _lazy_load_module(self, *args, **kwargs):
+        if self._handle is not None:
+            return
+        jit_options = self._jit_options
+        module = self._module
         if isinstance(module, str):
             # TODO: this option is only taken by the new library APIs, but we have
             # a bug that we can't easily support it just yet (NVIDIA/cuda-python#73).
             if jit_options is not None:
                 raise ValueError
-            module = module.encode()
             self._handle = handle_return(self._loader["file"](module))
         else:
             assert isinstance(module, bytes)
             if jit_options is None:
                 jit_options = {}
-            if backend == "new":
+            if self._backend_version == "new":
                 args = (
                     module,
                     list(jit_options.keys()),
@@ -141,15 +154,15 @@ def __init__(self, module, code_type, jit_options=None, *, symbol_mapping=None):
                     0,
                 )
             else:  # "old" backend
-                args = (module, len(jit_options), list(jit_options.keys()), list(jit_options.values()))
+                args = (
+                    module,
+                    len(jit_options),
+                    list(jit_options.keys()),
+                    list(jit_options.values()),
+                )
             self._handle = handle_return(self._loader["data"](*args))
 
-        self._code_type = code_type
-        self._module = module
-        self._sym_map = {} if symbol_mapping is None else symbol_mapping
-
-    # TODO: do we want to unload in a finalizer? Probably not..
-
+    @precondition(_lazy_load_module)
     def get_kernel(self, name):
         """Return the :obj:`Kernel` of a specified name from this object code.
 
@@ -168,6 +181,7 @@ def get_kernel(self, name):
             name = self._sym_map[name]
         except KeyError:
             name = name.encode()
+
         data = handle_return(self._loader["kernel"](self._handle, name))
         return Kernel._from_obj(data, self)
 

cuda_core/cuda/core/experimental/_system.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+from typing import Tuple
+
+from cuda import cuda, cudart
+from cuda.core.experimental._device import Device
+from cuda.core.experimental._utils import handle_return
+
+
+class System:
+    """Provide information about the cuda system.
+    This class is a singleton and should not be instantiated directly.
+    """
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if hasattr(self, "_initialized") and self._initialized:
+            return
+        self._initialized = True
+
+    @property
+    def driver_version(self) -> Tuple[int, int]:
+        """
+        Query the CUDA driver version.
+
+        Returns
+        -------
+        tuple of int
+            A 2-tuple of (major, minor) version numbers.
+        """
+        version = handle_return(cuda.cuDriverGetVersion())
+        major = version // 1000
+        minor = (version % 1000) // 10
+        return (major, minor)
+
+    @property
+    def num_devices(self) -> int:
+        """
+        Query the number of available GPUs.
+
+        Returns
+        -------
+        int
+            The number of available GPU devices.
+        """
+        return handle_return(cudart.cudaGetDeviceCount())
+
+    @property
+    def devices(self) -> tuple:
+        """
+        Query the available device instances.
+
+        Returns
+        -------
+        tuple of Device
+            A tuple containing instances of available devices.
+        """
+        total = self.num_devices
+        return tuple(Device(device_id) for device_id in range(total))

cuda_core/docs/source/api.rst

@@ -38,6 +38,17 @@ CUDA compilation toolchain
    LinkerOptions
 
 
+CUDA system information
+-----------------------
+
+.. autodata:: cuda.core.experimental.system.driver_version
+   :no-value:
+.. autodata:: cuda.core.experimental.system.num_devices
+   :no-value:
+.. autodata:: cuda.core.experimental.system.devices
+   :no-value:
+
+
 .. module:: cuda.core.experimental.utils
 
 Utility functions

cuda_core/docs/source/conf.py

@@ -91,3 +91,31 @@
 
 napoleon_google_docstring = False
 napoleon_numpy_docstring = True
+
+
+section_titles = ["Returns"]
+def autodoc_process_docstring(app, what, name, obj, options, lines):
+    if name.startswith("cuda.core.experimental.system"):
+        # patch the docstring (in lines) *in-place*. Should docstrings include section titles other than "Returns", 
+        # this will need to be modified to handle them.
+        attr = name.split(".")[-1]
+        from cuda.core.experimental._system import System
+
+        lines_new = getattr(System, attr).__doc__.split("\n")
+        formatted_lines = []
+        for line in lines_new:
+            title = line.strip()
+            if title in section_titles:
+                formatted_lines.append(line.replace(title, f".. rubric:: {title}"))
+            elif line.strip() == "-" * len(title):
+                formatted_lines.append(" " * len(title))
+            else:
+                formatted_lines.append(line)
+        n_pops = len(lines)
+        lines.extend(formatted_lines)
+        for _ in range(n_pops):
+            lines.pop(0)
+
+
+def setup(app):
+    app.connect("autodoc-process-docstring", autodoc_process_docstring)