async def on_complete(self, router):
tracer.add(
self.request_id, "router_recv_result", router_name=self.router_name
)
if self.batch_id is None:
await router.dequeue_request(self.backend_name, self.backend_worker)
elif self.idx_in_batch == 0:
await router.dequeue_request(self.backend_name, self.backend_worker)
def on_enqueue(self, endpoint, metadata=None):
tracer.add(
self.request_id,
"router_enqueue",
router_name=metadata.get("router_name", None)
if metadata is not None
else None,
)
# tracer.add_metadata(self.request_id, endpoint=endpoint)
if metadata:
assert isinstance(metadata, dict)
tracer.add_metadata(self.request_id, **metadata)
def on_assigned(self, backend_name, worker, batch_id, idx_in_batch):
self.backend_name = backend_name
self.backend_worker = worker
self.batch_id = batch_id
self.idx_in_batch = idx_in_batch
tracer.add(
self.request_id, "router_dequeue", router_name=self.router_name
)
tracer.add_metadata(
self.request_id,
backend=backend_name,
batch_id=batch_id,
router_name=self.router_name,
)
def on_worker_done(self):
tracer.add(self.request_id, "worker_done", router_name=self.router_name)
def on_worker_start(self):
tracer.add(
self.request_id, "worker_start", router_name=self.router_name
)