Minor changes

Author: codeSamuraii

README.md

@@ -38,8 +38,8 @@ Only the entry point needs `@trace`. Everything it calls -- `add()`, imports, cl
 Start an isolated worker (temporary venv with `--tmp`, cleaned up on exit):
 
 ```bash
-pyfuse worker --tmp --backend local://localhost:9748  # Same machine (async TCP)
-pyfuse worker --tmp --backend redis://localhost:6379  # Remote using Redis
+pyfuse worker --tmp --backend local://localhost:9748   # Same machine (async TCP)
+pyfuse worker --tmp --backend redis://localhost:6379   # Remote using Redis
 ```
 
 Run a script:

docs/QUICK_START.md

@@ -36,7 +36,7 @@ def hypotenuse(a: float, b: float) -> float:
 ### 2. Start a worker
 
 ```bash
-python -m pyfuse worker --backend redis://localhost:6379
+pyfuse worker --backend redis://localhost:6379
 ```
 
 The worker waits for tasks, reconstructs the function source on the fly, and executes it -- no prior knowledge of your code required.
@@ -250,13 +250,13 @@ result = await fetch_and_process.run("https://example.com")
 
 ```bash
 # 4 concurrent tasks
-python -m pyfuse worker --backend redis://localhost:6379 -c 4
+pyfuse worker --backend redis://localhost:6379 -c 4
 
 # Disable automatic pip installs
-python -m pyfuse worker --backend redis://localhost:6379 --no-auto-install
+pyfuse worker --backend redis://localhost:6379 --no-auto-install
 
 # Run in an isolated temporary venv (auto-cleaned on exit)
-python -m pyfuse worker --backend redis://localhost:6379 --tmp
+pyfuse worker --backend redis://localhost:6379 --tmp
 ```
 
 Or start a worker programmatically:
@@ -273,7 +273,7 @@ asyncio.run(pyfuse.serve("redis://localhost:6379", concurrency=4))
 The `run` command creates an isolated venv, auto-detects third-party dependencies from the script (including `install_package_as` blocks), installs them, and runs the script:
 
 ```bash
-python -m pyfuse run examples/remote_execution.py
+pyfuse run examples/remote_execution.py
 ```
 
 ## Backends
@@ -507,8 +507,8 @@ except TaskCancelled:
 
 ```bash
 # Remote execution
-python -m pyfuse worker --backend local://localhost:9748 --tmp   # Terminal 1
-python -m pyfuse run examples/remote_execution.py               # Terminal 2
+pyfuse worker --backend local://localhost:9748 --tmp   # Terminal 1
+pyfuse run examples/remote_execution.py               # Terminal 2
 ```
 
 The `--tmp` flag creates an isolated temporary venv for the worker, which is auto-cleaned on exit.

docs/TECHNICAL_OVERVIEW.md

@@ -7,7 +7,7 @@ pyfuse enables remote execution of Python functions without deploying code to wo
 ```
   Client                              Worker
   ──────                              ──────
-  @trace                              python -m pyfuse worker
+  @trace                              pyfuse worker
   await func.run(args)                ← listen for tasks
     │                                   │
     ├─ capture source + deps            ├─ deserialize graph
@@ -22,7 +22,7 @@ pyfuse enables remote execution of Python functions without deploying code to wo
 ```
 pyfuse/
     __init__.py              Public API: trace, connect, serve, serialize, reconstruct, ...
-    __main__.py              CLI: python -m pyfuse worker/run/info/serialize/reconstruct
+    __main__.py              CLI: pyfuse worker/run/info/serialize/reconstruct
     _venv.py                 Temporary virtual environment management (async)
     core/
         models.py            FunctionNode, ImportInfo dataclasses (incl. content hashing, class metadata)
@@ -55,7 +55,7 @@ pyfuse/
 3. **Submit** -- `await backend.submit(task_json)` sends the task to the transport layer.
 4. **Return** -- `.run()` awaits the result and returns the value directly. `.start()` returns a `Result` handle immediately.
 
-### Worker side: `await serve()` / `python -m pyfuse worker`
+### Worker side: `await serve()` / `pyfuse worker`
 
 1. **Listen** -- `async for task_json in backend.listen()` yields tasks as they arrive.
 2. **Cancellation check** -- If `await backend.is_cancelled(task_id)` returns ``True``, skip execution and log the cancellation.
@@ -513,20 +513,20 @@ When a module contains `from X import *`:
 
 ```bash
 # Start a worker
-python -m pyfuse worker --backend redis://localhost:6379
-python -m pyfuse worker --backend redis://localhost:6379 -c 4
-python -m pyfuse worker --backend redis://localhost:6379 --no-auto-install
-python -m pyfuse worker --backend redis://localhost:6379 --tmp   # isolated temp venv
+pyfuse worker --backend redis://localhost:6379
+pyfuse worker --backend redis://localhost:6379 -c 4
+pyfuse worker --backend redis://localhost:6379 --no-auto-install
+pyfuse worker --backend redis://localhost:6379 --tmp   # isolated temp venv
 
 # Run a script in a temporary venv (auto-detects and installs dependencies)
-python -m pyfuse run examples/script.py
+pyfuse run examples/script.py
 
 # Show configuration
-python -m pyfuse info
+pyfuse info
 
 # Serialize a function to JSON
-python -m pyfuse serialize mymodule:csv_to_json
+pyfuse serialize mymodule:csv_to_json
 
 # Reconstruct source from a graph file
-python -m pyfuse reconstruct graph.json csv_to_json
+pyfuse reconstruct graph.json csv_to_json
 ```

examples/cancellation.py

@@ -7,7 +7,7 @@
     pyfuse worker --backend redis://localhost:6379 --tmp
 
     # Terminal 2 -- run this script
-    python -m pyfuse run examples/cancellation.py
+    pyfuse run examples/cancellation.py
 """
 
 import asyncio

examples/large_module.py

@@ -8,7 +8,7 @@
 
 Usage:
     # Terminal 1 -- start a worker
-    python -m pyfuse worker --backend redis://localhost:6379
+    pyfuse worker --backend redis://localhost:6379
 
     # Terminal 2 -- run this script
     python examples/large_module.py

examples/package_installation.py

@@ -8,7 +8,7 @@
 
 Usage:
     # Terminal 1 -- start a worker
-    python -m pyfuse worker --backend redis://localhost:6379
+    pyfuse worker --backend redis://localhost:6379
 
     # Terminal 2 -- run this script
     python examples/package_installation.py

examples/progress_reporting.py

@@ -8,7 +8,7 @@
     pyfuse worker --backend redis://localhost:6379 --tmp
 
     # Terminal 2 -- run this script
-    python -m pyfuse run examples/progress_reporting.py
+    pyfuse run examples/progress_reporting.py
 """
 
 import asyncio

pyfuse/__main__.py

@@ -2,12 +2,12 @@
 
 Usage::
 
-    python -m pyfuse worker --backend redis://localhost:6379
-    python -m pyfuse worker --backend redis://localhost:6379 --tmp
-    python -m pyfuse run examples/script.py
-    python -m pyfuse info
-    python -m pyfuse serialize mymodule:csv_to_json
-    python -m pyfuse reconstruct graph.json csv_to_json
+    pyfuse worker --backend redis://localhost:6379
+    pyfuse worker --backend redis://localhost:6379 --tmp
+    pyfuse run examples/script.py
+    pyfuse info
+    pyfuse serialize mymodule:csv_to_json
+    pyfuse reconstruct graph.json csv_to_json
 """
 from __future__ import annotations
 

pyfuse/worker/remote.py

@@ -189,13 +189,27 @@ async def _heartbeat_loop(
     backend: Backend,
     task_id: str,
     cancel_event: asyncio.Event,
+    exec_task: asyncio.Task[Any] | None = None,
 ) -> None:
-    """Send periodic heartbeats until *cancel_event* is set."""
+    """Send periodic heartbeats and check for cancellation.
+
+    When *exec_task* is provided and the backend reports the task as
+    cancelled, the execution task is cancelled via
+    :meth:`asyncio.Task.cancel`, which raises :class:`CancelledError`
+    at the next ``await`` in async user functions.
+    """
     while not cancel_event.is_set():
         try:
             await backend.send_heartbeat(task_id)
         except Exception:
             pass  # best-effort
+        if exec_task is not None:
+            try:
+                if await backend.is_cancelled(task_id):
+                    exec_task.cancel()
+                    return
+            except Exception:
+                pass
         try:
             await asyncio.wait_for(cancel_event.wait(), timeout=_HEARTBEAT_INTERVAL)
         except asyncio.TimeoutError:
@@ -316,18 +330,25 @@ async def _handle_task(
         _log_task_result(task, envelope, 0, worker)
         return
 
-    cancel_event = asyncio.Event()
-    hb_task = asyncio.create_task(_heartbeat_loop(backend, task.task_id, cancel_event))
-
     # Set up rate-limited progress callback
     loop = asyncio.get_running_loop()
     progress_cb, flush = _make_progress_callback(backend, task.task_id, loop)
     token = _progress_callback.set(progress_cb)
 
+    # Run execution as a task so the heartbeat loop can cancel it
+    exec_task: asyncio.Task[Any] = asyncio.create_task(worker.run_with_policy(task))
+
+    cancel_event = asyncio.Event()
+    hb_task = asyncio.create_task(
+        _heartbeat_loop(backend, task.task_id, cancel_event, exec_task),
+    )
+
     t0 = time.monotonic()
     try:
-        result = await worker.run_with_policy(task)
+        result = await exec_task
         envelope = ResultEnvelope.success(task.task_id, result)
+    except asyncio.CancelledError:
+        envelope = ResultEnvelope.cancelled(task.task_id)
     except Exception as exc:
         logger.debug("Task %s failed", task.task_id, exc_info=True)
         envelope = ResultEnvelope.failure(task.task_id, exc)

tests/test_cancellation_progress.py

@@ -125,6 +125,22 @@ def _slow_store() -> tuple[Store, str]:
     return store, store.to_json()
 
 
+def _async_slow_store() -> tuple[Store, str]:
+    store = Store()
+    node = _node("slow_async", (
+        "async def slow_async(n):\n"
+        "    import asyncio\n"
+        "    total = 0\n"
+        "    for i in range(n):\n"
+        "        await asyncio.sleep(0.1)\n"
+        "        total += i\n"
+        "    return total"
+    ))
+    h = store.put(node)
+    store.set_ref("slow_async", h)
+    return store, store.to_json()
+
+
 def _progress_store() -> tuple[Store, str]:
     store = Store()
     node = _node("process", (
@@ -238,8 +254,42 @@ async def test_handle_task_skips_cancelled(self, backend: InMemoryBackend) -> No
         assert task.task_id not in backend.results
 
     @pytest.mark.asyncio
-    async def test_cancel_during_execution(self, backend: InMemoryBackend) -> None:
-        """Worker sends its result normally; cancel semantics are client-side."""
+    async def test_cancel_during_async_execution(
+        self, backend: InMemoryBackend, monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Cancelling during async execution interrupts the function."""
+        monkeypatch.setattr(_remote, "_HEARTBEAT_INTERVAL", 0.1)
+
+        _, json_str = _async_slow_store()
+        # 30 iterations × 0.1 s = 3 s without cancellation
+        task = Task(graph_json=json_str, function_name="slow_async", args=(30,))
+
+        worker = Worker(auto_install=False)
+
+        async def cancel_later() -> None:
+            await asyncio.sleep(0.15)
+            await backend.cancel_task(task.task_id)
+
+        t0 = time.monotonic()
+        asyncio.create_task(cancel_later())
+        await _handle_task(worker, backend, task.to_json())
+        elapsed = time.monotonic() - t0
+
+        # Should finish well before 3 s (the function was interrupted)
+        assert elapsed < 1.5
+
+        # Worker sends a cancelled envelope
+        assert task.task_id in backend.results
+        env = ResultEnvelope.from_json(backend.results[task.task_id])
+        assert env.status == "cancelled"
+
+    @pytest.mark.asyncio
+    async def test_cancel_during_sync_execution(
+        self, backend: InMemoryBackend, monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Cancelling during sync execution marks result as cancelled."""
+        monkeypatch.setattr(_remote, "_HEARTBEAT_INTERVAL", 0.1)
+
         _, json_str = _slow_store()
         task = Task(graph_json=json_str, function_name="slow", args=(5,))
 
@@ -252,10 +302,10 @@ async def cancel_later() -> None:
         asyncio.create_task(cancel_later())
         await _handle_task(worker, backend, task.to_json())
 
-        # Worker always sends its result; client reads cancel envelope first
+        # Executor thread can't be interrupted, but the coroutine is cancelled
         assert task.task_id in backend.results
         env = ResultEnvelope.from_json(backend.results[task.task_id])
-        assert env.status == "ok"
+        assert env.status == "cancelled"
 
 
 # ---------------------------------------------------------------------------