Use generator pattern (cross-process yield via spaces res_queue) — ZeroGPU forks process so streamer queue isn't shared

src/kpaa/llm/zerogpu_backend.py

@@ -109,39 +109,38 @@ class ZeroGPUBackend:
                 add_generation_prompt=True,
                 return_tensors="pt",
             )
-            # transformers 5.x 는 BatchEncoding(dict-like) 을, 4.x 는 Tensor 를
-            # 반환. model.generate(input_ids=...) 는 Tensor 만 받으므로 추출.
+            # transformers 5.x 는 BatchEncoding 을, 4.x 는 Tensor 를 반환.
             input_ids = _encoded["input_ids"] if hasattr(_encoded, "input_ids") else _encoded
 
-            from transformers import TextIteratorStreamer
-
-            streamer = TextIteratorStreamer(
-                tok,
-                skip_prompt=True,
-                skip_special_tokens=True,
-            )
-
+            # ZeroGPU 는 @spaces.GPU 함수를 fork() 자식 프로세스에서 실행한다.
+            # 자식 프로세스의 streamer.text_queue 는 부모 프로세스에서 접근 불가.
+            # → @spaces.GPU 를 generator 함수에 적용하면 spaces 가 res_queue 로
+            #    yield 값을 부모로 전송 (multiprocessing-safe). streamer 도 자식
+            #    프로세스 내에서 생성·소비.
             @_gpu_decorator(self._gpu_duration)
-            def _run_generate() -> None:
-                """ZeroGPU 점유 동안 generate 완료까지 블로킹.
+            def _run_generate_gen():
+                """ZeroGPU 자식 프로세스에서 토큰 generator.
 
-                내부에서 cuda 로 모델·입력을 옮기고, 별도 thread 로 generate
-                실행 → 메인 코루틴은 streamer 에서 토큰 빨아냄. 이 함수가
-                반환되어야 ZeroGPU 가 GPU 회수.
+                각 yield 는 spaces 의 res_queue 를 통해 부모 프로세스로 전달된다.
                 """
                 import torch as _t
-                from transformers import GenerationConfig
+                from transformers import GenerationConfig, TextIteratorStreamer
 
                 device = "cuda" if _t.cuda.is_available() else "cpu"
-                print(f"[kpaa.zerogpu] _run_generate start, device={device}", flush=True)
+                print(f"[kpaa.zerogpu] _gen start, device={device}", flush=True)
                 if device == "cuda":
                     model.to(device)
                     print(f"[kpaa.zerogpu] model moved to cuda", flush=True)
                 ids = input_ids.to(device)
                 print(f"[kpaa.zerogpu] input shape={tuple(ids.shape)}, max_new_tokens={opts.max_tokens}", flush=True)
 
-                # transformers 5.x 는 generate(temperature=...) 직접 호출을 무시하고
-                # GenerationConfig 객체로만 받음. 명시적으로 묶어서 전달.
+                # streamer 는 자식 프로세스 내에서 생성 — 자식 프로세스 thread 간 통신.
+                streamer = TextIteratorStreamer(
+                    tok,
+                    skip_prompt=True,
+                    skip_special_tokens=True,
+                )
+
                 gen_cfg = GenerationConfig(
                     max_new_tokens=opts.max_tokens,
                     do_sample=opts.temperature > 0.0,
@@ -161,10 +160,6 @@ class ZeroGPUBackend:
                         print(f"[kpaa.zerogpu] generate() raised: {type(e).__name__}: {e}", flush=True)
                         raise
                     finally:
-                        # transformers 5.x 에서 streamer.end() 가 generate 종료 시점에
-                        # 자동 호출되지 않는 케이스가 있음 → next(streamer) 가 영원히
-                        # 블록. 성공/실패 모두 명시적으로 end() 호출해 stop_signal
-                        # 을 큐에 넣어 consumer 의 StopIteration 을 보장.
                         try:
                             streamer.end()
                             print(f"[kpaa.zerogpu] streamer.end() called", flush=True)
@@ -173,43 +168,40 @@ class ZeroGPUBackend:
 
                 gen_thread = Thread(target=_generate_target, daemon=True)
                 gen_thread.start()
-                gen_thread.join()  # GPU 점유 중 generate 완료 대기
-                print(f"[kpaa.zerogpu] _run_generate end", flush=True)
 
-            # generate 호출 자체는 별도 thread (run_in_executor) — 그래야
-            # 본 코루틴이 streamer 에서 토큰을 비동기로 빨아낼 수 있음.
-            gen_future = loop.run_in_executor(None, _run_generate)
+                # streamer 에서 토큰 빨아내며 yield → spaces 가 res_queue 통해 부모 전송.
+                tok_n = 0
+                for token in streamer:
+                    if token:
+                        tok_n += 1
+                        if tok_n <= 3 or tok_n % 100 == 0:
+                            print(f"[kpaa.zerogpu] yield #{tok_n}: {token!r}", flush=True)
+                        yield token
 
-            _tok_count = [0]
+                gen_thread.join()
+                print(f"[kpaa.zerogpu] _gen end (total yielded={tok_n})", flush=True)
 
-            def _next_token() -> str | None:
+            # 부모: 자식의 generator 를 async iterate. spaces 가 cross-process 직렬화 처리.
+            def _safe_next(gen):
                 try:
-                    t = next(streamer)
-                    _tok_count[0] += 1
-                    if _tok_count[0] <= 3 or _tok_count[0] % 100 == 0:
-                        print(
-                            f"[kpaa.zerogpu] streamer #{_tok_count[0]}: {t!r}",
-                            flush=True,
-                        )
-                    return t
+                    return next(gen), False
                 except StopIteration:
-                    print(
-                        f"[kpaa.zerogpu] streamer exhausted, total={_tok_count[0]}",
-                        flush=True,
-                    )
-                    return None
+                    return None, True
 
+            gen = _run_generate_gen()
             try:
                 while True:
-                    token = await loop.run_in_executor(None, _next_token)
-                    if token is None:
+                    token, done = await loop.run_in_executor(None, _safe_next, gen)
+                    if done:
                         break
                     if token:
                         yield token
             finally:
-                # generate 가 아직 안 끝났으면 마무리 대기 (GPU 회수 보장).
-                if not gen_future.done():
-                    await gen_future
+                # gen 닫기 (자식 프로세스 종료 보장).
+                try:
+                    gen.close()
+                except Exception:
+                    pass
 
     async def close(self) -> None:
         # CUDA cache 는 ZeroGPU 가 자동 정리. transformers 는 GC.