More examples

Author: codeSamuraii

examples/async_db_client.py

@@ -0,0 +1,74 @@
+"""Async DB clients (Motor / asyncpg-style) on a pyfuse worker.
+
+Async clients are bound to the event loop they were created on.  Running
+them on a worker is fine *as long as* you create them inside the traced
+async function -- pyfuse's worker awaits the coroutine on its own loop,
+so a freshly-created client picks up the right loop automatically.
+
+This example uses Motor (async MongoDB driver).  The pattern is identical
+for ``asyncpg``, ``aioredis``, ``aiomysql``, etc.
+
+Usage:
+    docker run -d -p 27017:27017 --name pyfuse-mongo mongo:7
+    pyfuse worker --backend redis://localhost:6379 --tmp
+    python -m pyfuse run --tmp examples/async_db_client.py
+"""
+
+import asyncio
+from typing import Any
+
+import motor.motor_asyncio as motor
+
+import pyfuse
+from pyfuse import trace
+
+pyfuse.connect("redis://localhost:6379")
+
+MONGO_URL = "mongodb://localhost:27017"
+DB_NAME = "pyfuse_demo_async"
+
+
+@trace
+async def upsert_user(user_id: int, name: str, score: float) -> str:
+    """Async upsert -- exercises the worker's event loop end-to-end."""
+    client = motor.AsyncIOMotorClient(MONGO_URL, serverSelectionTimeoutMS=3000)
+    try:
+        db = client[DB_NAME]
+        await db.users.update_one(
+            {"_id": user_id},
+            {"$set": {"name": name, "score": score}},
+            upsert=True,
+        )
+        return f"upserted user {user_id}"
+    finally:
+        client.close()
+
+
+@trace
+async def top_users(limit: int = 5) -> list[dict[str, Any]]:
+    """Concurrent reads via asyncio.gather inside the worker loop."""
+    client = motor.AsyncIOMotorClient(MONGO_URL, serverSelectionTimeoutMS=3000)
+    try:
+        db = client[DB_NAME]
+        cursor = db.users.find({}, {"_id": 1, "name": 1, "score": 1})
+        cursor = cursor.sort("score", -1).limit(limit)
+        return [doc async for doc in cursor]
+    finally:
+        client.close()
+
+
+async def main() -> None:
+    # Fan-out a batch of upserts -- each runs as its own remote task,
+    # each with its own client / loop on the worker side.
+    users = [(i, f"user-{i}", float(i * 7 % 31)) for i in range(20)]
+    acks = await upsert_user.map(users)
+    print(f"{len(acks)} upserts acknowledged")
+
+    leaderboard = await top_users.run(5)
+    print("Leaderboard:")
+    for row in leaderboard:
+        print(f"  {row['name']:<10} score={row['score']}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/fastapi_offload.py

@@ -0,0 +1,110 @@
+"""Use pyfuse to offload heavy work from a FastAPI request handler.
+
+A common production pattern: the web app stays light and responsive, while
+CPU-bound or slow I/O work is shipped to a worker pool over Redis.  The
+request handler awaits the result and returns it.
+
+Endpoints:
+    GET  /                  -- healthcheck
+    POST /jobs/sync         -- await the worker, return result
+    POST /jobs/async        -- start the task, return task_id
+    GET  /jobs/{task_id}    -- poll status / result
+
+Usage:
+    pyfuse worker --backend redis://localhost:6379 --tmp
+    uvicorn examples.fastapi_offload:app --reload
+    curl -X POST localhost:8000/jobs/sync -d '{"text":"hello WORLD"}' \
+         -H 'content-type: application/json'
+"""
+
+import asyncio
+import hashlib
+from typing import Any
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+import pyfuse
+from pyfuse import trace, Result
+
+
+# --- traced workloads ------------------------------------------------------
+
+@trace
+def normalize(text: str) -> dict[str, Any]:
+    """Cheap demo workload -- runs on the worker."""
+    cleaned = " ".join(text.split()).lower()
+    return {
+        "length": len(cleaned),
+        "sha1": hashlib.sha1(cleaned.encode()).hexdigest(),
+        "preview": cleaned[:80],
+    }
+
+
+@trace(timeout=30, retries=2)
+def cpu_bound(n: int) -> int:
+    """Pretend-CPU work to demonstrate timeout / retry semantics."""
+    total = 0
+    for i in range(n):
+        total += (i * i) % 97
+    return total
+
+
+# --- FastAPI app -----------------------------------------------------------
+
+class JobRequest(BaseModel):
+    text: str = ""
+    n: int = 1_000_000
+
+
+app = FastAPI(title="pyfuse offload demo")
+_pending: dict[str, Result[Any]] = {}
+
+
+@app.on_event("startup")
+async def _startup() -> None:
+    # Connecting at startup avoids a per-request handshake.
+    pyfuse.connect("redis://localhost:6379")
+
+
+@app.on_event("shutdown")
+async def _shutdown() -> None:
+    await pyfuse.disconnect()
+
+
+@app.get("/")
+async def root() -> dict[str, str]:
+    return {"status": "ok"}
+
+
+@app.post("/jobs/sync")
+async def submit_sync(req: JobRequest) -> dict[str, Any]:
+    """Block the request until the worker returns."""
+    try:
+        return await asyncio.wait_for(normalize.run(req.text), timeout=10)
+    except asyncio.TimeoutError as exc:
+        raise HTTPException(status_code=504, detail="worker timeout") from exc
+
+
+@app.post("/jobs/async")
+async def submit_async(req: JobRequest) -> dict[str, str]:
+    """Fire-and-forget; return the task_id for polling."""
+    handle = await cpu_bound.start(req.n)
+    _pending[handle.task_id] = handle
+    return {"task_id": handle.task_id}
+
+
+@app.get("/jobs/{task_id}")
+async def poll(task_id: str) -> dict[str, Any]:
+    handle = _pending.get(task_id)
+    if handle is None:
+        raise HTTPException(status_code=404, detail="unknown task")
+    if not await handle.done():
+        return {"task_id": task_id, "status": "pending"}
+    try:
+        result = await handle
+    except Exception as exc:
+        return {"task_id": task_id, "status": "error", "detail": str(exc)}
+    finally:
+        _pending.pop(task_id, None)
+    return {"task_id": task_id, "status": "ok", "result": result}

examples/httpx_concurrent_scrape.py

@@ -0,0 +1,91 @@
+"""Async HTTP fan-out from a remote worker.
+
+Each traced call opens its own ``httpx.AsyncClient`` (HTTP/2, connection
+pool, async DNS) and probes a list of URLs concurrently.  This is a
+faithful test of the worker's event-loop integration: dozens of TCP
+sockets, TLS handshakes, and async iterators per task.
+
+The client is created *inside* the async function so it binds to the
+worker's running loop.  The ``async with`` block guarantees the pool is
+torn down before the task returns.
+
+Usage:
+    pyfuse worker --backend redis://localhost:6379 --tmp
+    python -m pyfuse run --tmp examples/httpx_concurrent_scrape.py
+"""
+
+import asyncio
+from typing import Any
+
+import httpx
+
+import pyfuse
+from pyfuse import trace, progress
+
+pyfuse.connect("redis://localhost:6379")
+
+
+URLS = [
+    "https://example.com",
+    "https://www.python.org",
+    "https://httpbin.org/get",
+    "https://httpbin.org/status/200",
+    "https://httpbin.org/status/404",
+    "https://httpbin.org/delay/1",
+    "https://httpbin.org/headers",
+    "https://httpbin.org/uuid",
+]
+
+
+@trace(timeout=30, retries=1)
+async def probe_many(urls: list[str], concurrency: int = 4) -> list[dict[str, Any]]:
+    """Hit every URL once, capture status / size / latency."""
+    sem = asyncio.Semaphore(concurrency)
+    results: list[dict[str, Any]] = []
+
+    async with httpx.AsyncClient(http2=True, timeout=10.0) as client:
+        async def probe(url: str) -> dict[str, Any]:
+            loop = asyncio.get_running_loop()
+            start = loop.time()
+            try:
+                async with sem:
+                    r = await client.get(url)
+                return {
+                    "url": url,
+                    "status": r.status_code,
+                    "bytes": len(r.content),
+                    "elapsed_ms": round((loop.time() - start) * 1000, 1),
+                }
+            except Exception as exc:
+                return {
+                    "url": url,
+                    "status": None,
+                    "error": type(exc).__name__,
+                    "elapsed_ms": round((loop.time() - start) * 1000, 1),
+                }
+
+        tasks = [asyncio.create_task(probe(u)) for u in urls]
+        for i, coro in enumerate(asyncio.as_completed(tasks), 1):
+            results.append(await coro)
+            progress(i, len(urls), message=f"{i}/{len(urls)} probed")
+
+    return results
+
+
+async def main() -> None:
+    handle = await probe_many.start(URLS, concurrency=4)
+    while not await handle.done():
+        p = await handle.progress()
+        if p is not None:
+            print(p)
+        await asyncio.sleep(0.3)
+
+    rows = await handle
+    print(f"\nProbed {len(rows)} URLs:")
+    for row in sorted(rows, key=lambda r: r["elapsed_ms"]):
+        status = row.get("status") or row.get("error")
+        print(f"  {row['elapsed_ms']:>7.1f}ms  {str(status):<5}  {row['url']}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/mongodb_client.py

@@ -0,0 +1,114 @@
+"""Real-world MongoDB usage with pyfuse.
+
+DB clients are stateful: connection pools, sockets, background heartbeat
+threads.  Two rules to keep in mind when running them on a remote worker:
+
+1. Open the client *inside* the traced function, not at module level.
+   pyfuse caches the reconstructed namespace per subgraph, so a module-level
+   client would survive across calls -- with stale sockets, wrong loop, etc.
+
+2. Close the client before returning.  Otherwise the worker leaks a pool
+   and a heartbeat thread per task.
+
+The example runs three independent tasks against a real MongoDB and uses
+``pymongo``'s aggregation framework so the round-trip is non-trivial.
+
+Usage:
+    # Start a Mongo + worker:
+    docker run -d -p 27017:27017 --name pyfuse-mongo mongo:7
+    pyfuse worker --backend redis://localhost:6379 --tmp
+
+    # Run the script (will pip-install pymongo on the worker):
+    python -m pyfuse run --tmp examples/mongodb_client.py
+"""
+
+import asyncio
+import random
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+import pymongo
+
+import pyfuse
+from pyfuse import trace
+
+pyfuse.connect("redis://localhost:6379")
+
+MONGO_URL = "mongodb://localhost:27017"
+DB_NAME = "pyfuse_demo"
+
+
+@contextmanager
+def mongo_db(url: str, db_name: str) -> Iterator[Any]:
+    """Open a client, yield the database, always close on exit."""
+    client = pymongo.MongoClient(url, serverSelectionTimeoutMS=3000)
+    try:
+        yield client[db_name]
+    finally:
+        client.close()
+
+
+@trace
+def seed_events(n: int, seed: int = 0) -> int:
+    """Insert *n* synthetic events, return the inserted count."""
+    rng = random.Random(seed)
+    levels = ["info", "warn", "error"]
+    docs = [
+        {
+            "level": rng.choice(levels),
+            "service": f"svc-{rng.randint(0, 4)}",
+            "latency_ms": round(rng.gauss(120, 40), 2),
+        }
+        for _ in range(n)
+    ]
+    with mongo_db(MONGO_URL, DB_NAME) as db:
+        db.events.drop()
+        result = db.events.insert_many(docs)
+        return len(result.inserted_ids)
+
+
+@trace
+def aggregate_by_service() -> list[dict[str, Any]]:
+    """Run an aggregation pipeline and return the result."""
+    pipeline = [
+        {"$match": {"level": {"$in": ["warn", "error"]}}},
+        {
+            "$group": {
+                "_id": "$service",
+                "count": {"$sum": 1},
+                "p95_latency": {"$max": "$latency_ms"},
+            }
+        },
+        {"$sort": {"count": -1}},
+    ]
+    with mongo_db(MONGO_URL, DB_NAME) as db:
+        return [
+            {"service": r["_id"], "count": r["count"], "p95": r["p95_latency"]}
+            for r in db.events.aggregate(pipeline)
+        ]
+
+
+@trace
+def purge_events() -> int:
+    """Drop the collection.  Returns the previous document count."""
+    with mongo_db(MONGO_URL, DB_NAME) as db:
+        before = db.events.count_documents({})
+        db.events.drop()
+        return before
+
+
+async def main() -> None:
+    inserted = await seed_events.run(2000, seed=42)
+    print(f"Inserted {inserted} events")
+
+    summary = await aggregate_by_service.run()
+    print("Top noisy services:")
+    for row in summary:
+        print(f"  {row['service']:<8} count={row['count']:<4} p95={row['p95']}ms")
+
+    removed = await purge_events.run()
+    print(f"Cleaned up {removed} events")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())