[RFC] Should Services know about Batching?

[Jack-Khuu]

Code Pointer: service.py

The idea behind batching is simple: you aggregate similar requests into a single payload for more efficient processing.
What I want to get discussion on is: should Services and Replicas handle batching or should they just blindly pass payloads along?

IMO there are currently 2 potential patterns (that can/should co-exist) for supporting batching without upstreaming to Services/Replicas

• Caller Batching: Agents expose an endpoint that explicitly expects a batched input
  • Caller + Agent operate on batched payloads
  • Con: Caller has to deal with batching
• Actor Batching: Agents expose an endpoint to enqueue requests and Agents writes separate logic to batch process internally
  • Caller operates on singletons
  • Agents handles aggregating and operates on batched payloads
  • Con: Agents are a tad more complex and will share similar boilerplate

---

Where Services/Replicas can potentially comes in is with replica routing and prebatching in the Actor Batching cases.

1. Callers call singleton endpoints
2. Services route batched requests to specific Replicas
   • Aware Routing: Services route payloads (w/ similarity criterion) and to specific Replicas
3. Replicas aggregate singleton requests
   • Prebatching: Replicas create an aggregated payload and passes it along to the Agent's Batched endpoint
4. Agents process aggregated payload

~ Note that the singleton endpoint called by the user, never reaches the Actor ~

In combination, this gives Forge the ability to generalize customizing the routing and batching process.
Beyond that, this has an added side effect of simplifying the Actor flow (execution and creating new ones)

[Current] Without Services/Replica handling batching

def NewActor(ForgeActor):
  ...
  @endpoint 
  async def add_request(self):
       # Enqueue

   async def run(self):
        # Process from internal Queue

   async def process_batched_request(self):
        # Execution

With Services/Replica handling batching

def NewActor(ForgeActor):
  ...
  @endpoint 
  async def process_batched_request(self):
       # Execution

  @classmethod
  def enbatch(cls, request, cummulative_batch):
       # Add request to batch struct

---

Welcome all feedback, but here are a few seeding questions

• If we plan to upstream Service/Replica, does this make them too RL specific?
• Does custom routing match the identity of Forge ("Users don't need to think about the network magic")?

[allenwang28]

interesting idea, this seems useful with some precedence, i.e. see: https://docs.ray.io/en/latest/serve/advanced-guides/dyn-req-batch.html

I think this could be added without too much overhead, but wondering how you're planning on using this in Forge?

[Jack-Khuu]

Oh didn't know that Ray did this

I was just thinking about adding batching to ReferenceModel and policy and this was a scope creep/generalization rabbit hole

[Jack-Khuu]

@allenwang28 Thoughts on

If we plan to upstream Service/Replica, does this make them too RL specific?

[allenwang28]

If we plan to upstream Service/Replica, does this make them too RL specific?

hmm, I don't think so tbh. Case in point, the Ray Serve example was created for serving, but we're finding that it has use cases in RL too

I think we're also too early for upstreaming into Monarch. When we get there, we will ideally have a good idea of all the ways we want to use services to guide the upstream implementation

[pbontrager]

I don’t think we can do batching in a generic way. Case in point, vllm’s batching system is going to be very different than a fixed batch system or one where the bs is data/memory dependent.

I love this RFC though because it is one of the bigger problems I noticed when we were working on the GRPO app. In my mind we need a couple of common patterns that people can use but I don’t think they can happen in the service runner itself.

[allenwang28]

I don’t think we can do batching in a generic way. Case in point, vllm’s batching system is going to be very different than a fixed batch system or one where the bs is data/memory dependent.

hmm I'm not sure that I would entangle these two statements - vLLM does break a lot of assumptions, but what other service do we expect to have custom batching requirements? E.g., Ray Serve's batching will not serve the needs of vLLM, but it is still useful for other use cases. A generic batcher with a timeout and batch size is not too difficult to implement in the current Service implementation IMO