AMOOOMA opened a new pull request, #40906:
URL: https://github.com/apache/beam/pull/40906
Split out of #40760, following review feedback to separate the generic
pieces from the vLLM-specific ones. This PR is a small generic RunInference
hook.
## Why
Runners that autoscale on `num_loaded_models_<model_identifier>` (Dataflow's
GPU autoscaling) look the metric up by the identifier RunInference annotates
the inference step with: the `model_identifier` passed to `RunInference`, or a
uuid it generates. `ModelManager` publishes under that identifier (and since
#40714 keeps it stable across side-input model updates). But a model handler
that publishes the metric itself, e.g. one that fronts an out-of-process model
server, has no way to learn that value.
## What
- `ModelHandler.set_model_identifier(model_identifier)`: a no-op by default.
`_RunInferenceDoFn.setup()` calls it on each worker, before the first
`load_model()`, with the same tag it annotates the step with and gives
`ModelManager`.
- Forwarded by the keyed, maybe-keyed, prebatched, preprocessing and
postprocessing wrappers (`KeyedModelHandler` forwards to every per-key handler).
No behavior change for existing handlers.
## Testing
`base_test.py`:
- An explicit `model_identifier` reaches the handler before `load_model()`.
- Without one, the generated uuid equals the `model_identifier` annotation
on `RunInference/BeamML_RunInference`.
- The identifier is forwarded through `KeyedModelHandler` and
pre/post-processing wrappers.
------------------------
Thank you for your contribution! Follow this checklist to help us
incorporate your contribution quickly and easily:
- [x] Replace these instructions with a description of your change and why
you have made it.
- [x] Mention the appropriate issue in your description (for example:
`addresses #123` or `fixes #456`).
- [ ] Update `CHANGES.md` with noteworthy changes.
--
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]
