feat: add subinterpreters support (#245)

Author: XuehaiPan

.github/workflows/build.yml

@@ -6,6 +6,8 @@ on:
       - main # allow to trigger the workflow with tag push event
   pull_request:
     types:
+      - labeled
+      - unlabeled
       - opened
       - synchronize
       - reopened
@@ -33,6 +35,11 @@ on:
           - build-only
           - build-and-publish
         required: true
+      nightly-pybind11:
+        description: "Use nightly pybind11"
+        type: boolean
+        required: false
+        default: false
 
 permissions:
   contents: read
@@ -85,6 +92,20 @@ jobs:
       - name: Print version
         run: python setup.py --version
 
+      - name: Use nightly pybind11
+        shell: bash
+        if: |
+          (github.event_name == 'workflow_dispatch' && github.event.inputs.nightly-pybind11 == 'true') ||
+          (
+            github.event_name == 'pull_request' &&
+            contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          )
+        run: |
+          python -m pip install --force-reinstall "$(python .github/workflows/set_setup_requires.py)"
+          echo "::group::pyproject.toml"
+          cat pyproject.toml
+          echo "::endgroup::"
+
       - name: Install dependencies
         run: python -m pip install --upgrade pip setuptools wheel build
 
@@ -246,6 +267,20 @@ jobs:
       - name: Print version
         run: python setup.py --version
 
+      - name: Use nightly pybind11
+        shell: bash
+        if: |
+          (github.event_name == 'workflow_dispatch' && github.event.inputs.nightly-pybind11 == 'true') ||
+          (
+            github.event_name == 'pull_request' &&
+            contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          )
+        run: |
+          python -m pip install --force-reinstall "$(python .github/workflows/set_setup_requires.py)"
+          echo "::group::pyproject.toml"
+          cat pyproject.toml
+          echo "::endgroup::"
+
       - name: Set CIBW_BUILD
         shell: bash
         run: |

.github/workflows/lint.yml

@@ -13,6 +13,12 @@ on:
       - reopened
   # Allow to trigger the workflow manually
   workflow_dispatch:
+    inputs:
+      nightly-pybind11:
+        description: "Use nightly pybind11"
+        type: boolean
+        required: false
+        default: false
 
 permissions:
   contents: read
@@ -72,8 +78,11 @@ jobs:
       - name: Install nightly pybind11
         shell: bash
         if: |
-          github.event_name == 'pull_request' &&
-          contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          (github.event_name == 'workflow_dispatch' && github.event.inputs.nightly-pybind11 == 'true') ||
+          (
+            github.event_name == 'pull_request' &&
+            contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          )
         run: |
           python -m pip install --force-reinstall "$(python .github/workflows/set_setup_requires.py)"
           echo "::group::pyproject.toml"

.github/workflows/tests-with-pydebug.yml

@@ -4,6 +4,9 @@ on:
   push:
     branches:
       - main
+  schedule:
+    # Run at 12:00 Asia/Shanghai (04:00 UTC) every three days with nightly pybind11
+    - cron: "0 4 */3 * *"
   pull_request:
     types:
       - labeled
@@ -24,6 +27,12 @@ on:
       - .github/workflows/tests-with-pydebug.yml
   # Allow to trigger the workflow manually
   workflow_dispatch:
+    inputs:
+      nightly-pybind11:
+        description: "Use nightly pybind11"
+        type: boolean
+        required: false
+        default: false
 
 permissions:
   contents: read
@@ -41,7 +50,6 @@ env:
   PYTHON: "python" # to be updated
   PYTHON_TAG: "py3" # to be updated
   PYTHON_VERSION: "3" # to be updated
-  pybind11_VERSION: "stable" # to be updated
   PYENV_ROOT: "~/.pyenv" # to be updated
   VENV_BIN_NAME: "bin" # to be updated
   COLUMNS: "100"
@@ -312,11 +320,14 @@ jobs:
       - name: Use nightly pybind11
         shell: bash
         if: |
-          github.event_name == 'pull_request' &&
-          contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          github.event_name == 'schedule' ||
+          (github.event_name == 'workflow_dispatch' && github.event.inputs.nightly-pybind11 == 'true') ||
+          (
+            github.event_name == 'pull_request' &&
+            contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          )
         run: |
           source "venv/${VENV_BIN_NAME}/activate"
-          ${{ env.PYTHON }} .github/workflows/set_setup_requires.py
           ${{ env.PYTHON }} -m pip install --force-reinstall "$(${{ env.PYTHON }} .github/workflows/set_setup_requires.py)"
           echo "::group::pyproject.toml"
           cat pyproject.toml
@@ -348,7 +359,13 @@ jobs:
             "--cov-report=xml:coverage-${{ env.PYTHON_TAG }}-${{ runner.os }}.xml"
             "--junit-xml=junit-${{ env.PYTHON_TAG }}-${{ runner.os }}.xml"
           )
-          make test PYTESTOPTS="${PYTESTOPTS[*]}"
+
+          if ${{ env.PYTHON }} -c 'import sys, optree; sys.exit(not optree._C.OPTREE_HAS_SUBINTERPRETER_SUPPORT)'; then
+            make test PYTESTOPTS="${PYTESTOPTS[*]} -k 'concurrent' --no-cov"
+            make test PYTESTOPTS="${PYTESTOPTS[*]} -k 'not subinterpreter'"
+          else
+            make test PYTESTOPTS="${PYTESTOPTS[*]}"
+          fi
 
           CORE_DUMP_FILES="$(
             find . -type d -path "./venv" -prune \

.github/workflows/tests.yml

@@ -25,6 +25,12 @@ on:
       - .github/workflows/tests.yml
   # Allow to trigger the workflow manually
   workflow_dispatch:
+    inputs:
+      nightly-pybind11:
+        description: "Use nightly pybind11"
+        type: boolean
+        required: false
+        default: false
 
 permissions:
   contents: read
@@ -42,7 +48,6 @@ env:
   PYTHONUNBUFFERED: "1"
   PYTHON: "python" # to be updated
   PYTHON_TAG: "py3" # to be updated
-  pybind11_VERSION: "stable" # to be updated
   VENV_BIN_NAME: "bin" # to be updated
   COLUMNS: "100"
   FORCE_COLOR: "1"
@@ -186,10 +191,12 @@ jobs:
       - name: Use nightly pybind11
         shell: bash
         if: |
-          github.event_name == 'pull_request' &&
-          contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          (github.event_name == 'workflow_dispatch' && github.event.inputs.nightly-pybind11 == 'true') ||
+          (
+            github.event_name == 'pull_request' &&
+            contains(github.event.pull_request.labels.*.name, 'test-with-nightly-pybind11')
+          )
         run: |
-          ${{ env.PYTHON }} .github/workflows/set_setup_requires.py
           ${{ env.PYTHON }} -m pip install --force-reinstall "$(${{ env.PYTHON }} .github/workflows/set_setup_requires.py)"
           echo "::group::pyproject.toml"
           cat pyproject.toml

.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
     hooks:
       - id: cpplint
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.14
+    rev: v0.15.2
     hooks:
       - id: ruff-check
         args: [--fix, --exit-non-zero-on-fix]

CHANGELOG.md

@@ -13,7 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
--
+- Add subinterpreters support for Python 3.14+ by [@XuehaiPan](https://github.com/XuehaiPan) in [#245](https://github.com/metaopt/optree/pull/245).
 
 ### Changed
 

include/optree/pymacros.h

@@ -17,6 +17,8 @@ limitations under the License.
 
 #pragma once
 
+#include <stdexcept>  // std::runtime_error
+
 #include <Python.h>
 
 #include <pybind11/pybind11.h>
@@ -32,6 +34,15 @@ limitations under the License.
 // NOLINTNEXTLINE[bugprone-macro-parentheses]
 #define NONZERO_OR_EMPTY(MACRO) ((MACRO + 0 != 0) || (0 - MACRO - 1 >= 0))
 
+#if !defined(PYPY_VERSION) && (PY_VERSION_HEX >= 0x030E0000 /* Python 3.14 */) &&                  \
+    (PYBIND11_VERSION_HEX >= 0x030002F0 /* pybind11 3.0.2 */) &&                                   \
+    (defined(PYBIND11_HAS_SUBINTERPRETER_SUPPORT) &&                                               \
+     NONZERO_OR_EMPTY(PYBIND11_HAS_SUBINTERPRETER_SUPPORT))
+#    define OPTREE_HAS_SUBINTERPRETER_SUPPORT 1
+#else
+#    undef OPTREE_HAS_SUBINTERPRETER_SUPPORT
+#endif
+
 namespace py = pybind11;
 
 #if !defined(Py_ALWAYS_INLINE)
@@ -59,3 +70,50 @@ inline constexpr Py_ALWAYS_INLINE bool Py_IsConstant(PyObject *x) noexcept {
     return Py_IsNone(x) || Py_IsTrue(x) || Py_IsFalse(x);
 }
 #define Py_IsConstant(x) Py_IsConstant(x)
+
+using interpid_t = decltype(PyInterpreterState_GetID(nullptr));
+
+#if defined(PYBIND11_HAS_SUBINTERPRETER_SUPPORT) &&                                                \
+    NONZERO_OR_EMPTY(PYBIND11_HAS_SUBINTERPRETER_SUPPORT)
+
+[[nodiscard]] inline bool IsCurrentPyInterpreterMain() {
+    return PyInterpreterState_Get() == PyInterpreterState_Main();
+}
+
+[[nodiscard]] inline interpid_t GetCurrentPyInterpreterID() {
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    if (PyErr_Occurred() != nullptr) [[unlikely]] {
+        throw py::error_already_set();
+    }
+    if (interp == nullptr) [[unlikely]] {
+        throw std::runtime_error("Failed to get the current Python interpreter state.");
+    }
+    const interpid_t interpid = PyInterpreterState_GetID(interp);
+    if (PyErr_Occurred() != nullptr) [[unlikely]] {
+        throw py::error_already_set();
+    }
+    return interpid;
+}
+
+[[nodiscard]] inline interpid_t GetMainPyInterpreterID() {
+    PyInterpreterState *interp = PyInterpreterState_Main();
+    if (PyErr_Occurred() != nullptr) [[unlikely]] {
+        throw py::error_already_set();
+    }
+    if (interp == nullptr) [[unlikely]] {
+        throw std::runtime_error("Failed to get the main Python interpreter state.");
+    }
+    const interpid_t interpid = PyInterpreterState_GetID(interp);
+    if (PyErr_Occurred() != nullptr) [[unlikely]] {
+        throw py::error_already_set();
+    }
+    return interpid;
+}
+
+#else
+
+[[nodiscard]] inline bool IsCurrentPyInterpreterMain() noexcept { return true; }
+[[nodiscard]] inline interpid_t GetCurrentPyInterpreterID() noexcept { return 0; }
+[[nodiscard]] inline interpid_t GetMainPyInterpreterID() noexcept { return 0; }
+
+#endif

include/optree/registry.h

@@ -23,12 +23,13 @@ limitations under the License.
 #include <string>         // std::string
 #include <unordered_map>  // std::unordered_map
 #include <unordered_set>  // std::unordered_set
-#include <utility>        // std::pair
+#include <utility>        // std::pair, std::make_pair
 
 #include <pybind11/pybind11.h>
 
 #include "optree/exceptions.h"
 #include "optree/hashing.h"
+#include "optree/pymacros.h"
 #include "optree/synchronization.h"
 
 namespace optree {
@@ -141,6 +142,52 @@ class PyTreeTypeRegistry {
         return count1;
     }
 
+    // Get the number of alive interpreters that have seen the registry.
+    [[nodiscard]] static inline Py_ALWAYS_INLINE ssize_t GetNumInterpretersAlive() {
+        const scoped_read_lock lock{sm_mutex};
+        return py::ssize_t_cast(sm_alive_interpids.size());
+    }
+
+    // Get the number of interpreters that have seen the registry.
+    [[nodiscard]] static inline Py_ALWAYS_INLINE ssize_t GetNumInterpretersSeen() {
+        const scoped_read_lock lock{sm_mutex};
+        return sm_num_interpreters_seen;
+    }
+
+    // Get the IDs of alive interpreters that have seen the registry.
+    [[nodiscard]] static inline Py_ALWAYS_INLINE std::unordered_set<interpid_t>
+    GetAliveInterpreterIDs() {
+        const scoped_read_lock lock{sm_mutex};
+        return sm_alive_interpids;
+    }
+
+    // Check if should preserve the insertion order of the dictionary keys during flattening.
+    [[nodiscard]] static inline Py_ALWAYS_INLINE bool IsDictInsertionOrdered(
+        const std::string &registry_namespace,
+        const bool &inherit_global_namespace = true) {
+        const scoped_read_lock lock{sm_dict_order_mutex};
+
+        const auto interpid = GetCurrentPyInterpreterID();
+        const auto &namespaces = sm_dict_insertion_ordered_namespaces;
+        return (namespaces.find({interpid, registry_namespace}) != namespaces.end()) ||
+               (inherit_global_namespace && namespaces.find({interpid, ""}) != namespaces.end());
+    }
+
+    // Set the namespace to preserve the insertion order of the dictionary keys during flattening.
+    static inline Py_ALWAYS_INLINE void SetDictInsertionOrdered(
+        const bool &mode,
+        const std::string &registry_namespace) {
+        const scoped_write_lock lock{sm_dict_order_mutex};
+
+        const auto interpid = GetCurrentPyInterpreterID();
+        const auto key = std::make_pair(interpid, registry_namespace);
+        if (mode) [[likely]] {
+            sm_dict_insertion_ordered_namespaces.insert(key);
+        } else [[unlikely]] {
+            sm_dict_insertion_ordered_namespaces.erase(key);
+        }
+    }
+
     friend void BuildModule(py::module_ &mod);  // NOLINT[runtime/references]
 
 private:
@@ -173,7 +220,16 @@ class PyTreeTypeRegistry {
     NamedRegistrationsMap m_named_registrations{};
     BuiltinsTypesSet m_builtins_types{};
 
+    // A set of namespaces that preserve the insertion order of the dictionary keys during
+    // flattening.
+    static inline std::unordered_set<std::pair<interpid_t, std::string>>
+        sm_dict_insertion_ordered_namespaces{};
+    static inline read_write_mutex sm_dict_order_mutex{};
+    friend class PyTreeSpec;
+
+    static inline std::unordered_set<interpid_t> sm_alive_interpids{};
     static inline read_write_mutex sm_mutex{};
+    static inline ssize_t sm_num_interpreters_seen = 0;
 };
 
 }  // namespace optree

include/optree/synchronization.h

@@ -62,15 +62,15 @@ using scoped_recursive_lock = std::scoped_lock<recursive_mutex>;
 #if (defined(__APPLE__) /* header <shared_mutex> is not available on macOS build target */ &&      \
      PY_VERSION_HEX < 0x030C00F0 /* Python 3.12.0 */)
 
-#    undef HAVE_READ_WRITE_LOCK
+#    undef OPTREE_HAS_READ_WRITE_LOCK
 
 using read_write_mutex = mutex;
 using scoped_read_lock = scoped_lock;
 using scoped_write_lock = scoped_lock;
 
 #else
 
-#    define HAVE_READ_WRITE_LOCK
+#    define OPTREE_HAS_READ_WRITE_LOCK 1
 
 #    include <shared_mutex>  // std::shared_mutex, std::shared_lock