Repository: volcengine/verl
File: verl/workers/rollout/sglang_rollout/async_sglang_server.py (Python)

Code before:
            )
        else:
            token_ids = output["output_ids"]
            log_probs = None

        routed_experts = None
        if self.config.enable_rollout_routing_replay:
            if self.config.skip_tokenizer_init:
                routed_experts = output.get("meta_info", {}).get("routed_experts", None)
            else:
                from sglang.srt.layers.moe.routed_experts_capturer import extract_routed_experts_from_meta_info

                hf_config = self.model_config.hf_config
                if not hasattr(hf_config, "num_hidden_layers") or not hasattr(hf_config, "num_experts_per_tok"):
                    raise AttributeError(
                        "enable_rollout_routing_replay is set, but hf_config is missing "
                        "'num_hidden_layers' or 'num_experts_per_tok'. This feature requires an MoE model "
                        "configuration that defines these attributes."
                    )
                routed_experts = extract_routed_experts_from_meta_info(output).reshape(
                    -1, hf_config.num_hidden_layers, hf_config.num_experts_per_tok
                )

        return TokenOutput(token_ids=token_ids, log_probs=log_probs, routed_experts=routed_experts)

    def _profile_args(self, **kwargs) -> dict[str, Any]:
        global_step = kwargs.pop("global_step", 0)
        config = self.profiler_controller.tool_config
        contents = self.profiler_controller.tool_config.contents
        save_path = os.path.join(
            self.config.profiler.save_path, f"rollout_step_{global_step}", f"agent_loop_replica_{self.replica_rank}"
        )

        profiler_tool = self.profiler_controller.config.tool
        activities: Optional[list[str]] = None
        if contents and profiler_tool:
            activities_tmp = []
            check_map = {
                "cpu": ("CPU", "torch"),
                "cuda|gpu": ("GPU", "torch"),
                "MEM": ("MEM", "torch_memory"),
            }
            for key, (act, tool) in check_map.items():
                if any(k in contents for k in key.split("|")):
                    activities_tmp.append(act)
                    if profiler_tool != tool:
                        raise ValueError(f"{act} profiling requires '{tool}' (got '{profiler_tool}')")
            for unsupported in ("CUDA_PROFILER", "RPD"):
                if unsupported in contents:
                    raise NotImplementedError(f"{unsupported} profiling is not supported")
            activities = activities_tmp if len(activities_tmp) > 0 else activities

Review comment (suggestion):
**Keep Profiler Logic Cohesive**: The profiler logic seems to be scattered across the main workflow. It would be better to encapsulate it to keep the main execution path clean and cohesive. Avoid adding too much profiling-specific code directly into the core rollout logic

Code after:

    async def wake_up(self):
        if self.rollout_mode == RolloutMode.HYBRID:
            # Call all workers to switch between trainer mode and rollout mode.
            await asyncio.gather(*[worker.wake_up.remote() for worker in self.workers])
        elif self.rollout_mode == RolloutMode.COLOCATED:
            # Directly call engine to wake up without sync weights.
            obj = ResumeMemoryOccupationReqInput(tags=["kv_cache", "weights"])
            await self.tokenizer_manager.resume_memory_occupation(obj, None)
            await self.tokenizer_manager.flush_cache()
        elif self.rollout_mode == RolloutMode.STANDALONE:
            logger.info("skip wake_up in standalone mode")

    async def sleep(self):
        if self.rollout_mode == RolloutMode.HYBRID:
            await asyncio.gather(*[worker.sleep.remote() for worker in self.workers])
        elif self.rollout_mode == RolloutMode.COLOCATED:
            obj = ReleaseMemoryOccupationReqInput(tags=["kv_cache", "weights"])
            await self.tokenizer_manager.release_memory_occupation(obj, None)
        elif self.rollout_mode == RolloutMode.STANDALONE:
            logger.info("skip sleep in standalone mode")

    async def clear_kv_cache(self):
        obj = ReleaseMemoryOccupationReqInput(tags=["kv_cache"])
        await self.tokenizer_manager.release_memory_occupation(obj, None)

    async def generate(
        self,
        prompt_ids: torch.Tensor,
        sampling_params: dict[str, Any],
        request_id: str,
        image_data: Optional[list[Any]] = None,
        video_data: Optional[list[Any]] = None,
    ) -> TokenOutput:
        """Generate sequence with token-in-token-out."""
        # TODO(@wuxibin): switch to `/generate` http endpoint once multi-modal support ready.
        max_possible_tokens = self.config.max_model_len - len(prompt_ids)

        if max_possible_tokens < 0:
            raise ValueError(
                f"Prompt length ({len(prompt_ids)}) exceeds the model's maximum context length "
                f"({self.config.max_model_len})."
            )

        if "max_new_tokens" in sampling_params:
            max_new_tokens = sampling_params.pop("max_new_tokens")
        elif "max_tokens" in sampling_params:
            # support vllm-style 'max_tokens' param
            max_new_tokens = sampling_params.pop("max_tokens")
        else:
            max_new_tokens = self.config.response_length + self.config.prompt_length - len(prompt_ids)