How do I trace an agent built with Pydantic AI and deployed to a Celery worker?

Rius traces a Pydantic AI agent in a Celery worker when each worker process calls rius.init() and Agent.instrument_all(True) after it forks, and shuts the client down before it exits. Pydantic AI emits its own spans for agent runs, model requests and tool calls, and Rius adds tokens and cost. Name the agent with service_name or RIUS_AGENT_NAME, and Pydantic AI names the spans after your agent, model and tools.

How it works

The Pydantic AI integration needs no extra and no wrappers. rius.init() installs the global OpenTelemetry tracer provider, and Agent.instrument_all(True) tells Pydantic AI to emit its spans through it. Because the order matters, call init first, then instrument_all, then run your agents. With no observe wrapper around the run, invoke_agent is the root span. Wrap the task with observe and it becomes a child of your span instead.

In the trace you get:

• A span named invoke_agent <agent name> for each agent run.
• A chat <model> generation for each model request, with tokens and cost.
• An execute_tool <tool name> span for each tool call.

Spans are queued in memory and exported in batches from a background thread, so tracing does not sit on your agent’s hot path. Rius computes cost server-side from the model and token counts, so you never send a cost value. The Integrations page lists every framework Rius covers.

Celery’s prefork pool forks its worker processes and ends them through os._exit, which skips Python’s exit hook. So start Rius in each child from the worker_process_init signal, and flush and shut it down from worker_process_shutdown. The reliability page covers the export pipeline.

Set it up

1. Install the SDK and Pydantic AI. No extra is needed.

   pip install glassflow-rius pydantic-ai

2. In each worker process, call rius.init() and then Agent.instrument_all(True), and give the service a name. Keep both out of module scope, because the parent process imports your tasks module too. Scope each run with rius.session() if one conversation spans several runs.

   import rius
   from celery import Celery
   from celery.signals import worker_process_init, worker_process_shutdown
   from pydantic_ai import Agent
    
   app = Celery("agents", broker="redis://localhost:6379/0")
   client = None
    
   @worker_process_init.connect
   def start_rius(**kwargs):
       global client
       client = rius.init(api_key="ri_...", service_name="my-agent")
       Agent.instrument_all(True)
    
   @worker_process_shutdown.connect
   def stop_rius(**kwargs):
       if client is not None:
           client.flush(timeout_millis=10_000)  # force an export
           client.shutdown()  # drain, stop the thread, send the stopped heartbeat
    
   support_agent = Agent(
       "openai:gpt-5-mini",
       instructions="Answer support questions; check accounts with the tool.",
   )
    
   @support_agent.tool_plain
   def account_status(email: str) -> str:
       """Fetch a customer's account standing by email."""
       return lookup(email)
    
   @app.task
   def handle_ticket(conversation_id: str, ticket: str) -> str:
       with rius.session(conversation_id):
           result = support_agent.run_sync(ticket)
       return result.output

   The threads, gevent and eventlet pools do not send worker_process_init, so with those pools call both once when the worker starts.

3. Name the agent. service_name sets the service.name of the agent. The name the agent is grouped under defaults to the service name, and you can set it separately with RIUS_AGENT_NAME:

   export RIUS_AGENT_NAME=support-agent

   Use a stable name. The Agents page finds agents from their traces and heartbeats, so there is nothing to register.

4. Flush at the process boundary, as the shutdown handler does, never inside a task, because flush() blocks until the export completes. It returns False on timeout or when the most recent export was rejected, such as a 401 on a bad API key. A worker process killed with SIGKILL, by a hard time limit or the out-of-memory killer, loses the spans still queued in the batch.

5. Verify. Run one request that calls a tool, then open the trace in the console. You should see the invoke_agent root, chat generations with model, tokens and cost, and execute_tool spans.

What you get

• A trace per agent run, with the model requests and tool calls as child spans named after your agent, model and tool.
• Tokens and cost on every model request, plus time to first token and reasoning tokens in the Pydantic AI runs.
• Conversations grouped by session.id when you wrap runs in rius.session().
• An agent on the Agents page under the name you chose, with heartbeats on by default.
• A clean stop signal. shutdown() sends the final stopped ping. A SIGKILL sends none, so the instance reports gone, which is how Rius shows a crash.
• One instance per worker process, each with its own heartbeat instance ID.

Pydantic AI controls what its spans carry, such as prompts, completions and tool arguments, through its InstrumentationSettings, on top of the SDK’s own privacy controls.

When something else fits

• If you run in a Lambda-style runtime that freezes the process between invocations, set heartbeat=False (or RIUS_HEARTBEAT=false) and call client.flush() at the end of every invocation.
• If you do not want the Rius SDK in the process, point any OpenTelemetry setup at the ingest endpoint instead, as the Pydantic AI integration page describes.

Related