Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support free-threaded Python #165

Merged
merged 14 commits into from
Nov 26, 2024
23 changes: 16 additions & 7 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,19 @@ jobs:

test:
needs: [resolve]
name: test rust-${{ matrix.rust-version }} on ${{ matrix.runs-on }}
name: test rust-${{ matrix.rust-version }} on ${{ matrix.runs-on }}, Python ${{ matrix.python-version }}
strategy:
fail-fast: false
matrix:
rust-version: [stable]
runs-on: [ubuntu, macos]
python-version: ["3.13"]
include:
- rust-version: ${{ needs.resolve.outputs.MSRV }}
runs-on: ubuntu
- rust-version: stable
runs-on: ubuntu
python-version: "3.13t"
- rust-version: nightly
runs-on: ubuntu

Expand All @@ -49,9 +53,9 @@ jobs:
- uses: actions/checkout@v4

- name: set up python
uses: actions/setup-python@v5
uses: quansight-labs/setup-python@v5
with:
python-version: "3.13"
python-version: ${{ matrix.python-version }}

- uses: dtolnay/rust-toolchain@master
with:
Expand All @@ -78,27 +82,32 @@ jobs:
token: ${{ secrets.CODECOV_TOKEN }}

test-python:
name: test jiter-python
name: test jiter-python ${{ matrix.python-version }}

runs-on: ubuntu-latest

strategy:
fail-fast: false
matrix:
python-version: ["3.13", "3.13t"]

env:
RUNS_ON: ubuntu-latest

steps:
- uses: actions/checkout@v4

- name: set up python
uses: actions/setup-python@v5
uses: quansight-labs/setup-python@v5
with:
python-version: "3.13"
python-version: ${{ matrix.python-version }}

- uses: dtolnay/rust-toolchain@stable
- uses: Swatinem/rust-cache@v2
- uses: taiki-e/install-action@cargo-llvm-cov

- name: Install python dependencies
run: make python-install
run: pip install -r crates/jiter-python/tests/requirements.txt

- name: Build jiter-python
run: |
Expand Down
8 changes: 4 additions & 4 deletions crates/jiter-python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,13 @@ mod jiter_python {
}

#[pyfunction]
pub fn cache_clear(py: Python<'_>) {
jiter::cache_clear(py);
pub fn cache_clear() {
jiter::cache_clear();
}

#[pyfunction]
pub fn cache_usage(py: Python<'_>) -> usize {
jiter::cache_usage(py)
pub fn cache_usage() -> usize {
jiter::cache_usage()
}

#[pymodule_init]
Expand Down
6 changes: 3 additions & 3 deletions crates/jiter/benches/python.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use jiter::{cache_clear, PythonParse, StringCacheMode};

fn python_parse_numeric(bench: &mut Bencher) {
Python::with_gil(|py| {
cache_clear(py);
cache_clear();
bench.iter(|| {
PythonParse::default()
.python_parse(
Expand All @@ -23,7 +23,7 @@ fn python_parse_numeric(bench: &mut Bencher) {

fn python_parse_other(bench: &mut Bencher) {
Python::with_gil(|py| {
cache_clear(py);
cache_clear();
bench.iter(|| {
PythonParse::default()
.python_parse(py, br#"["string", true, false, null]"#)
Expand All @@ -39,7 +39,7 @@ fn _python_parse_file(path: &str, bench: &mut Bencher, cache_mode: StringCacheMo
let json_data = contents.as_bytes();

Python::with_gil(|py| {
cache_clear(py);
cache_clear();
bench.iter(|| {
PythonParse {
cache_mode,
Expand Down
33 changes: 19 additions & 14 deletions crates/jiter/src/py_string_cache.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
use std::cell::RefCell;
use std::sync::{Mutex, OnceLock};

use ahash::random_state::RandomState;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::sync::{GILOnceCell, GILProtected};
use pyo3::types::{PyBool, PyString};

#[derive(Debug, Clone, Copy)]
Expand Down Expand Up @@ -86,28 +85,34 @@ impl StringMaybeCache for StringNoCache {
}
}

static STRING_CACHE: GILOnceCell<GILProtected<RefCell<PyStringCache>>> = GILOnceCell::new();
static STRING_CACHE: OnceLock<Mutex<PyStringCache>> = OnceLock::new();

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Despite the use of a mutex here, the single-threaded benchmark is not meaningfully impacted.

We can worry about multithreaded performance in the far future if highly parallel uses of jiter should arise; and if users hit a pathological case before we do anything fancy they can always turn off the string cache.

macro_rules! get_string_cache {
($py:ident) => {
STRING_CACHE
.get_or_init($py, || GILProtected::new(RefCell::new(PyStringCache::default())))
.get($py)
};
#[inline]
fn get_string_cache() -> &'static Mutex<PyStringCache> {
STRING_CACHE.get_or_init(|| Mutex::new(PyStringCache::default()))
}

pub fn cache_usage(py: Python) -> usize {
get_string_cache!(py).borrow().usage()
pub fn cache_usage() -> usize {
get_string_cache()
.lock()
.expect("no code panics with mutex locked")
.usage()
}

pub fn cache_clear(py: Python) {
get_string_cache!(py).borrow_mut().clear();
pub fn cache_clear() {
get_string_cache()
.lock()
.expect("no code panics with mutex locked")
.clear();
}

pub fn cached_py_string<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
// from tests, 0 and 1 character strings are faster not cached
if (2..64).contains(&s.len()) {
get_string_cache!(py).borrow_mut().get_or_insert(py, s, ascii_only)
get_string_cache()
.lock()
.expect("no code panics with mutex locked")
.get_or_insert(py, s, ascii_only)
} else {
pystring_fast_new(py, s, ascii_only)
}
Expand Down
Loading