Skip to content

Commit

Permalink
WIP: declare free-threaded support in pymodule macro
Browse files Browse the repository at this point in the history
  • Loading branch information
ngoldbaum committed Oct 1, 2024
1 parent ce18f79 commit 6834f3f
Show file tree
Hide file tree
Showing 11 changed files with 89 additions and 32 deletions.
8 changes: 0 additions & 8 deletions noxfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -654,14 +654,6 @@ def test_version_limits(session: nox.Session):
config_file.set("PyPy", "3.11")
_run_cargo(session, "check", env=env, expect_error=True)

# Python build with GIL disabled should fail building
config_file.set("CPython", "3.13", build_flags=["Py_GIL_DISABLED"])
_run_cargo(session, "check", env=env, expect_error=True)

# Python build with GIL disabled should pass with env flag on
env["UNSAFE_PYO3_BUILD_FREE_THREADED"] = "1"
_run_cargo(session, "check", env=env)


@nox.session(name="check-feature-powerset", venv_backend="none")
def check_feature_powerset(session: nox.Session):
Expand Down
8 changes: 0 additions & 8 deletions pyo3-ffi/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,14 +127,6 @@ fn ensure_gil_enabled(interpreter_config: &InterpreterConfig) -> Result<()> {
.0
.contains(&BuildFlag::Py_GIL_DISABLED)
.not();
ensure!(
gil_enabled || std::env::var("UNSAFE_PYO3_BUILD_FREE_THREADED").map_or(false, |os_str| os_str == "1"),
"the Python interpreter was built with the GIL disabled, which is not yet supported by PyO3\n\
= help: see https://github.com/PyO3/pyo3/issues/4265 for more information\n\
= help: please check if an updated version of PyO3 is available. Current version: {}\n\
= help: set UNSAFE_PYO3_BUILD_FREE_THREADED=1 to suppress this check and build anyway for free-threaded Python",
std::env::var("CARGO_PKG_VERSION").unwrap()
);
if !gil_enabled && interpreter_config.abi3 {
warn!(
"The free-threaded build of CPython does not yet support abi3 so the build artifacts will be version-specific."
Expand Down
22 changes: 17 additions & 5 deletions pyo3-ffi/src/moduleobject.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,17 +86,29 @@ impl Default for PyModuleDef_Slot {

pub const Py_mod_create: c_int = 1;
pub const Py_mod_exec: c_int = 2;
#[cfg(Py_3_12)]
#[cfg(all(not(Py_LIMITED_API), Py_3_12))]
pub const Py_mod_multiple_interpreters: c_int = 3;
#[cfg(all(not(Py_LIMITED_API), Py_3_13))]
pub const Py_mod_gil: c_int = 4;

#[cfg(Py_3_12)]
// skipped private _Py_mod_LAST_SLOT

#[cfg(all(not(Py_LIMITED_API), Py_3_12))]
pub const Py_MOD_MULTIPLE_INTERPRETERS_NOT_SUPPORTED: *mut c_void = 0 as *mut c_void;
#[cfg(Py_3_12)]
#[cfg(all(not(Py_LIMITED_API), Py_3_12))]
pub const Py_MOD_MULTIPLE_INTERPRETERS_SUPPORTED: *mut c_void = 1 as *mut c_void;
#[cfg(Py_3_12)]
#[cfg(all(not(Py_LIMITED_API), Py_3_12))]
pub const Py_MOD_PER_INTERPRETER_GIL_SUPPORTED: *mut c_void = 2 as *mut c_void;

// skipped non-limited _Py_mod_LAST_SLOT
#[cfg(all(not(Py_LIMITED_API), Py_3_13))]
pub const Py_MOD_GIL_USED: *mut c_void = 0 as *mut c_void;
#[cfg(all(not(Py_LIMITED_API), Py_3_13))]
pub const Py_MOD_GIL_NOT_USED: *mut c_void = 1 as *mut c_void;

#[cfg(all(not(Py_LIMITED_API), Py_GIL_DISABLED))]
extern "C" {
pub fn PyUnstable_Module_SetGIL(module: *mut PyObject, gil: *mut c_void);
}

#[repr(C)]
pub struct PyModuleDef {
Expand Down
4 changes: 3 additions & 1 deletion pyo3-macros-backend/src/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use syn::{
punctuated::Punctuated,
spanned::Spanned,
token::Comma,
Attribute, Expr, ExprPath, Ident, Index, LitStr, Member, Path, Result, Token,
Attribute, Expr, ExprPath, Ident, Index, LitBool, LitStr, Member, Path, Result, Token,
};

pub mod kw {
Expand Down Expand Up @@ -44,6 +44,7 @@ pub mod kw {
syn::custom_keyword!(transparent);
syn::custom_keyword!(unsendable);
syn::custom_keyword!(weakref);
syn::custom_keyword!(supports_free_threaded);
}

fn take_int(read: &mut &str, tracker: &mut usize) -> String {
Expand Down Expand Up @@ -308,6 +309,7 @@ pub type RenameAllAttribute = KeywordAttribute<kw::rename_all, RenamingRuleLitSt
pub type StrFormatterAttribute = OptionalKeywordAttribute<kw::str, StringFormatter>;
pub type TextSignatureAttribute = KeywordAttribute<kw::text_signature, TextSignatureAttributeValue>;
pub type SubmoduleAttribute = kw::submodule;
pub type FreeThreadedAttribute = KeywordAttribute<kw::supports_free_threaded, LitBool>;

impl<K: Parse + std::fmt::Debug, V: Parse> Parse for KeywordAttribute<K, V> {
fn parse(input: ParseStream<'_>) -> Result<Self> {
Expand Down
34 changes: 29 additions & 5 deletions pyo3-macros-backend/src/module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@

use crate::{
attributes::{
self, kw, take_attributes, take_pyo3_options, CrateAttribute, ModuleAttribute,
NameAttribute, SubmoduleAttribute,
self, kw, take_attributes, take_pyo3_options, CrateAttribute, FreeThreadedAttribute,
ModuleAttribute, NameAttribute, SubmoduleAttribute,
},
get_doc,
pyclass::PyClassPyO3Option,
Expand All @@ -29,6 +29,7 @@ pub struct PyModuleOptions {
name: Option<NameAttribute>,
module: Option<ModuleAttribute>,
submodule: Option<kw::submodule>,
supports_free_threaded: Option<FreeThreadedAttribute>,
}

impl Parse for PyModuleOptions {
Expand Down Expand Up @@ -72,6 +73,9 @@ impl PyModuleOptions {
submodule,
" (it is implicitly always specified for nested modules)"
),
PyModulePyO3Option::SupportsFreeThreaded(supports_free_threaded) => {
set_option!(supports_free_threaded)
}
}
}
Ok(())
Expand Down Expand Up @@ -344,7 +348,15 @@ pub fn pymodule_module_impl(
)
}
}};
let initialization = module_initialization(&name, ctx, module_def, options.submodule.is_some());
let initialization = module_initialization(
&name,
ctx,
module_def,
options.submodule.is_some(),
options
.supports_free_threaded
.is_some_and(|op| op.value.value),
);

Ok(quote!(
#(#attrs)*
Expand Down Expand Up @@ -383,7 +395,15 @@ pub fn pymodule_function_impl(
let vis = &function.vis;
let doc = get_doc(&function.attrs, None, ctx);

let initialization = module_initialization(&name, ctx, quote! { MakeDef::make_def() }, false);
let initialization = module_initialization(
&name,
ctx,
quote! { MakeDef::make_def() },
false,
options
.supports_free_threaded
.is_some_and(|op| op.value.value),
);

// Module function called with optional Python<'_> marker as first arg, followed by the module.
let mut module_args = Vec::new();
Expand Down Expand Up @@ -428,6 +448,7 @@ fn module_initialization(
ctx: &Ctx,
module_def: TokenStream,
is_submodule: bool,
supports_free_threaded: bool,
) -> TokenStream {
let Ctx { pyo3_path, .. } = ctx;
let pyinit_symbol = format!("PyInit_{}", name);
Expand All @@ -449,7 +470,7 @@ fn module_initialization(
#[doc(hidden)]
#[export_name = #pyinit_symbol]
pub unsafe extern "C" fn __pyo3_init() -> *mut #pyo3_path::ffi::PyObject {
#pyo3_path::impl_::trampoline::module_init(|py| _PYO3_DEF.make_module(py))
#pyo3_path::impl_::trampoline::module_init(|py| _PYO3_DEF.make_module(py, #supports_free_threaded))
}
});
}
Expand Down Expand Up @@ -596,6 +617,7 @@ enum PyModulePyO3Option {
Crate(CrateAttribute),
Name(NameAttribute),
Module(ModuleAttribute),
SupportsFreeThreaded(FreeThreadedAttribute),
}

impl Parse for PyModulePyO3Option {
Expand All @@ -609,6 +631,8 @@ impl Parse for PyModulePyO3Option {
input.parse().map(PyModulePyO3Option::Module)
} else if lookahead.peek(attributes::kw::submodule) {
input.parse().map(PyModulePyO3Option::Submodule)
} else if lookahead.peek(attributes::kw::supports_free_threaded) {
input.parse().map(PyModulePyO3Option::SupportsFreeThreaded)
} else {
Err(lookahead.error())
}
Expand Down
12 changes: 12 additions & 0 deletions pytests/src/free_threaded_mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
use pyo3::prelude::*;

#[pyfunction]
fn add_two(x: usize) -> usize {
x + 2
}

#[pymodule]
pub fn free_threaded_mod(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(add_two, m)?)?;
Ok(())
}
2 changes: 2 additions & 0 deletions pytests/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ pub mod comparisons;
pub mod datetime;
pub mod dict_iter;
pub mod enums;
pub mod free_threaded_mod;
pub mod misc;
pub mod objstore;
pub mod othermod;
Expand Down Expand Up @@ -35,6 +36,7 @@ fn pyo3_pytests(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pymodule!(pyfunctions::pyfunctions))?;
m.add_wrapped(wrap_pymodule!(sequence::sequence))?;
m.add_wrapped(wrap_pymodule!(subclassing::subclassing))?;
m.add_wrapped(wrap_pymodule!(free_threaded_mod::free_threaded_mod))?;

// Inserting to sys.modules allows importing submodules nicely from Python
// e.g. import pyo3_pytests.buf_and_str as bas
Expand Down
1 change: 1 addition & 0 deletions pytests/tests/test_free_threaded.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from pyo3_pytests import free_threaded_mod
26 changes: 22 additions & 4 deletions src/impl_/pymodule.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use portable_atomic::{AtomicI64, Ordering};
not(all(windows, Py_LIMITED_API, not(Py_3_10))),
target_has_atomic = "64",
))]
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};

#[cfg(not(any(PyPy, GraalPy)))]
use crate::exceptions::PyImportError;
Expand All @@ -41,6 +41,8 @@ pub struct ModuleDef {
interpreter: AtomicI64,
/// Initialized module object, cached to avoid reinitialization.
module: GILOnceCell<Py<PyModule>>,
/// Whether or not the module supports running without the GIL
supports_free_threaded: AtomicBool,
}

/// Wrapper to enable initializer to be used in const fns.
Expand Down Expand Up @@ -85,10 +87,15 @@ impl ModuleDef {
))]
interpreter: AtomicI64::new(-1),
module: GILOnceCell::new(),
supports_free_threaded: AtomicBool::new(false),
}
}
/// Builds a module using user given initializer. Used for [`#[pymodule]`][crate::pymodule].
pub fn make_module(&'static self, py: Python<'_>) -> PyResult<Py<PyModule>> {
pub fn make_module(
&'static self,
py: Python<'_>,
supports_free_threaded: bool,
) -> PyResult<Py<PyModule>> {
// Check the interpreter ID has not changed, since we currently have no way to guarantee
// that static data is not reused across interpreters.
//
Expand Down Expand Up @@ -134,6 +141,11 @@ impl ModuleDef {
ffi::PyModule_Create(self.ffi_def.get()),
)?
};
if supports_free_threaded {
unsafe {
ffi::PyUnstable_Module_SetGIL(module.as_ptr(), ffi::Py_MOD_GIL_NOT_USED)
};
}
self.initializer.0(module.bind(py))?;
Ok(module)
})
Expand Down Expand Up @@ -190,7 +202,13 @@ impl PyAddToModule for PyMethodDef {
/// For adding a module to a module.
impl PyAddToModule for ModuleDef {
fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_submodule(self.make_module(module.py())?.bind(module.py()))
module.add_submodule(
self.make_module(
module.py(),
self.supports_free_threaded.load(Ordering::Relaxed),
)?
.bind(module.py()),
)
}
}

Expand Down Expand Up @@ -223,7 +241,7 @@ mod tests {
)
};
Python::with_gil(|py| {
let module = MODULE_DEF.make_module(py).unwrap().into_bound(py);
let module = MODULE_DEF.make_module(py, false).unwrap().into_bound(py);
assert_eq!(
module
.getattr("__name__")
Expand Down
3 changes: 2 additions & 1 deletion src/macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,8 @@ macro_rules! wrap_pymodule {
&|py| {
use $module as wrapped_pymodule;
wrapped_pymodule::_PYO3_DEF
.make_module(py)
// corrected in add_wrapped later based on the parent module's settings
.make_module(py, false)
.expect("failed to wrap pymodule")
}
};
Expand Down
1 change: 1 addition & 0 deletions src/types/module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -488,6 +488,7 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
T: IntoPyCallbackOutput<PyObject>,
{
fn inner(module: &Bound<'_, PyModule>, object: Bound<'_, PyAny>) -> PyResult<()> {
if object.is_instance_of::<PyModule>() {}
let name = object.getattr(__name__(module.py()))?;
module.add(name.downcast_into::<PyString>()?, object)
}
Expand Down

0 comments on commit 6834f3f

Please sign in to comment.