From 230c3032e5e8a660f1f158a8fff3511ef91485bb Mon Sep 17 00:00:00 2001 From: Michael Pilosov Date: Tue, 21 Apr 2026 20:16:33 -0600 Subject: [PATCH] parallelism limits --- flows/embedding_flow.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/flows/embedding_flow.py b/flows/embedding_flow.py index 0cfe553..103c460 100644 --- a/flows/embedding_flow.py +++ b/flows/embedding_flow.py @@ -7,6 +7,18 @@ import sys os.environ.setdefault("PREFECT_API_URL", "http://localhost:4200/api") os.environ.setdefault("DO_NOT_TRACK", "1") +# Pin per-process thread pools to 1 so Ray's worker parallelism doesn't +# multiply against BLAS/numba/etc. thread pools — otherwise 4 workers × N +# cores → thrash. Must be set before numpy/numba/sklearn import, since +# those libs latch onto these env vars at import time. Ray manages OMP +# per-task-CPU but does NOT manage NUMBA_NUM_THREADS, which is what +# PaCMAP/UMAP use for their optimization loops. +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("MKL_NUM_THREADS", "1") +os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") +os.environ.setdefault("NUMEXPR_NUM_THREADS", "1") +os.environ.setdefault("NUMBA_NUM_THREADS", "1") + from datetime import timedelta import math from pathlib import Path