How do I trace a Python AI agent running in a Celery worker?

Rius traces a Python agent in a Celery worker by calling rius.init once per worker process and wrapping your agent function, with spans exported in the background. Each worker process shows up as its own instance of one agent. On a clean shutdown the SDK drains pending spans, and you add an explicit flush only where a process is killed hard or frozen.

How it works

• One init per process. Call rius.init() once, as early as possible, in the code every worker process loads. The OpenTelemetry global provider is write-once, so a second global init() logs a warning and returns the existing client unchanged. To reconfigure, call client.shutdown() first.
• Export runs off the hot path. Spans are queued in memory and sent in batches from a background thread, so a slow network never stalls a task. Failures are logged and never raise into your code.
• Forked workers are handled. When worker processes are forked, as with gunicorn or multiprocessing, the exporter re-creates its background thread in each child. The heartbeat re-arms with a new instance ID, so you see one instance per worker.
• Integrations turn on by themselves. With a global init(), every installed integration is enabled, so OpenAI, Anthropic, LangChain and the other libraries you use record model calls with tokens and cost. See Integrations for the full list. A scoped client (set_global=False) does not auto-instrument unless you pass instruments explicitly.
• Heartbeat. The SDK pings the platform for the life of the process, which powers the Agents view. It needs an API key or an Authorization header.

Set it up

1. Install the SDK and the extra for the library your agent calls. Python 3.10 or newer is required. The installation page covers every option.

   pip install "glassflow-rius[openai]"

2. Initialize once, in a module each worker process imports before it runs tasks. Set service_name explicitly and keep it the same across replicas and deploys. Put release numbers in service_version, not in the name.

   import rius
    
   client = rius.init(
       api_key="ri_...",
       service_name="my-agent",
   )

3. Wrap the function that runs your agent with rius.observe, and call it from your task. Your model calls inside it are recorded by the integration, and the wrapper gives them one parent span.

   @rius.observe
   def run_agent(query: str) -> str:
       ...  # your agent code
    
   @app.task
   def answer(query: str) -> str:
       return run_agent(query)

4. Decide how your processes end.

   • Long-running worker that stops cleanly: Celery’s default prefork pool ends its child processes without running exit hooks, so the automatic drain at interpreter exit never runs there. Connect worker_process_shutdown to client.shutdown():

     from celery.signals import worker_process_shutdown
      
     @worker_process_shutdown.connect
     def _rius_shutdown(**kwargs):
         client.shutdown()

   • Worker that may be killed hard: only an unclean shutdown can lose the last batch. Lowering OTEL_BSP_SCHEDULE_DELAY (default 5000 ms) shortens the window.

   • Short-lived process, or a runtime that freezes the process between jobs: a frozen background thread exports nothing and the exit hook may never run, so flush at the end of each job or invocation.

   client.flush(timeout_millis=30_000)   # end of each job

   Call client.shutdown() once, when the worker process stops, never per job: it stops the export thread and releases init(), so later tasks in that process go untraced.

   Flush at process or job boundaries only. flush() blocks until the export completes, so never call it in a hot path. It returns False on timeout or when the latest export was rejected, such as a bad API key.

5. Check the settings that matter in a worker:

   • heartbeat (RIUS_HEARTBEAT): leave it on for long-lived workers. Set it to false if the runtime freezes the process between jobs, because pings stop and the instance flaps between alive and stale or gone.
   • heartbeat_interval (RIUS_HEARTBEAT_INTERVAL): keep it at 15 seconds or below, or a live instance flaps to stale.
   • OTEL_BSP_MAX_QUEUE_SIZE (default 2048): raise it for a very chatty agent before you lower the sample rate. When the queue is full, new spans are dropped with a warning.
   • session_id (RIUS_SESSION_ID): a process-wide session id. Do not set it in a worker that handles many users.

   The full list of arguments is in Reliability and the installation page.

What you get

• A trace for each agent run, with a span for each model call, tokens and a list-price cost estimate in US dollars.
• One agent under a stable name, with each worker process listed as an instance when heartbeats are on and credentials are set.
• Spans exported without blocking your tasks, and a clear rule for when to flush so the last batch is not lost.
• Failures in tracing that are logged and never break a task.

When something else fits

• If your runtime freezes the process between jobs, turn the heartbeat off and flush at the end of each job, as described in Reliability.
• If another OpenTelemetry SDK already owns the global tracer provider in your worker, set RIUS_BRIDGE_FOREIGN_PROVIDER so its spans reach Rius too, or use a scoped client for a separate pipeline.

Related