cancellation and progress

Author: codeSamuraii

README.md

@@ -81,6 +81,9 @@ Common mappings (`cv2` -> `opencv-python`, `PIL` -> `Pillow`, etc.) are built in
 - **Third-party package auto-install** -- Workers install missing packages via pip before execution.
 - **Async-native** -- The entire I/O layer is built on `asyncio`. `.run()`, `.start()`, `.map()`, `await result`, and `asyncio.gather` all work out of the box.
 - **Heartbeat & stall detection** -- Workers send periodic heartbeats. Clients raise `TaskStalled` when a worker stops responding.
+- **Task cancellation** -- `await future.cancel()` cancels pending or in-progress tasks.
+- **Progress reporting** -- Call `pyfuse.progress(75.0)` or `pyfuse.progress(3, 10)` inside tasks; query with `await future.progress()`.
+- **Graceful shutdown** -- Workers finish in-progress tasks before stopping. Second Ctrl+C force-quits.
 - **Class methods** -- `self.method()` and `cls.method()` dependencies are detected. Entire class hierarchies (including `super()`), class-level attributes, decorators (`@dataclass`, etc.), and metaclass keywords are reconstructed.
 - **Retry and timeout** -- `@trace(timeout=30, retries=3)` with exponential backoff.
 - **Batch submission** -- `await func.map([(a1, b1), (a2, b2)])` submits and awaits multiple tasks.
@@ -98,6 +101,8 @@ pyfuse run examples/remote_execution.py
 - **[`examples/remote_execution.py`](examples/remote_execution.py)** -- Remote execution with auto-discovered dependencies
 - **[`examples/async_execution.py`](examples/async_execution.py)** -- Async: `.run()`, `.start()`, `.map()`, `asyncio.gather`
 - **[`examples/package_installation.py`](examples/package_installation.py)** -- Auto-installing third-party packages on workers
+- **[`examples/progress_reporting.py`](examples/progress_reporting.py)** -- Real-time progress tracking from long-running tasks
+- **[`examples/cancellation.py`](examples/cancellation.py)** -- Cancelling pending or in-progress tasks
 - **[`examples/large_module.py`](examples/large_module.py)** -- Stress test: 47 functions across 7 files, one `@trace`
 
 ## Documentation

docs/CONTEXT.md

@@ -18,7 +18,8 @@ pyfuse/
 │   ├── task.py              # Task dataclass: serializable envelope (graph + args + options)
 │   ├── models.py            # FunctionNode and ImportInfo dataclasses, content hashing
 │   ├── version.py           # _VERSION = "0.4.0"
-│   └── errors.py            # Error, WorkerError, RemoteError, DependencyError, TaskStalled
+│   ├── errors.py            # Error, WorkerError, RemoteError, DependencyError, TaskStalled, TaskCancelled
+│   └── progress.py          # ProgressInfo dataclass, progress() function, context variable
 ├── graph/
 │   ├── decorator.py         # @trace: marks functions, adds .run()/.start()/.map()
 │   ├── graph.py             # Graph class: registration, auto-discovery, serialization
@@ -54,9 +55,9 @@ Data models and error types. `FunctionNode` represents a function in the graph,
 
 Handles remote execution. Built entirely on `asyncio`:
 
-- **`worker.py`**: `Worker` class reconstructs functions from serialized stores, caches compiled namespaces by subgraph content hash, and executes with retry/timeout policies. Async user functions are awaited directly; sync user functions run in `loop.run_in_executor()`. Timeouts use `asyncio.wait_for()`.
-- **`remote.py`**: Orchestrates the connection lifecycle, worker event loop (`asyncio.TaskGroup` + `asyncio.Semaphore` for bounded concurrency), and heartbeat tasks (`asyncio.create_task`).
-- **`result.py`**: `Result` is an awaitable future returned by `.start()`. Simple async polling loop for stall detection.
+- **`worker.py`**: `Worker` class reconstructs functions from serialized stores, caches compiled namespaces by subgraph content hash, and executes with retry/timeout policies. Async user functions are awaited directly; sync user functions run in `loop.run_in_executor()` with explicit context propagation via `contextvars.copy_context()`. Timeouts use `asyncio.wait_for()`.
+- **`remote.py`**: Orchestrates the connection lifecycle, worker event loop (`asyncio.Semaphore` for bounded concurrency), heartbeat tasks (`asyncio.create_task`), progress injection, cancellation checking, and graceful shutdown via signal handling.
+- **`result.py`**: `Result` is an awaitable future returned by `.start()`. Simple async polling loop for stall detection. Supports `cancel()` and `progress()` methods.
 - **`deps.py`**: Package installation via `asyncio.create_subprocess_exec`.
 - **`backends/`**: All backend methods are `async def`. `listen()` and `subscribe_results()` are async generators.
 
@@ -99,8 +100,11 @@ await Worker.run(task)
 - **Decorator stripping**: `@trace` lines are removed from captured source so reconstructed code doesn't depend on pyfuse.
 - **Backend auto-detection**: `connect()` picks Redis or local TCP backend based on URL scheme. Falls back to `PYFUSE_BACKEND` env var.
 - **Worker caching**: Keyed by SHA-256 of all reachable content hashes (sorted + joined). Same code from different clients = cache hit.
-- **Async-native I/O**: All backend methods, worker execution, result handling, pip installation, and subprocess management use `asyncio`. Sync user functions run in `loop.run_in_executor()` to avoid blocking the event loop.
+- **Async-native I/O**: All backend methods, worker execution, result handling, pip installation, and subprocess management use `asyncio`. Sync user functions run in `loop.run_in_executor()` with explicit `contextvars.copy_context()` to propagate progress callbacks.
 - **Heartbeat**: Workers send heartbeats via `asyncio.create_task`. Client-side stall detection tracks when heartbeat *values* last changed using local monotonic clock (no cross-machine timestamp comparison).
+- **Task cancellation**: Cooperative via backend flags. Workers check before execution; clients store a "cancelled" result envelope. `TaskCancelled` is raised on await.
+- **Progress reporting**: Uses `contextvars.ContextVar` for the progress callback. Sync functions get context propagated via explicit copy. Progress updates are fire-and-forget async tasks.
+- **Graceful shutdown**: Signal-based (`SIGINT`/`SIGTERM`). First signal stops the listener and waits for in-flight tasks. Second signal cancels all tasks.
 
 ## Serialization format (v0.4.0)
 
@@ -146,7 +150,16 @@ results = await func.map([(a1, b1), ...]) # batch submit + await all (returns va
 result = await future                     # shorthand for await future.result()
 result = await future.result(timeout=10, stall_timeout=10.0)  # with options
 await future.done()                       # non-blocking check
-await future.status()                     # "pending", "success", or "error"
+await future.status()                     # "pending", "success", "error", or "cancelled"
+
+# Cancellation
+await future.cancel()                     # cancel task; raises TaskCancelled when awaited
+
+# Progress reporting
+pyfuse.progress(75.0)                     # report percentage (no-op locally)
+pyfuse.progress(3, 10, message="step 3")  # report current/total with message
+p = await future.progress()               # get latest ProgressInfo (or None)
+if p: print(f"{p.current}/{p.total} {p.percent:.0f}%")
 
 # Serialization (sync -- pure CPU)
 pyfuse.serialize(func)                    # -> JSON string
@@ -166,7 +179,7 @@ pytest                    # run all tests
 pytest tests/test_api.py  # specific module
 ```
 
-15 test modules covering: API surface, AST analysis, async features (Result.result, await, .run(), .start(), .map(), gather, heartbeat, stall detection), auto-discovery (including metaclass keywords, class attributes, class decorators, `__init_subclass__`), dependency management, graph operations, integration scenarios, local backend (async-native TCP), remote execution, runtime tracing (including closure capture of non-traced functions, lambdas, constructor expressions, pickle fallback), store operations, stress tests (47 functions across 7 files), task serialization, temp venv management, and worker caching/execution.
+16 test modules covering: API surface, AST analysis, async features (Result.result, await, .run(), .start(), .map(), gather, heartbeat, stall detection), auto-discovery (including metaclass keywords, class attributes, class decorators, `__init_subclass__`), cancellation and progress reporting, dependency management, graph operations, integration scenarios, local backend (async-native TCP), remote execution, runtime tracing (including closure capture of non-traced functions, lambdas, constructor expressions, pickle fallback), store operations, stress tests (47 functions across 7 files), task serialization, temp venv management, and worker caching/execution.
 
 All async tests use `pytest-asyncio` with `asyncio_mode = "auto"`.
 
@@ -185,7 +198,7 @@ pytest                          # test suite
 - `analyzer.py` is the core of static analysis (~365 lines). Changes here affect what gets captured.
 - `tracing.py` uses `contextvars.ContextVar` for thread/async safety. The `_runtime_deps` dict is guarded by `threading.Lock`.
 - The `Task` wire format keeps `graph` as a JSON string (not nested object) to keep the envelope flat.
-- Backend implementations must satisfy the `Backend` ABC in `backends/base.py`. All methods are `async def`. New methods (`notify_result`, `subscribe_results`, `get_heartbeats`) are non-abstract with safe defaults -- custom backends don't break.
+- Backend implementations must satisfy the `Backend` ABC in `backends/base.py`. All methods are `async def`. New methods (`notify_result`, `subscribe_results`, `get_heartbeats`, `cancel_task`, `is_cancelled`, `send_progress`, `get_progress`) are non-abstract with safe defaults -- custom backends don't break.
 - `install_package_as()` is a no-op at runtime; the AST analyzer in `decorator.py`/`analyzer.py` detects the `with` block pattern and tags `ImportInfo` objects with the package name.
 - `_capture_closure()` in `graph.py` uses a multi-tier strategy: repr validation → traced functions → lambdas (source extraction) → non-traced user functions (auto-registration) → constructor expressions (defaultdict/Counter/deque) → pickle fallback → warning. Returns function objects for auto-registration.
 - `_set_class_metadata()` in `graph.py` captures class-level attributes and decorators from the class source AST. Called from both `_auto_register_class` and `_discover_self_call_deps` to handle both constructor-discovered and directly-traced method classes.

docs/QUICK_START.md

@@ -118,6 +118,57 @@ result = await to_yaml.run({"key": "value"})
 
 Common mappings like `cv2` -> `opencv-python` and `PIL` -> `Pillow` are built in.
 
+## Task cancellation
+
+Cancel a pending or in-progress task with `await future.cancel()`:
+
+```python
+from pyfuse import TaskCancelled
+
+future = await slow_task.start(data)
+await asyncio.sleep(1)
+await future.cancel()
+
+try:
+    result = await future
+except TaskCancelled:
+    print("Task was cancelled")
+```
+
+If the worker hasn't started execution yet, it skips the task entirely. If execution is already in progress, the client receives `TaskCancelled` when awaiting the result.
+
+## Progress reporting
+
+Long-running tasks can report progress back to the client:
+
+```python
+from pyfuse import trace, progress
+
+@trace
+def process_batch(items: list[str]) -> list[str]:
+    results = []
+    for i, item in enumerate(items):
+        results.append(transform(item))
+        progress(i + 1, len(items), message=f"Processing {item}")
+    return results
+```
+
+Query progress from the client:
+
+```python
+future = await process_batch.start(items)
+
+while not await future.done():
+    p = await future.progress()
+    if p is not None:
+        print(f"{p.current}/{p.total} ({p.percent:.0f}%) - {p.message}")
+    await asyncio.sleep(0.5)
+
+result = await future
+```
+
+`progress()` is a silent no-op when called outside a worker, so the function works unchanged locally.
+
 ## Async API
 
 pyfuse is async-native. All remote execution methods are coroutines.
@@ -259,7 +310,14 @@ future = await hypotenuse.start(3.0, 4.0)
 
 # Check status without blocking
 await future.done()      # True / False
-await future.status()    # "pending", "success", or "error"
+await future.status()    # "pending", "success", "error", or "cancelled"
+
+# Check progress (from pyfuse.progress() calls)
+p = await future.progress()    # ProgressInfo or None
+if p: print(f"{p.current}/{p.total}")
+
+# Cancel a task
+await future.cancel()    # raises TaskCancelled when awaited
 
 # Await the result
 result = await future                          # shorthand
@@ -428,7 +486,7 @@ merged = json.dumps({
 ## Error handling
 
 ```python
-from pyfuse import trace, Error, RemoteError
+from pyfuse import trace, Error, RemoteError, TaskCancelled
 
 # Tracing errors
 try:
@@ -441,6 +499,8 @@ try:
     result = await future.result()
 except RemoteError as e:
     print(e)  # includes remote traceback
+except TaskCancelled:
+    print("Task was cancelled")
 ```
 
 ## Running the examples

docs/TECHNICAL_OVERVIEW.md

@@ -58,12 +58,13 @@ pyfuse/
 ### Worker side: `await serve()` / `python -m pyfuse worker`
 
 1. **Listen** -- `async for task_json in backend.listen()` yields tasks as they arrive.
-2. **Deserialize** -- Parse the JSON graph into a `Store`.
-3. **Cache check** -- Compute a subgraph key (SHA-256 of all reachable content hashes). If cached, skip to step 6.
-4. **Install dependencies** -- Extract third-party imports, install missing packages via `asyncio.create_subprocess_exec`.
-5. **Reconstruct** -- Produce a self-contained Python script from the store. `compile()` and `exec()` it into a fresh namespace.
-6. **Execute** -- Call the function with the provided arguments. Async functions are awaited directly; sync functions run in an executor. Apply retry/timeout policies via `asyncio.wait_for`.
-7. **Send result** -- Wrap the return value (or exception traceback) in a `ResultEnvelope` and send it back.
+2. **Cancellation check** -- If `await backend.is_cancelled(task_id)` returns ``True``, skip execution and log the cancellation.
+3. **Deserialize** -- Parse the JSON graph into a `Store`.
+4. **Cache check** -- Compute a subgraph key (SHA-256 of all reachable content hashes). If cached, skip to step 7.
+5. **Install dependencies** -- Extract third-party imports, install missing packages via `asyncio.create_subprocess_exec`.
+6. **Reconstruct** -- Produce a self-contained Python script from the store. `compile()` and `exec()` it into a fresh namespace.
+7. **Execute** -- Call the function with the provided arguments. Async functions are awaited directly; sync functions run in an executor with explicit context propagation (for progress reporting). Apply retry/timeout policies via `asyncio.wait_for`.
+8. **Send result** -- Wrap the return value (or exception traceback) in a `ResultEnvelope` and send it back. If cancelled during execution, skip result delivery.
 
 ### Client side: `await future` / `await future.result()`
 
@@ -89,6 +90,10 @@ The `Backend` ABC defines an async transport interface:
 | `async send_heartbeat(task_id)` | Signal active processing (no-op default) |
 | `async get_heartbeat(task_id)` | Get last heartbeat timestamp |
 | `async get_heartbeats(task_ids)` | Batch heartbeat fetch (default loops over `get_heartbeat`) |
+| `async cancel_task(task_id)` | Mark a task as cancelled (no-op default) |
+| `async is_cancelled(task_id)` | Check cancellation flag (default returns `False`) |
+| `async send_progress(task_id, json)` | Store latest progress data (no-op default) |
+| `async get_progress(task_id)` | Get latest progress JSON (default returns `None`) |
 | `async notify_result(task_id)` | Push notification that result is ready (no-op default) |
 | `async subscribe_results()` | Async iterator yielding task_ids on result arrival |
 | `async close()` | Release resources |
@@ -101,6 +106,8 @@ Uses `redis.asyncio.Redis` with `RPUSH`/`BLPOP` patterns. Keys:
 - `pyfuse:tasks` -- task queue
 - `pyfuse:result:{task_id}` -- per-task result (TTL: 300s)
 - `pyfuse:heartbeat:{task_id}` -- worker heartbeat timestamp (TTL: 30s)
+- `pyfuse:cancel:{task_id}` -- cancellation flag (TTL: 3600s)
+- `pyfuse:progress:{task_id}` -- latest progress JSON (TTL: 300s)
 - `pyfuse:notify` -- Pub/Sub channel for result notifications
 
 Result notifications use Redis Pub/Sub (`PUBLISH`/`SUBSCRIBE`). Batch heartbeat fetching uses `MGET` for efficiency.
@@ -283,9 +290,11 @@ Error (includes remote traceback):
 | Method / Property | Description |
 |------------------|-------------|
 | `await result` | Shorthand for `await result.result()` |
-| `await result.result(timeout, stall_timeout=10.0)` | Await with options; raises `RemoteError` on failure, `TaskStalled` on stall |
+| `await result.result(timeout, stall_timeout=10.0)` | Await with options; raises `RemoteError` on failure, `TaskStalled` on stall, `TaskCancelled` on cancel |
+| `await result.cancel()` | Cancel the task; awaiting raises `TaskCancelled` |
+| `await result.progress()` | Latest `ProgressInfo` (or `None`) |
 | `await result.done()` | Non-blocking check |
-| `await result.status()` | `"pending"`, `"success"`, or `"error"` |
+| `await result.status()` | `"pending"`, `"success"`, `"error"`, or `"cancelled"` |
 | `.task_id` | The task identifier |
 
 ## Serialization format
@@ -417,6 +426,54 @@ Stall detection only triggers after at least one heartbeat has been observed, av
 |--------|----------------|
 | `await result.result()` | On by default (`stall_timeout=10.0`). Disable with `stall_timeout=None` |
 
+## Task cancellation
+
+Tasks can be cancelled via `await result.cancel()`. Cancellation is cooperative:
+
+1. **Client** calls `cancel()`, which sets a cancellation flag in the backend and stores a ``"cancelled"`` result envelope.
+2. **Worker** checks `is_cancelled()` before starting execution. If cancelled, the task is skipped entirely (no result sent, since the cancel already stored one).
+3. **During execution**: if `cancel()` is called while a task is running, the execution continues. When the worker finishes, it checks `is_cancelled()` again and discards the result if cancelled.
+4. **Client** receives `TaskCancelled` when awaiting a cancelled task.
+
+### Backend storage
+
+| Backend | Cancellation storage |
+|---------|---------------------|
+| Redis | `SET pyfuse:cancel:{task_id} 1 EX 3600` |
+| Local | In-memory `set()` in the broker |
+
+## Progress reporting
+
+Tasks can report progress via `pyfuse.progress(percent)` or `pyfuse.progress(current, total)`.
+
+### How it works
+
+1. **Context variable**: A `contextvars.ContextVar` holds the progress callback. The worker sets it before executing each task.
+2. **Sync function support**: `Worker.run()` explicitly propagates context variables to executor threads via `contextvars.copy_context().run()`.
+3. **Rate-limited sends**: The progress callback rate-limits backend sends to one per 50 ms. Intermediate updates are stored locally. A ``flush()`` coroutine sends the final state after execution completes.
+4. **No-op locally**: When called outside a worker, `progress()` is a silent no-op (context variable is ``None``).
+
+### Backend storage
+
+| Backend | Progress storage |
+|---------|-----------------|
+| Redis | `SET pyfuse:progress:{task_id} <json> EX 300` |
+| Local | In-memory `dict` in the broker |
+
+## Graceful shutdown
+
+Workers support graceful shutdown via signal handling:
+
+1. **First SIGINT/SIGTERM**: Sets a shutdown event, stops accepting new tasks, and waits for in-progress tasks to complete.
+2. **Second SIGINT/SIGTERM**: Cancels all in-progress tasks immediately and exits.
+
+Signal handlers are installed via `loop.add_signal_handler()` (Unix). On Windows, falls back to `KeyboardInterrupt` handling. The worker logs shutdown progress:
+
+```
+12:30:00 INFO    Graceful shutdown: waiting for 2 task(s) to complete... (Ctrl+C to force quit)
+12:30:02 INFO    Worker stopped.
+```
+
 ## Thread and task safety
 
 - The runtime call stack uses `contextvars.ContextVar`, providing per-thread isolation in sync code and per-task isolation in async code.