@service(
    dynamo={
        "enabled": True,
        "namespace": "dynamo",
    },
    resources={"gpu": 1, "cpu": "10", "memory": "20Gi"},
    workers=1,
)
class VllmWorker:
    """vLLM Worker 服务"""

    prefill_worker = depends(PrefillWorker)  # 可选依赖

@async_on_start
async def async_init(self):
    """服务启动时的异步初始化"""

    # 1. 构建 vLLM 引擎客户端
    self._engine_context = build_async_engine_client_from_engine_args(
        self.engine_args
    )
    self.engine_client = await self._engine_context.__aenter__()

    # 2. 配置 KV 感知路由
    if self.engine_args.router == "kv":
        self.engine_client.set_metrics_publisher(self.metrics_publisher)
        self.metrics_publisher.publish(
            request_active_slots=0,
            kv_active_blocks=0,
            # ...
        )

    # 3. 注册 NIXL 元数据
    if self.engine_args.remote_prefill:
        metadata = self.engine_client.nixl_metadata
        metadata_store = NixlMetadataStore("dynamo", runtime)
        await metadata_store.put(metadata.engine_id, metadata)

    # 4. 配置分离式路由器
    if self.engine_args.conditional_disagg:
        self.disaggregated_router = PyDisaggregatedRouter(
            runtime,
            self.model_name,
            max_local_prefill_length=self.engine_args.max_local_prefill_length,
        )

@dynamo_endpoint()
async def generate(self, request: vLLMGenerateRequest):
    """生成端点 - 处理推理请求"""

@dynamo_endpoint()
async def generate(self, request: vLLMGenerateRequest):
    # 1. 决定是否远程 prefill
    if self.disaggregated_router is not None:
        disagg_router_decision = self.disaggregated_router.prefill_remote(
            len(request.engine_prompt["prompt_token_ids"]),
            request.prefix_hit_rate
        )
    else:
        disagg_router_decision = True

    # 2. 配置远程 prefill 参数
    if self.do_remote_prefill and disagg_router_decision:
        remote_prefill_params = RemotePrefillParams(
            is_remote_prefill=True,
            remote_prefill_request_callback=self.get_remote_prefill_request_callback(),
        )
    else:
        remote_prefill_params = None

    # 3. 调用引擎生成
    async for response in self.engine_client.generate(
        prompt=request.engine_prompt,
        sampling_params=request.sampling_params,
        request_id=request.request_id,
        remote_prefill_params=remote_prefill_params,
    ):
        yield MyRequestOutput(
            request_id=response.request_id,
            outputs=response.outputs,
            finished=response.finished,
        ).model_dump_json()

class KvMetricsPublisher:
    def publish(
        self,
        request_active_slots: int,
        request_total_slots: int,
        kv_active_blocks: int,
        kv_total_blocks: int,
        num_requests_waiting: int,
        gpu_cache_usage_perc: float,
        gpu_prefix_cache_hit_rate: float,
    ):
        # 发布指标到 NATS

VllmWorker:
  model: deepseek-ai/DeepSeek-R1-Distill-Llama-8B
  enable-prefix-caching: true
  tensor-parallel-size: 1
  router: kv
  remote-prefill: true
  conditional-disagg: true
  max-local-prefill-length: 512
  ServiceArgs:
    workers: 4
    resources:
      gpu: 1
      cpu: 10
      memory: 20Gi