File size: 946 Bytes
bb1a40e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
# early_env.py — import this BEFORE numpy/torch/cv2/etc.
import os, re
def _clean_int_env(name: str, default: str | None = None):
val = os.environ.get(name)
if val is None:
if default is not None:
os.environ[name] = default
return
if not re.fullmatch(r"\d+", str(val).strip()):
if default is None:
os.environ.pop(name, None)
else:
os.environ[name] = default
# Keep thread env sane for Spaces (prevents libgomp errors)
_clean_int_env("OMP_NUM_THREADS", "2")
_clean_int_env("MKL_NUM_THREADS", "2")
_clean_int_env("OPENBLAS_NUM_THREADS", "2")
_clean_int_env("NUMEXPR_NUM_THREADS", "2")
# Configure PyTorch threading EARLY (before any parallel work)
try:
import torch
if hasattr(torch, "set_num_interop_threads"):
torch.set_num_interop_threads(2)
if hasattr(torch, "set_num_threads"):
torch.set_num_threads(2)
except Exception:
pass
|