From 3c1a72c396b3d50c7f6c0f47902871c498768017 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 29 Apr 2024 22:25:01 -0700
Subject: [PATCH] Only use functions in the limited API (#5871)

This PR removes usage of the only method in raft's Cython that is not part of the Python limited API. Contributes to https://github.com/rapidsai/build-planning/issues/42

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: https://github.com/rapidsai/cuml/pull/5871
---
 python/CMakeLists.txt                        |  1 -
 python/cuml/ensemble/randomforest_shared.pyx | 21 ++++++++++++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index 52872fd0d5..036d91eef1 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -41,7 +41,6 @@ option(SINGLEGPU "Disable all mnmg components and comms libraries" OFF)
 set(CUML_RAFT_CLONE_ON_PIN OFF)
 
 
-
 # todo: use CMAKE_MESSAGE_CONTEXT for prefix for logging.
 # https://github.com/rapidsai/cuml/issues/4843
 message(VERBOSE "CUML_PY: Build only cuML CPU Python components.: ${CUML_CPU}")
diff --git a/python/cuml/ensemble/randomforest_shared.pyx b/python/cuml/ensemble/randomforest_shared.pyx
index 4e8c86341b..a164bbcd02 100644
--- a/python/cuml/ensemble/randomforest_shared.pyx
+++ b/python/cuml/ensemble/randomforest_shared.pyx
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,8 +16,9 @@
 
 # distutils: language = c++
 
+from cpython.buffer cimport PyObject_GetBuffer, PyBuffer_Release, PyBUF_FULL_RO
+
 from libcpp.vector cimport vector
-from cpython.object cimport PyObject
 from libc.stdint cimport uintptr_t
 from libcpp.memory cimport unique_ptr
 from typing import Dict, List, Union
@@ -37,9 +38,6 @@ cdef extern from "treelite/tree.h" namespace "treelite":
         @staticmethod
         unique_ptr[Model] DeserializeFromPyBuffer(const vector[TreelitePyBufferFrame] &) except +
 
-cdef extern from "Python.h":
-    Py_buffer* PyMemoryView_GET_BUFFER(PyObject* mview)
-
 cdef class PyBufferFrameWrapper:
     cdef TreelitePyBufferFrame _handle
     cdef Py_ssize_t shape[1]
@@ -92,18 +90,25 @@ def get_frames(model: uintptr_t) -> List[memoryview]:
 def init_from_frames(frames: List[np.ndarray],
                      format_str: List[str], itemsize: List[int]) -> uintptr_t:
     cdef vector[TreelitePyBufferFrame] cpp_frames
+    # Need to keep track of the buffers to release them later.
+    cdef vector[Py_buffer] buffers
     cdef Py_buffer* buf
     cdef TreelitePyBufferFrame cpp_frame
     format_bytes = [s.encode('utf-8') for s in format_str]
     for i, frame in enumerate(frames):
-        x = memoryview(frame)
-        buf = PyMemoryView_GET_BUFFER(<PyObject*>x)
+        buffers.emplace_back()
+        buf = &buffers.back()
+        PyObject_GetBuffer(frame, buf, PyBUF_FULL_RO)
         cpp_frame.buf = buf.buf
         cpp_frame.format = format_bytes[i]
         cpp_frame.itemsize = itemsize[i]
         cpp_frame.nitem = buf.len // itemsize[i]
         cpp_frames.push_back(cpp_frame)
-    return <uintptr_t> _init_from_frames(cpp_frames)
+    output = <uintptr_t> _init_from_frames(cpp_frames)
+    cdef int j
+    for j in range(buffers.size()):
+        PyBuffer_Release(&buffers[j])
+    return output
 
 
 def treelite_serialize(