Use public pylibcudf APIs in graph_primtypes (#5036)

mroeschke · web-flow · commit 5da5bbe33e04 · 2025-05-07T15:13:08.000Z
`graph_primtypes` currently uses private cuDF Python APIs to construct a cuDF Python column from a pointer to an RMM buffer. After rapidsai/cudf#18502 is merged, this can be constructed using public pylibcudf methods. Authors: - Matthew Roeschke (https://github.com/mroeschke) Approvers: - Rick Ratzel (https://github.com/rlratzel) - James Lamb (https://github.com/jameslamb) URL: #5036
diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
@@ -59,6 +59,7 @@ DEPENDENCIES=(
   libcugraph-tests
   libraft
   librmm
+  pylibcudf
   pylibcugraph
   pylibwholegraph
   pylibraft
diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml
@@ -45,6 +45,7 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
+- pylibcudf==25.6.*,>=0.0.0a0
 - pylibraft==25.6.*,>=0.0.0a0
 - pylibwholegraph==25.6.*,>=0.0.0a0
 - pytest
diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -45,6 +45,7 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
+- pylibcudf==25.6.*,>=0.0.0a0
 - pylibraft==25.6.*,>=0.0.0a0
 - pylibwholegraph==25.6.*,>=0.0.0a0
 - pytest
diff --git a/conda/environments/all_cuda-128_arch-aarch64.yaml b/conda/environments/all_cuda-128_arch-aarch64.yaml
@@ -50,6 +50,7 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
+- pylibcudf==25.6.*,>=0.0.0a0
 - pylibraft==25.6.*,>=0.0.0a0
 - pylibwholegraph==25.6.*,>=0.0.0a0
 - pytest
diff --git a/conda/environments/all_cuda-128_arch-x86_64.yaml b/conda/environments/all_cuda-128_arch-x86_64.yaml
@@ -50,6 +50,7 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
+- pylibcudf==25.6.*,>=0.0.0a0
 - pylibraft==25.6.*,>=0.0.0a0
 - pylibwholegraph==25.6.*,>=0.0.0a0
 - pytest
diff --git a/conda/recipes/cugraph/recipe.yaml b/conda/recipes/cugraph/recipe.yaml
@@ -60,6 +60,7 @@ requirements:
     - cython >=3.0.3
     - libcugraph =${{ version }}
     - pip
+    - pylibcudf =${{ minor_version }}
     - pylibraft =${{ minor_version }}
     - python =${{ py_version }}
     - raft-dask =${{ minor_version }}
@@ -80,6 +81,7 @@ requirements:
     - dask-cudf =${{ minor_version }}
     - fsspec>=0.6.0
     - libcugraph =${{ version }}
+    - pylibcudf =${{ minor_version }}
     - pylibcugraph =${{ version }}
     - pylibraft =${{ minor_version }}
     - python
diff --git a/dependencies.yaml b/dependencies.yaml
@@ -18,6 +18,7 @@ files:
       - depends_on_libcudf
       - depends_on_libraft
       - depends_on_librmm
+      - depends_on_pylibcudf
       - depends_on_pylibraft
       - depends_on_pylibwholegraph
       - depends_on_pytorch
@@ -122,6 +123,7 @@ files:
       - depends_on_dask_cuda
       - depends_on_dask_cudf
       - depends_on_libcugraph
+      - depends_on_pylibcudf
       - depends_on_pylibcugraph
       - depends_on_pylibraft
       - depends_on_raft_dask
@@ -750,6 +752,31 @@ dependencies:
               - rmm-cu11==25.6.*,>=0.0.0a0
           - {matrix: null, packages: [*rmm_unsuffixed]}
 
+  depends_on_pylibcudf:
+    common:
+      - output_types: conda
+        packages:
+          - &pylibcudf_unsuffixed pylibcudf==25.6.*,>=0.0.0a0
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+              cuda_suffixed: "true"
+            packages:
+              - pylibcudf-cu12==25.6.*,>=0.0.0a0
+          - matrix:
+              cuda: "11.*"
+              cuda_suffixed: "true"
+            packages:
+              - pylibcudf-cu11==25.6.*,>=0.0.0a0
+          - {matrix: null, packages: [*pylibcudf_unsuffixed]}
+
   depends_on_cudf:
     common:
       - output_types: conda
diff --git a/python/cugraph/cugraph/structure/graph_primtypes.pxd b/python/cugraph/cugraph/structure/graph_primtypes.pxd
@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# Copyright (c) 2019-2025, NVIDIA CORPORATION.
 # 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
@@ -154,8 +154,17 @@ ctypedef GraphCOOView[int,int,double] GraphCOOViewDouble
 ctypedef GraphCSRView[int,int,float] GraphCSRViewFloat
 ctypedef GraphCSRView[int,int,double] GraphCSRViewDouble
 
-cdef move_device_buffer_to_column(unique_ptr[device_buffer] device_buffer_unique_ptr, dtype)
-cdef move_device_buffer_to_series(unique_ptr[device_buffer] device_buffer_unique_ptr, dtype, series_name)
+cdef move_device_buffer_to_column(
+    unique_ptr[device_buffer] device_buffer_unique_ptr,
+    dtype,
+    size_t itemsize,
+)
+cdef move_device_buffer_to_series(
+    unique_ptr[device_buffer] device_buffer_unique_ptr,
+    dtype,
+    size_t itemsize,
+    series_name,
+)
 cdef coo_to_df(GraphCOOPtrType graph)
 cdef csr_to_series(GraphCSRPtrType graph)
 cdef GraphCOOViewFloat get_coo_float_graph_view(input_graph, bool weighted=*)
diff --git a/python/cugraph/cugraph/structure/graph_primtypes.pyx b/python/cugraph/cugraph/structure/graph_primtypes.pyx
@@ -16,61 +16,93 @@
 # cython: embedsignature = True
 # cython: language_level = 3
 
-import numpy as np
 from libc.stdint cimport uintptr_t
 from libcpp.utility cimport move
 
 from rmm.pylibrmm.device_buffer cimport DeviceBuffer
-from cudf.core.buffer import as_buffer
+import pylibcudf
 import cudf
 
 
 cdef move_device_buffer_to_column(
-    unique_ptr[device_buffer] device_buffer_unique_ptr, dtype):
+    unique_ptr[device_buffer] device_buffer_unique_ptr,
+    dtype,
+    size_t itemsize,
+):
     """
     Transfers ownership of device_buffer_unique_ptr to a cuDF buffer which is
     used to construct a cudf column object, which is then returned. If the
     intermediate buffer is empty, the device_buffer_unique_ptr is still
     transfered but None is returned.
     """
-    buff = DeviceBuffer.c_from_unique_ptr(move(device_buffer_unique_ptr))
-    buff = as_buffer(buff)
-    if buff.nbytes != 0:
-        column = cudf.core.column.build_column(buff, dtype=cudf.dtype(dtype))
-        return column
+    cdef size_t buff_size = device_buffer_unique_ptr.get().size()
+    cdef DeviceBuffer buff = DeviceBuffer.c_from_unique_ptr(move(device_buffer_unique_ptr))
+    cdef size_t col_size = buff_size // itemsize
+    result_column = pylibcudf.Column.from_rmm_buffer(
+        buff,
+        dtype,
+        col_size,
+        [],
+    )
+    if buff_size != 0:
+        return result_column
     return None
 
 
 cdef move_device_buffer_to_series(
-    unique_ptr[device_buffer] device_buffer_unique_ptr, dtype, series_name):
+    unique_ptr[device_buffer] device_buffer_unique_ptr,
+    dtype,
+    size_t itemsize,
+    series_name
+):
     """
     Transfers ownership of device_buffer_unique_ptr to a cuDF buffer which is
     used to construct a cudf.Series object with name series_name, which is then
     returned. If the intermediate buffer is empty, the device_buffer_unique_ptr
     is still transfered but None is returned.
     """
-    column = move_device_buffer_to_column(move(device_buffer_unique_ptr), dtype)
+    column = move_device_buffer_to_column(
+        move(device_buffer_unique_ptr),
+        dtype,
+        itemsize,
+    )
     if column is not None:
-        series = cudf.Series._from_data({series_name: column})
-        return series
+        return cudf.Series.from_pylibcudf(column, metadata={"name": series_name})
     return None
 
 
 cdef coo_to_df(GraphCOOPtrType graph):
     # FIXME: this function assumes columns named "src" and "dst" and can only
     # be used for SG graphs due to that assumption.
     contents = move(graph.get()[0].release())
-    src = move_device_buffer_to_column(move(contents.src_indices), "int32")
-    dst = move_device_buffer_to_column(move(contents.dst_indices), "int32")
+    src = move_device_buffer_to_series(
+        move(contents.src_indices),
+        pylibcudf.DataType(pylibcudf.TypeId.INT32),
+        4,
+        None,
+    )
+    dst = move_device_buffer_to_series(
+        move(contents.dst_indices),
+        pylibcudf.DataType(pylibcudf.TypeId.INT32),
+        4,
+        None,
+    )
 
     if GraphCOOPtrType is GraphCOOPtrFloat:
-        weight_type = "float32"
+        weight_type = pylibcudf.DataType(pylibcudf.TypeId.FLOAT32)
+        itemsize = 4
     elif GraphCOOPtrType is GraphCOOPtrDouble:
-        weight_type = "float64"
+        weight_type = pylibcudf.DataType(pylibcudf.TypeId.FLOAT64)
+        itemsize = 8
     else:
         raise TypeError("Invalid GraphCOOPtrType")
 
-    wgt = move_device_buffer_to_column(move(contents.edge_data), weight_type)
+    wgt = move_device_buffer_to_series(
+        move(contents.edge_data),
+        weight_type,
+        itemsize,
+        None,
+    )
 
     df = cudf.DataFrame()
     df['src'] = src
@@ -83,20 +115,34 @@ cdef coo_to_df(GraphCOOPtrType graph):
 
 cdef csr_to_series(GraphCSRPtrType graph):
     contents = move(graph.get()[0].release())
-    csr_offsets = move_device_buffer_to_series(move(contents.offsets),
-                                               "int32", "csr_offsets")
-    csr_indices = move_device_buffer_to_series(move(contents.indices),
-                                               "int32", "csr_indices")
+    csr_offsets = move_device_buffer_to_series(
+        move(contents.offsets),
+        pylibcudf.DataType(pylibcudf.TypeId.INT32),
+        4,
+        "csr_offsets"
+    )
+    csr_indices = move_device_buffer_to_series(
+        move(contents.indices),
+        pylibcudf.DataType(pylibcudf.TypeId.INT32),
+        4,
+        "csr_indices"
+    )
 
     if GraphCSRPtrType is GraphCSRPtrFloat:
-        weight_type = "float32"
+        weight_type = pylibcudf.DataType(pylibcudf.TypeId.FLOAT32)
+        itemsize = 4
     elif GraphCSRPtrType is GraphCSRPtrDouble:
-        weight_type = "float64"
+        weight_type = pylibcudf.DataType(pylibcudf.TypeId.FLOAT64)
+        itemsize = 8
     else:
         raise TypeError("Invalid GraphCSRPtrType")
 
-    csr_weights = move_device_buffer_to_series(move(contents.edge_data),
-                                               weight_type, "csr_weights")
+    csr_weights = move_device_buffer_to_series(
+        move(contents.edge_data),
+        weight_type,
+        itemsize,
+        "csr_weights"
+    )
 
     return (csr_offsets, csr_indices, csr_weights)
 
diff --git a/python/cugraph/pyproject.toml b/python/cugraph/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
     "libcugraph==25.6.*,>=0.0.0a0",
     "numba>=0.59.1,<0.61.0a0",
     "numpy>=1.23,<3.0a0",
+    "pylibcudf==25.6.*,>=0.0.0a0",
     "pylibcugraph==25.6.*,>=0.0.0a0",
     "pylibraft==25.6.*,>=0.0.0a0",
     "raft-dask==25.6.*,>=0.0.0a0",
