Skip to content

Commit 0302bf5

Browse files
authored
Implemented MLTransform generate vocab Dataflow benchmark (#38215)
* Implemented MLTransform generate vocab Dataflow benchmark * applied anomaly transform test for reshuffled ZScore order * remove redundant money nanos fallback * resolved comments
1 parent 54f7c7d commit 0302bf5

12 files changed

Lines changed: 791 additions & 2 deletions

File tree

.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ jobs:
9494
${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_VLLM_Gemma_Batch.txt
9595
${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Table_Row_Inference_Batch.txt
9696
${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Table_Row_Inference_Stream.txt
97+
${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Generate_Vocab_Batch.txt
9798
# The env variables are created and populated in the test-arguments-action as "<github.job>_test_arguments_<argument_file_paths_index>"
9899
- name: get current time
99100
run: echo "NOW_UTC=$(date '+%m%d%H%M%S' --utc)" >> $GITHUB_ENV
@@ -214,3 +215,14 @@ jobs:
214215
-PpythonVersion=3.10 \
215216
-PloadTest.requirementsTxtFile=apache_beam/ml/inference/table_row_inference_requirements.txt \
216217
'-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_10 }} --autoscaling_algorithm=THROUGHPUT_BASED --max_num_workers=20 --metrics_table=result_table_row_inference_stream --influx_measurement=result_table_row_inference_stream --mode=streaming --input_subscription=projects/apache-beam-testing/subscriptions/table_row_inference_benchmark --window_size_sec=60 --trigger_interval_sec=30 --timeout_ms=900000 --output_table=apache-beam-testing:beam_run_inference.result_table_row_inference_stream_outputs --job_name=benchmark-tests-table-row-inference-stream-${{env.NOW_UTC}}'
218+
- name: run MLTransform Generate Vocab Batch
219+
uses: ./.github/actions/gradle-command-self-hosted-action
220+
timeout-minutes: 180
221+
with:
222+
gradle-command: :sdks:python:apache_beam:testing:load_tests:run
223+
arguments: |
224+
-PloadTest.mainClass=apache_beam.testing.benchmarks.inference.mltransform_generate_vocab_benchmark \
225+
-Prunner=DataflowRunner \
226+
-PpythonVersion=3.10 \
227+
-PloadTest.requirementsTxtFile=apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt \
228+
'-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_11 }} --job_name=benchmark-tests-mltransform-generate-vocab-batch-${{env.NOW_UTC}}'
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one or more
2+
# contributor license agreements. See the NOTICE file distributed with
3+
# this work for additional information regarding copyright ownership.
4+
# The ASF licenses this file to You under the Apache License, Version 2.0
5+
# (the "License"); you may not use this file except in compliance with
6+
# the License. You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
17+
--project=apache-beam-testing
18+
--region=us-central1
19+
--runner=DataflowRunner
20+
--temp_location=gs://temp-storage-for-perf-tests/loadtests
21+
--staging_location=gs://temp-storage-for-perf-tests/loadtests
22+
--machine_type=n1-standard-4
23+
--disk_size_gb=100
24+
--num_workers=8
25+
--max_num_workers=16
26+
--autoscaling_algorithm=THROUGHPUT_BASED
27+
--worker_zone=us-central1-b
28+
--sdk_location=container
29+
--requirements_file=apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
30+
--input_options={}
31+
--publish_to_big_query=true
32+
--metrics_dataset=beam_run_inference
33+
--metrics_table=mltransform_generate_vocab_batch
34+
--influx_measurement=mltransform_generate_vocab_batch
35+
--input_file=gs://apache-beam-ml/testing/inputs/sentences_50k.txt
36+
--output_vocab=gs://temp-storage-for-perf-tests/mltransform/vocab_outputs/mltransform_generate_vocab_batch
37+
--columns=text
38+
--vocab_size=50000
39+
--min_frequency=1
40+
--lowercase=true
41+
--input_expand_factor=1
42+

.test-infra/tools/refresh_looker_metrics.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@
4444
("85", ["268", "269", "270", "271", "272"]), # PyTorch Sentiment Batch DistilBERT base uncased
4545
("86", ["284", "285", "286", "287", "288"]), # VLLM Batch Gemma
4646
("96", ["270", "304", "305", "353", "354"]), # Table Row Inference Sklearn Batch
47-
("106", ["355", "356", "357", "358", "359"]) # Table Row Inference Sklearn Streaming
47+
("106", ["355", "356", "357", "358", "359"]), # Table Row Inference Sklearn Streaming
48+
("107", ["360", "361", "362", "363", "364"]), # MLTransform Generate Vocab Batch
4849
]
4950

5051
def get_look(id: str) -> models.Look:
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
<!--
2+
Licensed under the Apache License, Version 2.0 (the "License");
3+
you may not use this file except in compliance with the License.
4+
You may obtain a copy of the License at
5+
6+
http://www.apache.org/licenses/LICENSE-2.0
7+
8+
Unless required by applicable law or agreed to in writing, software
9+
distributed under the License is distributed on an "AS IS" BASIS,
10+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
See the License for the specific language governing permissions and
12+
limitations under the License.
13+
-->
14+
15+
# MLTransform Examples
16+
17+
This directory contains Apache Beam examples for MLTransform pipelines.
18+
19+
## MLTransform - Generate Vocab (Batch only)
20+
21+
`mltransform_generate_vocab.py` builds a vocabulary artifact from batch input
22+
rows using `MLTransform` + `ComputeAndApplyVocabulary`.
23+
24+
### What it does
25+
26+
1. Reads input rows from JSONL (`--input_file`) or BigQuery (`--input_table`).
27+
2. Extracts specified columns (`--columns`).
28+
3. Normalizes and combines text values (`trim`, optional lowercasing).
29+
4. Runs `ComputeAndApplyVocabulary` with top-k and min-frequency constraints
30+
using space-delimited token splitting.
31+
5. Writes the vocabulary as one token per line.
32+
33+
### Required arguments
34+
35+
- `--output_vocab`
36+
- `--columns`
37+
- and one of:
38+
- `--input_file`
39+
- `--input_table`
40+
41+
### Optional arguments
42+
43+
- `--vocab_size` (default: `50000`)
44+
- `--min_frequency` (default: `1`)
45+
- `--lowercase` (default: `true`)
46+
- `--input_expand_factor` (default: `1`, useful for perf/load testing)
47+
48+
### Local batch example
49+
50+
```sh
51+
python -m apache_beam.examples.ml_transform.mltransform_generate_vocab \
52+
--input_file=/tmp/input.jsonl \
53+
--output_vocab=/tmp/vocab.txt \
54+
--columns=text,category \
55+
--vocab_size=5 \
56+
--min_frequency=1 \
57+
--lowercase=true \
58+
--input_expand_factor=1 \
59+
--runner=DirectRunner
60+
```
61+
62+
### Input format
63+
64+
JSONL input with object rows, for example:
65+
66+
```json
67+
{"id":"1","text":"Beam beam ML pipeline"}
68+
{"id":"2","text":"Beam pipeline dataflow"}
69+
{"id":"3","text":"ML transform beam"}
70+
{"id":"4","text":"vocab vocab vocab test"}
71+
{"id":"5","text":"rare_token_once"}
72+
{"id":"6","text":""}
73+
{"id":"7","text":null}
74+
```
75+
76+
The integration tests in `mltransform_generate_vocab_test.py` generate this
77+
sample data programmatically.
78+
79+
### Output format
80+
81+
One token per line:
82+
83+
1. tokens follow the vocabulary order produced by `ComputeAndApplyVocabulary`.
84+
85+
Example output:
86+
87+
```txt
88+
beam
89+
ml
90+
```
91+
92+
For this sample and config:
93+
94+
```sh
95+
--columns=text --min_frequency=2 --vocab_size=3
96+
```
97+
98+
the expected output is:
99+
100+
```txt
101+
beam
102+
vocab
103+
ml
104+
```
105+
106+
### Additional test datasets
107+
108+
Test data for happy path and null/empty/missing columns is generated inline in
109+
`mltransform_generate_vocab_test.py`.
110+
111+
### Performance testing pattern
112+
113+
- Small local files: functional correctness and output-stability tests.
114+
- Large GCS files (or moderate file + `--input_expand_factor`): throughput/cost
115+
benchmarking on Dataflow.
116+

0 commit comments

Comments
 (0)