[GR-50282] Add support for native PickleBuffer as_buffer

PullRequest: graalpython/3254

Author: qunaibit

graalpython/com.oracle.graal.python.cext/src/capi.c

@@ -96,6 +96,14 @@ typedef struct {
     PyObject *mapping;
 } mappingproxyobject;
 
+typedef struct {
+    PyObject_HEAD
+    /* The view exported by the original object */
+    Py_buffer view;
+    PyObject *weakreflist;
+} PyPickleBufferObject;
+
+
 static void object_dealloc(PyObject *self) {
     Py_TYPE(self)->tp_free(self);
 }
@@ -278,6 +286,16 @@ void bytearray_releasebuffer(PyByteArrayObject *obj, Py_buffer *view);
 int memory_getbuf(PyMemoryViewObject *self, Py_buffer *view, int flags);
 void memory_releasebuf(PyMemoryViewObject *self, Py_buffer *view);
 
+/* PICKLEBUFFER */
+static int
+picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
+{
+    PyObject *self_view_obj = GraalPyTruffle_PickleBuffer_viewobj(self);
+    return PyObject_GetBuffer(self_view_obj, view, flags);
+}
+
+static void empty_releasebuf(PyObject *self, Py_buffer *view) {}
+
 static void initialize_bufferprocs() {
     static PyBufferProcs bytes_as_buffer = {
         (getbufferproc)bytes_buffer_getbuffer,       /* bf_getbuffer */
@@ -305,6 +323,12 @@ static void initialize_bufferprocs() {
     array_as_buffer.bf_releasebuffer = GraalPyTruffle_Array_releasebuffer,
     Arraytype.tp_as_buffer = &array_as_buffer;
     GraalPy_set_PyTypeObject_tp_as_buffer(&Arraytype, &array_as_buffer);
+
+    static PyBufferProcs picklebuf_as_buffer;
+    picklebuf_as_buffer.bf_getbuffer = (getbufferproc)picklebuf_getbuf,
+    picklebuf_as_buffer.bf_releasebuffer = empty_releasebuf,
+    PyPickleBuffer_Type.tp_as_buffer = &picklebuf_as_buffer;
+    GraalPy_set_PyTypeObject_tp_as_buffer(&PyPickleBuffer_Type, &picklebuf_as_buffer);
 }
 
 int is_builtin_type(PyTypeObject *tp) {

graalpython/com.oracle.graal.python.cext/src/capi.h

@@ -342,6 +342,8 @@ PY_TRUFFLE_TYPE(UnionType_Type, 		"_ctypes.UnionType",		&PyType_Type, sizeof(PyO
 PY_TRUFFLE_TYPE(PyCPointerType_Type,	"PyCPointerType", 		    &PyType_Type, sizeof(PyObject)) \
 PY_TRUFFLE_TYPE(PyCArrayType_Type,		"PyCArrayType", 			&PyType_Type, sizeof(PyObject)) \
 PY_TRUFFLE_TYPE(PyCoro_Type, 			"coroutine", 				&PyType_Type, sizeof(PyCoroObject)) \
+/* PyPickleBufferObject (PyObject_HEAD + Py_buffer + PyObject*) is defined within Objects/picklebufobject.c, so its not exposed. */ \
+PY_TRUFFLE_TYPE(PyPickleBuffer_Type, 	"_pickle.PickleBuffer",     &PyType_Type, sizeof(PyPickleBufferObject)) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(_PyAIterWrapper_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(_PyAsyncGenASend_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(_PyAsyncGenAThrow_Type) \
@@ -385,7 +387,6 @@ PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictItems_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictIter_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictKeys_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyODictValues_Type) \
-PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyPickleBuffer_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PyRangeIter_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PySetIter_Type) \
 PY_TRUFFLE_TYPE_UNIMPLEMENTED(PySortWrapper_Type) \

graalpython/com.oracle.graal.python.cext/src/obmalloc.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * The Universal Permissive License (UPL), Version 1.0
@@ -110,7 +110,7 @@ int PyTruffle_AllocMemory(size_t size) {
 
 void PyTruffle_FreeMemory(size_t size) {
     if (PyTruffle_AllocatedMemory < size) {
-        PyTruffle_Log(PY_TRUFFLE_LOG_INFO, "PyTruffle_FreeMemory: freed memory size (%lu) is larger than allocated memory size (%lu)\n", size, PyTruffle_AllocMemory);
+        PyTruffle_Log(PY_TRUFFLE_LOG_INFO, "PyTruffle_FreeMemory: freed memory size (%lu) is larger than allocated memory size (%lu)\n", size, PyTruffle_AllocatedMemory);
         PyTruffle_AllocatedMemory = size;
     }
     PyTruffle_AllocatedMemory -= size;

graalpython/com.oracle.graal.python.test/src/tests/cpyext/test_bytes.py

@@ -582,6 +582,45 @@ def test_create_from_buffer_exception(self):
         self.assertRaises(ValueError, bytes, TestType())
         self.assertRaises(ValueError, bytearray, TestType())
 
+    def test_tp_as_buffer_pickle(self):
+        TestAsBufferPickle = CPyExtType(
+            "TestAsBufferPickle",
+            """
+            static PyObject* get_bytes(PyObject* self, PyObject *obj) {
+                Py_buffer view;
+                if (PyObject_GetBuffer(obj, &view, PyBUF_FULL_RO) != 0)
+                    return NULL;
+                PyObject* bytes = PyBytes_FromStringAndSize(view.buf, view.len);
+                PyBuffer_Release(&view);
+                return bytes;
+            }
+            """,
+            tp_methods='{"get_bytes", (PyCFunction)get_bytes, METH_O | METH_CLASS, ""}',
+        )
+        TestType = CPyExtType(
+            "TestMemoryViewBufferPickle",
+            """
+            char buf[] = {1,2,3,4};
+            int getbuffer(TestMemoryViewBufferPickleObject *self, Py_buffer *view, int flags) {
+                return PyBuffer_FillInfo(view, (PyObject*)self, buf, 4, 1, flags);
+            }
+            void releasebuffer(TestMemoryViewBufferPickleObject *self, Py_buffer *view) {}
+            
+            static PyBufferProcs as_buffer = {
+                (getbufferproc)getbuffer,
+                (releasebufferproc)releasebuffer,
+            };
+            """,
+            tp_as_buffer='&as_buffer',
+        )
+        obj = TestType()
+        b = bytes([1,2,3,4]) # same as `buf[] = {1,2,3,4};`
+        import pickle
+        b2 = TestAsBufferPickle.get_bytes(pickle.PickleBuffer(memoryview(obj)))
+        b1 = TestAsBufferPickle.get_bytes(pickle.PickleBuffer(bytearray([1,2,3,4])))
+        assert b == b1
+        assert b == b2
+
 
 class TestNativeSubclass(unittest.TestCase):
     def test_builtins(self):

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/cext/PythonCextPickleBufferBuiltins.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2023, 2024, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * The Universal Permissive License (UPL), Version 1.0
+ *
+ * Subject to the condition set forth below, permission is hereby granted to any
+ * person obtaining a copy of this software, associated documentation and/or
+ * data (collectively the "Software"), free of charge and under any and all
+ * copyright rights in the Software, and any and all patent rights owned or
+ * freely licensable by each licensor hereunder covering either (i) the
+ * unmodified Software as contributed to or provided by such licensor, or (ii)
+ * the Larger Works (as defined below), to deal in both
+ *
+ * (a) the Software, and
+ *
+ * (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+ * one is included with the Software each a "Larger Work" to which the Software
+ * is contributed by such licensors),
+ *
+ * without restriction, including without limitation the rights to copy, create
+ * derivative works of, display, perform, and distribute the Software and make,
+ * use, sell, offer for sale, import, export, have made, and have sold the
+ * Software and the Larger Work(s), and to sublicense the foregoing rights on
+ * either these or other terms.
+ *
+ * This license is subject to the following condition:
+ *
+ * The above copyright notice and either this complete permission notice or at a
+ * minimum a reference to the UPL must be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+package com.oracle.graal.python.builtins.modules.cext;
+
+import static com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiCallPath.Ignored;
+import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyObject;
+import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyObjectBorrowed;
+import static com.oracle.graal.python.runtime.exception.PythonErrorType.ValueError;
+
+import com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiBuiltin;
+import com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiUnaryBuiltinNode;
+import com.oracle.graal.python.builtins.modules.pickle.PPickleBuffer;
+import com.oracle.graal.python.builtins.objects.buffer.PythonBufferAccessLibrary;
+import com.oracle.graal.python.nodes.ErrorMessages;
+import com.oracle.graal.python.nodes.PRaiseNode;
+import com.oracle.truffle.api.dsl.Bind;
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.library.CachedLibrary;
+import com.oracle.truffle.api.nodes.Node;
+
+public final class PythonCextPickleBufferBuiltins {
+
+    @CApiBuiltin(ret = PyObjectBorrowed, args = {PyObject}, call = Ignored)
+    abstract static class PyTruffle_PickleBuffer_viewobj extends CApiUnaryBuiltinNode {
+        @Specialization
+        static Object getviewobj(PPickleBuffer object,
+                        @Bind("this") Node inliningTarget,
+                        @CachedLibrary(limit = "3") PythonBufferAccessLibrary bufferLib,
+                        @Cached PRaiseNode.Lazy raiseNode) {
+            Object owner = null;
+            if (object.getView() != null) {
+                owner = bufferLib.getOwner(object.getView());
+            }
+            if (owner == null) {
+                throw raiseNode.get(inliningTarget).raise(ValueError, ErrorMessages.OP_FORBIDDEN_ON_OBJECT, "PickleBuffer");
+            }
+            return owner;
+        }
+
+    }
+}

graalpython/lib-graalpython/patches/Cython/Cython-0.29.32.patch

@@ -202,6 +202,19 @@ index b2c67dc..ad80ca4 100644
      }
  #endif
      else if (PyMethod_Check(method)) {
+diff --git a/Cython/Utility/Exceptions.c b/Cython/Utility/Exceptions.c
+index 2cd4b60..2c92132 100644
+--- a/Cython/Utility/Exceptions.c
++++ b/Cython/Utility/Exceptions.c
+@@ -661,7 +661,7 @@ static int __Pyx_CLineForTraceback(CYTHON_NCP_UNUSED PyThreadState *tstate, int
+ 
+     __Pyx_ErrFetchInState(tstate, &ptype, &pvalue, &ptraceback);
+ 
+-#if CYTHON_COMPILING_IN_CPYTHON
++#if 0
+     cython_runtime_dict = _PyObject_GetDictPtr(${cython_runtime_cname});
+     if (likely(cython_runtime_dict)) {
+         __PYX_PY_DICT_LOOKUP_IF_MODIFIED(
 diff --git a/Cython/Utility/ModuleSetupCode.c b/Cython/Utility/ModuleSetupCode.c
 index f8bf885..94d0ca6 100644
 --- a/Cython/Utility/ModuleSetupCode.c

graalpython/lib-graalpython/patches/msgpack/metadata.toml

@@ -0,0 +1,3 @@
+[[rules]]
+version = '== 1.0.8'
+patch = 'msgpack-1.0.8.patch'

graalpython/lib-graalpython/patches/msgpack/msgpack-1.0.8.patch

@@ -0,0 +1,13 @@
+diff --git a/pyproject.toml b/pyproject.toml
+index f9af967..1f5a002 100644
+--- a/pyproject.toml
++++ b/pyproject.toml
+@@ -2,7 +2,7 @@
+ requires = [
+     # Also declared in requirements.txt, if updating here please also update
+     # there
+-    "Cython~=3.0.8",
++    "Cython < 3",
+     "setuptools >= 35.0.2",
+ ]
+ build-backend = "setuptools.build_meta"

graalpython/lib-graalpython/patches/ray/ray-2.9.1.patch

@@ -123,6 +123,36 @@ index bee0ae6..ee0e3ea 100644
              "actor_method_names": json.dumps(list(actor_method_names)),
          }
 
+diff --git a/python/ray/_private/resource_spec.py b/python/ray/_private/resource_spec.py
+index 69780aade9..3be8cbbc4c 100644
+--- a/python/ray/_private/resource_spec.py
++++ b/python/ray/_private/resource_spec.py
+@@ -160,9 +160,22 @@ class ResourceSpec(
+         if is_head:
+             resources[HEAD_NODE_RESOURCE_NAME] = 1.0
+ 
+-        num_cpus = self.num_cpus
+-        if num_cpus is None:
+-            num_cpus = ray._private.utils.get_num_cpus()
++        try:
++            system_num_cpus = __graalpython__.get_max_process_count()
++            factor = min(system_num_cpus, 4)
++            num_cpus = system_num_cpus // factor
++            if self.num_cpus:
++                num_cpus = min(self.num_cpus, num_cpus)
++                if num_cpus != self.num_cpus:
++                    logger.warning(
++                        f"GraalPy: user requested {self.num_cpus} for number"
++                        "of CPUs, but GraalPy can only support up to {num_cpus} CPUs"
++                    )
++        except:
++            # We are not running on GraalPy
++            num_cpus = self.num_cpus
++            if num_cpus is None:
++                num_cpus = ray._private.utils.get_num_cpus()
+ 
+         num_gpus = 0
+         for (
 diff --git a/python/ray/_private/services.py b/python/ray/_private/services.py
 index 867b748..40b3fb2 100644
 --- a/python/ray/_private/services.py
@@ -272,10 +302,10 @@ index a34a39c..51a1a17 100644
      if BAZEL_ARGS:
 diff --git a/ray_build_backend.py b/ray_build_backend.py
 new file mode 100644
-index 0000000..9e7f3fd
+index 0000000000..10f532ae9e
 --- /dev/null
 +++ b/ray_build_backend.py
-@@ -0,0 +1,117 @@
+@@ -0,0 +1,125 @@
 +# We need a whole custom build backend just because the setup.py is in a subdir
 +import os
 +import re
@@ -385,6 +415,14 @@ index 0000000..9e7f3fd
 +    env = os.environ.copy()
 +    env['SKIP_THIRDPARTY_INSTALL'] = 'true'
 +    try:
++        system_num_cpus = __graalpython__.get_max_process_count()
++        env['BAZEL_LIMIT_CPUS'] = '%d' % system_num_cpus
++    except:
++        # do not limit bazel
++        pass
++    try:
++        if os.path.isdir(Path('dist')):
++            shutil.rmtree(Path('dist'))
 +        subprocess.run([sys.executable, 'setup.py', 'bdist_wheel'], env=env, check=True)
 +        wheels = list(Path('dist').glob('*.whl'))
 +        assert len(wheels) == 1, f"Expected 1 wheel, found {len(wheels)}"

mx.graalpython/mx_graalpython.py

@@ -2128,7 +2128,7 @@ def _python_checkpatchfiles():
             # Puts whole license into the field. It's BSD 3-Clause
             'pythran-0.12.0.patch',
             'pythran-0.13.patch',
-            # Puts whole license into the field. It's BSD 3-Clause
+            'pyzmq.patch',
             'jupyter_server.patch',
             # Empty license field. It's BSD-3-Clause
             'prompt_toolkit.patch',