[data][llm] Allow tokenized_prompt without prompt in vLLMEngineStage

[jeffreywang-anyscale]

Description

When using vLLMEngineProcessor, vLLMEngineStage only requires either prompt or tokenized_prompt to be present. However, the current implementation raises a validation error when prompt is missing.

This PR updates the validation logic in vLLMEngineStage to allow inputs where the prompt column is absent, as long as tokenized_prompt is provided. If both prompt and tokenized_prompt are present, the existing behavior is preserved where tokenized_prompt is preferred and tokens are passed directly to the engine.

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request correctly updates the validation logic in vLLMEngineStage to allow either prompt or tokenized_prompt as input, which aligns with the behavior of vLLMEngineProcessor. The changes are logical and are accompanied by good test coverage, including both positive and negative test cases.

I have one suggestion to refactor the validate_inputs method in vllm_engine_stage.py to improve maintainability by avoiding temporary state modification. Otherwise, the changes look good.

[gemini-code-assist]

The current implementation of validate_inputs temporarily modifies the instance state (self.expected_input_keys). While this works and is protected by a try...finally block, it can be fragile and is generally not a good practice as it can lead to subtle bugs, especially if the class is used in a concurrent environment in the future.

A cleaner approach would be to reimplement the validation logic without modifying instance state. This would involve checking for prompt or tokenized_prompt and then checking for the other expected keys, similar to what the superclass does. This makes the method self-contained and easier to reason about.

    def validate_inputs(self, inputs: List[Dict[str, Any]]):
        """Validate the inputs to make sure the required keys are present.
        
        Overrides base class to handle the requirement for prompt/tokenized_prompt.

        Args:
            inputs: The inputs.

        Raises:
            ValueError: If the required keys are not found.
        """
        # All expected keys except for prompt/tokenized_prompt, which are handled specially.
        other_expected_keys = self.expected_input_keys - {"prompt", "tokenized_prompt"}

        for inp in inputs:
            input_keys = set(inp.keys())

            if self.IDX_IN_BATCH_COLUMN in input_keys:
                raise ValueError(
                    f"The input column {self.IDX_IN_BATCH_COLUMN} is reserved "
                    "for internal use."
                )

            if "prompt" not in input_keys and "tokenized_prompt" not in input_keys:
                raise ValueError(
                    "Either 'prompt' (text) or 'tokenized_prompt' (tokens) "
                    f"must be provided. Input keys: {input_keys}"
                )

            # Check for other required keys.
            missing_required = other_expected_keys - input_keys
            if missing_required:
                raise ValueError(
                    f"Required input keys {missing_required} not found at the input of "
                    f"{self.__class__.__name__}. Input keys: {input_keys}"
                )

[jeffreywang-anyscale]

Open to suggestions -- at least one of prompt or tokenized_prompt should be provided. Is marking both as required the right heuristic?

[kouroshHakha]

what happens if we don't validate here? And let the engine fail? What would be the error message coming from the engine? It won't fail silently, right?

The reason that I am interested in not validating like this at all, is that we are adding too much extra complexity for a little gain on input key validation in a case where we want either/or type of expectation.

[jeffreywang-anyscale]

In fact, we are validating that prompt should exist if tokenized_prompt does not present here. From the assertion, users may be confused that prompt must be provided. Although it's not immediately clear that only 1 of prompt or tokenized prompt should be provided, I agree that the additional validation complexity isn’t worthwhile.

Note: Even if the existing validation is removed, the engine will raise vllm.v1.engine.exceptions.EngineGenerateError.

[jeffreywang-anyscale]

Let's put both of them under get_optional_input_keys then.

[cursor]

Error message uses internal field name instead of user-facing name

Low Severity

The error message says "Either 'prompt' or 'prompt_token_ids' must be provided" but users provide data using the field name tokenized_prompt, not prompt_token_ids. The _prepare_llm_request method maps the user-facing tokenized_prompt field to the internal prompt_token_ids field when creating vLLMEngineRequest. This mismatch between the error message and the user-facing API could confuse users who receive this validation error.

 

[kouroshHakha]

code looks good.

[jeffreywang-anyscale]

Looking into premerge failures 🤔