runtime.py

"""
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


@contextlib.contextmanager
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()