Add guards to .run_every()

Author: codeSamuraii

examples/stress_task.py

@@ -0,0 +1,112 @@
+"""CPU and memory stress task for cloud_poc.
+
+Submits a task that burns CPU for a given duration while repeatedly
+allocating and releasing chunks of memory.
+
+Usage:
+    BROKER_URL="http://localhost:8000/api/v1/broker?api_key=<key>" \
+        python stress_task.py [cpu_seconds] [mem_mib_per_iter] [mem_iters]
+
+Arguments (positional, all optional):
+    cpu_seconds       Seconds to keep the CPU busy     (default: 10)
+    mem_mib_per_iter  MiB to allocate per iteration    (default: 64)
+    mem_iters         Number of allocation iterations  (default: 8)
+"""
+
+import asyncio
+import os
+import sys
+
+import offwork
+from offwork import progress
+
+broker_url = os.environ.get("BROKER_URL")
+if not broker_url:
+    print("error: BROKER_URL is not set", file=sys.stderr)
+    sys.exit(1)
+
+offwork.connect(broker_url)
+
+
+def _burn_cpu(seconds: float) -> int:
+    """Busy-loop for *seconds* wall time. Returns iteration count."""
+    import time
+    import hashlib
+
+    deadline = time.monotonic() + seconds
+    data = b"offwork-stress" * 256  # ~3.5 KB
+    iters = 0
+    while time.monotonic() < deadline:
+        hashlib.sha256(data).digest()
+        iters += 1
+    return iters
+
+
+def _alloc_mib(mib: float) -> int:
+    """Allocate *mib* MiB, touch every page, return checksum byte."""
+    size = int(mib * 1024 * 1024)
+    chunk = bytearray(size)
+    # Touch every 4096-byte page so the OS actually faults the memory in.
+    for i in range(0, size, 4096):
+        chunk[i] = (i // 4096) & 0xFF
+    return chunk[-1]
+
+
+@offwork.task
+def stress(cpu_seconds: float, mem_mib_per_iter: float, mem_iters: int) -> dict:
+    """Stress the CPU and memory on the worker.
+
+    Args:
+        cpu_seconds:      Wall-clock seconds to spend hashing.
+        mem_mib_per_iter: MiB to allocate and touch per memory iteration.
+        mem_iters:        How many times to allocate / release that block.
+
+    Returns a summary dict with hash iteration count and total memory touched.
+    """
+    import time
+
+    start = time.monotonic()
+    total_steps = mem_iters  # N CPU phase + N mem phases
+
+    checksums = []
+    for i in range(mem_iters):
+        # --- CPU phase ---
+        progress(i, total_steps, message=f"Iter {i + 1}/{mem_iters}")
+        hash_iters = _burn_cpu(cpu_seconds)
+        checksums.append(_alloc_mib(mem_mib_per_iter))
+
+    elapsed = time.monotonic() - start
+    total_mib = mem_mib_per_iter * mem_iters
+    return {
+        "elapsed_seconds": round(elapsed, 2),
+        "cpu_seconds_requested": cpu_seconds,
+        "hash_iterations": hash_iters,
+        "memory_iterations": mem_iters,
+        "mem_mib_per_iter": mem_mib_per_iter,
+        "total_mem_touched_mib": total_mib,
+        "checksum": sum(checksums) & 0xFF,
+    }
+
+
+async def main() -> None:
+    cpu_seconds = float(sys.argv[1]) if len(sys.argv) > 1 else 10.0
+    mem_mib_per_iter = float(sys.argv[2]) if len(sys.argv) > 2 else 64.0
+    mem_iters = int(sys.argv[3]) if len(sys.argv) > 3 else 8
+
+    print(f"submitting stress task: cpu={cpu_seconds}s  mem={mem_mib_per_iter}MiB x {mem_iters} iters")
+    future = await stress.start(cpu_seconds, mem_mib_per_iter, mem_iters)
+    print(f"task id: {future.task_id}")
+
+    while not await future.done():
+        p = await future.progress()
+        if p is not None:
+            print(f"  [{p.current}/{p.total}] {p.message}")
+        await asyncio.sleep(2.0)
+
+    result = await future
+    print("\nresult:")
+    for k, v in result.items():
+        print(f"  {k}: {v}")
+
+
+asyncio.run(main())

offwork/core/envelope.py

@@ -171,6 +171,8 @@ def verify_task_envelope(
         retry_delay=data.get("retry_delay", 1.0),
         scheduled_at=data.get("scheduled_at"),
         recur_interval=data.get("recur_interval"),
+        recur_deadline=data.get("recur_deadline"),
+        recur_remaining=data.get("recur_remaining"),
         schedule_id=data.get("schedule_id"),
         throttle=data.get("throttle"),
     )

offwork/core/task.py

@@ -418,6 +418,8 @@ class Task:
     retry_delay: float = 1.0
     scheduled_at: float | None = None
     recur_interval: float | None = None
+    recur_deadline: float | None = None
+    recur_remaining: int | None = None
     schedule_id: str | None = None
     throttle: float | None = None
 
@@ -442,6 +444,10 @@ def _to_dict(self) -> dict[str, Any]:
             d["scheduled_at"] = self.scheduled_at
         if self.recur_interval is not None:
             d["recur_interval"] = self.recur_interval
+        if self.recur_deadline is not None:
+            d["recur_deadline"] = self.recur_deadline
+        if self.recur_remaining is not None:
+            d["recur_remaining"] = self.recur_remaining
         if self.schedule_id is not None:
             d["schedule_id"] = self.schedule_id
         if self.throttle is not None:
@@ -467,6 +473,8 @@ def from_json(cls, json_str: str | bytes) -> Self:
             retry_delay=data.get("retry_delay", 1.0),
             scheduled_at=data.get("scheduled_at"),
             recur_interval=data.get("recur_interval"),
+            recur_deadline=data.get("recur_deadline"),
+            recur_remaining=data.get("recur_remaining"),
             schedule_id=data.get("schedule_id"),
             throttle=data.get("throttle"),
         )

offwork/graph/tracing.py

@@ -137,6 +137,8 @@ async def run_every(
         frequency: Any,
         *args: Any,
         _start_at: Any = None,
+        run_for: Any = None,
+        max_runs: int | None = None,
         backend: str | Backend | None = None,
         **kwargs: Any,
     ) -> object:
@@ -146,9 +148,19 @@ async def run_every(
         start_ts: float | None = None
         if _start_at is not None:
             start_ts = _start_at.timestamp() if isinstance(_start_at, datetime) else float(_start_at)
+        if run_for is None and max_runs is None:
+            run_for = timedelta(hours=1)
+        run_for_seconds: float | None = None
+        if run_for is not None:
+            run_for_seconds = run_for.total_seconds() if isinstance(run_for, timedelta) else float(run_for)
+            if run_for_seconds <= 0:
+                raise ValueError(f"run_for must be positive, got {run_for}")
+        if max_runs is not None and max_runs <= 0:
+            raise ValueError(f"max_runs must be positive, got {max_runs}")
         return await submit_recurring(
             func, wrapper, *args,
             _backend=backend, _interval=interval, _start_at=start_ts,
+            _run_for=run_for_seconds, _max_runs=max_runs,
             **kwargs,
         )
 

offwork/typing.py

@@ -38,6 +38,8 @@ async def run_every(
         frequency: timedelta | float,
         *args: Any,
         _start_at: datetime | None = ...,
+        run_for: timedelta | float | None = ...,
+        max_runs: int | None = ...,
         **kwargs: Any,
     ) -> ScheduleHandle: ...
 

offwork/worker/remote.py

@@ -280,6 +280,8 @@ async def submit_recurring(
     _root_token: bytes | None = None,
     _interval: float = 0,
     _start_at: float | None = None,
+    _run_for: float | None = None,
+    _max_runs: int | None = None,
     **kwargs: Any,
 ) -> ScheduleHandle:
     """Submit a recurring task and return a :class:`ScheduleHandle`."""
@@ -302,6 +304,7 @@ async def submit_recurring(
 
     schedule_id = uuid.uuid4().hex[:12]
     scheduled_at = _start_at or time.time()
+    recur_deadline = scheduled_at + _run_for if _run_for is not None else None
 
     opts = getattr(wrapper, "__offwork_options__", {})
     task = Task(
@@ -315,6 +318,8 @@ async def submit_recurring(
         throttle=opts.get("throttle"),
         scheduled_at=scheduled_at,
         recur_interval=_interval,
+        recur_deadline=recur_deadline,
+        recur_remaining=_max_runs,
         schedule_id=schedule_id,
     )
 
@@ -634,7 +639,18 @@ async def _handle_task(
 
     # Re-enqueue recurring task — worker re-signs with its own identity.
     if task.recur_interval is not None and task.schedule_id is not None:
-        if not await backend.is_schedule_cancelled(task.schedule_id):
+        next_at = time.time() + task.recur_interval
+        remaining = task.recur_remaining
+        deadline_exceeded = task.recur_deadline is not None and next_at > task.recur_deadline
+        runs_exhausted = remaining is not None and remaining <= 1
+        if deadline_exceeded or runs_exhausted:
+            await backend.cancel_schedule(task.schedule_id)
+            logger.info(
+                "Recurring schedule %s exhausted (%s)",
+                task.schedule_id,
+                "deadline" if deadline_exceeded else "max_runs",
+            )
+        elif not await backend.is_schedule_cancelled(task.schedule_id):
             next_task = Task(
                 graph_json=task.graph_json,
                 function_name=task.function_name,
@@ -644,8 +660,10 @@ async def _handle_task(
                 retries=task.retries,
                 retry_delay=task.retry_delay,
                 throttle=task.throttle,
-                scheduled_at=time.time() + task.recur_interval,
+                scheduled_at=next_at,
                 recur_interval=task.recur_interval,
+                recur_deadline=task.recur_deadline,
+                recur_remaining=remaining - 1 if remaining is not None else None,
                 schedule_id=task.schedule_id,
             )
             await backend.submit(_encode_task(next_task, root_token))

offwork/worker/result.py

@@ -214,22 +214,52 @@ def __await__(self) -> Generator[Any, None, Any]:
 
     # -- cancellation ----------------------------------------------------------
 
-    async def cancel(self) -> None:
+    async def cancel(self, wait: bool | float = False) -> bool:
         """Cancel the task.
 
-        Marks the task as cancelled in the backend.  If the worker
-        hasn't started execution yet, it will skip the task.  If
-        execution is already in progress, it will continue but the
-        client will receive a :class:`TaskCancelled` error.
+        Marks the task as cancelled in the backend.  The worker
+        observes the flag via its heartbeat loop and aborts execution
+        cooperatively.
 
-        Awaiting the result after cancellation raises
-        :class:`TaskCancelled`.
+        Parameters
+        ----------
+        wait
+            If ``False`` (default), return immediately after signalling
+            cancellation. If ``True``, block until the worker confirms
+            (default 30s timeout). If a number, wait that many seconds
+            for confirmation.
+
+        Returns
+        -------
+        bool
+            ``True`` if cancellation was confirmed by the worker (or
+            ``wait=False``).  ``False`` if the wait timed out.
         """
         await self._backend.cancel_task(self._task_id)
-        await self._backend.send_result(
-            self._task_id,
-            ResultEnvelope.cancelled(self._task_id).to_json(),
-        )
+        if wait is False:
+            # Pre-seed a cancelled envelope so a client that never awaits
+            # confirmation still gets TaskCancelled when it reads.
+            await self._backend.send_result(
+                self._task_id,
+                ResultEnvelope.cancelled(self._task_id).to_json(),
+            )
+            return True
+        timeout = 30.0 if wait is True else float(wait)
+        deadline = time.monotonic() + timeout
+        while True:
+            raw = await self._backend.try_get_result(self._task_id)
+            if raw is not None:
+                self._envelope = ResultEnvelope.from_json(raw)
+                return True
+            if time.monotonic() >= deadline:
+                # Fall back: seed a cancelled envelope so subsequent reads
+                # don't hang forever, and tell the caller we timed out.
+                await self._backend.send_result(
+                    self._task_id,
+                    ResultEnvelope.cancelled(self._task_id).to_json(),
+                )
+                return False
+            await asyncio.sleep(0.5)
 
     # -- progress --------------------------------------------------------------
 

offwork/worker/schedule.py

@@ -1,5 +1,8 @@
 """Schedule handle for recurring task management."""
 
+import asyncio
+import time
+
 from offwork.worker.backends.base import Backend
 
 
@@ -14,13 +17,37 @@ def __init__(self, schedule_id: str, backend: Backend) -> None:
     def schedule_id(self) -> str:
         return self._schedule_id
 
-    async def cancel(self) -> None:
+    async def cancel(self, wait: bool | float = False) -> bool:
         """Cancel this recurring schedule.
 
-        The worker will stop re-enqueuing new occurrences after the
+        The worker stops re-enqueuing new occurrences after the
         current one completes.
+
+        Parameters
+        ----------
+        wait
+            If ``False`` (default), return immediately. If ``True``,
+            block until the backend confirms the schedule is marked
+            cancelled (default 30s timeout). If a number, wait that
+            many seconds.
+
+        Returns
+        -------
+        bool
+            ``True`` if cancellation was acknowledged (or
+            ``wait=False``). ``False`` on timeout.
         """
         await self._backend.cancel_schedule(self._schedule_id)
+        if wait is False:
+            return True
+        timeout = 30.0 if wait is True else float(wait)
+        deadline = time.monotonic() + timeout
+        while True:
+            if await self._backend.is_schedule_cancelled(self._schedule_id):
+                return True
+            if time.monotonic() >= deadline:
+                return False
+            await asyncio.sleep(0.5)
 
     def __repr__(self) -> str:
         return f"ScheduleHandle(schedule_id={self._schedule_id!r})"