Spaces:
Running
Running
unknown
Claude Opus 5.5
Fix 'Already borrowed': every fast tokenizer gets its own lock, so jobs running at the same time can share a model
634fc07 | """ | |
| Shared CPU scheduling for CPU-only hardware (e.g. HF "CPU Upgrade": 8 vCPU / 32 GB). | |
| * At most CPU_JOBS model calls run at once (default cores // 2); further calls wait for a slot. | |
| * PyTorch threads adapt to the load: a job running alone uses ALL cores, two jobs get half each, etc. | |
| (Before, every job was capped at cores // CPU_JOBS threads even when the server was idle.) | |
| * quantize(model): int8 dynamic quantisation of the Linear layers for CPU inference (≈2x faster, | |
| near-identical output). Disable with CPU_INT8=0. | |
| """ | |
| import contextlib | |
| import os | |
| import sys | |
| import threading | |
| import torch | |
| # A log line must never fail a job. The logs carry "→", "°C", emoji and customers' text; on a console that is not | |
| # UTF-8 (Windows, a container with a C/POSIX locale) print() raised UnicodeEncodeError and the finished job was | |
| # reported as "Processing failed". Characters the console cannot show are now written as escapes instead. | |
| for _stream in (sys.stdout, sys.stderr): | |
| try: | |
| _stream.reconfigure(errors="backslashreplace") | |
| except Exception: # replaced / closed stream: nothing to configure | |
| pass | |
| HAS_CUDA = torch.cuda.is_available() | |
| # The checking models (meaning, fluency, clarity, vocabulary) ran on the CPU even on a GPU Space: with several | |
| # versions of every sentence checked on every pass, they were most of the waiting time. On a dedicated GPU they now | |
| # run there too. Not on ZeroGPU (the GPU may only be used inside a GPU job there). HELPERS_GPU=0 = keep on CPU. | |
| _ZERO = os.getenv("SPACES_ZERO_GPU", "").strip().lower() in ("1", "true", "yes") | |
| HELPERS_GPU = HAS_CUDA and not _ZERO and os.getenv("HELPERS_GPU", "1") != "0" | |
| def to_gpu(model): | |
| """Move a small checking model to the GPU when there is a dedicated one (it stays on the CPU otherwise).""" | |
| if HELPERS_GPU: | |
| try: | |
| return model.to("cuda") | |
| except Exception as err: # out of GPU memory etc.: the CPU still works | |
| print(f"[runtime] model kept on CPU: {err}", flush=True) | |
| return model | |
| def device_of(model): | |
| return next(model.parameters()).device | |
| CORES = os.cpu_count() or 2 | |
| CPU_JOBS = max(1, int(os.getenv("CPU_JOBS", str(max(1, CORES // 2))))) | |
| cpu_slots = threading.BoundedSemaphore(CPU_JOBS) | |
| _active = 0 | |
| _lock = threading.Lock() | |
| if not HAS_CUDA: | |
| try: | |
| torch.set_num_interop_threads(1) | |
| except Exception: | |
| pass | |
| def _retune(): | |
| try: | |
| torch.set_num_threads(max(1, CORES // max(1, _active))) | |
| except Exception: | |
| pass | |
| def cpu_job(): | |
| """Hold a CPU slot for one model call; threads = cores ÷ jobs running right now.""" | |
| global _active | |
| cpu_slots.acquire() | |
| with _lock: | |
| _active += 1 | |
| _retune() | |
| try: | |
| yield | |
| finally: | |
| with _lock: | |
| _active -= 1 | |
| _retune() | |
| cpu_slots.release() | |
| def quantize(model): | |
| """int8 dynamic quantisation for CPU (falls back to the original model on any problem).""" | |
| if HAS_CUDA or os.getenv("CPU_INT8", "1") == "0": | |
| return model | |
| try: | |
| q = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8) | |
| print("[runtime] int8 CPU quantisation enabled", flush=True) | |
| return q | |
| except Exception as err: | |
| print(f"[runtime] quantisation skipped: {err}", flush=True) | |
| return model | |
| # --------------------------------------------------------------------------------------------------------------- | |
| # Fast tokenizers are not thread-safe. Hugging Face's Rust tokenizer is re-configured (padding / truncation) on | |
| # every call, so two jobs tokenizing with the same model at the same moment fail with "Already borrowed" | |
| # (live Space, 2026-10-19). Every fast tokenizer now gets its own lock: one call at a time per tokenizer, any number | |
| # of tokenizers in parallel. A call takes milliseconds, so jobs do not notice the wait. | |
| def _guard_tokenizers(): | |
| try: | |
| from transformers import PreTrainedTokenizerFast as T | |
| except Exception: # transformers not installed: nothing to guard | |
| return | |
| if getattr(T, "_cs_guarded", False): | |
| return | |
| make = threading.Lock() | |
| def lock_of(tok): | |
| lk = tok.__dict__.get("_cs_lock") | |
| if lk is None: | |
| with make: | |
| lk = tok.__dict__.setdefault("_cs_lock", threading.RLock()) | |
| return lk | |
| def guard(name): | |
| fn = getattr(T, name, None) | |
| if fn is None: | |
| return | |
| def wrapped(self, *a, **k): | |
| with lock_of(self): | |
| return fn(self, *a, **k) | |
| wrapped.__name__, wrapped.__doc__ = fn.__name__, fn.__doc__ | |
| setattr(T, name, wrapped) | |
| for name in ("_batch_encode_plus", "_encode_plus", "_decode", "set_truncation_and_padding"): | |
| guard(name) | |
| T._cs_guarded = True | |
| _guard_tokenizers() | |