diff --git a/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml b/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml index 195ce693a388..88ae683dfe0c 100644 --- a/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml +++ b/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml @@ -96,6 +96,9 @@ jobs: ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Table_Row_Inference_Stream.txt ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Generate_Vocab_Batch.txt ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_One_Hot_Encoding_Batch.txt + ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_GPU_Batch.txt + ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Text_Embedding_Batch.txt + ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_CPU_Batch.txt # The env variables are created and populated in the test-arguments-action as "_test_arguments_" - name: get current time run: echo "NOW_UTC=$(date '+%m%d%H%M%S' --utc)" >> $GITHUB_ENV @@ -238,3 +241,36 @@ jobs: -PpythonVersion=3.10 \ -PloadTest.requirementsTxtFile=apache_beam/examples/ml_transform/mltransform_one_hot_encoding_requirements.txt \ '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_12 }} --autoscaling_algorithm=NONE --metrics_table=mltransform_one_hot_encoding_batch --influx_measurement=mltransform_one_hot_encoding_batch --job_name=benchmark-tests-mltransform-one-hot-encoding-batch-${{env.NOW_UTC}} --output_file=gs://temp-storage-for-end-to-end-tests/mltransform/one_hot_output_${{env.NOW_UTC}} --artifact_location=gs://temp-storage-for-end-to-end-tests/mltransform/artifacts_${{env.NOW_UTC}}' + - name: run MLTransform Image Embedding GPU Batch + uses: ./.github/actions/gradle-command-self-hosted-action + timeout-minutes: 180 + with: + gradle-command: :sdks:python:apache_beam:testing:load_tests:run + arguments: | + -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.mltransform_image_embedding_benchmark \ + -Prunner=DataflowRunner \ + -PpythonVersion=3.10 \ + -PloadTest.requirementsTxtFile=apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt \ + '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_13 }} --artifact_location=gs://temp-storage-for-perf-tests/loadtests/mltransform_image_embedding_gpu/artifacts-${{env.NOW_UTC}} --job_name=benchmark-tests-mltransform-image-embedding-gpu-batch-${{env.NOW_UTC}}' + - name: run MLTransform Text Embedding Batch + uses: ./.github/actions/gradle-command-self-hosted-action + timeout-minutes: 180 + with: + gradle-command: :sdks:python:apache_beam:testing:load_tests:run + arguments: | + -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.mltransform_text_embedding_benchmark \ + -Prunner=DataflowRunner \ + -PpythonVersion=3.10 \ + -PloadTest.requirementsTxtFile=apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt \ + '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_14 }} --output=gs://temp-storage-for-perf-tests/loadtests/mltransform_text_embedding/results-${{env.NOW_UTC}} --artifact_location=gs://temp-storage-for-perf-tests/loadtests/mltransform_text_embedding/artifacts-${{env.NOW_UTC}} --job_name=benchmark-tests-mltransform-text-embedding-batch-${{env.NOW_UTC}}' + - name: run MLTransform Image Embedding CPU Batch + uses: ./.github/actions/gradle-command-self-hosted-action + timeout-minutes: 180 + with: + gradle-command: :sdks:python:apache_beam:testing:load_tests:run + arguments: | + -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.mltransform_image_embedding_benchmark \ + -Prunner=DataflowRunner \ + -PpythonVersion=3.10 \ + -PloadTest.requirementsTxtFile=apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt \ + '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_15 }} --artifact_location=gs://temp-storage-for-perf-tests/loadtests/mltransform_image_embedding_cpu/artifacts-${{env.NOW_UTC}} --job_name=benchmark-tests-mltransform-image-embedding-cpu-batch-${{env.NOW_UTC}}' diff --git a/.github/workflows/beam_PostCommit_Java_DataflowV2.yml b/.github/workflows/beam_PostCommit_Java_DataflowV2.yml index 5c38e94ed7dd..1cbc9fd25a89 100644 --- a/.github/workflows/beam_PostCommit_Java_DataflowV2.yml +++ b/.github/workflows/beam_PostCommit_Java_DataflowV2.yml @@ -54,7 +54,7 @@ jobs: beam_PostCommit_Java_DataflowV2: name: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) runs-on: [self-hosted, ubuntu-24.04, main] - timeout-minutes: 360 + timeout-minutes: 480 strategy: matrix: job_name: [beam_PostCommit_Java_DataflowV2] @@ -91,4 +91,4 @@ jobs: commit: '${{ env.prsha || env.GITHUB_SHA }}' comment_mode: ${{ github.event_name == 'issue_comment' && 'always' || 'off' }} files: '**/build/test-results/**/*.xml' - large_files: true \ No newline at end of file + large_files: true diff --git a/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_CPU_Batch.txt b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_CPU_Batch.txt new file mode 100644 index 000000000000..37c47ca25542 --- /dev/null +++ b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_CPU_Batch.txt @@ -0,0 +1,42 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +--runner=DataflowRunner +--project=apache-beam-testing +--region=us-central1 +--num_workers=5 +--max_num_workers=20 +--disk_size_gb=100 +--autoscaling_algorithm=THROUGHPUT_BASED +--staging_location=gs://temp-storage-for-perf-tests/loadtests +--temp_location=gs://temp-storage-for-perf-tests/loadtests +--requirements_file=apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt +--publish_to_big_query=true +--metrics_dataset=beam_run_inference +--metrics_table=mltransform_image_embedding_cpu_batch +--input_options={} +--influx_measurement=mltransform_image_embedding_cpu_batch +--mode=batch +--input=gs://apache-beam-ml/testing/inputs/openimage_50k_benchmark.txt +--artifact_location=gs://temp-storage-for-perf-tests/loadtests/mltransform_image_embedding_cpu/artifacts +--output_table=apache-beam-testing:beam_run_inference.result_mltransform_image_embedding_cpu_batch +--pretrained_model_name=clip-ViT-B-32 +--device=CPU +--min_batch_size=8 +--max_batch_size=64 +--embedding_min_ram=16GB +--dataflow_service_options=enable_prime +--experiments=use_runner_v2 +--timeout_ms=3600000 diff --git a/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_GPU_Batch.txt b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_GPU_Batch.txt new file mode 100644 index 000000000000..274dc43f9416 --- /dev/null +++ b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_GPU_Batch.txt @@ -0,0 +1,44 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +--runner=DataflowRunner +--project=apache-beam-testing +--region=us-central1 +--num_workers=5 +--max_num_workers=20 +--disk_size_gb=100 +--autoscaling_algorithm=THROUGHPUT_BASED +--staging_location=gs://temp-storage-for-perf-tests/loadtests +--temp_location=gs://temp-storage-for-perf-tests/loadtests +--requirements_file=apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt +--publish_to_big_query=true +--metrics_dataset=beam_run_inference +--metrics_table=mltransform_image_embedding_gpu_batch +--input_options={} +--influx_measurement=mltransform_image_embedding_gpu_batch +--mode=batch +--input=gs://apache-beam-ml/testing/inputs/openimage_50k_benchmark.txt +--artifact_location=gs://temp-storage-for-perf-tests/loadtests/mltransform_image_embedding_gpu/artifacts +--output_table=apache-beam-testing:beam_run_inference.result_mltransform_image_embedding_gpu_batch +--pretrained_model_name=clip-ViT-B-32 +--device=GPU +--min_batch_size=8 +--max_batch_size=64 +--embedding_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver +--sdk_location=container +--sdk_container_image=us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest +--dataflow_service_options=enable_prime +--experiments=use_runner_v2 +--timeout_ms=3600000 diff --git a/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Text_Embedding_Batch.txt b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Text_Embedding_Batch.txt new file mode 100644 index 000000000000..1835dc260025 --- /dev/null +++ b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Text_Embedding_Batch.txt @@ -0,0 +1,42 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +--runner=DataflowRunner +--project=apache-beam-testing +--region=us-central1 +--machine_type=n1-standard-4 +--num_workers=10 +--max_num_workers=50 +--disk_size_gb=100 +--autoscaling_algorithm=THROUGHPUT_BASED +--staging_location=gs://temp-storage-for-perf-tests/loadtests +--temp_location=gs://temp-storage-for-perf-tests/loadtests +--requirements_file=apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt +--publish_to_big_query=true +--metrics_dataset=beam_run_inference +--metrics_table=mltransform_text_embedding_batch +--input_options={} +--influx_measurement=mltransform_text_embedding_batch +--input_file=gs://apache-beam-ml/testing/inputs/sentences_50k.txt +--output=gs://temp-storage-for-perf-tests/loadtests/mltransform_text_embedding/results +--artifact_location=gs://temp-storage-for-perf-tests/loadtests/mltransform_text_embedding/artifacts +--model_name=sentence-transformers/all-MiniLM-L6-v2 +--min_batch_size=16 +--max_batch_size=128 +--model_batch_size=32 +--device=CPU +--sdk_location=container +--sdk_container_image=us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest +--experiments=use_runner_v2 diff --git a/.test-infra/tools/refresh_looker_metrics.py b/.test-infra/tools/refresh_looker_metrics.py index c8d66f4a4bd3..69798e6f24e8 100644 --- a/.test-infra/tools/refresh_looker_metrics.py +++ b/.test-infra/tools/refresh_looker_metrics.py @@ -46,7 +46,10 @@ ("96", ["270", "304", "305", "353", "354"]), # Table Row Inference Sklearn Batch ("106", ["355", "356", "357", "358", "359"]), # Table Row Inference Sklearn Streaming ("107", ["360", "361", "362", "363", "364"]), # MLTransform Generate Vocab Batch - ("108", ["365", "366", "367", "368", "369"]) # MLTransform One-Hot Encoding Batch + ("108", ["365", "366", "367", "368", "369"]), # MLTransform One-Hot Encoding Batch + ("109", ["375", "376", "377", "378", "379"]), # MLTransform Text Embedding Batch + ("110", ["380", "381", "382", "383", "385"]), # MLTransform Image Embedding GPU Batch + ("111", ["370", "371", "372", "373", "374"]), # MLTransform Image Embedding CPU Batch ] def get_look(id: str) -> models.Look: diff --git a/CHANGES.md b/CHANGES.md index 8bc74b958bad..ca3e0a4c4252 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -72,6 +72,7 @@ ## New Features / Improvements +* (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). * Dataflow Runner v2 has been renamed to Dataflow Portable Runner. Please refer to Dataflow [public documentation](https://docs.cloud.google.com/dataflow/docs/runner-v2) on when to enable Portable Runner.([#39000](https://github.com/apache/beam/issues/39000)). * (Java) Enabled state tag encoding v2 by default for new Dataflow Streaming Engine jobs. It can be disabled by passing `--experiments=disable_streaming_engine_state_tag_encoding_v2` or `--updateCompatibilityVersion=2.74.0` pipeline option. Note that the tag encoding version cannot change during a job update. Jobs using tag encoding v2 (enabled by default for new jobs on 2.75.0+) cannot be downgraded to Beam versions prior to 2.73.0, as only versions 2.73.0 and later support tag encoding v2. ([#38705](https://github.com/apache/beam/issues/38705)). * (Python) Added instrumentation to support off-the-shelf profiling agents when launching Python SDK Harness ([#38853](https://github.com/apache/beam/issues/38853)). diff --git a/sdks/go.mod b/sdks/go.mod index a597578df528..aadeaf57ccf3 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -33,10 +33,10 @@ require ( cloud.google.com/go/spanner v1.92.0 cloud.google.com/go/storage v1.63.0 github.com/aws/aws-sdk-go-v2 v1.42.1 - github.com/aws/aws-sdk-go-v2/config v1.32.27 - github.com/aws/aws-sdk-go-v2/credentials v1.19.26 - github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.30 - github.com/aws/aws-sdk-go-v2/service/s3 v1.104.2 + github.com/aws/aws-sdk-go-v2/config v1.32.28 + github.com/aws/aws-sdk-go-v2/credentials v1.19.27 + github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.31 + github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0 github.com/aws/smithy-go v1.27.3 github.com/docker/go-connections v0.7.0 // indirect github.com/dustin/go-humanize v1.0.1 @@ -59,7 +59,7 @@ require ( golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.21.0 golang.org/x/sys v0.46.0 - golang.org/x/text v0.38.0 + golang.org/x/text v0.39.0 google.golang.org/api v0.287.0 google.golang.org/genproto v0.0.0-20260523011958-0a33c5d7ca68 google.golang.org/grpc v1.82.0 @@ -91,7 +91,7 @@ require ( github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.57.0 // indirect github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op // indirect github.com/apache/arrow/go/v15 v15.0.2 // indirect - github.com/aws/aws-sdk-go-v2/service/signin v1.2.2 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.3.0 // indirect github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect github.com/containerd/log v0.1.0 // indirect @@ -134,7 +134,7 @@ require ( go.opentelemetry.io/otel/sdk/metric v1.44.0 // indirect go.opentelemetry.io/otel/trace v1.44.0 // indirect go.shabbyrobe.org/gocovmerge v0.0.0-20230507111327-fa4f82cfbf4d // indirect - golang.org/x/telemetry v0.0.0-20260519152614-eab6ae52b5e2 // indirect + golang.org/x/telemetry v0.0.0-20260625142307-59b4966ccb57 // indirect golang.org/x/time v0.15.0 // indirect ) @@ -157,9 +157,9 @@ require ( github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 // indirect github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 // indirect - github.com/aws/aws-sdk-go-v2/service/sso v1.31.5 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.8 // indirect - github.com/aws/aws-sdk-go-v2/service/sts v1.43.5 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.32.0 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.44.0 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect @@ -200,8 +200,8 @@ require ( github.com/zeebo/xxh3 v1.1.0 // indirect go.opencensus.io v0.24.0 // indirect golang.org/x/crypto v0.53.0 // indirect - golang.org/x/mod v0.36.0 // indirect - golang.org/x/tools v0.45.0 // indirect + golang.org/x/mod v0.37.0 // indirect + golang.org/x/tools v0.47.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260622175928-b703f567277d // indirect diff --git a/sdks/go.sum b/sdks/go.sum index ea12281ba155..f5865534d2f9 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -207,20 +207,20 @@ github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHz github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14/go.mod h1:zwM6veDkhGgQFqkBy+uT28AAYpLu+uFMlPl+rCg/73E= github.com/aws/aws-sdk-go-v2/config v1.15.3/go.mod h1:9YL3v07Xc/ohTsxFXzan9ZpFpdTOFl4X65BAKYaz8jg= github.com/aws/aws-sdk-go-v2/config v1.25.3/go.mod h1:tAByZy03nH5jcq0vZmkcVoo6tRzRHEwSFx3QW4NmDw8= -github.com/aws/aws-sdk-go-v2/config v1.32.27 h1:SJwJ9Q4kM7v5QVSYYyXj3znRr6lNyZEhSgAXmXXcVbI= -github.com/aws/aws-sdk-go-v2/config v1.32.27/go.mod h1:uBfrzTRedDmB2u+b6+UlaKJy2O6VSH5un2jP24t/KvQ= +github.com/aws/aws-sdk-go-v2/config v1.32.28 h1:qY6afygxK5c2PPU3Sz8W6yB5W44RF1vnmPdBwViDN+Y= +github.com/aws/aws-sdk-go-v2/config v1.32.28/go.mod h1:WeS/wN1IDs8YC+BxTrFz9ZyJ1rufRBQfirOcDusEpmQ= github.com/aws/aws-sdk-go-v2/credentials v1.11.2/go.mod h1:j8YsY9TXTm31k4eFhspiQicfXPLZ0gYXA50i4gxPE8g= github.com/aws/aws-sdk-go-v2/credentials v1.16.2/go.mod h1:sDdvGhXrSVT5yzBDR7qXz+rhbpiMpUYfF3vJ01QSdrc= -github.com/aws/aws-sdk-go-v2/credentials v1.19.26 h1:Si8kk1kyJnuJWCEgiwpBtTdtgSdR7i611596NnC0YIQ= -github.com/aws/aws-sdk-go-v2/credentials v1.19.26/go.mod h1:lBckz+W9SAdNtSDw3pYgQUJDJFcBBWry0GSzw+bK0TY= +github.com/aws/aws-sdk-go-v2/credentials v1.19.27 h1:cFksKkdaBGGmpe6XJpvrxFNWkbXY5/gwFqZNB2O9WCM= +github.com/aws/aws-sdk-go-v2/credentials v1.19.27/go.mod h1:20CoObBgNhFfl8/ggDQu2IZmItxDhkLcWSy4C3alDPI= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.3/go.mod h1:uk1vhHHERfSVCUnqSqz8O48LBYDSC+k6brng09jcMOk= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.14.4/go.mod h1:t4i+yGHMCcUNIX1x7YVYa6bH/Do7civ5I6cG/6PMfyA= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQu0EtP23pP553Cf9KgsDV4= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30/go.mod h1:/3AOgy4K17Dm4ucMZVC/MJkzy5kmfKUcINRHZyo0koQ= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.3/go.mod h1:0dHuD2HZZSiwfJSy1FO5bX1hQ1TxVV1QXXjpn3XUE44= github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.14.0/go.mod h1:UcgIwJ9KHquYxs6Q5skC9qXjhYMK+JASDYcXQ4X7JZE= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.30 h1:qzaF2qXmjqBcLbtcpIxOddoDH5eKIRdmOiKt7BatgDc= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.30/go.mod h1:5LlJNctj1Y7l59FPonG4LNjK5DvorTKIS+pAVuWAh3s= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.31 h1:7lcHy74J5ajNAuTvwZYS9Tw4Ept/eLCavGEnykacT8o= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.31/go.mod h1:7eILDitqxXgdStzJvYvCBuVL9k5ftKcPQ6F4ZRUV3+A= github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.9/go.mod h1:AnVH5pvai0pAF4lXRq0bmhbes1u9R8wTE+g+183bZNM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.3/go.mod h1:7sGSz1JCKHWWBHq98m6sMtWQikmYPpxjqOydDemiVoM= github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 h1:xM/Is9cKMHa8Jj8zkvWhvrFkZsXJV9E+BB4g0HW0duQ= @@ -253,25 +253,25 @@ github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31/go.mod h1:I/1+z0 github.com/aws/aws-sdk-go-v2/service/kms v1.16.3/go.mod h1:QuiHPBqlOFCi4LqdSskYYAWpQlx3PKmohy+rE2F+o5g= github.com/aws/aws-sdk-go-v2/service/s3 v1.26.3/go.mod h1:g1qvDuRsJY+XghsV6zg00Z4KJ7DtFFCx8fJD2a491Ak= github.com/aws/aws-sdk-go-v2/service/s3 v1.43.0/go.mod h1:NXRKkiRF+erX2hnybnVU660cYT5/KChRD4iUgJ97cI8= -github.com/aws/aws-sdk-go-v2/service/s3 v1.104.2 h1:bAY6O/TDv1HQnvylh9E247IyIKsUWUt2G965S7qX110= -github.com/aws/aws-sdk-go-v2/service/s3 v1.104.2/go.mod h1:zdmCoFO/dSI7GlrwsPqFJI+WlFnSU4Tc8TJnlXrM1Do= +github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0 h1:XptwLL+UHXgafYMIHTy59IRovLbhz3znkxY2uS/pbXU= +github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0/go.mod h1:zdmCoFO/dSI7GlrwsPqFJI+WlFnSU4Tc8TJnlXrM1Do= github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.15.4/go.mod h1:PJc8s+lxyU8rrre0/4a0pn2wgwiDvOEzoOjcJUBr67o= -github.com/aws/aws-sdk-go-v2/service/signin v1.2.2 h1:69JEZSDTQ+UNbTWQJCZMmbpQb5sfc79KUt0O7Pyfjmo= -github.com/aws/aws-sdk-go-v2/service/signin v1.2.2/go.mod h1:mxC0nT/C8wMMS97DemZPzvUZxvIt+2Iq+eS3JdFZGgg= +github.com/aws/aws-sdk-go-v2/service/signin v1.3.0 h1:i0+tbB9QBnzL5NrF2WR/zk8q2s+1N+RaDYr2627E8UI= +github.com/aws/aws-sdk-go-v2/service/signin v1.3.0/go.mod h1:mxC0nT/C8wMMS97DemZPzvUZxvIt+2Iq+eS3JdFZGgg= github.com/aws/aws-sdk-go-v2/service/sns v1.17.4/go.mod h1:kElt+uCcXxcqFyc+bQqZPFD9DME/eC6oHBXvFzQ9Bcw= github.com/aws/aws-sdk-go-v2/service/sqs v1.18.3/go.mod h1:skmQo0UPvsjsuYYSYMVmrPc1HWCbHUJyrCEp+ZaLzqM= github.com/aws/aws-sdk-go-v2/service/ssm v1.24.1/go.mod h1:NR/xoKjdbRJ+qx0pMR4mI+N/H1I1ynHwXnO6FowXJc0= github.com/aws/aws-sdk-go-v2/service/sso v1.11.3/go.mod h1:7UQ/e69kU7LDPtY40OyoHYgRmgfGM4mgsLYtcObdveU= github.com/aws/aws-sdk-go-v2/service/sso v1.17.2/go.mod h1:/pE21vno3q1h4bbhUOEi+6Zu/aT26UK2WKkDXd+TssQ= -github.com/aws/aws-sdk-go-v2/service/sso v1.31.5 h1:xlK3Tdc8FO7Tq1k0+hL+otF33glj+dE+qeM5iINiDvU= -github.com/aws/aws-sdk-go-v2/service/sso v1.31.5/go.mod h1:u8af9Nqkmqnr96f7v9nHqzZT9XBwbXEkTiqT4ROuJSE= +github.com/aws/aws-sdk-go-v2/service/sso v1.32.0 h1:qjMmry/cBDee1E/2gyvel0uRYCi3mwRZ2hf6N+GAodo= +github.com/aws/aws-sdk-go-v2/service/sso v1.32.0/go.mod h1:u8af9Nqkmqnr96f7v9nHqzZT9XBwbXEkTiqT4ROuJSE= github.com/aws/aws-sdk-go-v2/service/ssooidc v1.20.0/go.mod h1:dWqm5G767qwKPuayKfzm4rjzFmVjiBFbOJrpSPnAMDs= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.8 h1:yX1IbiBfC7SdEgDwIGnRaZyPPDRbQPDOJxl8102PcGk= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.8/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.0 h1:fpOlDPI55HdszaxapEGk6HsGosOUaM2YPWJpjMgp8UI= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.0/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= github.com/aws/aws-sdk-go-v2/service/sts v1.16.3/go.mod h1:bfBj0iVmsUyUg4weDB4NxktD9rDGeKSVWnjTnwbx9b8= github.com/aws/aws-sdk-go-v2/service/sts v1.25.3/go.mod h1:4EqRHDCKP78hq3zOnmFXu5k0j4bXbRFfCh/zQ6KnEfQ= -github.com/aws/aws-sdk-go-v2/service/sts v1.43.5 h1:T3ANO8QWDbzQD8f4+UaX+fvJlyGnOFMKLbW+NGBHg04= -github.com/aws/aws-sdk-go-v2/service/sts v1.43.5/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= +github.com/aws/aws-sdk-go-v2/service/sts v1.44.0 h1:bLZ0PolJ8J+HkJHztcXORUpHXBye2U8298lCEMi6ZCU= +github.com/aws/aws-sdk-go-v2/service/sts v1.44.0/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= github.com/aws/smithy-go v1.11.2/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM= github.com/aws/smithy-go v1.17.0/go.mod h1:NukqUGpCZIILqqiV0NIjeFh24kd/FAa4beRb6nbIUPE= github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= @@ -979,8 +979,8 @@ golang.org/x/mod v0.5.0/go.mod h1:5OXOZSfqPIIbmVBIIKWRFfZjPR0E5r58TLhUjH0a2Ro= golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/mod v0.10.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= -golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= -golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= +golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= +golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190108225652-1e06a53dbb7e/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -1162,8 +1162,8 @@ golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/telemetry v0.0.0-20260519152614-eab6ae52b5e2 h1:2EucmYlcIsc8Y6aLj+kX90Y00hmjqLNlw935kc13R2k= -golang.org/x/telemetry v0.0.0-20260519152614-eab6ae52b5e2/go.mod h1:Eqhaxk/wZsWEH8CRxLwj6xzEJbz7k1EFGqx7nyCoabE= +golang.org/x/telemetry v0.0.0-20260625142307-59b4966ccb57 h1:nwGZBCt+FnXUrGsj5vjzAsEmkcaFvd82BbOjECiFYZc= +golang.org/x/telemetry v0.0.0-20260625142307-59b4966ccb57/go.mod h1:3AWMyWHS+caVoiEXpiq6+tzKA40J4vQT3MYr80ZtQpc= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= @@ -1188,8 +1188,8 @@ golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= -golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= -golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= +golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus= +golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= @@ -1264,8 +1264,8 @@ golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= golang.org/x/tools v0.8.0/go.mod h1:JxBZ99ISMI5ViVkT1tr6tdNmXeTrcpVSD3vZ1RsRdN4= -golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= -golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= +golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= +golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= golang.org/x/xerrors v0.0.0-20190410155217-1f06c39b4373/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsIOWriteBatchesTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsIOWriteBatchesTest.java index e92720bfb5a5..90ea28d60447 100644 --- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsIOWriteBatchesTest.java +++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/sqs/SqsIOWriteBatchesTest.java @@ -21,8 +21,10 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.concurrent.CompletableFuture.completedFuture; import static java.util.concurrent.CompletableFuture.supplyAsync; +import static java.util.stream.Collectors.groupingBy; import static java.util.stream.Collectors.toList; import static java.util.stream.IntStream.range; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -30,6 +32,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.atLeast; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -38,6 +41,7 @@ import java.util.Arrays; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; @@ -263,10 +267,16 @@ public void testWriteBatchesWithTimeout() { p.run().waitUntilFinish(); - SendMessageBatchRequestEntry[] entries = entries(range(0, 5)); - // due to added delay, batches are timed out on arrival of every 3rd msg - verify(sqs).sendMessageBatch(request("queue", entries[0], entries[1], entries[2])); - verify(sqs).sendMessageBatch(request("queue", entries[3], entries[4])); + // Nominally batches time out on arrival of every 3rd message ([0,1,2], [3,4]). The exact + // grouping depends on wall clock time and is unreliable on loaded machines (#38946), so + // verify timing-independent invariants instead: expired batches are flushed on append, so + // with >= 100ms between messages no batch can exceed 3 entries before exceeding the 150ms + // timeout, forcing at least 2 batches for 5 messages. + Map> requests = captureBatchRequests(atLeast(2)); + assertThat(requests.keySet()).containsExactly("queue"); + assertMessageBodies(requestsFor(requests, "queue"), range(0, 5)); + assertThat(requestsFor(requests, "queue")) + .allSatisfy(req -> assertThat(req.entries().size()).isLessThanOrEqualTo(3)); } @Test @@ -285,11 +295,16 @@ public void testWriteBatchesWithStrictTimeout() { p.run().waitUntilFinish(); - SendMessageBatchRequestEntry[] entries = entries(range(0, 5)); - // using strict timeouts batches, batches are timed out by a separate thread - verify(sqs).sendMessageBatch(request("queue", entries[0], entries[1])); - verify(sqs).sendMessageBatch(request("queue", entries[2], entries[3])); - verify(sqs).sendMessageBatch(request("queue", entries[4])); + // Nominally the separate timeout thread flushes [0,1], [2,3], [4]. The exact grouping + // depends on wall clock time and is unreliable on loaded machines (#38946), so verify + // timing-independent invariants instead. Expired batches are also flushed on append + // (independently of the timeout thread), so no batch can exceed 3 entries and 5 messages + // require at least 2 batches. + Map> requests = captureBatchRequests(atLeast(2)); + assertThat(requests.keySet()).containsExactly("queue"); + assertMessageBodies(requestsFor(requests, "queue"), range(0, 5)); + assertThat(requestsFor(requests, "queue")) + .allSatisfy(req -> assertThat(req.entries().size()).isLessThanOrEqualTo(3)); } @Test @@ -337,11 +352,18 @@ public void testWriteBatchesToDynamicWithTimeout() { p.run().waitUntilFinish(); - SendMessageBatchRequestEntry[] entries = entries(range(0, 5)); - // due to added delay, dynamic batches are timed out on arrival of every 2nd msg (per batch) - verify(sqs).sendMessageBatch(request("even", entries[0], entries[2])); - verify(sqs).sendMessageBatch(request("uneven", entries[1], entries[3])); - verify(sqs).sendMessageBatch(request("even", entries[4])); + // Nominally dynamic batches are timed out on arrival of every 2nd message per batch + // ([0,2], [1,3], [4]). The exact grouping depends on wall clock time and is unreliable on + // loaded machines (#38946), so verify timing-independent invariants instead: per queue, + // consecutive messages are >= 200ms apart, so expired batches are flushed on append once + // reaching 2 entries, requiring at least 3 batches overall. + Map> requests = captureBatchRequests(atLeast(3)); + assertThat(requests.keySet()).containsExactlyInAnyOrder("even", "uneven"); + assertMessageBodies(requestsFor(requests, "even"), range(0, 5).filter(i -> i % 2 == 0)); + assertMessageBodies(requestsFor(requests, "uneven"), range(0, 5).filter(i -> i % 2 == 1)); + requests.values().stream() + .flatMap(List::stream) + .forEach(req -> assertThat(req.entries().size()).isLessThanOrEqualTo(2)); } @Test @@ -360,14 +382,19 @@ public void testWriteBatchesToDynamicWithStrictTimeout() { p.run().waitUntilFinish(); - SendMessageBatchRequestEntry[] entries = entries(range(0, 5)); - // using strict timeouts batches, batches are timed out by a separate thread before any 2nd - // entry - verify(sqs).sendMessageBatch(request("even", entries[0])); - verify(sqs).sendMessageBatch(request("uneven", entries[1])); - verify(sqs).sendMessageBatch(request("even", entries[2])); - verify(sqs).sendMessageBatch(request("uneven", entries[3])); - verify(sqs).sendMessageBatch(request("even", entries[4])); + // Nominally the separate timeout thread flushes every batch before a 2nd entry arrives + // (5 singleton batches). The exact grouping depends on wall clock time and is unreliable + // on loaded machines (#38946), so verify timing-independent invariants instead. Expired + // batches are also flushed on append (independently of the timeout thread) and per queue + // consecutive messages are >= 200ms apart, so no batch can exceed 2 entries and at least + // 3 batches are required overall. + Map> requests = captureBatchRequests(atLeast(3)); + assertThat(requests.keySet()).containsExactlyInAnyOrder("even", "uneven"); + assertMessageBodies(requestsFor(requests, "even"), range(0, 5).filter(i -> i % 2 == 0)); + assertMessageBodies(requestsFor(requests, "uneven"), range(0, 5).filter(i -> i % 2 == 1)); + requests.values().stream() + .flatMap(List::stream) + .forEach(req -> assertThat(req.entries().size()).isLessThanOrEqualTo(2)); } @Test @@ -406,6 +433,32 @@ private SendMessageBatchRequest anyRequest() { return any(); } + /** Captures all batch requests, verifying the given mode, and groups them by queue url. */ + private Map> captureBatchRequests( + org.mockito.verification.VerificationMode mode) { + ArgumentCaptor captor = + ArgumentCaptor.forClass(SendMessageBatchRequest.class); + verify(sqs, mode).sendMessageBatch(captor.capture()); + return captor.getAllValues().stream().collect(groupingBy(SendMessageBatchRequest::queueUrl)); + } + + /** Returns the (non-null) batch requests for a queue, failing if there were none. */ + private List requestsFor( + Map> requests, String queue) { + return checkStateNotNull(requests.get(queue), "no batch requests for queue %s", queue); + } + + /** Asserts that the requests contain exactly the expected message bodies, each exactly once. */ + private void assertMessageBodies(List requests, IntStream expectedMsgs) { + assertThat( + requests.stream() + .flatMap(req -> req.entries().stream()) + .map(SendMessageBatchRequestEntry::messageBody) + .collect(toList())) + .containsExactlyInAnyOrderElementsOf( + expectedMsgs.mapToObj(Integer::toString).collect(toList())); + } + private SendMessageBatchRequest request(String queue, SendMessageBatchRequestEntry... entries) { return SendMessageBatchRequest.builder() .queueUrl(queue) diff --git a/sdks/python/.yapfignore b/sdks/python/.yapfignore index ca33e35e8fea..d37beaaac266 100644 --- a/sdks/python/.yapfignore +++ b/sdks/python/.yapfignore @@ -21,6 +21,7 @@ apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_messages.py apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py apache_beam/io/gcp/internal/clients/storage/storage_v1_messages.py apache_beam/coders/proto2_coder_test_messages_pb2.py +apache_beam/io/components/rate_limit_pb2.py apache_beam/portability/api/* # Avoid excessive wrapping. diff --git a/sdks/python/apache_beam/coders/coder_impl.pxd b/sdks/python/apache_beam/coders/coder_impl.pxd index 89730c17fe1d..e64177e6fd34 100644 --- a/sdks/python/apache_beam/coders/coder_impl.pxd +++ b/sdks/python/apache_beam/coders/coder_impl.pxd @@ -135,6 +135,7 @@ cdef class TimestampCoderImpl(StreamCoderImpl): cdef list small_ints cdef class VarIntCoderImpl(StreamCoderImpl): + @cython.locals(ivalue=libc.stdint.int64_t) cpdef bytes encode(self, value) diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py index e6b1fb6c37cd..0bded25e05d2 100644 --- a/sdks/python/apache_beam/coders/coder_impl.py +++ b/sdks/python/apache_beam/coders/coder_impl.py @@ -90,20 +90,6 @@ is_compiled = False fits_in_64_bits = lambda x: -(1 << 63) <= x <= (1 << 63) - 1 - -def _as_signed_int64(value): - # type: (int) -> int - - """Folds a uint64 to the signed int64 with the same bits (and VarInt - encoding), which the Cython int64_t stream params accept. Values outside the - 64-bit range raise here so the pure-Python path matches Cython's overflow.""" - if (1 << 63) <= value < (1 << 64): - return int(value) - (1 << 64) - if not fits_in_64_bits(value): - raise OverflowError("%d is out of range for a 64-bit integer." % value) - return value - - if TYPE_CHECKING or SLOW_STREAM: from .slow_stream import ByteCountingOutputStream from .slow_stream import InputStream as create_InputStream @@ -1058,14 +1044,12 @@ class VarIntCoderImpl(StreamCoderImpl): def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None try: - # Fold uint64 values into signed int64 so Cython doesn't overflow. - out.write_var_int64(_as_signed_int64(value)) + out.write_var_int64(value) except OverflowError as e: raise OverflowError( f"Integer value '{value}' is out of the encodable range for " - f"VarIntCoder. This coder is limited to 64-bit integers: the " - f"signed range -(2**63) to 2**63 - 1, plus unsigned values up to " - f"2**64 - 1 which share the same wire encoding. " + f"VarIntCoder. This coder is limited to values that fit " + f"within a 64-bit signed integer (-(2**63) to 2**63 - 1). " f"Original error: {e}") from e def decode_from_stream(self, in_stream, nested): @@ -1073,11 +1057,9 @@ def decode_from_stream(self, in_stream, nested): return in_stream.read_var_int64() def encode(self, value): - # Compare as a Python object: a uint64 value overflows the int64_t cast - # the compiled fast path used to do here. Non-small values (including - # uint64) fall through to encode_to_stream, which folds them. - if 0 <= value < len(small_ints): - return small_ints[value] + ivalue = value # type cast + if 0 <= ivalue < len(small_ints): + return small_ints[ivalue] return StreamCoderImpl.encode(self, value) def decode(self, encoded): @@ -1091,11 +1073,11 @@ def estimate_size(self, value, nested=False): # type: (Any, bool) -> int # Note that VarInts are encoded the same way regardless of nesting. try: - return get_varint_size(_as_signed_int64(value)) + return get_varint_size(value) except OverflowError as e: raise OverflowError( f"Cannot estimate size for integer value '{value}'. " - f"Value is out of the range for VarIntCoder (64-bit integer). " + f"Value is out of the range for VarIntCoder (64-bit signed integer). " f"Original error: {e}") from e diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py index 939e44aff468..422d494b61c7 100644 --- a/sdks/python/apache_beam/coders/coders_test_common.py +++ b/sdks/python/apache_beam/coders/coders_test_common.py @@ -437,24 +437,6 @@ def test_varint_coder(self): for k in range(0, int(math.log(MAX_64_BIT_INT))) ]) - def test_varint_coder_uint64(self): - # uint64 values [2**63, 2**64) must encode like the signed int64 with the - # same bits instead of overflowing Cython's int64_t. Decoding is signed, - # matching Java's VarIntCoder. - coder = coders.VarIntCoder() - impl = coder.get_impl() - for v in [1 << 63, (1 << 63) + 12345, (1 << 64) - 1]: - signed_twin = v - (1 << 64) - encoded = coder.encode(v) - self.assertEqual(encoded, coder.encode(signed_twin)) - self.assertEqual(impl.estimate_size(v), len(encoded)) - self.assertEqual(coder.decode(encoded), signed_twin) - - # Values outside the 64-bit range stay out of range on both paths. - for v in [1 << 64, (1 << 70), -(1 << 63) - 1]: - with self.assertRaises(OverflowError): - coder.encode(v) - def test_varint32_coder(self): # Small ints. self.check_coder(coders.VarInt32Coder(), *range(-10, 10)) diff --git a/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding.py b/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding.py new file mode 100644 index 000000000000..c0c7b180a175 --- /dev/null +++ b/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding.py @@ -0,0 +1,263 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Batch image embedding pipeline using MLTransform. + +The pipeline reads image URIs from a text file, decodes images with Pillow, +generates SentenceTransformers image embeddings through MLTransform, and writes +results to BigQuery using batch file loads. +""" + +import argparse +import hashlib +import io +import logging +import time +from collections.abc import Iterable +from typing import Any + +import apache_beam as beam +from apache_beam.io.filesystems import FileSystems +from apache_beam.ml.transforms.base import MLTransform +from apache_beam.ml.transforms.embeddings.huggingface import SentenceTransformerEmbeddings +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.options.pipeline_options import SetupOptions +from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.runners.runner import PipelineResult +from PIL import Image + +IMAGE_COLUMN = 'image' +IMAGE_ID_COLUMN = 'image_id' +IMAGE_URI_COLUMN = 'image_uri' + +DEFAULT_IMAGE_MODEL_NAME = 'clip-ViT-B-32' +DEFAULT_ACCELERATOR = 'type:nvidia-tesla-t4;count:1;install-nvidia-driver' +DEFAULT_EMBEDDING_MIN_RAM = '16GB' + +OUTPUT_TABLE_SCHEMA = { + 'fields': [ + { + 'name': 'image_id', 'type': 'STRING' + }, + { + 'name': 'image_uri', 'type': 'STRING' + }, + { + 'name': 'model_name', 'type': 'STRING' + }, + { + 'name': 'embedding', 'type': 'FLOAT64', 'mode': 'REPEATED' + }, + { + 'name': 'embedding_dim', 'type': 'INT64' + }, + { + 'name': 'infer_ms', 'type': 'INT64' + }, + ] +} + + +def now_millis() -> int: + return int(time.time() * 1000) + + +def sha1_hex(value: str) -> str: + return hashlib.sha1(value.encode('utf-8')).hexdigest() + + +def filter_empty_uri(uri: str) -> Iterable[str]: + uri = uri.strip() + if uri: + yield uri + + +def load_image_from_uri(uri: str) -> bytes: + with FileSystems.open(uri) as file: + return file.read() + + +def decode_pil(image_bytes: bytes) -> Image.Image: + with Image.open(io.BytesIO(image_bytes)) as image: + image = image.convert('RGB') + image.load() + return image + + +class ReadImage(beam.DoFn): + def process(self, uri: str) -> Iterable[dict[str, Any]]: + image_id = sha1_hex(uri) + try: + yield { + IMAGE_ID_COLUMN: image_id, + IMAGE_URI_COLUMN: uri, + IMAGE_COLUMN: decode_pil(load_image_from_uri(uri)), + } + except Exception as exc: + logging.warning( + 'Failed to read or decode image %s (%s): %s', image_id, uri, exc) + + +def _as_dict(row: Any) -> dict[str, Any]: + if hasattr(row, 'as_dict'): + return row.as_dict() + if hasattr(row, '_asdict'): + return row._asdict() + return dict(row) + + +def embedding_to_list(value: Any) -> list[float]: + if hasattr(value, 'tolist'): + return value.tolist() + return [float(item) for item in value] + + +class FormatImageEmbeddingOutput(beam.DoFn): + def __init__(self, model_name: str): + self.model_name = model_name + + def process(self, row: Any) -> Iterable[dict[str, Any]]: + row = _as_dict(row) + embedding = embedding_to_list(row[IMAGE_COLUMN]) + yield { + IMAGE_ID_COLUMN: row[IMAGE_ID_COLUMN], + IMAGE_URI_COLUMN: row[IMAGE_URI_COLUMN], + 'model_name': self.model_name, + 'embedding': embedding, + 'embedding_dim': len(embedding), + 'infer_ms': now_millis(), + } + + +def _str_to_bool(value: str) -> bool: + if value.lower() == 'true': + return True + if value.lower() == 'false': + return False + raise argparse.ArgumentTypeError( + f'"true" or "false" expected, got "{value}" instead.') + + +def parse_known_args(argv): + parser = argparse.ArgumentParser() + parser.add_argument('--mode', default='batch', choices=['batch']) + parser.add_argument( + '--input', required=True, help='Path to a text file with image URIs.') + parser.add_argument( + '--output_table', required=True, help='BigQuery table for embeddings.') + parser.add_argument( + '--publish_to_big_query', + type=_str_to_bool, + default=True, + help='Whether to write embedding rows to BigQuery.') + parser.add_argument( + '--artifact_location', + required=True, + help='Path where MLTransform artifacts are written.') + parser.add_argument( + '--pretrained_model_name', + default=DEFAULT_IMAGE_MODEL_NAME, + help='SentenceTransformers image model name.') + parser.add_argument( + '--device', + default='CPU', + choices=['CPU', 'GPU'], + help='Device used by SentenceTransformers on the worker.') + parser.add_argument( + '--min_batch_size', + type=int, + default=8, + help='Minimum Beam inference batch size.') + parser.add_argument( + '--max_batch_size', + type=int, + default=64, + help='Maximum Beam inference batch size.') + parser.add_argument( + '--embedding_accelerator', + default=DEFAULT_ACCELERATOR, + help='GPU accelerator resource hint for the MLTransform embedding step.') + parser.add_argument( + '--embedding_min_ram', + default=DEFAULT_EMBEDDING_MIN_RAM, + help='CPU right-fitting min RAM resource hint for the embedding step.') + return parser.parse_known_args(argv) + + +def run( + argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: + known_args, pipeline_args = parse_known_args(argv) + pipeline_options = PipelineOptions(pipeline_args) + pipeline_options.view_as(SetupOptions).save_main_session = save_main_session + pipeline_options.view_as(StandardOptions).streaming = False + + device = 'cuda' if known_args.device == 'GPU' else 'cpu' + embedding_transform = SentenceTransformerEmbeddings( + model_name=known_args.pretrained_model_name, + columns=[IMAGE_COLUMN], + image_model=True, + min_batch_size=known_args.min_batch_size, + max_batch_size=known_args.max_batch_size, + load_model_args={'device': device}, + inference_args={ + 'convert_to_numpy': True, + 'show_progress_bar': False, + }) + + ml_transform = MLTransform( + write_artifact_location=known_args.artifact_location).with_transform( + embedding_transform) + if known_args.device == 'GPU' and known_args.embedding_accelerator: + ml_transform = ml_transform.with_resource_hints( + accelerator=known_args.embedding_accelerator) + elif known_args.embedding_min_ram: + ml_transform = ml_transform.with_resource_hints( + min_ram=known_args.embedding_min_ram) + + pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) + rows = ( + pipeline + | 'ReadImageURIs' >> beam.io.ReadFromText(known_args.input) + | 'FilterEmptyURIs' >> beam.FlatMap(filter_empty_uri) + | 'ReshuffleBeforeEmbedding' >> beam.Reshuffle() + | 'ReadImages' >> beam.ParDo(ReadImage())) + results = ( + rows + | 'MLTransformImageEmbeddings' >> ml_transform + | 'FormatOutput' >> beam.ParDo( + FormatImageEmbeddingOutput( + model_name=known_args.pretrained_model_name))) + + if known_args.publish_to_big_query: + _ = ( + results + | 'WriteToBigQuery' >> beam.io.WriteToBigQuery( + known_args.output_table, + schema=OUTPUT_TABLE_SCHEMA, + write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, + create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, + method=beam.io.WriteToBigQuery.Method.FILE_LOADS)) + + result = pipeline.run() + if not test_pipeline: + result.wait_until_finish() + return result + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + run() diff --git a/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding_test.py b/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding_test.py new file mode 100644 index 000000000000..2dd8a46c043e --- /dev/null +++ b/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding_test.py @@ -0,0 +1,105 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import io +import unittest + +import apache_beam as beam +from apache_beam.io.gcp import bigquery_tools +from PIL import Image + +from apache_beam.examples.ml_transform import mltransform_image_embedding + + +class FakeArray: + def __init__(self, value): + self.value = value + + def tolist(self): + return self.value + + +class MLTransformImageEmbeddingTest(unittest.TestCase): + def test_filter_empty_uri(self): + self.assertEqual( + list(mltransform_image_embedding.filter_empty_uri(' ')), []) + self.assertEqual( + list(mltransform_image_embedding.filter_empty_uri(' gs://bucket/img ')), + ['gs://bucket/img']) + + def test_sha1_hex(self): + self.assertEqual( + mltransform_image_embedding.sha1_hex('gs://bucket/img.jpg'), + 'f3557ebc47344e0189d39dcad3642c8e8afbc10a') + + def test_decode_pil_returns_rgb_image(self): + image = Image.new('RGBA', (2, 2), color=(255, 0, 0, 255)) + image_bytes = io.BytesIO() + image.save(image_bytes, format='PNG') + + decoded = mltransform_image_embedding.decode_pil(image_bytes.getvalue()) + + self.assertEqual(decoded.mode, 'RGB') + self.assertEqual(decoded.size, (2, 2)) + + def test_embedding_to_list(self): + self.assertEqual( + mltransform_image_embedding.embedding_to_list(FakeArray([1, 2.5])), + [1, 2.5]) + + def test_format_output_with_dict(self): + row = { + 'image_id': 'abc', + 'image_uri': 'gs://bucket/img.jpg', + 'image': FakeArray([0.1, 0.2, 0.3]), + } + + output = next( + mltransform_image_embedding.FormatImageEmbeddingOutput( + 'clip-ViT-B-32').process(row)) + + self.assertEqual(output['image_id'], 'abc') + self.assertEqual(output['image_uri'], 'gs://bucket/img.jpg') + self.assertEqual(output['model_name'], 'clip-ViT-B-32') + self.assertEqual(output['embedding'], [0.1, 0.2, 0.3]) + self.assertEqual(output['embedding_dim'], 3) + self.assertIn('infer_ms', output) + + def test_output_table_schema_marks_embedding_as_repeated(self): + schema = bigquery_tools.get_dict_table_schema( + mltransform_image_embedding.OUTPUT_TABLE_SCHEMA) + embedding_field = next( + field for field in schema['fields'] if field['name'] == 'embedding') + self.assertEqual(embedding_field['type'], 'FLOAT64') + self.assertEqual(embedding_field['mode'], 'REPEATED') + + def test_format_output_with_beam_row(self): + row = beam.Row( + image_id='abc', + image_uri='gs://bucket/img.jpg', + image=FakeArray([0.1, 0.2]), + ) + + output = next( + mltransform_image_embedding.FormatImageEmbeddingOutput( + 'clip-ViT-B-32').process(row)) + + self.assertEqual(output['embedding_dim'], 2) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding.py b/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding.py new file mode 100644 index 000000000000..b692ac322e09 --- /dev/null +++ b/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding.py @@ -0,0 +1,186 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Batch text embedding pipeline using MLTransform. + +The pipeline reads text lines, generates sentence-transformer embeddings through +MLTransform, and writes JSONL output to a sharded text sink. +""" + +import argparse +import hashlib +import json +import logging +from collections.abc import Iterable +from typing import Any + +import apache_beam as beam +from apache_beam.ml.transforms.base import MLTransform +from apache_beam.ml.transforms.embeddings.huggingface import SentenceTransformerEmbeddings +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.options.pipeline_options import SetupOptions +from apache_beam.runners.runner import PipelineResult + +ID_COLUMN = 'id' +RAW_TEXT_COLUMN = 'raw_text' +TEXT_COLUMN = 'text' + +DEFAULT_MODEL_NAME = 'sentence-transformers/all-MiniLM-L6-v2' + + +def _str_to_bool(value: str) -> bool: + if value.lower() == 'true': + return True + if value.lower() == 'false': + return False + raise argparse.ArgumentTypeError( + f'"true" or "false" expected, got "{value}" instead.') + + +def parse_known_args(argv): + parser = argparse.ArgumentParser() + parser.add_argument( + '--input', required=True, help='Path to the input text file.') + parser.add_argument( + '--output', + required=True, + help='Output prefix for sharded JSONL embedding results.') + parser.add_argument( + '--artifact_location', + required=True, + help='Path where MLTransform artifacts are written.') + parser.add_argument( + '--model_name', + default=DEFAULT_MODEL_NAME, + help='SentenceTransformers model name.') + parser.add_argument( + '--min_batch_size', + type=int, + default=16, + help='Minimum Beam inference batch size.') + parser.add_argument( + '--max_batch_size', + type=int, + default=128, + help='Maximum Beam inference batch size.') + parser.add_argument( + '--model_batch_size', + type=int, + default=32, + help='Batch size passed to SentenceTransformer.encode.') + parser.add_argument( + '--device', + default='CPU', + choices=['CPU', 'GPU'], + help='Device used by SentenceTransformers on the worker.') + parser.add_argument( + '--large_model', + type=_str_to_bool, + default=False, + help='Whether RunInference should share the model across processes.') + return parser.parse_known_args(argv) + + +def text_to_record(line: str) -> Iterable[dict[str, str]]: + text = line.strip() + if not text: + return + + yield { + ID_COLUMN: hashlib.sha256(text.encode('utf-8')).hexdigest(), + RAW_TEXT_COLUMN: text, + TEXT_COLUMN: text, + } + + +def _as_dict(row: Any) -> dict[str, Any]: + if hasattr(row, 'as_dict'): + return row.as_dict() + if hasattr(row, '_asdict'): + return row._asdict() + return dict(row) + + +def embedding_to_list(value: Any) -> list[float]: + if hasattr(value, 'tolist'): + return value.tolist() + return [float(item) for item in value] + + +class FormatEmbeddingOutput(beam.DoFn): + def __init__(self, model_name: str): + self.model_name = model_name + + def process(self, row: Any) -> Iterable[str]: + row = _as_dict(row) + embedding = embedding_to_list(row[TEXT_COLUMN]) + yield json.dumps({ + ID_COLUMN: row[ID_COLUMN], + 'model_name': self.model_name, + RAW_TEXT_COLUMN: row[RAW_TEXT_COLUMN], + 'embedding': embedding, + 'embedding_dim': len(embedding) if isinstance(embedding, list) else 0, + }, + sort_keys=True) + + +def run( + argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: + known_args, pipeline_args = parse_known_args(argv) + pipeline_options = PipelineOptions(pipeline_args) + pipeline_options.view_as(SetupOptions).save_main_session = save_main_session + + device = 'cuda' if known_args.device == 'GPU' else 'cpu' + embedding_transform = SentenceTransformerEmbeddings( + model_name=known_args.model_name, + columns=[TEXT_COLUMN], + min_batch_size=known_args.min_batch_size, + max_batch_size=known_args.max_batch_size, + large_model=known_args.large_model, + load_model_args={'device': device}, + inference_args={ + 'batch_size': known_args.model_batch_size, + 'convert_to_numpy': True, + 'show_progress_bar': False, + }) + + pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) + records = ( + pipeline + | 'ReadTextLines' >> beam.io.ReadFromText(known_args.input) + | 'ToEmbeddingRecords' >> beam.FlatMap(text_to_record)) + embedded = ( + records + | 'MLTransformTextEmbeddings' >> MLTransform( + write_artifact_location=known_args.artifact_location).with_transform( + embedding_transform)) + _ = ( + embedded + | 'FormatOutput' >> beam.ParDo( + FormatEmbeddingOutput(model_name=known_args.model_name)) + | 'WriteOutput' >> beam.io.WriteToText( + known_args.output, file_name_suffix='.jsonl')) + + result = pipeline.run() + if not test_pipeline: + result.wait_until_finish() + return result + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + run() diff --git a/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding_test.py b/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding_test.py new file mode 100644 index 000000000000..dc22cc42b680 --- /dev/null +++ b/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding_test.py @@ -0,0 +1,94 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import hashlib +import json +import unittest + +import apache_beam as beam + +from apache_beam.examples.ml_transform import mltransform_text_embedding + + +class FakeArray: + def __init__(self, value): + self.value = value + + def tolist(self): + return self.value + + +class MLTransformTextEmbeddingTest(unittest.TestCase): + def test_text_to_record_skips_empty_lines(self): + self.assertEqual(list(mltransform_text_embedding.text_to_record(' ')), []) + + def test_text_to_record_preserves_text_and_adds_stable_id(self): + records = list( + mltransform_text_embedding.text_to_record(' Apache Beam is fun ')) + + self.assertEqual(len(records), 1) + expected_id = hashlib.sha256(b'Apache Beam is fun').hexdigest() + self.assertEqual( + records[0], + { + 'id': expected_id, + 'raw_text': 'Apache Beam is fun', + 'text': 'Apache Beam is fun', + }) + + def test_embedding_to_list(self): + self.assertEqual( + mltransform_text_embedding.embedding_to_list(FakeArray([1, 2.5])), + [1, 2.5]) + + def test_format_embedding_output_with_dict(self): + row = { + 'id': 'abc', + 'raw_text': 'hello', + 'text': FakeArray([0.1, 0.2, 0.3]), + } + + output = next( + mltransform_text_embedding.FormatEmbeddingOutput('test-model').process( + row)) + + self.assertEqual( + json.loads(output), + { + 'id': 'abc', + 'raw_text': 'hello', + 'model_name': 'test-model', + 'embedding': [0.1, 0.2, 0.3], + 'embedding_dim': 3, + }) + + def test_format_embedding_output_with_beam_row(self): + row = beam.Row( + id='abc', + raw_text='hello', + text=FakeArray([0.1, 0.2]), + ) + + output = next( + mltransform_text_embedding.FormatEmbeddingOutput('test-model').process( + row)) + + self.assertEqual(json.loads(output)['embedding_dim'], 2) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/io/components/rate_limit.proto b/sdks/python/apache_beam/io/components/rate_limit.proto new file mode 100644 index 000000000000..dccf65f931ab --- /dev/null +++ b/sdks/python/apache_beam/io/components/rate_limit.proto @@ -0,0 +1,64 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Minimal, self-contained subset of the Envoy Rate Limit Service protocol, +// used by EnvoyRateLimiter in rate_limiter.py. Only the fields the client +// touches are declared. Field numbers MUST match Envoy's rls.proto and +// ratelimit.proto so this is wire-compatible with a real Envoy RLS server +// (protobuf carries only field numbers and types on the wire, not message or +// package names). This lets Beam use its standard protobuf/grpcio stack +// instead of pulling in envoy-data-plane and its betterproto dependency. +// +// Regenerate rate_limit_pb2.py after editing this file: +// cd sdks/python +// python -m grpc_tools.protoc -I. --python_out=. \ +// apache_beam/io/components/rate_limit.proto +// then prepend the Apache license header and remove the generated +// `runtime_version` guard so the module stays compatible with the full +// protobuf runtime range Beam supports (see setup.py). + +syntax = "proto3"; + +package apache_beam.io.components.ratelimit; + +import "google/protobuf/duration.proto"; + +message RateLimitDescriptor { + message Entry { + string key = 1; + string value = 2; + } + repeated Entry entries = 1; +} + +message RateLimitRequest { + string domain = 1; + repeated RateLimitDescriptor descriptors = 2; + uint32 hits_addend = 3; +} + +message RateLimitResponse { + enum Code { + UNKNOWN = 0; + OK = 1; + OVER_LIMIT = 2; + } + message DescriptorStatus { + Code code = 1; + google.protobuf.Duration duration_until_reset = 4; + } + Code overall_code = 1; + repeated DescriptorStatus statuses = 2; +} diff --git a/sdks/python/apache_beam/io/components/rate_limit_pb2.py b/sdks/python/apache_beam/io/components/rate_limit_pb2.py new file mode 100644 index 000000000000..1a7af8278575 --- /dev/null +++ b/sdks/python/apache_beam/io/components/rate_limit_pb2.py @@ -0,0 +1,63 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: apache_beam/io/components/rate_limit.proto +# +# Regenerate with (from sdks/python): +# python -m grpc_tools.protoc -I. --python_out=. \ +# apache_beam/io/components/rate_limit.proto +# then re-apply this license header and delete the generated +# `runtime_version` guard so the module works across the full protobuf +# runtime range Beam supports (see setup.py). + +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + +from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n*apache_beam/io/components/rate_limit.proto\x12#apache_beam.io.components.ratelimit\x1a\x1egoogle/protobuf/duration.proto\"\x8b\x01\n\x13RateLimitDescriptor\x12O\n\x07\x65ntries\x18\x01 \x03(\x0b\x32>.apache_beam.io.components.ratelimit.RateLimitDescriptor.Entry\x1a#\n\x05\x45ntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t\"\x86\x01\n\x10RateLimitRequest\x12\x0e\n\x06\x64omain\x18\x01 \x01(\t\x12M\n\x0b\x64\x65scriptors\x18\x02 \x03(\x0b\x32\x38.apache_beam.io.components.ratelimit.RateLimitDescriptor\x12\x13\n\x0bhits_addend\x18\x03 \x01(\r\"\x87\x03\n\x11RateLimitResponse\x12Q\n\x0coverall_code\x18\x01 \x01(\x0e\x32;.apache_beam.io.components.ratelimit.RateLimitResponse.Code\x12Y\n\x08statuses\x18\x02 \x03(\x0b\x32G.apache_beam.io.components.ratelimit.RateLimitResponse.DescriptorStatus\x1a\x96\x01\n\x10\x44\x65scriptorStatus\x12I\n\x04\x63ode\x18\x01 \x01(\x0e\x32;.apache_beam.io.components.ratelimit.RateLimitResponse.Code\x12\x37\n\x14\x64uration_until_reset\x18\x04 \x01(\x0b\x32\x19.google.protobuf.Duration\"+\n\x04\x43ode\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x06\n\x02OK\x10\x01\x12\x0e\n\nOVER_LIMIT\x10\x02\x62\x06proto3' +) + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages( + DESCRIPTOR, 'apache_beam.io.components.rate_limit_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_RATELIMITDESCRIPTOR']._serialized_start = 116 + _globals['_RATELIMITDESCRIPTOR']._serialized_end = 255 + _globals['_RATELIMITDESCRIPTOR_ENTRY']._serialized_start = 220 + _globals['_RATELIMITDESCRIPTOR_ENTRY']._serialized_end = 255 + _globals['_RATELIMITREQUEST']._serialized_start = 258 + _globals['_RATELIMITREQUEST']._serialized_end = 392 + _globals['_RATELIMITRESPONSE']._serialized_start = 395 + _globals['_RATELIMITRESPONSE']._serialized_end = 786 + _globals['_RATELIMITRESPONSE_DESCRIPTORSTATUS']._serialized_start = 591 + _globals['_RATELIMITRESPONSE_DESCRIPTORSTATUS']._serialized_end = 741 + _globals['_RATELIMITRESPONSE_CODE']._serialized_start = 743 + _globals['_RATELIMITRESPONSE_CODE']._serialized_end = 786 +# @@protoc_insertion_point(module_scope) diff --git a/sdks/python/apache_beam/io/components/rate_limiter.py b/sdks/python/apache_beam/io/components/rate_limiter.py index 4d6bf2f17a8c..8e6779737951 100644 --- a/sdks/python/apache_beam/io/components/rate_limiter.py +++ b/sdks/python/apache_beam/io/components/rate_limiter.py @@ -27,13 +27,9 @@ import time import grpc -from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptor -from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptorEntry -from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitRequest -from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponse -from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseCode from apache_beam.io.components import adaptive_throttler +from apache_beam.io.components import rate_limit_pb2 from apache_beam.metrics import Metrics _LOGGER = logging.getLogger(__name__) @@ -122,19 +118,18 @@ def __init__( self._lock = threading.Lock() class RateLimitServiceStub(object): - """ - Wrapper for gRPC stub to be compatible with envoy_data_plane messages. - - The envoy-data-plane package uses 'betterproto' which generates async stubs - for 'grpclib'. As Beam uses standard synchronous 'grpcio', - RateLimitServiceStub is a bridge class to use the betterproto Message types - (RateLimitRequest) with a standard grpcio Channel. + """ + Minimal gRPC stub for the Envoy Rate Limit Service ShouldRateLimit method. + + The method path is fixed by the Envoy RLS proto, so we bind it by hand + against a standard synchronous grpcio Channel using the protobuf message + types from rate_limit_pb2. """ def __init__(self, channel): self.ShouldRateLimit = channel.unary_unary( '/envoy.service.ratelimit.v3.RateLimitService/ShouldRateLimit', - request_serializer=RateLimitRequest.SerializeToString, - response_deserializer=RateLimitResponse.FromString, + request_serializer=rate_limit_pb2.RateLimitRequest.SerializeToString, + response_deserializer=rate_limit_pb2.RateLimitResponse.FromString, ) def init_connection(self): @@ -171,10 +166,11 @@ def allow(self, hits_added: int = 1) -> bool: for d in self.descriptors: entries = [] for k, v in d.items(): - entries.append(RateLimitDescriptorEntry(key=k, value=v)) - proto_descriptors.append(RateLimitDescriptor(entries=entries)) + entries.append(rate_limit_pb2.RateLimitDescriptor.Entry(key=k, value=v)) + proto_descriptors.append( + rate_limit_pb2.RateLimitDescriptor(entries=entries)) - request = RateLimitRequest( + request = rate_limit_pb2.RateLimitRequest( domain=self.domain, descriptors=proto_descriptors, hits_addend=hits_added) @@ -205,11 +201,11 @@ def allow(self, hits_added: int = 1) -> bool: e) time.sleep(_RPC_RETRY_DELAY_SECONDS) - if response.overall_code == RateLimitResponseCode.OK: + if response.overall_code == rate_limit_pb2.RateLimitResponse.OK: self.requests_allowed.inc() throttled = True break - elif response.overall_code == RateLimitResponseCode.OVER_LIMIT: + elif response.overall_code == rate_limit_pb2.RateLimitResponse.OVER_LIMIT: self.requests_throttled.inc() # Ratelimit exceeded, sleep for duration until reset and retry # multiple rules can be set in the RLS config, so we need to find the @@ -217,10 +213,9 @@ def allow(self, hits_added: int = 1) -> bool: sleep_s = 0.0 if response.statuses: for status in response.statuses: - if status.code == RateLimitResponseCode.OVER_LIMIT: + if status.code == rate_limit_pb2.RateLimitResponse.OVER_LIMIT: dur = status.duration_until_reset - # duration_until_reset is converted to timedelta by betterproto - val = dur.total_seconds() + val = dur.ToTimedelta().total_seconds() if val > sleep_s: sleep_s = val diff --git a/sdks/python/apache_beam/io/components/rate_limiter_test.py b/sdks/python/apache_beam/io/components/rate_limiter_test.py index 24d30a1c5c93..490d3a04e2d5 100644 --- a/sdks/python/apache_beam/io/components/rate_limiter_test.py +++ b/sdks/python/apache_beam/io/components/rate_limiter_test.py @@ -16,16 +16,16 @@ # import unittest -from datetime import timedelta from unittest import mock import grpc -from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponse -from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseCode -from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseDescriptorStatus +from google.protobuf.duration_pb2 import Duration +from apache_beam.io.components import rate_limit_pb2 from apache_beam.io.components import rate_limiter +RateLimitResponse = rate_limit_pb2.RateLimitResponse + class EnvoyRateLimiterTest(unittest.TestCase): def setUp(self): @@ -45,7 +45,7 @@ def setUp(self): def test_allow_success(self, mock_channel): # Mock successful OK response mock_stub = mock.Mock() - mock_response = RateLimitResponse(overall_code=RateLimitResponseCode.OK) + mock_response = RateLimitResponse(overall_code=RateLimitResponse.OK) mock_stub.ShouldRateLimit.return_value = mock_response # Inject mock stub @@ -60,8 +60,7 @@ def test_allow_success(self, mock_channel): def test_allow_over_limit_retries_exceeded(self, mock_channel): # Mock OVER_LIMIT response mock_stub = mock.Mock() - mock_response = RateLimitResponse( - overall_code=RateLimitResponseCode.OVER_LIMIT) + mock_response = RateLimitResponse(overall_code=RateLimitResponse.OVER_LIMIT) mock_stub.ShouldRateLimit.return_value = mock_response self.limiter._stub = mock_stub @@ -86,7 +85,7 @@ def test_allow_over_limit_retries_exceeded(self, mock_channel): def test_allow_rpc_error_retry(self, mock_channel): # Mock RpcError then Success mock_stub = mock.Mock() - mock_response = RateLimitResponse(overall_code=RateLimitResponseCode.OK) + mock_response = RateLimitResponse(overall_code=RateLimitResponse.OK) # Side effect: Error, Error, Success error = grpc.RpcError() @@ -123,11 +122,11 @@ def test_extract_duration_from_response(self, mock_random, mock_channel): mock_stub = mock.Mock() # Valid until 5 seconds - status = RateLimitResponseDescriptorStatus( - code=RateLimitResponseCode.OVER_LIMIT, - duration_until_reset=timedelta(seconds=5)) + status = RateLimitResponse.DescriptorStatus( + code=RateLimitResponse.OVER_LIMIT, + duration_until_reset=Duration(seconds=5)) mock_response = RateLimitResponse( - overall_code=RateLimitResponseCode.OVER_LIMIT, statuses=[status]) + overall_code=RateLimitResponse.OVER_LIMIT, statuses=[status]) mock_stub.ShouldRateLimit.return_value = mock_response self.limiter._stub = mock_stub @@ -139,5 +138,44 @@ def test_extract_duration_from_response(self, mock_random, mock_channel): mock_sleep.assert_called_with(5.0) +class RateLimitWireFormatTest(unittest.TestCase): + """Pins the on-the-wire layout of the vendored rate_limit_pb2 messages. + + Wire compatibility with a real Envoy Rate Limit Service depends solely on + field numbers and types (protobuf carries neither message nor package names + on the wire), so these must stay in lockstep with Envoy's rls.proto and + ratelimit.proto. The mock-based tests above would pass even if a field were + renumbered; these golden-byte assertions fail if that ever happens. + """ + def test_request_wire_layout(self): + request = rate_limit_pb2.RateLimitRequest( + domain='d', + descriptors=[ + rate_limit_pb2.RateLimitDescriptor( + entries=[ + rate_limit_pb2.RateLimitDescriptor.Entry( + key='k', value='v') + ]) + ], + hits_addend=1) + # domain=1 (LEN "d"); descriptors=2 (LEN {entries=1 (LEN {key=1 "k", + # value=2 "v"})}); hits_addend=3 (VARINT 1). + self.assertEqual( + request.SerializeToString().hex(), '0a016412080a060a016b1201761801') + + def test_descriptor_status_wire_layout(self): + status = rate_limit_pb2.RateLimitResponse.DescriptorStatus( + code=rate_limit_pb2.RateLimitResponse.OVER_LIMIT, + duration_until_reset=Duration(seconds=5)) + # code=1 (VARINT OVER_LIMIT=2); duration_until_reset=4 (LEN + # Duration{seconds=1 (VARINT 5)}). + self.assertEqual(status.SerializeToString().hex(), '080222020805') + + def test_response_code_enum_values(self): + self.assertEqual(int(rate_limit_pb2.RateLimitResponse.UNKNOWN), 0) + self.assertEqual(int(rate_limit_pb2.RateLimitResponse.OK), 1) + self.assertEqual(int(rate_limit_pb2.RateLimitResponse.OVER_LIMIT), 2) + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt b/sdks/python/apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt new file mode 100644 index 000000000000..77019383dd77 --- /dev/null +++ b/sdks/python/apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt @@ -0,0 +1,20 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +sentence-transformers~=5.0.0 +Pillow>=9.0.0 +google-cloud-monitoring>=2.27.0 diff --git a/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_image_embedding_benchmark.py b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_image_embedding_benchmark.py new file mode 100644 index 000000000000..2fb207707ee6 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_image_embedding_benchmark.py @@ -0,0 +1,128 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import logging +from datetime import datetime +from typing import Optional + +from apache_beam.examples.ml_transform import mltransform_image_embedding +from apache_beam.options.pipeline_options import DebugOptions +from apache_beam.options.pipeline_options import GoogleCloudOptions +from apache_beam.options.pipeline_options import SetupOptions +from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.options.pipeline_options import WorkerOptions +from apache_beam.testing.load_tests import dataflow_cost_consts as costs +from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark +from apache_beam.testing.load_tests.load_test import LoadTestOptions + + +class MLTransformImageEmbeddingOptions( + LoadTestOptions, + StandardOptions, + GoogleCloudOptions, + WorkerOptions, + DebugOptions, + SetupOptions, +): + @classmethod + def _add_argparse_args(cls, parser): + parser.add_argument('--mode', default='batch') + parser.add_argument('--input', default='') + parser.add_argument('--input_file', default='') + parser.add_argument('--output_table', default='') + parser.add_argument('--artifact_location', default='') + parser.add_argument( + '--pretrained_model_name', + default=mltransform_image_embedding.DEFAULT_IMAGE_MODEL_NAME) + parser.add_argument('--device', default='CPU') + parser.add_argument('--min_batch_size', type=int, default=8) + parser.add_argument('--max_batch_size', type=int, default=64) + parser.add_argument( + '--embedding_accelerator', + default=mltransform_image_embedding.DEFAULT_ACCELERATOR) + parser.add_argument( + '--embedding_min_ram', + default=mltransform_image_embedding.DEFAULT_EMBEDDING_MIN_RAM) + + +class MLTransformImageEmbeddingBenchmarkTest(DataflowCostBenchmark): + options_class = MLTransformImageEmbeddingOptions + + def __init__(self): + self.metrics_namespace = 'BeamML_MLTransform' + super().__init__( + metrics_namespace=self.metrics_namespace, + pcollection='FormatOutput.out0') + self.opts = self.pipeline.get_pipeline_options().view_as( + MLTransformImageEmbeddingOptions) + if self.opts.device == 'GPU': + self.gpu = costs.Accelerator.T4 + self.multi_worker_pool = True + + def _get_throughput_metrics( + self, + project: str, + job_id: str, + start_time: datetime, + end_time: datetime, + pcollection_name: Optional[str] = None, + ) -> dict[str, float]: + return self._get_throughput_metrics_with_pcollection_fallback( + project, + job_id, + start_time, + end_time, + pcollection_candidates=[ + self.pcollection, + 'MLTransformImageEmbeddings.out0', + 'MLTransformImageEmbeddings/RunInference.out0', + 'MLTransformImageEmbeddings/RunInference/' + 'BeamML_RunInference_Postprocess-0.out0', + 'WriteToBigQuery/BigQueryBatchFileLoads/TriggerLoadJobs.out0', + ], + pcollection_name=pcollection_name) + + def test(self): + input_path = self.opts.input or self.opts.input_file + if not input_path: + raise RuntimeError('Please provide --input or --input_file.') + if not self.opts.output_table: + raise RuntimeError('Please provide --output_table.') + if not self.opts.artifact_location: + raise RuntimeError('Please provide --artifact_location.') + + extra_opts = { + 'mode': self.opts.mode, + 'input': input_path, + 'output_table': self.opts.output_table, + 'artifact_location': self.opts.artifact_location, + 'pretrained_model_name': self.opts.pretrained_model_name, + 'device': self.opts.device, + 'min_batch_size': self.opts.min_batch_size, + 'max_batch_size': self.opts.max_batch_size, + 'embedding_accelerator': self.opts.embedding_accelerator, + 'embedding_min_ram': self.opts.embedding_min_ram, + } + + self.result = mltransform_image_embedding.run( + self.pipeline.get_full_options_as_args(**extra_opts), + test_pipeline=self.pipeline) + + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO) + MLTransformImageEmbeddingBenchmarkTest().run() diff --git a/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_one_hot_encoding_benchmark.py b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_one_hot_encoding_benchmark.py index e80fca633352..3998a492c907 100644 --- a/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_one_hot_encoding_benchmark.py +++ b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_one_hot_encoding_benchmark.py @@ -24,9 +24,8 @@ """ import logging - -from google.cloud import monitoring_v3 -from google.protobuf.duration_pb2 import Duration +from datetime import datetime +from typing import Optional from apache_beam.examples.ml_transform import mltransform_one_hot_encoding from apache_beam.options.pipeline_options import DebugOptions @@ -128,68 +127,22 @@ def _get_throughput_metrics( self, project: str, job_id: str, - start_time: str, - end_time: str, - pcollection_name: str | None = None, + start_time: datetime, + end_time: datetime, + pcollection_name: Optional[str] = None, ) -> dict[str, float]: """Get throughput metrics with runner-v2-friendly fallbacks.""" - candidate_pcollections = [] - if pcollection_name: - candidate_pcollections.append(pcollection_name) - candidate_pcollections.extend([ - self.pcollection, - 'MLTransform.out0', - 'FormatOutput.out0', - ]) - - # Deduplicate while preserving order. - seen = set() - unique_candidates = [] - for name in candidate_pcollections: - if name and name not in seen: - seen.add(name) - unique_candidates.append(name) - - for name in unique_candidates: - metrics = super()._get_throughput_metrics( - project, job_id, start_time, end_time, pcollection_name=name) - if (metrics.get('AvgThroughputBytes', 0) > 0 or - metrics.get('AvgThroughputElements', 0) > 0): - return metrics - - # Final fallback: aggregate job-level throughput without pcollection label. - interval = monitoring_v3.TimeInterval( - start_time=start_time, end_time=end_time) - aggregation = monitoring_v3.Aggregation( - alignment_period=Duration(seconds=60), - per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_MEAN) - requests = { - "Bytes": monitoring_v3.ListTimeSeriesRequest( - name=f"projects/{project}", - filter=( - 'metric.type="dataflow.googleapis.com/job/estimated_byte_count" ' - f'AND metric.labels.job_id="{job_id}"'), - interval=interval, - aggregation=aggregation), - "Elements": monitoring_v3.ListTimeSeriesRequest( - name=f"projects/{project}", - filter=( - 'metric.type="dataflow.googleapis.com/job/element_count" ' - f'AND metric.labels.job_id="{job_id}"'), - interval=interval, - aggregation=aggregation), - } - - fallback_metrics = {} - for key, req in requests.items(): - time_series = self.monitoring_client.list_time_series(request=req) - values = [ - point.value.double_value for series in time_series - for point in series.points - ] - fallback_metrics[f"AvgThroughput{key}"] = ( - sum(values) / len(values) if values else 0.0) - return fallback_metrics + return self._get_throughput_metrics_with_pcollection_fallback( + project, + job_id, + start_time, + end_time, + pcollection_candidates=[ + self.pcollection, + 'MLTransform.out0', + 'FormatOutput.out0', + ], + pcollection_name=pcollection_name) if __name__ == '__main__': diff --git a/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_text_embedding_benchmark.py b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_text_embedding_benchmark.py new file mode 100644 index 000000000000..a03a2578fd4a --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_text_embedding_benchmark.py @@ -0,0 +1,117 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import logging +from datetime import datetime +from typing import Optional + +from apache_beam.examples.ml_transform import mltransform_text_embedding +from apache_beam.options.pipeline_options import DebugOptions +from apache_beam.options.pipeline_options import GoogleCloudOptions +from apache_beam.options.pipeline_options import SetupOptions +from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.options.pipeline_options import WorkerOptions +from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark +from apache_beam.testing.load_tests.load_test import LoadTestOptions + + +class MLTransformTextEmbeddingOptions( + LoadTestOptions, + StandardOptions, + GoogleCloudOptions, + WorkerOptions, + DebugOptions, + SetupOptions, +): + @classmethod + def _add_argparse_args(cls, parser): + parser.add_argument('--input', default='') + parser.add_argument('--input_file', default='') + parser.add_argument('--output', default='') + parser.add_argument('--artifact_location', default='') + parser.add_argument( + '--model_name', default=mltransform_text_embedding.DEFAULT_MODEL_NAME) + parser.add_argument('--min_batch_size', type=int, default=16) + parser.add_argument('--max_batch_size', type=int, default=128) + parser.add_argument('--model_batch_size', type=int, default=32) + parser.add_argument('--device', default='CPU') + parser.add_argument('--large_model', default='false') + + +class MLTransformTextEmbeddingBenchmarkTest(DataflowCostBenchmark): + options_class = MLTransformTextEmbeddingOptions + + def __init__(self): + self.metrics_namespace = 'BeamML_MLTransform' + super().__init__( + metrics_namespace=self.metrics_namespace, + pcollection='FormatOutput.out0') + self.opts = self.pipeline.get_pipeline_options().view_as( + MLTransformTextEmbeddingOptions) + + def _get_throughput_metrics( + self, + project: str, + job_id: str, + start_time: datetime, + end_time: datetime, + pcollection_name: Optional[str] = None, + ) -> dict[str, float]: + return self._get_throughput_metrics_with_pcollection_fallback( + project, + job_id, + start_time, + end_time, + pcollection_candidates=[ + self.pcollection, + 'MLTransformTextEmbeddings.out0', + 'MLTransformTextEmbeddings/RunInference.out0', + 'MLTransformTextEmbeddings/RunInference/' + 'BeamML_RunInference_Postprocess-0.out0', + 'WriteOutput/Write/WriteImpl/FinalizeWrite.out0', + ], + pcollection_name=pcollection_name) + + def test(self): + input_path = self.opts.input or self.opts.input_file + if not input_path: + raise RuntimeError('Please provide --input or --input_file.') + if not self.opts.output: + raise RuntimeError('Please provide --output.') + if not self.opts.artifact_location: + raise RuntimeError('Please provide --artifact_location.') + + extra_opts = { + 'input': input_path, + 'output': self.opts.output, + 'artifact_location': self.opts.artifact_location, + 'model_name': self.opts.model_name, + 'min_batch_size': self.opts.min_batch_size, + 'max_batch_size': self.opts.max_batch_size, + 'model_batch_size': self.opts.model_batch_size, + 'device': self.opts.device, + 'large_model': self.opts.large_model, + } + + self.result = mltransform_text_embedding.run( + self.pipeline.get_full_options_as_args(**extra_opts), + test_pipeline=self.pipeline) + + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO) + MLTransformTextEmbeddingBenchmarkTest().run() diff --git a/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py b/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py index 4d412eb228ec..8c64026c016c 100644 --- a/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py +++ b/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py @@ -62,11 +62,13 @@ def __init__( is_streaming: bool = False, gpu: Optional[costs.Accelerator] = None, pcollection: str = 'ProcessOutput.out0', - subscription: Optional[str] = None): + subscription: Optional[str] = None, + multi_worker_pool: bool = False): self.is_streaming = is_streaming self.gpu = gpu self.pcollection = pcollection self.subscription = subscription + self.multi_worker_pool = multi_worker_pool super().__init__(metrics_namespace=metrics_namespace) self.dataflow_client = DataflowApplicationClient( self.pipeline.get_pipeline_options()) @@ -140,22 +142,32 @@ def _process_metrics_list(self, system_metrics[metric.name] = entry.committed or 0.0 return system_metrics + def _get_message_text(self, message) -> Optional[str]: + return getattr(message, 'messageText', None) or getattr( + message, 'message_text', None) + def _get_worker_time_interval( self, job_id: str) -> tuple[Optional[datetime], Optional[datetime]]: """Extracts worker start and stop times from job messages.""" + if self.multi_worker_pool: + return self._get_multi_worker_pool_time_interval(job_id) + return self._get_single_worker_pool_time_interval(job_id) + + def _get_single_worker_pool_time_interval( + self, job_id: str) -> tuple[Optional[datetime], Optional[datetime]]: start_time, end_time = None, None page_token = None all_messages = [] last_message_time = None while True: messages, page_token = self.dataflow_client.list_messages( - job_id=job_id, - start_time=None, - end_time=None, - page_token=page_token, - minimum_importance='JOB_MESSAGE_DEBUG') + job_id=job_id, + start_time=None, + end_time=None, + page_token=page_token, + minimum_importance='JOB_MESSAGE_DEBUG') for message in messages: - text = message.message_text + text = self._get_message_text(message) if getattr(message, 'time', None): last_message_time = message.time if text: @@ -182,6 +194,51 @@ def _get_worker_time_interval( len(all_messages)) return start_time, end_time + def _get_multi_worker_pool_time_interval( + self, job_id: str) -> tuple[Optional[datetime], Optional[datetime]]: + """Extract worker interval across multiple pools (e.g. CPU + GPU).""" + start_time, end_time = None, None + page_token = None + message_count = 0 + last_message_time = None + while True: + messages, page_token = self.dataflow_client.list_messages( + job_id=job_id, + start_time=None, + end_time=None, + page_token=page_token, + minimum_importance='JOB_MESSAGE_DEBUG') + for message in messages: + message_count += 1 + text = self._get_message_text(message) + if getattr(message, 'time', None): + last_message_time = message.time + if not text: + continue + if self.WORKER_START_PATTERN.search(text): + if start_time is None or message.time < start_time: + start_time = message.time + logging.info('Matched WORKER_START_PATTERN: %r', text) + if self.WORKER_STOP_PATTERN.search(text): + if end_time is None or message.time > end_time: + end_time = message.time + logging.info('Matched WORKER_STOP_PATTERN: %r', text) + if not page_token: + break + if start_time and not end_time and last_message_time: + end_time = last_message_time + logging.info( + 'Using last job message time as end_time for multi-pool job: %s', + end_time) + if not start_time or not end_time: + logging.warning( + 'Could not determine multi-pool worker time interval. ' + 'start_time=%s, end_time=%s, total messages=%d', + start_time, + end_time, + message_count) + return start_time, end_time + def _get_throughput_metrics( self, project: str, @@ -255,6 +312,78 @@ def _point_numeric_value(point) -> float: return metrics + def _get_job_level_throughput_metrics( + self, + project: str, + job_id: str, + start_time: datetime, + end_time: datetime, + ) -> dict[str, float]: + """Query Cloud Monitoring for job-level throughput (no PCollection label).""" + interval = monitoring_v3.TimeInterval( + start_time=start_time, end_time=end_time) + aggregation = monitoring_v3.Aggregation( + alignment_period=Duration(seconds=60), + per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_MEAN) + requests = { + 'Bytes': monitoring_v3.ListTimeSeriesRequest( + name=f'projects/{project}', + filter=( + 'metric.type="dataflow.googleapis.com/job/estimated_byte_count" ' + f'AND metric.labels.job_id="{job_id}"'), + interval=interval, + aggregation=aggregation), + 'Elements': monitoring_v3.ListTimeSeriesRequest( + name=f'projects/{project}', + filter=( + 'metric.type="dataflow.googleapis.com/job/element_count" ' + f'AND metric.labels.job_id="{job_id}"'), + interval=interval, + aggregation=aggregation), + } + metrics = {} + for key, request in requests.items(): + values = [ + point.value.double_value + for series in self.monitoring_client.list_time_series( + request=request) for point in series.points + ] + metrics[f'AvgThroughput{key}'] = sum(values) / len( + values) if values else 0.0 + return metrics + + def _get_throughput_metrics_with_pcollection_fallback( + self, + project: str, + job_id: str, + start_time: datetime, + end_time: datetime, + pcollection_candidates: list[str], + pcollection_name: Optional[str] = None, + ) -> dict[str, float]: + """Try PCollection throughput metrics, then fall back to job-level metrics.""" + candidates = [] + if pcollection_name: + candidates.append(pcollection_name) + for candidate in pcollection_candidates: + if candidate and candidate not in candidates: + candidates.append(candidate) + + for candidate in candidates: + metrics = DataflowCostBenchmark._get_throughput_metrics( + self, project, job_id, start_time, end_time, candidate) + if (metrics.get('AvgThroughputBytes', 0) > 0 or + metrics.get('AvgThroughputElements', 0) > 0): + logging.info('Using throughput metrics for PCollection %s', candidate) + return metrics + + logging.warning( + 'No PCollection-level throughput metrics found for candidates %s. ' + 'Falling back to job-level Dataflow throughput metrics.', + candidates) + return self._get_job_level_throughput_metrics( + project, job_id, start_time, end_time) + def _get_streaming_throughput_metrics( self, project: str, start_time: datetime, end_time: datetime) -> dict[str, float]: diff --git a/sdks/python/container/ml/py310/base_image_requirements.txt b/sdks/python/container/ml/py310/base_image_requirements.txt index bd851ba73543..bec06ec1befb 100644 --- a/sdks/python/container/ml/py310/base_image_requirements.txt +++ b/sdks/python/container/ml/py310/base_image_requirements.txt @@ -35,7 +35,6 @@ attrs==26.1.0 backports.tarfile==1.2.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b7 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -52,7 +51,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy-data-plane==0.2.6 exceptiongroup==1.3.1 execnet==2.1.2 fastavro==1.12.2 diff --git a/sdks/python/container/ml/py310/gpu_image_requirements.txt b/sdks/python/container/ml/py310/gpu_image_requirements.txt index a9846379404d..4f9e02edf772 100644 --- a/sdks/python/container/ml/py310/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py310/gpu_image_requirements.txt @@ -37,7 +37,6 @@ attrs==26.1.0 backports.tarfile==1.2.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b7 blake3==1.0.9 bs4==0.0.2 build==1.5.0 @@ -67,7 +66,6 @@ docker==7.1.0 docstring_parser==0.18.0 einops==0.8.2 email-validator==2.3.0 -envoy-data-plane==0.2.6 exceptiongroup==1.3.1 execnet==2.1.2 fastapi==0.138.1 diff --git a/sdks/python/container/ml/py311/base_image_requirements.txt b/sdks/python/container/ml/py311/base_image_requirements.txt index 28278e54d349..3f9a18099fa2 100644 --- a/sdks/python/container/ml/py311/base_image_requirements.txt +++ b/sdks/python/container/ml/py311/base_image_requirements.txt @@ -34,7 +34,6 @@ attrs==26.1.0 backports.tarfile==1.2.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -51,7 +50,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/container/ml/py311/gpu_image_requirements.txt b/sdks/python/container/ml/py311/gpu_image_requirements.txt index 88fc7aad84c8..7cb148c46f98 100644 --- a/sdks/python/container/ml/py311/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py311/gpu_image_requirements.txt @@ -36,7 +36,6 @@ attrs==26.1.0 backports.tarfile==1.2.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 blake3==1.0.9 bs4==0.0.2 build==1.5.0 @@ -66,7 +65,6 @@ docker==7.1.0 docstring_parser==0.18.0 einops==0.8.2 email-validator==2.3.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastapi==0.138.1 fastapi-cli==0.0.27 diff --git a/sdks/python/container/ml/py312/base_image_requirements.txt b/sdks/python/container/ml/py312/base_image_requirements.txt index 7e7c5728a4a5..265696dc7e8d 100644 --- a/sdks/python/container/ml/py312/base_image_requirements.txt +++ b/sdks/python/container/ml/py312/base_image_requirements.txt @@ -33,7 +33,6 @@ astunparse==1.6.3 attrs==26.1.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -50,7 +49,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/container/ml/py312/gpu_image_requirements.txt b/sdks/python/container/ml/py312/gpu_image_requirements.txt index f563d4457dcd..7cc83b10ca0f 100644 --- a/sdks/python/container/ml/py312/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py312/gpu_image_requirements.txt @@ -35,7 +35,6 @@ astunparse==1.6.3 attrs==26.1.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 blake3==1.0.9 bs4==0.0.2 build==1.5.0 @@ -65,7 +64,6 @@ docker==7.1.0 docstring_parser==0.18.0 einops==0.8.2 email-validator==2.3.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastapi==0.138.1 fastapi-cli==0.0.27 diff --git a/sdks/python/container/ml/py313/base_image_requirements.txt b/sdks/python/container/ml/py313/base_image_requirements.txt index 390de543ed5f..67f459b4fb2f 100644 --- a/sdks/python/container/ml/py313/base_image_requirements.txt +++ b/sdks/python/container/ml/py313/base_image_requirements.txt @@ -33,7 +33,6 @@ astunparse==1.6.3 attrs==26.1.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -50,7 +49,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/container/py310/base_image_requirements.txt b/sdks/python/container/py310/base_image_requirements.txt index 5c4bac1c2b2a..85e7615cfdba 100644 --- a/sdks/python/container/py310/base_image_requirements.txt +++ b/sdks/python/container/py310/base_image_requirements.txt @@ -33,7 +33,6 @@ attrs==26.1.0 backports.tarfile==1.2.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b7 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -50,7 +49,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy-data-plane==0.2.6 exceptiongroup==1.3.1 execnet==2.1.2 fastavro==1.12.2 diff --git a/sdks/python/container/py311/base_image_requirements.txt b/sdks/python/container/py311/base_image_requirements.txt index 3d9abe3dd9bc..b9ede50e41de 100644 --- a/sdks/python/container/py311/base_image_requirements.txt +++ b/sdks/python/container/py311/base_image_requirements.txt @@ -32,7 +32,6 @@ attrs==26.1.0 backports.tarfile==1.2.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -49,7 +48,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/container/py312/base_image_requirements.txt b/sdks/python/container/py312/base_image_requirements.txt index d6324c02562b..4edcb6421100 100644 --- a/sdks/python/container/py312/base_image_requirements.txt +++ b/sdks/python/container/py312/base_image_requirements.txt @@ -31,7 +31,6 @@ asn1crypto==1.5.1 attrs==26.1.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -48,7 +47,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/container/py313/base_image_requirements.txt b/sdks/python/container/py313/base_image_requirements.txt index 0cdce630cadc..a9728cd5e106 100644 --- a/sdks/python/container/py313/base_image_requirements.txt +++ b/sdks/python/container/py313/base_image_requirements.txt @@ -31,7 +31,6 @@ asn1crypto==1.5.1 attrs==26.1.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -48,7 +47,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/container/py314/base_image_requirements.txt b/sdks/python/container/py314/base_image_requirements.txt index 66a95d3e331b..1598d940e93d 100644 --- a/sdks/python/container/py314/base_image_requirements.txt +++ b/sdks/python/container/py314/base_image_requirements.txt @@ -31,7 +31,6 @@ asn1crypto==1.5.1 attrs==26.1.0 beartype==0.22.9 beautifulsoup4==4.15.0 -betterproto==2.0.0b6 bs4==0.0.2 build==1.5.0 cachetools==6.2.6 @@ -48,7 +47,6 @@ distro==1.9.0 dnspython==2.8.0 docker==7.1.0 docstring_parser==0.18.0 -envoy_data_plane==1.0.3 execnet==2.1.2 fastavro==1.12.2 fasteners==0.20 diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index 41c7d6d60d78..a3a59e878c29 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -82,6 +82,7 @@ skip = [ "storage_v1_client.py", "storage_v1_messages.py", "proto2_coder_test_messages_pb2.py", + "rate_limit_pb2.py", "cloudbuild_v1_client.py", "cloudbuild_v1_messages.py", "boto3_client.py", diff --git a/sdks/python/scripts/run_lint.sh b/sdks/python/scripts/run_lint.sh index 58d75cd7bace..068c838952dd 100755 --- a/sdks/python/scripts/run_lint.sh +++ b/sdks/python/scripts/run_lint.sh @@ -57,6 +57,7 @@ EXCLUDED_GENERATED_FILES=( "apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py" "apache_beam/io/gcp/internal/clients/storage/storage_v1_messages.py" "apache_beam/coders/proto2_coder_test_messages_pb2.py" +"apache_beam/io/components/rate_limit_pb2.py" "apache_beam/runners/dataflow/internal/clients/cloudbuild/cloudbuild_v1_client.py" "apache_beam/runners/dataflow/internal/clients/cloudbuild/cloudbuild_v1_messages.py" "apache_beam/io/aws/clients/s3/boto3_client.py" diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 9a566b15d1d1..e764f935bea6 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -418,10 +418,6 @@ def get_portability_package_data(): ext_modules=extensions, install_requires=[ 'cryptography>=39.0.0,<49.0.0', - 'envoy-data-plane>=1.0.3,<2; python_version >= "3.11"', - # Newer version only work on Python 3.11. Versions 0.3 <= ver < 1.x - # conflict with other GCP dependencies. - 'envoy-data-plane<0.3.0; python_version < "3.11"', 'fastavro>=0.23.6,<2', 'fasteners>=0.3,<1.0', 'grpcio>=1.33.1,<2,!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*', # pylint: disable=line-too-long diff --git a/website/www/site/content/en/performance/_index.md b/website/www/site/content/en/performance/_index.md index a0eaba2aa0ec..347b7c1110c3 100644 --- a/website/www/site/content/en/performance/_index.md +++ b/website/www/site/content/en/performance/_index.md @@ -61,3 +61,6 @@ See the following pages for performance measures recorded when running various B - [Table Row Inference Sklearn Batch](/performance/tablerowinference) - [MLTransform Generate Vocab (batch)](/performance/mltransformvocab) - [MLTransform One-Hot Encoding](/performance/mltransformonehot) +- [MLTransform Text Embedding](/performance/mltransform-text-embedding) +- [MLTransform Image Embedding GPU](/performance/mltransform-image-embedding-gpu) +- [MLTransform Image Embedding CPU](/performance/mltransform-image-embedding-cpu) diff --git a/website/www/site/content/en/performance/mltransform-image-embedding-cpu/_index.md b/website/www/site/content/en/performance/mltransform-image-embedding-cpu/_index.md new file mode 100644 index 000000000000..6780093b244a --- /dev/null +++ b/website/www/site/content/en/performance/mltransform-image-embedding-cpu/_index.md @@ -0,0 +1,48 @@ +--- +title: "MLTransform Image Embedding CPU Performance" +--- + + + +# MLTransform Image Embedding CPU Performance + +**Model**: Sentence Transformers — clip-ViT-B-32 (image) +**Accelerator**: CPU with Dataflow Prime right-fitting (16 GB min RAM) +**Host**: Dataflow Prime with throughput-based autoscaling + +This batch pipeline reads image URIs from GCS, decodes images with Pillow, +generates image embeddings through `MLTransform` with +`SentenceTransformerEmbeddings(image_model=True)`, and writes results to +BigQuery using batch file loads. + +See the [glossary](/performance/glossary) for definitions. + +Full pipeline implementation is available +[here](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding.py). + +## What is the estimated cost to run the pipeline? + +{{< performance_looks io="mltransform-image-embedding-cpu" read_or_write="write" section="cost" >}} + +## How has various metrics changed when running the pipeline for different Beam SDK versions? + +{{< performance_looks io="mltransform-image-embedding-cpu" read_or_write="write" section="version" >}} + +## How has various metrics changed over time when running the pipeline? + +{{< performance_looks io="mltransform-image-embedding-cpu" read_or_write="write" section="date" >}} + +See also [MLTransform Image Embedding GPU](/performance/mltransform-image-embedding-gpu) +for the Tesla T4 GPU variant of this pipeline. diff --git a/website/www/site/content/en/performance/mltransform-image-embedding-gpu/_index.md b/website/www/site/content/en/performance/mltransform-image-embedding-gpu/_index.md new file mode 100644 index 000000000000..f5f5f1fee50c --- /dev/null +++ b/website/www/site/content/en/performance/mltransform-image-embedding-gpu/_index.md @@ -0,0 +1,48 @@ +--- +title: "MLTransform Image Embedding GPU Performance" +--- + + + +# MLTransform Image Embedding GPU Performance + +**Model**: Sentence Transformers — clip-ViT-B-32 (image) +**Accelerator**: NVIDIA Tesla T4 GPU +**Host**: Dataflow Prime with throughput-based autoscaling + +This batch pipeline reads image URIs from GCS, decodes images with Pillow, +generates image embeddings through `MLTransform` with +`SentenceTransformerEmbeddings(image_model=True)`, and writes results to +BigQuery using batch file loads. + +See the [glossary](/performance/glossary) for definitions. + +Full pipeline implementation is available +[here](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding.py). + +## What is the estimated cost to run the pipeline? + +{{< performance_looks io="mltransform-image-embedding-gpu" read_or_write="write" section="cost" >}} + +## How has various metrics changed when running the pipeline for different Beam SDK versions? + +{{< performance_looks io="mltransform-image-embedding-gpu" read_or_write="write" section="version" >}} + +## How has various metrics changed over time when running the pipeline? + +{{< performance_looks io="mltransform-image-embedding-gpu" read_or_write="write" section="date" >}} + +See also [MLTransform Image Embedding CPU](/performance/mltransform-image-embedding-cpu) +for the CPU right-fitting variant of this pipeline. diff --git a/website/www/site/content/en/performance/mltransform-text-embedding/_index.md b/website/www/site/content/en/performance/mltransform-text-embedding/_index.md new file mode 100644 index 000000000000..051fb88b6c08 --- /dev/null +++ b/website/www/site/content/en/performance/mltransform-text-embedding/_index.md @@ -0,0 +1,44 @@ +--- +title: "MLTransform Text Embedding Performance" +--- + + + +# MLTransform Text Embedding Performance + +**Model**: Sentence Transformers — all-MiniLM-L6-v2 +**Accelerator**: CPU only +**Host**: 10 × n1-standard-4 (4 vCPUs, 15 GB RAM) + +This batch pipeline reads text lines from GCS, generates sentence embeddings +through `MLTransform` with `SentenceTransformerEmbeddings`, and writes JSONL +results to GCS. + +See the [glossary](/performance/glossary) for definitions. + +Full pipeline implementation is available +[here](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding.py). + +## What is the estimated cost to run the pipeline? + +{{< performance_looks io="mltransform-text-embedding" read_or_write="write" section="cost" >}} + +## How has various metrics changed when running the pipeline for different Beam SDK versions? + +{{< performance_looks io="mltransform-text-embedding" read_or_write="write" section="version" >}} + +## How has various metrics changed over time when running the pipeline? + +{{< performance_looks io="mltransform-text-embedding" read_or_write="write" section="date" >}} diff --git a/website/www/site/data/performance.yaml b/website/www/site/data/performance.yaml index 42fb5bee92f5..c27dcc22f237 100644 --- a/website/www/site/data/performance.yaml +++ b/website/www/site/data/performance.yaml @@ -315,4 +315,51 @@ looks: title: AvgThroughputBytesPerSec by Version - id: Cn5FsXkdy2ZXCxCJshSCxcjsTW3TXf3c title: AvgThroughputElementsPerSec by Version - + mltransform-text-embedding: + write: + folder: 109 + cost: + - id: wm4mr82Q6MGcJZbvQF32MQB6vYMh3pBM + title: RunTime and EstimatedCost + date: + - id: dcCym4PGpqg3mTXsTXS5tfmTcJCVNTBD + title: AvgThroughputBytesPerSec by Date + - id: Q5sKJqB53Jkhy5TKCybfQgXjJnhp7pRQ + title: AvgThroughputElementsPerSec by Date + version: + - id: qdpHqnkzRPb2zqJr8sjCpW84KSQkhDHJ + title: AvgThroughputBytesPerSec by Version + - id: HcWgFVVct3xq2WjCMmY7Q8KcHGms7KDW + title: AvgThroughputElementsPerSec by Version + mltransform-image-embedding-gpu: + write: + folder: 110 + cost: + - id: zQZzvdcy3srprGq5qxqmRRN4GRHpdRDX + title: RunTime and EstimatedCost + date: + - id: X2wdRXcgsFG7sxfSQykhypFDkYtRjN2J + title: AvgThroughputBytesPerSec by Date + - id: zppHCwdN5qrDpyn2WQpJd8YmfQCPw4dk + title: AvgThroughputElementsPerSec by Date + version: + - id: FZXH7JKyFNWmrkgcFQZYxH7wywZgQpYF + title: AvgThroughputBytesPerSec by Version + - id: FQ3QYHgZTxvh3zPp3BcsGnmJx7yxfVzc + title: AvgThroughputElementsPerSec by Version + mltransform-image-embedding-cpu: + write: + folder: 111 + cost: + - id: dxnGs2xD2vNH4Y6HbFMDd5Hc2d5cqVWc + title: RunTime and EstimatedCost + date: + - id: tkwdB46WT8YDkR4Q37jJ2p7y5fGchmTC + title: AvgThroughputBytesPerSec by Date + - id: x3f5pRDvjKxVPy5Y6GjghhqSddYxprn3 + title: AvgThroughputElementsPerSec by Date + version: + - id: ZPHYbRdBMBSvd7RGS8wF8C8d8hysQmbH + title: AvgThroughputBytesPerSec by Version + - id: zjMRJssC6tNrtKPBgs8yPW8jHkbk8wjk + title: AvgThroughputElementsPerSec by Version