Constrain native math library threading to fix XGBoost stalling in containers

Found the actual difference between "works locally, hangs on the
Space": locally numpy uses Apple's Accelerate framework, unaffected by
OMP_NUM_THREADS; on Linux it uses OpenBLAS, governed by its own
OPENBLAS_NUM_THREADS. Neither our env-var setup nor the fitted
XGBClassifier's n_jobs=None was actually constraining thread count on
Linux, so a single prediction could spawn as many native threads as
the *host* reports — potentially far more than a cgroup-limited
container's actual CPU quota, which starves/stalls the process.

- config.py now also sets OPENBLAS_NUM_THREADS, MKL_NUM_THREADS,
  NUMEXPR_NUM_THREADS, VECLIB_MAXIMUM_THREADS, BLIS_NUM_THREADS.
- These only take effect if set before the native libraries are first
  loaded, so reordered imports in protac_splitter.py and the app
  script to import protac_splitter.config before pandas/rdkit/datasets
  (and, in the app, before gradio, which pulls those in transitively).
- GraphEdgeClassifier.load() now forces the fitted classifier's
  n_jobs to 1 explicitly, since that default is baked into the
  pickled model at training time and env vars alone don't override
  XGBoost's own internal thread pool.

Author: ribesstefano

protac_splitter/config.py

@@ -6,9 +6,20 @@
 
 load_dotenv()
 
-# XGBoost with OpenMP can deadlock on macOS ARM64 when spawning multiple threads
-# during model deserialization. setdefault respects any existing user override.
-os.environ.setdefault("OMP_NUM_THREADS", "1")
+# Native math libraries (OpenMP, OpenBLAS, MKL, Apple Accelerate) each read their own
+# thread-count env var and default to "use every core the process can see" — which on
+# a cgroup-limited container can wildly exceed the actual CPU quota and turn a
+# millisecond-scale single-row prediction into a multi-minute stall. setdefault
+# respects any existing user override.
+for _threads_var in (
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "NUMEXPR_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+    "BLIS_NUM_THREADS",
+):
+    os.environ.setdefault(_threads_var, "1")
 
 
 def get_cache_dir() -> Path:

protac_splitter/graphs/edge_classifier.py

@@ -137,7 +137,16 @@ def save(self, path: Union[str, Path]):
 
     @classmethod
     def load(cls, path: Union[str, Path]) -> "GraphEdgeClassifier":
-        return joblib.load(str(path))
+        model = joblib.load(str(path))
+        # The fitted estimator's n_jobs (None = auto-detect all cores) is baked in at
+        # training time. On cgroup-limited containers where the reported core count
+        # doesn't match the actual CPU quota, that over-provisions native threads and
+        # can make a single prediction take orders of magnitude longer than it should.
+        # A single row gains nothing from parallelism anyway, so force it to 1.
+        clf = model.pipeline.named_steps.get("clf")
+        if clf is not None and hasattr(clf, "set_params"):
+            clf.set_params(n_jobs=1)
+        return model
 
     @staticmethod
     def extract_graph_features(

protac_splitter/protac_splitter.py

@@ -1,15 +1,21 @@
 import hashlib
 import logging
 import warnings
-import requests
 from pathlib import Path
 from typing import Union, Optional, Dict, List, Literal
 
+# Import first, before numpy-linked packages (rdkit/datasets/pandas below): this sets
+# thread-count env vars for native math libraries (OpenMP/OpenBLAS/MKL/Accelerate),
+# which some of them only read once, at first load. Importing it after would be too
+# late and leave those libraries free to over-provision threads on cgroup-limited
+# containers.
+from protac_splitter.config import get_cache_dir, get_hf_token
+
+import requests
 from rdkit import Chem
 from datasets import Dataset
 import pandas as pd
 
-from protac_splitter.config import get_cache_dir, get_hf_token
 from protac_splitter.chemoinformatics import canonize
 from protac_splitter.evaluation import split_prediction
 from protac_splitter.fixing_functions import fix_prediction

scripts/protac_splitter_app.py

@@ -24,16 +24,20 @@
 from pathlib import Path
 from typing import Union
 
+# Import protac_splitter before gradio/pandas/rdkit below: it sets thread-count env
+# vars for native math libraries (see protac_splitter/config.py), which some of them
+# only read once, at first load — importing it after gradio (which pulls in numpy
+# transitively) would be too late.
+from protac_splitter import split_protac
+from protac_splitter.config import get_cache_dir
+from protac_splitter.evaluation import split_prediction
+
 from PIL import Image
 import gradio as gr
 import pandas as pd
 from rdkit import Chem
 from rdkit.Chem import Draw
 
-from protac_splitter import split_protac
-from protac_splitter.config import get_cache_dir
-from protac_splitter.evaluation import split_prediction
-
 # HF Spaces sets SPACE_ID automatically; cap parallelism on the (limited) free tier.
 IS_HF_SPACE = os.environ.get("SPACE_ID") is not None
 MAX_NUM_PROC = 2 if IS_HF_SPACE else 8