feat(tts): enhance TTS request handling with concurrency limits and pacing for improved performance

app/main.py

@@ -85,6 +85,7 @@ def subject_code(label: object) -> str:
     return next((code for code, name in SUBJECT_LABELS.items() if name == label), "")
 MODEL_ID = os.environ.get("MODEL_ID", "facebook/wav2vec2-xlsr-53-espeak-cv-ft")
 LOCAL_TZ = ZoneInfo("Asia/Jerusalem")  # a day of practice is the students' calendar day
+WARM_WORKERS = 3  # Sawtak requests the cache-warm sweep runs at once (tts.MAX_CONCURRENT and its pace apply)
 MAX_UPLOAD = 2_000_000  # bytes; 15 s of opus is ~100 KB
 MAX_BACKUP_UPLOAD = 250_000_000
 EMAIL = re.compile(r"[^@\s,;<>]+@[^@\s,;<>]+\.[^@\s,;<>]+")
@@ -2665,23 +2666,27 @@ def create_app(scorer: Scorer | None = None, tts: SawtakTTS | None = None,
         remote: dict[str, int] = {}
         pending: list[tuple[Path | bytes, str, int]] = []
         oldest = [0.0]  # when the oldest pending file was queued: a batch waits at most 30 s for its 25th file
+        uploads = threading.Lock()  # guards remote, pending and oldest: the workers queue at the same time
 
         def flush() -> None:
-            if not pending:
+            with uploads:
+                batch = pending[:]
+                pending.clear()
+            if not batch:
                 return
             try:
-                mirror.upload([(src, dest) for src, dest, _size in pending])
-                remote.update({dest: size for _src, dest, size in pending})
+                mirror.upload([(src, dest) for src, dest, _size in batch])
+                with uploads:
+                    remote.update({dest: size for _src, dest, size in batch})
                 with app.state.tts_warm_lock:
-                    app.state.tts_warm["uploaded"] += len(pending)
+                    app.state.tts_warm["uploaded"] += len(batch)
             except Exception as e:
-                log.warning("tts bucket upload of %d files failed: %s", len(pending), e)
+                log.warning("tts bucket upload of %d files failed: %s", len(batch), e)
                 with app.state.tts_warm_lock:
-                    app.state.tts_warm["upload_failed"] += len(pending)
+                    app.state.tts_warm["upload_failed"] += len(batch)
                     app.state.tts_warm["upload_error"] = str(e)[:300]
-            pending.clear()
 
-        def queue(src: Path | bytes, rel: str) -> None:
+        def queue_upload(src: Path | bytes, rel: str) -> None:
             if mirror is None:
                 return
             try:
@@ -2689,14 +2694,48 @@ def create_app(scorer: Scorer | None = None, tts: SawtakTTS | None = None,
             except OSError:
                 return
             dest = mirror.remote_path(rel)
-            if remote.get(dest) == size or any(d == dest for _s, d, _z in pending):
-                return
-            if not pending:
-                oldest[0] = time.monotonic()
-            pending.append((src, dest, size))
-            if len(pending) >= 25 or time.monotonic() - oldest[0] > 30:
+            with uploads:
+                if remote.get(dest) == size or any(d == dest for _s, d, _z in pending):
+                    return
+                if not pending:
+                    oldest[0] = time.monotonic()
+                pending.append((src, dest, size))
+                due = len(pending) >= 25 or time.monotonic() - oldest[0] > 30
+            if due:
                 flush()
 
+        def bump(**counts: int) -> None:
+            with app.state.tts_warm_lock:
+                for name, n in counts.items():
+                    app.state.tts_warm[name] += n
+
+        halted = threading.Event()  # a phrase failed for good: no worker starts another
+
+        def generate(text: str) -> bool:
+            """One phrase, retried on a rate limit, a provider 5xx or a dropped connection. False: stopped."""
+            for attempt in range(5):
+                if stop.is_set() or halted.is_set():
+                    return False
+                try:
+                    data, _used = tts.wav_for_voice(text, voice)
+                    if mirror is not None:  # the disk copy if it was kept, else the bytes themselves
+                        queue_upload(tts.cache_file(text, voice) or data, tts.cache_key(text, voice))
+                    return True
+                except TTSUnavailable as e:
+                    if not e.transient or attempt == 4:
+                        raise
+                    delay = max(.2, min(e.retry_after or 2 ** (attempt + 1), 60))
+                    if e.status == 429 and hasattr(tts, "hold"):
+                        tts.hold(delay)  # the other workers wait too, instead of hitting the limit again
+                    with app.state.tts_warm_lock:
+                        app.state.tts_warm["waiting"] = round(delay, 1)
+                    interrupted = stop.wait(delay)  # a stop request ends the wait at once
+                    with app.state.tts_warm_lock:
+                        app.state.tts_warm["waiting"] = 0
+                    if interrupted:
+                        return False
+            return False
+
         def run() -> None:
             if mirror is not None:
                 try:
@@ -2705,57 +2744,56 @@ def create_app(scorer: Scorer | None = None, tts: SawtakTTS | None = None,
                     log.warning("tts bucket listing failed: %s", e)
                     with app.state.tts_warm_lock:
                         app.state.tts_warm["upload_error"] = f"listing failed: {str(e)[:300]}"
+            # Already cached: skip without reading it (the cache can sit on a bucket mount, where reading each WAV is
+            # a network fetch); the upload pass at the end still sends it if the bucket lacks it.
+            todo = []
             for text in jobs:
                 if stop.is_set():
                     break
-                with app.state.tts_warm_lock:
-                    app.state.tts_warm["current"] = text
                 try:
-                    # Already cached: skip without reading it. The cache can sit on a bucket mount, where reading
-                    # each WAV is a network fetch; the upload pass at the end still sends it if the bucket lacks it.
-                    if tts.cached(text, voice):
-                        with app.state.tts_warm_lock:
-                            app.state.tts_warm["skipped"] += 1
-                            app.state.tts_warm["done"] += 1
-                        continue
-                    was_cached = interrupted = False
-                    for attempt in range(5):
-                        try:
-                            data, _used = tts.wav_for_voice(text, voice)  # serialised: one request in this replica
-                            if mirror is not None:  # the disk copy if it was kept, else the bytes themselves
-                                queue(tts.cache_file(text, voice) or data, tts.cache_key(text, voice))
-                            break
-                        except TTSUnavailable as e:  # rate limit, provider 5xx or a dropped connection: retry
-                            if not e.transient or attempt == 4:
-                                raise
-                            delay = max(.2, min(e.retry_after or 2 ** (attempt + 1), 60))
-                            with app.state.tts_warm_lock:
-                                app.state.tts_warm["waiting"] = round(delay, 1)
-                            interrupted = stop.wait(delay)  # a stop request ends the wait at once
-                            with app.state.tts_warm_lock:
-                                app.state.tts_warm["waiting"] = 0
-                            if interrupted:
-                                break
-                    if interrupted:
-                        break
+                    cached = tts.cached(text, voice)
+                except OSError:
+                    cached = False
+                if cached:
+                    bump(skipped=1, done=1)
+                else:
+                    todo.append(text)
+            work = iter(todo)
+            take = threading.Lock()
+
+            def worker() -> None:
+                while not (stop.is_set() or halted.is_set()):
+                    with take:
+                        text = next(work, None)
+                    if text is None:
+                        return
                     with app.state.tts_warm_lock:
-                        app.state.tts_warm["skipped" if was_cached else "generated"] += 1
-                        app.state.tts_warm["done"] += 1
-                except Exception as e:
-                    log.warning("tts cache warm stopped at %r: %s", text, e)
-                    with app.state.tts_warm_lock:
-                        app.state.tts_warm["failed"] += 1
-                        app.state.tts_warm["error"] = str(e)
-                    break
-            if mirror is not None and stop.is_set():
-                flush()  # what was generated before the stop still reaches the bucket; the full pass waits for next run
+                        app.state.tts_warm["current"] = text
+                    try:
+                        if generate(text):
+                            bump(generated=1, done=1)
+                    except Exception as e:
+                        log.warning("tts cache warm stopped at %r: %s", text, e)
+                        halted.set()
+                        with app.state.tts_warm_lock:
+                            app.state.tts_warm["failed"] += 1
+                            app.state.tts_warm["error"] = str(e)
+
+            workers = [threading.Thread(target=worker, name=f"ihki-tts-warm-{i}", daemon=True)
+                       for i in range(WARM_WORKERS)]
+            for w in workers:
+                w.start()
+            for w in workers:
+                w.join()
+            if mirror is not None and (stop.is_set() or halted.is_set()):
+                flush()  # what was generated still reaches the bucket; the full pass waits for the next run
             elif mirror is not None:
                 # also everything cached earlier: other voices, playback misses, the legacy flat layout
                 with app.state.tts_warm_lock:
                     app.state.tts_warm["current"] = "uploading cached files"
                 try:
                     for path, rel in tts.cached_files():
-                        queue(path, rel)
+                        queue_upload(path, rel)
                 except OSError as e:
                     log.warning("tts cache listing failed: %s", e)
                 flush()

app/tts.py

@@ -34,7 +34,10 @@ CACHE_VERSION = "raef-pointed-enhanced-provider-default-v1"
 TIMEOUT = 60  # per socket operation: a response that keeps trickling never hits it, hence the two caps below
 MAX_SECONDS = 90  # one whole generation, download included
 MAX_AUDIO_SECONDS = 60  # the longest phrase is ~240 characters; Sawtak's broken answer is minutes of silence
-MIN_REQUEST_INTERVAL = .2  # spacing between Sawtak requests; the account allows 60/minute (docs/limits), 429s are retried
+# The account allows 60 requests a minute and 5 at once (GET /v1/me). Request starts are spaced to at most 45 a minute
+# for this server, and at most 3 are in flight, so student playback and other clients keep some room. 429s are retried.
+MIN_REQUEST_INTERVAL = 60 / 45
+MAX_CONCURRENT = 3
 # The single Sawtak voice used throughout the application.
 RAEF_VOICE_ID = "273d81b2-991b-5b76-8525-eeacc6a3d759"
 VOICES = {"רַאאִף": RAEF_VOICE_ID}
@@ -87,8 +90,11 @@ class SawtakTTS:
         self.api_key, self.voice_id = api_key, RAEF_VOICE_ID
         self.voice_name = "רַאאִף"
         self.cache_dir = cache_dir or Path.home() / ".cache" / "ihki-tts"
-        self._lock = threading.Lock()  # one generation at a time; also guards the voice lookup
-        self._last_request = 0.0
+        self._lock = threading.Lock()  # guards the counters and clearing the cache
+        self._pace = threading.Lock()  # hands out request start times, MIN_REQUEST_INTERVAL apart
+        self._next_start = 0.0
+        self._in_flight = threading.Semaphore(MAX_CONCURRENT)
+        self._key_locks: dict[str, threading.Lock] = {}  # one generation per recording at a time
         self._cache_hits = self._cache_misses = 0
         self.mirror: BucketMirror | None = None  # set by load_tts when TTS_BUCKET is configured
         # With a mirror, new recordings are written here (local disk) and uploaded through the Bucket API, instead
@@ -175,12 +181,26 @@ class SawtakTTS:
         """Has this voice said `text` before? Then wav() answers from disk, without the provider."""
         return any(c.exists() for c in self._candidates(text, voice))
 
+    def hold(self, seconds: float) -> None:
+        """No new request starts for `seconds` (after a 429, so the other requests in flight back off too)."""
+        with self._pace:
+            self._next_start = max(self._next_start, time.monotonic() + seconds)
+
+    def _key_lock(self, key: str) -> threading.Lock:
+        with self._lock:
+            return self._key_locks.setdefault(key, threading.Lock())
+
     def _generate(self, text: str, voice: str) -> bytes:
-        """Generate one uncached WAV recording."""
-        wait = MIN_REQUEST_INTERVAL - (time.monotonic() - self._last_request)
-        if wait > 0:
-            time.sleep(wait)
-        self._last_request = time.monotonic()
+        """Generate one uncached WAV recording: paced, and at most MAX_CONCURRENT at once."""
+        with self._pace:
+            start = max(time.monotonic(), self._next_start)
+            self._next_start = start + MIN_REQUEST_INTERVAL
+        if start > time.monotonic():
+            time.sleep(start - time.monotonic())
+        with self._in_flight:
+            return self._generate_now(text, voice)
+
+    def _generate_now(self, text: str, voice: str) -> bytes:
         # Sawtak currently returns 320 seconds of silence whenever `temperature` is present,
         # including for 0 and 0.1. Omitting it uses the provider's deterministic supported path.
         with self._request("/audio/speech", {"model": MODEL, "voice": voice, "input": text,
@@ -244,11 +264,12 @@ class SawtakTTS:
 
     def _wav(self, text: str, voice: str, upload: bool = False) -> tuple[bytes, str]:
         """Generate or read `text` for an already validated voice id. `upload`: send a new recording to the bucket
-        in the background (playback); the cache-warm sweep uploads its own in batches."""
-        with self._lock:
+        in the background (playback); the cache-warm sweep uploads its own in batches. Different recordings can be
+        generated at the same time (see _generate); the same one waits for the first request and then reads it."""
+        with self._key_lock(self.cache_key(text, voice)):
             found = self.cache_file(text, voice) if self.staging_dir else None
             if found is not None:  # with a mirror the cache folder may be a slow mount: no migration there
-                self._cache_hits += 1
+                self._count(hit=True)
                 return found.read_bytes(), voice
             path = self._path(text, voice)
             legacy = self._legacy_path(path)
@@ -260,9 +281,9 @@ class SawtakTTS:
                     log.warning("tts cache migration failed: %s: %s", type(e).__name__, e)
                     path = legacy
             if path.exists():
-                self._cache_hits += 1
+                self._count(hit=True)
                 return path.read_bytes(), voice
-            self._cache_misses += 1
+            self._count(hit=False)
             # The caller supplies the phrase's Palestinian vowel marks. Sawtak preserves existing
             # marks; enhancement fills remaining pronunciation cues using the selected voice's dialect.
             data = self._generate(text, voice)
@@ -279,6 +300,13 @@ class SawtakTTS:
                 self._upload_soon(data, self.cache_key(text, voice))
             return data, voice
 
+    def _count(self, hit: bool) -> None:
+        with self._lock:
+            if hit:
+                self._cache_hits += 1
+            else:
+                self._cache_misses += 1
+
     def _upload_soon(self, data: bytes, rel: str) -> None:
         """One playback recording to the bucket, off the request: the student does not wait for it."""
         mirror = self.mirror

tests/test_scoring.py

@@ -2299,7 +2299,7 @@ def test_the_warm_sweep_stops_on_request_and_keeps_what_it_made(monkeypatch):
     from app.main import create_app
     from app.tts import RAEF_VOICE_ID
 
-    phrases = [{"variants": [(word, "a")]} for word in ("واحد", "اثنين", "ثلاثة", "اربعة")]
+    phrases = [{"variants": [(word, "a")]} for word in ("واحد", "اثنين", "ثلاثة", "اربعة", "خمسة")]
     monkeypatch.setattr(app.main, "PHRASES", phrases)
     monkeypatch.setattr(app.main, "load_auth", lambda: None)
     first_started, release = threading.Event(), threading.Event()
@@ -2308,20 +2308,23 @@ def test_the_warm_sweep_stops_on_request_and_keeps_what_it_made(monkeypatch):
         def __init__(self): self.calls = []
         def cached(self, text, voice): return False
         def wav_for_voice(self, text, voice):
-            self.calls.append(text); first_started.set(); release.wait(5); return b"RIFF", voice
+            self.calls.append(text)
+            if len(self.calls) == app.main.WARM_WORKERS: first_started.set()
+            release.wait(5); return b"RIFF", voice
         def cache_stats(self): return {"files": 0, "bytes": 0, "hits": 0, "misses": 0, "hit_rate": None}
 
     tts = SlowTTS()
     with TestClient(create_app(scorer_for(emissions(["w"])), tts, tutor=object())) as c:
-        assert c.post("/api/tts-cache/warm", json={"voice_id": RAEF_VOICE_ID}).json()["total"] == 4
+        assert c.post("/api/tts-cache/warm", json={"voice_id": RAEF_VOICE_ID}).json()["total"] == 5
         assert first_started.wait(5)
         assert c.post("/api/tts-cache/warm/stop").status_code == 200
-        release.set()  # the phrase in progress finishes and is kept; no further phrase starts
+        release.set()  # the phrases in progress (one per worker) finish and are kept; no further phrase starts
         for _ in range(200):
             state = c.get("/api/tts-cache").json()["warm"]
             if not state["running"]: break
             time.sleep(.01)
-    assert state["stopped"] is True and state["done"] == state["generated"] == 1 and len(tts.calls) == 1
+    workers = app.main.WARM_WORKERS
+    assert state["stopped"] is True and state["done"] == state["generated"] == workers == len(tts.calls) < 5
 
 
 def test_with_a_bucket_new_recordings_are_staged_locally_not_written_to_the_mount(monkeypatch, tmp_path):
@@ -2394,3 +2397,38 @@ def test_an_offline_copy_loads_back_into_the_cache_and_the_bucket(monkeypatch, t
         if len(uploaded) == 2: break
         time.sleep(.01)
     assert sorted(uploaded) == sorted([f"ihki-tts/{new}", f"ihki-tts/{legacy}"])
+
+
+def test_the_sweep_runs_sawtak_requests_in_parallel_within_the_pace(monkeypatch, tmp_path):
+    import threading
+    import time
+
+    import app.tts as tts_module
+    from app.tts import RAEF_VOICE_ID, SawtakTTS
+
+    monkeypatch.setattr(tts_module, "MIN_REQUEST_INTERVAL", .05)
+    tts = SawtakTTS("key", cache_dir=tmp_path)
+    active, peak, starts, guard = [0], [0], [], threading.Lock()
+
+    def slow(text, voice):
+        with guard:
+            active[0] += 1; peak[0] = max(peak[0], active[0]); starts.append(time.monotonic())
+        time.sleep(.3)
+        with guard:
+            active[0] -= 1
+        return b"RIFF" + text.encode()
+    monkeypatch.setattr(tts, "_generate_now", slow)
+    threads = [threading.Thread(target=tts.wav_for_voice, args=(f"نص {i}", RAEF_VOICE_ID)) for i in range(6)]
+    t0 = time.monotonic()
+    for t in threads: t.start()
+    for t in threads: t.join()
+    assert peak[0] == tts_module.MAX_CONCURRENT                      # several at once, never more than allowed
+    gaps = [b - a for a, b in zip(sorted(starts), sorted(starts)[1:])]
+    assert min(gaps) >= .045 and time.monotonic() - t0 < 6 * .3      # paced, yet faster than one at a time
+    # the same phrase asked twice at once is generated once
+    calls = []
+    monkeypatch.setattr(tts, "_generate_now", lambda text, voice: (calls.append(text), time.sleep(.1), b"RIFF")[2])
+    pair = [threading.Thread(target=tts.wav_for_voice, args=("مرة", RAEF_VOICE_ID)) for _ in range(2)]
+    for t in pair: t.start()
+    for t in pair: t.join()
+    assert calls == ["مرة"]