Upload 2 files

cancel_registry.py

@@ -0,0 +1,79 @@
+"""
+Cancel Registry for OCR Jobs
+Manages cancellation flags for cooperative job cancellation
+"""
+import asyncio
+from typing import Optional
+
+# Global registry: job_id -> asyncio.Event
+CANCEL_FLAGS: dict[str, asyncio.Event] = {}
+_cancel_lock = asyncio.Lock()
+
+
+def new_cancel_flag(job_id: str) -> asyncio.Event:
+    """Create a new cancellation flag for a job"""
+    async def _create():
+        async with _cancel_lock:
+            ev = asyncio.Event()
+            CANCEL_FLAGS[job_id] = ev
+            return ev
+    # Run in event loop if available, otherwise create synchronously
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            ev = asyncio.Event()
+            asyncio.create_task(_create_sync(job_id, ev))
+            return ev
+        else:
+            return loop.run_until_complete(_create())
+    except RuntimeError:
+        # No event loop, create directly
+        ev = asyncio.Event()
+        CANCEL_FLAGS[job_id] = ev
+        return ev
+
+
+async def _create_sync(job_id: str, ev: asyncio.Event):
+    """Helper to register event"""
+    async with _cancel_lock:
+        CANCEL_FLAGS[job_id] = ev
+
+
+def get_cancel_flag(job_id: str) -> Optional[asyncio.Event]:
+    """Get cancellation flag for a job"""
+    return CANCEL_FLAGS.get(job_id)
+
+
+def cancel_job(job_id: str) -> bool:
+    """Cancel a job by setting its flag. Returns True if job exists."""
+    ev = CANCEL_FLAGS.get(job_id)
+    if ev:
+        ev.set()
+        return True
+    return False
+
+
+async def is_cancelled(job_id: str) -> bool:
+    """Check if a job has been cancelled"""
+    ev = CANCEL_FLAGS.get(job_id)
+    if ev:
+        return ev.is_set()
+    return False
+
+
+def remove_cancel_flag(job_id: str):
+    """Remove cancellation flag (cleanup after job completes)"""
+    async def _remove():
+        async with _cancel_lock:
+            CANCEL_FLAGS.pop(job_id, None)
+    
+    try:
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            asyncio.create_task(_remove())
+        else:
+            loop.run_until_complete(_remove())
+    except RuntimeError:
+        # No event loop, remove directly
+        CANCEL_FLAGS.pop(job_id, None)
+

ocr_service.py

@@ -133,6 +133,36 @@ _jobs: dict[str, dict] = {}  # job_id -> {status, progress, result, error, cance
 _jobs_lock = asyncio.Lock()
 _cancellation_tokens: dict[str, asyncio.Event] = {}  # job_id -> cancellation event
 
+# Import cancel registry
+try:
+    from cancel_registry import cancel_job, get_cancel_flag, new_cancel_flag, remove_cancel_flag, is_cancelled
+except ImportError:
+    # Fallback if cancel_registry not available
+    def cancel_job(job_id: str): return False
+    def get_cancel_flag(job_id: str): return _cancellation_tokens.get(job_id)
+    def new_cancel_flag(job_id: str): return _cancellation_tokens.setdefault(job_id, asyncio.Event())
+    def remove_cancel_flag(job_id: str): pass
+    async def is_cancelled(job_id: str): return False
+
+# StoppingCriteria for generation (if transformers supports it)
+try:
+    from transformers import StoppingCriteria, StoppingCriteriaList
+    _STOPPING_CRITERIA_AVAILABLE = True
+except ImportError:
+    _STOPPING_CRITERIA_AVAILABLE = False
+    StoppingCriteria = None
+    StoppingCriteriaList = None
+
+
+class CancelCriterion(StoppingCriteria):
+    """Stopping criteria that checks a cancellation flag"""
+    def __init__(self, cancel_flag: asyncio.Event):
+        self.cancel_flag = cancel_flag
+    
+    def __call__(self, input_ids, scores, **kwargs):
+        """Return True to stop generation immediately"""
+        return self.cancel_flag.is_set()
+
 
 def _download_and_patch_model_locally(model_id: str, revision: str) -> str:
     """
@@ -338,8 +368,17 @@ async def run_deepseek_ocr(
         # Note: We can't interrupt inference mid-process, but we can check before/after
         torch = _get_torch()
         with torch.inference_mode():
+            # Check cancellation one more time right before inference (critical point)
+            if job_id:
+                async with _jobs_lock:
+                    cancel_event = _cancellation_tokens.get(job_id)
+                    if cancel_event and cancel_event.is_set():
+                        raise asyncio.CancelledError(f"Job {job_id} was cancelled")
+            
             # Estimate inference takes ~80% of time (10-90%)
             # We'll update progress during post-processing
+            # Note: This is a blocking call - once it starts, it runs to completion
+            # The cancellation will be checked immediately after it returns
             result = model.infer(
                 tokenizer,
                 prompt=prompt,
@@ -351,6 +390,13 @@ async def run_deepseek_ocr(
                 save_results=False,
                 test_compress=False,
             )
+            
+            # Check cancellation immediately after inference completes
+            if job_id:
+                async with _jobs_lock:
+                    cancel_event = _cancellation_tokens.get(job_id)
+                    if cancel_event and cancel_event.is_set():
+                        raise asyncio.CancelledError(f"Job {job_id} was cancelled during inference")
         
         # Check for cancellation after inference
         if job_id:
@@ -402,6 +448,13 @@ async def run_deepseek_ocr(
                             if cancel_event and cancel_event.is_set():
                                 break
                     
+                    # Check cancellation right before each field detection
+                    if job_id:
+                        async with _jobs_lock:
+                            cancel_event = _cancellation_tokens.get(job_id)
+                            if cancel_event and cancel_event.is_set():
+                                raise asyncio.CancelledError(f"Job {job_id} was cancelled during field detection")
+                    
                     # Run locator query for this field
                     with torch.inference_mode():
                         locator_result = model.infer(
@@ -416,6 +469,13 @@ async def run_deepseek_ocr(
                             test_compress=False,
                         )
                     
+                    # Check cancellation immediately after locator inference
+                    if job_id:
+                        async with _jobs_lock:
+                            cancel_event = _cancellation_tokens.get(job_id)
+                            if cancel_event and cancel_event.is_set():
+                                raise asyncio.CancelledError(f"Job {job_id} was cancelled after field detection")
+                    
                     # Parse locator boxes from result
                     locator_text = locator_result if isinstance(locator_result, str) else str(locator_result)
                     locator_boxes = _parse_locator_boxes(locator_text, field_name)
@@ -988,6 +1048,7 @@ async def ocr_page(
                     _jobs[job_id]["status"] = "cancelled"
                     _jobs[job_id]["message"] = "Job was cancelled"
                 _cancellation_tokens.pop(job_id, None)
+                remove_cancel_flag(job_id)  # Cleanup cancel registry
             raise HTTPException(status_code=499, detail="Job was cancelled")
         except Exception as e:
             # Log the error and update job status
@@ -1120,15 +1181,37 @@ async def run_ocr_job_async(job_id: str, file: UploadFile, bus):
 
 
 @app.get("/progress/{job_id}")
-async def get_progress_stream(job_id: str):
-    """SSE stream for real-time OCR progress updates"""
+async def get_progress_stream(job_id: str, request: Request):
+    """SSE stream for real-time OCR progress updates with client disconnect detection"""
     try:
         from progress_bus import bus
     except ImportError:
         raise HTTPException(status_code=503, detail="SSE streaming not available")
     
+    async def gen_with_disconnect_check():
+        """Generator that checks for client disconnect and auto-cancels"""
+        try:
+            async for event in bus.stream(job_id):
+                # Check if client disconnected
+                if await request.is_disconnected():
+                    # Auto-cancel job on disconnect (optional but recommended)
+                    cancel_job(job_id)
+                    if job_id in _cancellation_tokens:
+                        _cancellation_tokens[job_id].set()
+                    async with _jobs_lock:
+                        if job_id in _jobs:
+                            _jobs[job_id]["status"] = "cancelled"
+                            _jobs[job_id]["message"] = "Client disconnected"
+                    break
+                yield event
+        except asyncio.CancelledError:
+            # Stream was cancelled
+            cancel_job(job_id)
+            if job_id in _cancellation_tokens:
+                _cancellation_tokens[job_id].set()
+    
     return StreamingResponse(
-        bus.stream(job_id),
+        gen_with_disconnect_check(),
         media_type="text/event-stream",
         headers={
             "Cache-Control": "no-cache",
@@ -1156,31 +1239,35 @@ async def get_job_status(job_id: str):
 
 
 @app.post("/jobs/{job_id}/cancel")
-async def cancel_job(job_id: str):
-    """Cancel a running OCR job"""
+async def cancel_job_endpoint(job_id: str):
+    """Cancel a running OCR job (cooperative cancellation with StoppingCriteria)"""
     async with _jobs_lock:
         if job_id not in _jobs:
             raise HTTPException(status_code=404, detail="Job not found")
         
         job = _jobs[job_id]
+        
+        # Already finished?
         if job["status"] in ("completed", "failed", "cancelled"):
-            return {"message": f"Job already {job['status']}"}
+            return {"ok": True, "message": f"Job already {job['status']}", "job_id": job_id}
         
-        # Set cancellation flag
+        # Set cancellation flag (use cancel_registry for consistency)
+        success = cancel_job(job_id)
         if job_id in _cancellation_tokens:
             _cancellation_tokens[job_id].set()
         
-        job["status"] = "cancelling"
+        job["status"] = "cancelled"
         job["message"] = "Cancellation requested..."
+        job["progress"] = job.get("progress", 0.0)
         
         # Send cancellation to SSE stream
         try:
             from progress_bus import bus
-            await bus.error(job_id, "Cancellation requested")
+            await bus.error(job_id, "Job cancelled by user")
         except ImportError:
             pass
         
-        return {"message": "Cancellation requested", "job_id": job_id}
+        return {"ok": True, "message": "Cancellation requested", "job_id": job_id}
 
 
 @app.post("/split")