From 904f8898ab28e9f86b575fdc71262e45627cee9c Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Tue, 7 Jul 2026 14:57:28 +0400 Subject: [PATCH 1/9] Fix PostCommit Java DataflowV2 workflow (#39232) --- .github/workflows/beam_PostCommit_Java_DataflowV2.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From 4bf5a6afb5067653df37d586a42de4ac9cd592b9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 07:12:44 -0400 Subject: [PATCH 2/9] Bump github.com/aws/aws-sdk-go-v2/service/s3 in /sdks (#39227) Bumps [github.com/aws/aws-sdk-go-v2/service/s3](https://github.com/aws/aws-sdk-go-v2) from 1.104.2 to 1.105.0. - [Release notes](https://github.com/aws/aws-sdk-go-v2/releases) - [Commits](https://github.com/aws/aws-sdk-go-v2/compare/service/s3/v1.104.2...service/s3/v1.105.0) --- updated-dependencies: - dependency-name: github.com/aws/aws-sdk-go-v2/service/s3 dependency-version: 1.105.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- sdks/go.mod | 2 +- sdks/go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index a597578df528..9ffb77db1533 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -36,7 +36,7 @@ require ( 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/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 diff --git a/sdks/go.sum b/sdks/go.sum index ea12281ba155..21b26922de21 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -253,8 +253,8 @@ 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= From 4cde7ed933aeb3a74e903eb88a56d9550d0956d7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 07:13:20 -0400 Subject: [PATCH 3/9] Bump github.com/aws/aws-sdk-go-v2/config in /sdks (#39229) Bumps [github.com/aws/aws-sdk-go-v2/config](https://github.com/aws/aws-sdk-go-v2) from 1.32.27 to 1.32.28. - [Release notes](https://github.com/aws/aws-sdk-go-v2/releases) - [Commits](https://github.com/aws/aws-sdk-go-v2/compare/config/v1.32.27...config/v1.32.28) --- updated-dependencies: - dependency-name: github.com/aws/aws-sdk-go-v2/config dependency-version: 1.32.28 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- sdks/go.mod | 12 ++++++------ sdks/go.sum | 24 ++++++++++++------------ 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index 9ffb77db1533..1ed27ede6001 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -33,8 +33,8 @@ 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/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.30 github.com/aws/aws-sdk-go-v2/service/s3 v1.105.0 github.com/aws/smithy-go v1.27.3 @@ -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 @@ -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 diff --git a/sdks/go.sum b/sdks/go.sum index 21b26922de21..0cbeaa2174e4 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -207,12 +207,12 @@ 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= @@ -256,22 +256,22 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.43.0/go.mod h1:NXRKkiRF+erX2hnybnVU66 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= From 7502ae2d5293eb64f70255ee5d212bbb65765c8c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 07:14:00 -0400 Subject: [PATCH 4/9] Bump golang.org/x/text from 0.38.0 to 0.39.0 in /sdks (#39230) Bumps [golang.org/x/text](https://github.com/golang/text) from 0.38.0 to 0.39.0. - [Release notes](https://github.com/golang/text/releases) - [Commits](https://github.com/golang/text/compare/v0.38.0...v0.39.0) --- updated-dependencies: - dependency-name: golang.org/x/text dependency-version: 0.39.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- sdks/go.mod | 8 ++++---- sdks/go.sum | 16 ++++++++-------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index 1ed27ede6001..65c989c2a52f 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -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 @@ -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 ) @@ -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 0cbeaa2174e4..6afd872ebea0 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -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= From f35fa826f44fe0a18813bf49c039af601f04d3d1 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 08:32:26 -0400 Subject: [PATCH 5/9] Bump github.com/aws/aws-sdk-go-v2/feature/s3/manager in /sdks (#39231) Bumps [github.com/aws/aws-sdk-go-v2/feature/s3/manager](https://github.com/aws/aws-sdk-go-v2) from 1.22.30 to 1.22.31. - [Release notes](https://github.com/aws/aws-sdk-go-v2/releases) - [Commits](https://github.com/aws/aws-sdk-go-v2/compare/feature/s3/manager/v1.22.30...feature/s3/manager/v1.22.31) --- updated-dependencies: - dependency-name: github.com/aws/aws-sdk-go-v2/feature/s3/manager dependency-version: 1.22.31 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Derrick Williams --- sdks/go.mod | 2 +- sdks/go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/go.mod b/sdks/go.mod index 65c989c2a52f..aadeaf57ccf3 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -35,7 +35,7 @@ require ( github.com/aws/aws-sdk-go-v2 v1.42.1 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.30 + 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 diff --git a/sdks/go.sum b/sdks/go.sum index 6afd872ebea0..f5865534d2f9 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -219,8 +219,8 @@ github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQ 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= From bdc26675c91de1e7ca63f85d45b4c1fbdc97021f Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:53:36 +0300 Subject: [PATCH 6/9] Drop envoy-data-plane/betterproto dependency from Python SDK (#39213) EnvoyRateLimiter only needed a handful of protobuf message classes from envoy-data-plane, which pulls in the outdated betterproto==2.0.0b6 beta (a protobuf reimplementation) plus grpclib and transitive deps. That pin is the last blocker stopping Apache Airflow from un-suspending its Beam provider (apache/airflow#66952), and it forced a Python-version split in setup.py. The dependency was already fought rather than used: the RateLimitServiceStub bridge exists solely because betterproto emits async grpclib stubs that don't work with Beam's synchronous grpcio, and the wire is plain protobuf over grpcio regardless. Replace it with a minimal, self-contained rate_limit.proto compiled to a checked-in rate_limit_pb2.py (following the existing proto2_coder_test_messages_pb2.py precedent). Field numbers match Envoy's rls.proto/ratelimit.proto, so it stays wire-compatible with a real RLS server. This removes the conflict permanently for every downstream, deletes the py<3.11 split, and drops betterproto + grpclib from every container image. New wire-format tests pin the field numbers/enum values so a renumbering can't silently break live rate limiting (the mock-based tests would not). Note: the container base_image_requirements.txt files had the two direct packages removed to stay consistent with setup.py; a full `generatePythonRequirementsAll` regeneration should follow in CI to also prune now-orphaned transitives (grpclib, h2, multidict, ...) and refresh pins. Part of #37854 Unblocks apache/airflow#66952 Co-authored-by: Claude Opus 4.8 (1M context) --- CHANGES.md | 1 + sdks/python/.yapfignore | 1 + .../io/components/rate_limit.proto | 64 +++++++++++++++++++ .../io/components/rate_limit_pb2.py | 63 ++++++++++++++++++ .../apache_beam/io/components/rate_limiter.py | 39 +++++------ .../io/components/rate_limiter_test.py | 62 ++++++++++++++---- .../ml/py310/base_image_requirements.txt | 2 - .../ml/py310/gpu_image_requirements.txt | 2 - .../ml/py311/base_image_requirements.txt | 2 - .../ml/py311/gpu_image_requirements.txt | 2 - .../ml/py312/base_image_requirements.txt | 2 - .../ml/py312/gpu_image_requirements.txt | 2 - .../ml/py313/base_image_requirements.txt | 2 - .../py310/base_image_requirements.txt | 2 - .../py311/base_image_requirements.txt | 2 - .../py312/base_image_requirements.txt | 2 - .../py313/base_image_requirements.txt | 2 - .../py314/base_image_requirements.txt | 2 - sdks/python/pyproject.toml | 1 + sdks/python/scripts/run_lint.sh | 1 + sdks/python/setup.py | 4 -- 21 files changed, 198 insertions(+), 62 deletions(-) create mode 100644 sdks/python/apache_beam/io/components/rate_limit.proto create mode 100644 sdks/python/apache_beam/io/components/rate_limit_pb2.py 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/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/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/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 From c6e13b60fa49b84c71c9fc554ba395895b14d517 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 7 Jul 2026 10:18:50 -0400 Subject: [PATCH 7/9] Revert "Fix Python VarIntCoder OverflowError on uint64 values (#39047)" (#39218) This reverts commit e87f50e3a9a0d5650cfb53af805fa8691981d277. --- sdks/python/apache_beam/coders/coder_impl.pxd | 1 + sdks/python/apache_beam/coders/coder_impl.py | 34 +++++-------------- .../apache_beam/coders/coders_test_common.py | 18 ---------- 3 files changed, 9 insertions(+), 44 deletions(-) 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)) From 20494736bb92d57166ea031a9fb225a87ba1318f Mon Sep 17 00:00:00 2001 From: Abdelrahman Ibrahim Date: Tue, 7 Jul 2026 17:22:07 +0300 Subject: [PATCH 8/9] Add MLTransform embedding benchmarks (#38917) * Add MLTransform embedding benchmarks * resolved code review comments * Refactor DataflowCostBenchmark worker interval and throughput helpers --- ...m_Inference_Python_Benchmarks_Dataflow.yml | 36 +++ ..._MLTransform_Image_Embedding_CPU_Batch.txt | 42 +++ ..._MLTransform_Image_Embedding_GPU_Batch.txt | 44 +++ ...aflow_MLTransform_Text_Embedding_Batch.txt | 42 +++ .test-infra/tools/refresh_looker_metrics.py | 5 +- .../mltransform_image_embedding.py | 263 ++++++++++++++++++ .../mltransform_image_embedding_test.py | 105 +++++++ .../mltransform_text_embedding.py | 186 +++++++++++++ .../mltransform_text_embedding_test.py | 94 +++++++ ...transform_embedding_tests_requirements.txt | 20 ++ .../mltransform_image_embedding_benchmark.py | 128 +++++++++ .../mltransform_one_hot_encoding_benchmark.py | 79 ++---- .../mltransform_text_embedding_benchmark.py | 117 ++++++++ .../load_tests/dataflow_cost_benchmark.py | 143 +++++++++- .../www/site/content/en/performance/_index.md | 3 + .../mltransform-image-embedding-cpu/_index.md | 48 ++++ .../mltransform-image-embedding-gpu/_index.md | 48 ++++ .../mltransform-text-embedding/_index.md | 44 +++ website/www/site/data/performance.yaml | 49 +++- 19 files changed, 1424 insertions(+), 72 deletions(-) create mode 100644 .github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_CPU_Batch.txt create mode 100644 .github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Image_Embedding_GPU_Batch.txt create mode 100644 .github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Text_Embedding_Batch.txt create mode 100644 sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding.py create mode 100644 sdks/python/apache_beam/examples/ml_transform/mltransform_image_embedding_test.py create mode 100644 sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding.py create mode 100644 sdks/python/apache_beam/examples/ml_transform/mltransform_text_embedding_test.py create mode 100644 sdks/python/apache_beam/ml/transforms/mltransform_embedding_tests_requirements.txt create mode 100644 sdks/python/apache_beam/testing/benchmarks/inference/mltransform_image_embedding_benchmark.py create mode 100644 sdks/python/apache_beam/testing/benchmarks/inference/mltransform_text_embedding_benchmark.py create mode 100644 website/www/site/content/en/performance/mltransform-image-embedding-cpu/_index.md create mode 100644 website/www/site/content/en/performance/mltransform-image-embedding-gpu/_index.md create mode 100644 website/www/site/content/en/performance/mltransform-text-embedding/_index.md 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/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/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/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/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 From d1c1a681d83ada3edf51bf8642e2ff6a2811a045 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Tue, 7 Jul 2026 16:41:37 +0200 Subject: [PATCH 9/9] [aws2] Make SqsIOWriteBatchesTest timeout tests timing-independent (#38971) The four timeout-related tests asserted the exact grouping of messages into SendMessageBatch calls. Those groupings depend on wall-clock timing (the per-message Thread.sleep delay racing the configured batch timeout), so on loaded CI runners batches form differently and the strict verify(...).sendMessageBatch(request(exact entries)) checks fail with Mockito ArgumentsAreDifferent. Rewrite the assertions to verify timing-independent invariants instead: all expected message bodies are sent exactly once, no batch exceeds the size implied by the timeout cadence, and at least the minimum number of batches is produced. This still exercises the timeout-driven flushing (both synchronous and the strict separate-thread variant) without depending on exact wall-clock behavior. Fixes #38946 --- .../io/aws2/sqs/SqsIOWriteBatchesTest.java | 97 ++++++++++++++----- 1 file changed, 75 insertions(+), 22 deletions(-) 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)