| """ZeroGPU-safe decorator. |
| |
| spaces.GPU attaches a device for the duration of the call and enforces a |
| wall-clock ceiling. Wrapping the whole pipeline in one call blows that ceiling |
| on anything longer than a short clip, so each GPU stage gets its own call and |
| all CPU work stays outside. Off ZeroGPU this is a no-op. |
| |
| Adds over the bare decorator: |
| * the worker traceback is printed and appended to DWCT_GPU_TRACE, because the |
| ZeroGPU wrapper otherwise reduces a failure to a bare exception class name; |
| """ |
| from __future__ import annotations |
|
|
| import functools |
| import os |
| import sys |
| import traceback |
|
|
| TRACE_FILE = os.environ.get("DWCT_GPU_TRACE", "/tmp/dwct_gpu_trace.txt") |
|
|
| try: |
| import spaces |
|
|
| _GPU = spaces.GPU |
| HAVE_ZEROGPU = True |
| except Exception: |
| HAVE_ZEROGPU = False |
|
|
|
|
| def _loud(fn): |
| @functools.wraps(fn) |
| def inner(*a, **k): |
| try: |
| return fn(*a, **k) |
| except BaseException as exc: |
| text = "[gpu] %s raised %s: %s\n%s" % ( |
| fn.__name__, type(exc).__name__, exc, traceback.format_exc()) |
| print(text, file=sys.stderr) |
| try: |
| with open(TRACE_FILE, "a") as fh: |
| fh.write(text + "\n" + "=" * 60 + "\n") |
| except Exception: |
| pass |
| sys.stderr.flush() |
| sys.stdout.flush() |
| raise |
|
|
| return inner |
|
|
|
|
| def gpu(duration: int = 120): |
| if HAVE_ZEROGPU: |
| def deco(fn): |
| return _GPU(duration=int(duration))(_loud(fn)) |
| return deco |
|
|
| def deco(fn): |
| return _loud(fn) |
|
|
| return deco |
|
|
