damccorm commented on code in PR #40908:
URL: https://github.com/apache/beam/pull/40908#discussion_r4564213567
##########
sdks/python/apache_beam/ml/inference/vllm_inference.py:
##########
@@ -370,11 +406,14 @@ def __init__(
dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
- max_batch_duration_secs: Optional[int] = None,
+ max_batch_duration_secs: Optional[float] = None,
max_batch_weight: Optional[int] = None,
element_size_fn: Optional[Callable[[Any], int]] = None,
batch_length_fn: Optional[Callable[[Any], int]] = None,
- batch_bucket_boundaries: Optional[list[int]] = None):
+ batch_bucket_boundaries: Optional[list[int]] = None,
+ prefix_batching_config: Optional[PrefixBatchingConfig] = None,
+ enable_autoscaling_metric: bool = True,
Review Comment:
Is there any reason a user would set this to false? Or should we just
default to having it on?
##########
sdks/python/apache_beam/ml/inference/vllm_inference.py:
##########
@@ -425,15 +465,31 @@ def __init__(
batch_bucket_boundaries: optional. a sorted list of positive boundary
values for length-aware batching buckets. Cannot be used together with
max_batch_weight.
+ prefix_batching_config: optional :class:`PrefixBatchingConfig`. When set,
+ prompts sharing a leading prefix are routed to the same worker and
+ batched together to maximize vLLM prefix cache hits (see
+ ``apache_beam.ml.inference.prefix_batching``).
+ enable_autoscaling_metric: if True (default), publish the
+ ``num_loaded_models_<tag>`` worker capacity signal from vLLM's request
+ metrics so Dataflow's GPU autoscaler scales workers with load. Requires
+ ``RunInference(..., use_model_manager=False)`` (the default); not
Review Comment:
Doesn't need to come in this pr, but I'm wondering if we should allow
ModelHandlers to have some variable or lifecycle method such that we throw in
cases where this is misconfigured. WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
