Set default max threads in Python

fdintino · Oct 22, 2024 · fea9ba6 · fea9ba6
1 parent 5f07715
commit fea9ba6
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 63 deletions.
diff --git a/src/PIL/AvifImagePlugin.py b/src/PIL/AvifImagePlugin.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from io import BytesIO
 from typing import IO
 
@@ -46,6 +47,15 @@ def _accept(prefix: bytes) -> bool | str:
     return False
 
 
+def _get_default_max_threads():
+    if DEFAULT_MAX_THREADS:
+        return DEFAULT_MAX_THREADS
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+    else:
+        return os.cpu_count() or 1
+
+
 class AvifImageFile(ImageFile.ImageFile):
     format = "AVIF"
     format_description = "AVIF image"
@@ -64,7 +74,10 @@ def _open(self) -> None:
             raise SyntaxError(msg)
 
         self._decoder = _avif.AvifDecoder(
-            self.fp.read(), DECODE_CODEC_CHOICE, CHROMA_UPSAMPLING, DEFAULT_MAX_THREADS
+            self.fp.read(),
+            DECODE_CODEC_CHOICE,
+            CHROMA_UPSAMPLING,
+            _get_default_max_threads(),
         )
 
         # Get info from decoder
@@ -140,7 +153,7 @@ def _save(
     duration = info.get("duration", 0)
     subsampling = info.get("subsampling", "4:2:0")
     speed = info.get("speed", 6)
-    max_threads = info.get("max_threads", DEFAULT_MAX_THREADS)
+    max_threads = info.get("max_threads", _get_default_max_threads())
     codec = info.get("codec", "auto")
     range_ = info.get("range", "full")
     tile_rows_log2 = info.get("tile_rows", 0)

diff --git a/src/_avif.c b/src/_avif.c
@@ -44,54 +44,6 @@ typedef struct {
 
 static PyTypeObject AvifDecoder_Type;
 
-static int default_max_threads = 0;
-
-static void
-init_max_threads(void) {
-    PyObject *os = NULL;
-    PyObject *n = NULL;
-    long num_cpus;
-
-    os = PyImport_ImportModule("os");
-    if (os == NULL) {
-        goto error;
-    }
-
-    if (PyObject_HasAttrString(os, "sched_getaffinity")) {
-        n = PyObject_CallMethod(os, "sched_getaffinity", "i", 0);
-        if (n == NULL) {
-            goto error;
-        }
-        num_cpus = PySet_Size(n);
-    } else {
-        n = PyObject_CallMethod(os, "cpu_count", NULL);
-        if (n == NULL) {
-            goto error;
-        }
-        num_cpus = PyLong_AsLong(n);
-    }
-
-    if (num_cpus < 1) {
-        goto error;
-    }
-
-    default_max_threads = (int)num_cpus;
-
-done:
-    Py_XDECREF(os);
-    Py_XDECREF(n);
-    return;
-
-error:
-    if (PyErr_Occurred()) {
-        PyErr_Clear();
-    }
-    PyErr_WarnEx(
-        PyExc_RuntimeWarning, "could not get cpu count: using max_threads=1", 1
-    );
-    goto done;
-}
-
 static int
 normalize_quantize_value(int qvalue) {
     if (qvalue < AVIF_QUANTIZER_BEST_QUALITY) {
@@ -438,13 +390,6 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
 
         encoder = avifEncoderCreate();
 
-        if (max_threads == 0) {
-            if (default_max_threads == 0) {
-                init_max_threads();
-            }
-            max_threads = default_max_threads;
-        }
-
         int is_aom_encode = strcmp(codec, "aom") == 0 ||
                             (strcmp(codec, "auto") == 0 &&
                              _codec_available("aom", AVIF_CODEC_FLAG_CAN_ENCODE));
@@ -785,12 +730,6 @@ AvifDecoderNew(PyObject *self_, PyObject *args) {
 
     self->decoder = avifDecoderCreate();
 #if AVIF_VERSION >= 80400
-    if (max_threads == 0) {
-        if (default_max_threads == 0) {
-            init_max_threads();
-        }
-        max_threads = default_max_threads;
-    }
     self->decoder->maxThreads = max_threads;
 #endif
 #if AVIF_VERSION >= 90200