Spaces:
Running on Zero
Running on Zero
File size: 2,219 Bytes
6d1acb0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Session-scoped cooperative cancellation, safe across the ZeroGPU worker boundary.
The SDK owns its worker. This app creates no inference thread or GPU generator.
The cancellation marker lives on the shared filesystem and is unlinked only after
the decorated call has returned. A new request cannot reuse the previous marker.
"""
from dataclasses import dataclass
import hashlib
import logging
from pathlib import Path
import secrets
import threading
import time
log = logging.getLogger("zonos.lifecycle")
class Cancelled(Exception):
pass
class AlreadyRunning(Exception):
pass
def check_cancelled(marker):
if Path(marker).exists():
raise Cancelled()
@dataclass(frozen=True)
class Job:
id: str
owner: str
marker: Path
started: float
class Jobs:
def __init__(self, directory):
self.directory = Path(directory)
self.directory.mkdir(parents=True, exist_ok=True)
self._lock = threading.Lock()
self._active = {}
def begin(self, owner):
if not isinstance(owner, str) or not owner:
raise ValueError("Missing session identity")
key = hashlib.sha256(owner.encode()).hexdigest()
with self._lock:
if key in self._active:
raise AlreadyRunning()
identifier = secrets.token_hex(24)
job = Job(identifier, key, self.directory / (identifier + ".cancel"), time.monotonic())
self._active[key] = job
log.info("request_started request=%s", identifier)
return job
def cancel(self, owner):
key = hashlib.sha256(owner.encode()).hexdigest()
with self._lock:
job = self._active.get(key)
if job is None:
return False
job.marker.touch(exist_ok=True)
log.info("cancellation_requested request=%s", job.id)
return True
def finish(self, job):
with self._lock:
if self._active.get(job.owner) == job:
del self._active[job.owner]
job.marker.unlink(missing_ok=True)
log.info("request_finished request=%s", job.id)
def active(self):
with self._lock:
return len(self._active)
|