AMOOOMA opened a new pull request, #40671:
URL: https://github.com/apache/beam/pull/40671
## Summary
This PR addresses two related issues with `model_identifier` and
`ModelManager` metrics during runtime side-input model updates in
`RunInference`:
1. **Scope `model_identifier` transform annotation strictly to
`BeamML_RunInference`**:
- Previously, `RunInference.annotations()` attached `model_identifier` to
the outer `RunInference` composite `PTransform`, which caused all nested
sub-transforms (such as `WatchFilePattern` side-input polling transforms,
`BatchElements`, preprocess/postprocess transforms, etc.) to inherit the
`model_identifier` annotation. Runners that group steps by `model_identifier`
to look up `num_loaded_models_<model_identifier>` would then fail to resolve
the metric on steps that never execute `_RunInferenceDoFn`.
- Remove `model_identifier` from `RunInference.annotations()` and instead
attach `transform_annotations(model_identifier=self._model_tag)` directly
around `BeamML_RunInference` in `RunInference.expand()`.
2. **Preserve `num_loaded_models_<model_identifier>` metric across runtime
side-input model updates**:
- When `RunInference` updates its model via side input (`ModelMetadata`),
`_RunInferenceDoFn._load_model` acquires the new model under the updated
`model_id` tag (`self._cur_tag`), while the `BeamML_RunInference` transform
annotation remains fixed at construction time (`self._model_tag` /
`model_identifier`).
- Pass `metric_tag=self._model_tag` through `_SharedModelWrapper` to
`ModelManager.acquire_model(..., metric_tag=...)`, and track
`self._metric_tag_to_active_tag[metric_tag] = tag` in `ModelManager` so
`_update_model_count_metric()` and `_clear_all_model_metrics()` continue
reporting `num_loaded_models_<metric_tag>` alongside
`num_loaded_models_<active_tag>`.
## Testing
- Added
`test_run_inference_only_attaches_model_identifier_to_inference_step` in
`base_test.py`.
- Updated
`test_run_inference_dofn_model_update_acquires_and_releases_matching_tag` in
`base_test.py` to verify `metric_tag` preservation across model updates.
- Added `test_model_manager_metrics_with_metric_tag` in
`model_manager_test.py`.
- Updated `test_monitoring_outlet_clears_model_identifier_annotation` in
`vertex_ai_model_monitoring_v2_test.py`.
--
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]
