diff --git a/.github/actions/setup-k8s-access/action.yml b/.github/actions/setup-k8s-access/action.yml index cb00c853738b..264aa77cfc88 100644 --- a/.github/actions/setup-k8s-access/action.yml +++ b/.github/actions/setup-k8s-access/action.yml @@ -68,7 +68,7 @@ runs: run: | kubectl config set-context --current --namespace=${{ steps.replace_namespace.outputs.TEST_NAMESPACE }} - name: Post cleanup - uses: pyTooling/Actions/with-post-step@v4.2.2 + uses: pytooling/actions/with-post-step@42e17fae05f224e5ac3d79d021a4e3577878efe5 # v4.2.2 with: main: echo "Post Cleanup" post: | diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json index 60eb7272a936..327fe1987c50 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 4 + "modification": 5 } \ No newline at end of file diff --git a/.github/trigger_files/beam_PreCommit_Flink_Container.json b/.github/trigger_files/beam_PreCommit_Flink_Container.json index bbdc3a3910ef..62ae7886c573 100644 --- a/.github/trigger_files/beam_PreCommit_Flink_Container.json +++ b/.github/trigger_files/beam_PreCommit_Flink_Container.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 3 + "modification": 4 } diff --git a/.github/workflows/README.md b/.github/workflows/README.md index dcc885feec9d..41287872ee25 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -58,6 +58,7 @@ On top of normal Actions workflow steps, all new CI workflows (excluding release 4) A set of GitHub token permissions 5) Concurrency Groups 6) Comment Triggering Support +7) Allowed Actions (See below) Each of these is described in more detail below. ## Name and Phrase @@ -165,6 +166,14 @@ In order to make it easier for non-committers to interact with workflows, workfl **Note:** Comment triggering is found not scalable ([#28909](https://github.com/apache/beam/issues/28909)) and is currently limited to a subset of suites. For more information see the [Running Workflows Manually](#running-workflows-manually) section. +## Allowed Actions + +GitHub Actions under the trusted namespaces `actions/*`, `github/*`, and `apache/*` are automatically approved. Any third-party actions outside of these namespaces must be explicitly approved and added to the ASF-wide allowlist before use. + +Additionally, third-party actions must be pinned to a specific Git commit SHA (e.g., `zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02`) instead of a version tag or branch name to comply with ASF security policy. + +For details on the allowlist and how to request new actions, see the [ASF Infrastructure Actions Allowlist Check Documentation](https://github.com/apache/infrastructure-actions/blob/main/allowlist-check/README.md). + # Testing new workflows or workflow updates ## Testing New Workflows diff --git a/.github/workflows/beam_Infrastructure_UsersPermissions.yml b/.github/workflows/beam_Infrastructure_UsersPermissions.yml index 0e806034b5ca..b188e5ff875b 100644 --- a/.github/workflows/beam_Infrastructure_UsersPermissions.yml +++ b/.github/workflows/beam_Infrastructure_UsersPermissions.yml @@ -52,7 +52,7 @@ jobs: - name: Setup gcloud uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db - name: Install Terraform - uses: hashicorp/setup-terraform@v4 + uses: hashicorp/setup-terraform@dfe3c3f87815947d99a8997f908cb6525fc44e9e # v4.0.1 with: terraform_version: 1.12.2 - name: Initialize Terraform diff --git a/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml b/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml index 383db7e137fd..d0a5deb8c4ab 100644 --- a/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml +++ b/.github/workflows/beam_PostCommit_Java_IO_Performance_Tests.yml @@ -64,7 +64,7 @@ jobs: matrix: job_name: ["beam_PostCommit_Java_IO_Performance_Tests"] job_phrase: ["Run Java PostCommit IO Performance Tests"] - test_case: ["GCSPerformanceTest", "BigTablePerformanceTest", "BigQueryStorageApiStreamingPerformanceTest"] + test_case: ["IcebergPerformanceTest", "GCSPerformanceTest", "BigTablePerformanceTest", "BigQueryStorageApiStreamingPerformanceTest"] steps: - uses: actions/checkout@v7 with: diff --git a/.github/workflows/beam_PostCommit_Python_Arm.yml b/.github/workflows/beam_PostCommit_Python_Arm.yml index 0cb03894b255..961dfbb2c53a 100644 --- a/.github/workflows/beam_PostCommit_Python_Arm.yml +++ b/.github/workflows/beam_PostCommit_Python_Arm.yml @@ -54,7 +54,8 @@ jobs: beam_PostCommit_Python_Arm: name: ${{ matrix.job_name }} ${{ matrix.python_version }} runs-on: [self-hosted, ubuntu-24.04, main] - timeout-minutes: 480 + # Python 3.14 multiarch container builds take much longer than other versions. + timeout-minutes: ${{ matrix.python_version == '3.14' && 600 || 480 }} strategy: fail-fast: false matrix: diff --git a/.github/workflows/beam_PreCommit_GHA.yml b/.github/workflows/beam_PreCommit_GHA.yml index 7ad6b5e67583..fa3d2c9605b4 100644 --- a/.github/workflows/beam_PreCommit_GHA.yml +++ b/.github/workflows/beam_PreCommit_GHA.yml @@ -80,15 +80,19 @@ jobs: comment_phrase: ${{ matrix.job_phrase }} github_token: ${{ secrets.GITHUB_TOKEN }} github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) - - name: Run zizmor - uses: zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02 # v0.6.0 - with: - advanced-security: true - name: Setup environment uses: ./.github/actions/setup-environment-action with: java-version: default go-version: default + python-version: default + # For details on the ASF GHA allowlist, see: https://github.com/apache/infrastructure-actions/blob/main/allowlist-check/README.md + - name: Validate GHA Allowlist + uses: apache/infrastructure-actions/allowlist-check@main # zizmor: ignore[unpinned-uses] + - name: Run zizmor + uses: zizmorcore/zizmor-action@6599ee8b7a49aef6a770f63d261d214911a7ce02 # v0.6.0 + with: + advanced-security: true - name: run GHA PreCommit script uses: ./.github/actions/gradle-command-self-hosted-action with: diff --git a/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml b/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml index c37ddf7e4f3d..6bf05e99bd23 100644 --- a/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml +++ b/.github/workflows/beam_Publish_Beam_SDK_Snapshots.yml @@ -62,12 +62,10 @@ jobs: - "python:container:py311:docker" - "python:container:py312:docker" - "python:container:py313:docker" - - "python:container:py314:docker" - "python:container:distroless:py310:docker" - "python:container:distroless:py311:docker" - "python:container:distroless:py312:docker" - "python:container:distroless:py313:docker" - - "python:container:distroless:py314:docker" - "python:container:ml:py310:docker" - "python:container:ml:py311:docker" - "python:container:ml:py312:docker" @@ -90,8 +88,16 @@ jobs: # This is needed to run pipelines that use the default environment at HEAD, for example, when a # pipeline uses an expansion service built from HEAD. run: | - BEAM_VERSION_LINE=$(cat gradle.properties | grep "sdk_version") - echo "BEAM_VERSION=${BEAM_VERSION_LINE#*sdk_version=}" >> $GITHUB_ENV + BEAM_VERSION=$(grep -E '^sdk_version=' gradle.properties | cut -d= -f2-) + if [[ ! "${BEAM_VERSION}" =~ ^[0-9A-Za-z._-]+$ ]]; then + echo "Invalid sdk_version in gradle.properties: ${BEAM_VERSION}" + exit 1 + fi + { + echo "BEAM_VERSION<> "$GITHUB_ENV" - name: Set latest tag only on master branch if: github.ref == 'refs/heads/master' run: echo "LATEST_TAG=,latest" >> $GITHUB_ENV @@ -127,3 +133,75 @@ jobs: -Pcontainer-architecture-list=arm64,amd64 \ -Ppush-containers \ -Pdocker-pull-licenses + + beam_Publish_Beam_SDK_Snapshots_py314: + needs: beam_Publish_Beam_SDK_Snapshots + if: | + always() && + !cancelled() && + (github.event_name == 'workflow_dispatch' || + (github.event_name == 'schedule' && github.repository == 'apache/beam')) + runs-on: [self-hosted, ubuntu-24.04, main] + timeout-minutes: 300 + name: ${{ matrix.job_name }} (${{ matrix.container_task }}) + strategy: + fail-fast: false + matrix: + job_name: ["beam_Publish_Beam_SDK_Snapshots"] + job_phrase: ["N/A"] + container_task: + - "python:container:py314:docker" + - "python:container:distroless:py314:docker" + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + - name: Free Disk Space (Ubuntu) + uses: jlumbroso/free-disk-space@v1.3.1 + - name: Setup repository + uses: ./.github/actions/setup-action + with: + comment_phrase: ${{ matrix.job_phrase }} + github_token: ${{ secrets.GITHUB_TOKEN }} + github_job: ${{ matrix.job_name }} (${{ matrix.container_task }}) + - name: Find Beam Version + run: | + BEAM_VERSION=$(grep -E '^sdk_version=' gradle.properties | cut -d= -f2-) + if [[ ! "${BEAM_VERSION}" =~ ^[0-9A-Za-z._-]+$ ]]; then + echo "Invalid sdk_version in gradle.properties: ${BEAM_VERSION}" + exit 1 + fi + { + echo "BEAM_VERSION<> "$GITHUB_ENV" + - name: Set latest tag only on master branch + if: github.ref == 'refs/heads/master' + run: echo "LATEST_TAG=,latest" >> $GITHUB_ENV + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c + - name: Authenticate on GCP + uses: google-github-actions/auth@7c6bc770dae815cd3e89ee6cdf493a5fab2cc093 + with: + service_account: ${{ secrets.GCP_SA_EMAIL }} + credentials_json: ${{ secrets.GCP_SA_KEY }} + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db + - name: GCloud Docker credential helper + run: | + gcloud auth configure-docker ${{ env.docker_registry }} + - name: Setup Python environment + uses: ./.github/actions/setup-environment-action + with: + python-version: default + - name: run Publish Beam SDK Snapshots script + uses: ./.github/actions/gradle-command-self-hosted-action + with: + gradle-command: :sdks:${{ matrix.container_task }} + arguments: | + -Pdocker-repository-root=gcr.io/apache-beam-testing/beam-sdk \ + -Pdocker-tag-list=${{ github.sha }},${BEAM_VERSION}${LATEST_TAG} \ + -Pcontainer-architecture-list=arm64,amd64 \ + -Ppush-containers \ + -Pdocker-pull-licenses diff --git a/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml b/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml index d7e012e9f036..6636f70197e6 100644 --- a/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml +++ b/.github/workflows/beam_Python_ValidatesContainer_Dataflow_ARM.yml @@ -62,6 +62,8 @@ jobs: startsWith(github.event.comment.body, 'Run Python ValidatesContainer Dataflow ARM') runs-on: ubuntu-22.04 + # Python 3.14 multiarch container builds compile many deps from source on arm64. + timeout-minutes: ${{ matrix.python_version == '3.14' && 300 || 120 }} steps: - uses: actions/checkout@v7 with: diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml index 6236888521a5..02b223943fea 100644 --- a/.github/workflows/build_wheels.yml +++ b/.github/workflows/build_wheels.yml @@ -37,7 +37,8 @@ concurrency: cancel-in-progress: true env: - GCP_PATH: "gs://${{ secrets.GCP_PYTHON_WHEELS_BUCKET }}/${GITHUB_REF##*/}/${GITHUB_SHA}-${GITHUB_RUN_ID}/" + # Use github.* context; workflow env is not bash-evaluated. + GCP_PATH: "gs://${{ secrets.GCP_PYTHON_WHEELS_BUCKET }}/${{ github.ref_name }}/${{ github.sha }}-${{ github.run_id }}/" jobs: diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index be534203eaed..7072b9a8da4a 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -36,7 +36,7 @@ jobs: matrix: ${{ steps.set-matrix.outputs.matrix }} steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v7 - uses: dorny/paths-filter@fbd0ab8f3e69293af611ebaee6363fc25e6d187d id: filter with: diff --git a/.github/workflows/update_python_dependencies.yml b/.github/workflows/update_python_dependencies.yml index 9da85090c4c7..f8630ebd0aec 100644 --- a/.github/workflows/update_python_dependencies.yml +++ b/.github/workflows/update_python_dependencies.yml @@ -51,6 +51,7 @@ jobs: runs-on: [self-hosted, ubuntu-24.04, main] needs: set-properties name: Update Python Dependencies + timeout-minutes: 180 steps: - name: Checkout code uses: actions/checkout@v7 diff --git a/.test-infra/tools/refresh_looker_metrics.py b/.test-infra/tools/refresh_looker_metrics.py index 351d8998f8f6..b0800905e140 100644 --- a/.test-infra/tools/refresh_looker_metrics.py +++ b/.test-infra/tools/refresh_looker_metrics.py @@ -34,6 +34,8 @@ ("33", ["21", "70", "116", "69", "115"]), # BigTableIO_Write ("34", ["22", "56", "96", "55", "95"]), # TextIO_Read ("35", ["23", "64", "110", "63", "109"]), # TextIO_Write + ("113", ["386", "388", "390", "392", "394"]), # IcebergIO_Read + ("114", ["387", "389", "391", "393", "395"]), # IcebergIO_Write ("75", ["258", "259", "260", "261", "262"]), # TensorFlow MNIST ("76", ["233", "234", "235", "236", "237"]), # PyTorch BERT base uncased ("77", ["238", "239", "240", "241", "242"]), # PyTorch BERT large uncased diff --git a/CHANGES.md b/CHANGES.md index d39a62f1a780..87fdaa8d0216 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -75,6 +75,7 @@ ## Breaking Changes * (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)). +* (Java) Added `DRAINING` and `DRAINED` states to `PipelineResult`, including runner state mappings and Dataflow update handling ([#39020](https://github.com/apache/beam/issues/39020)). ## Deprecations @@ -84,7 +85,7 @@ * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)). * Improved Java pipeline performance by avoiding repeated `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)). -* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Python) Fixed a memory leak in Python SDK caused by storing exceptions with potentially large stack frames in a cache ([#39406](https://github.com/apache/beam/issues/39406)). ## Security Fixes @@ -138,6 +139,7 @@ ## Known Issues * (Java) Projects using the Flink runner with Flink 2.1 or later alongside libraries requiring `org.lz4:lz4-java` (e.g., Kafka clients) may encounter a Gradle capability conflict, because Flink 2.1+ ships `at.yawk.lz4:lz4-java` which declares the same capability. To resolve, add a `capabilitiesResolution` rule to your `build.gradle` that selects `at.yawk.lz4:lz4-java` ([#38947](https://github.com/apache/beam/issues/38947)). +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.74.0] - 2026-06-02 @@ -179,6 +181,10 @@ * Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). * Added `max_batch_duration_secs` passthrough support in Python Enrichment BigQuery and CloudSQL handlers so batching duration can be forwarded to `BatchElements` ([#38243](https://github.com/apache/beam/issues/38243)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.73.0] - 2026-04-29 ## Highlights @@ -215,6 +221,10 @@ * Fixed [CVE-2023-46604](https://www.cve.org/CVERecord?id=CVE-2023-46604) (CVSS 10.0) and [CVE-2022-41678](https://www.cve.org/CVERecord?id=CVE-2022-41678) by upgrading ActiveMQ from 5.14.5 to 5.19.2 (Java) ([#37943](https://github.com/apache/beam/issues/37943)). * Fixed [CVE-2024-1597](https://www.cve.org/CVERecord?id=CVE-2024-1597), [CVE-2022-31197](https://www.cve.org/CVERecord?id=CVE-2022-31197), and [CVE-2022-21724](https://www.cve.org/CVERecord?id=CVE-2022-21724) by upgrading PostgreSQL JDBC Driver from 42.2.16 to 42.6.2 (Java) ([#37942](https://github.com/apache/beam/issues/37942)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.72.0] - 2026-03-30 ## Highlights @@ -248,6 +258,10 @@ * Fixed [CVE-2024-28397](https://www.cve.org/CVERecord?id=CVE-2024-28397) by switching from js2py to pythonmonkey (Yaml) ([#37560](https://github.com/apache/beam/issues/37560)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.71.0] - 2026-01-22 ## I/Os @@ -269,6 +283,7 @@ ## Known Issues +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.70.0] - 2025-12-16 @@ -291,6 +306,10 @@ Now Beam has full support for Milvus integration including Milvus enrichment and * (Python) Python 3.9 reached EOL in October 2025 and support for the language version has been removed. ([#36665](https://github.com/apache/beam/issues/36665)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.69.0] - 2025-10-28 ## Highlights @@ -350,6 +369,10 @@ Now Beam has full support for Milvus integration including Milvus enrichment and ([#36141](https://github.com/apache/beam/issues/36141)). * Fixed Spanner Change Stream reading stuck issue due to watermark of partition moving backwards ([#36470](https://github.com/apache/beam/issues/36470)). +## Known Issues + +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). + # [2.68.0] - 2025-09-22 ## Highlights @@ -404,6 +427,7 @@ Now Beam has full support for Milvus integration including Milvus enrichment and ## Known Issues * ([#36470](https://github.com/apache/beam/issues/36470)). Spanner Change Stream reading stuck issue due to watermark of partition moving backwards. This issue exists in 2.67.0 and 2.68.0. To mitigate the issue, either use old version 2.66.0 or go to 2.69.0. +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.67.0] - 2025-08-12 @@ -452,6 +476,7 @@ Now Beam has full support for Milvus integration including Milvus enrichment and * ([#35666](https://github.com/apache/beam/issues/35666)). YAML Flatten incorrectly drops fields when input PCollections' schema are different. This issue exists for all versions since 2.52.0. * ([#36470](https://github.com/apache/beam/issues/36470)). Spanner Change Stream reading stuck issue due to watermark of partition moving backwards. This issue exists in 2.67.0 and 2.68.0. To mitigate the issue, either use old version 2.66.0 or go to 2.69.0. +* (Python) Long-running Python pipelines might experience memory growth and periodic OOMs ([#39406](https://github.com/apache/beam/issues/39406)). # [2.66.0] - 2025-07-01 diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy index 6963f96d7313..574e34118905 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamDockerPlugin.groovy @@ -61,6 +61,7 @@ class BeamDockerPlugin implements Plugin { boolean push = false String builder = null String target = null + String compression = 'zstd' File resolvedDockerfile = null File resolvedDockerComposeTemplate = null @@ -72,6 +73,9 @@ class BeamDockerPlugin implements Plugin { DockerExtension(Project project) { this.project = project this.copySpec = project.copySpec() + if (project.hasProperty('docker-compression')) { + this.compression = project.property('docker-compression') + } } void resolvePathsAndValidate() { @@ -235,13 +239,23 @@ class BeamDockerPlugin implements Plugin { if (!ext.platform.isEmpty()) { buildCommandLine.addAll('--platform', String.join(',', ext.platform)) } - if (ext.load) { - buildCommandLine.add '--load' + if (ext.load && ext.push) { + throw new Exception("cannot combine 'push' and 'load' options") } - if (ext.push) { - buildCommandLine.add '--push' + if (ext.compression != null && !ext.compression.isEmpty()) { + if (ext.push) { + buildCommandLine.add "--output=type=registry,compression=${ext.compression},force-compression=true,oci-mediatypes=true" + } else if (ext.load) { + buildCommandLine.add '--load' + } else { + buildCommandLine.add "--output=type=image,compression=${ext.compression},force-compression=true,oci-mediatypes=true" + } + } else { if (ext.load) { - throw new Exception("cannot combine 'push' and 'load' options") + buildCommandLine.add '--load' + } + if (ext.push) { + buildCommandLine.add '--push' } } if (ext.builder != null) { diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy index d73f2e7a2bae..70a22e397724 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy @@ -613,7 +613,7 @@ class BeamModulePlugin implements Plugin { def autoservice_version = "1.0.1" def aws_java_sdk2_version = "2.20.162" def cassandra_driver_version = "3.10.2" - def cdap_version = "6.5.1" + def cdap_version = "6.11.4" def checkerframework_version = "3.42.0" def classgraph_version = "4.8.162" def delta_lake_version = "4.2.0" diff --git a/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py b/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py index fc83c8d443cc..e9f07daef101 100644 --- a/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py +++ b/examples/notebooks/beam-ml/rag_usecase/opensearch_connector.py @@ -88,8 +88,8 @@ def __init__(self, """ self.host = host self.port = port - self.username = username | os.getenv("OPENSEARCH_USERNAME") - self.password = password | os.getenv("OPENSEARCH_PASSWORD") + self.username = username or os.getenv("OPENSEARCH_USERNAME") + self.password = password or os.getenv("OPENSEARCH_PASSWORD") self._batch_size = batch_size if not self.username or not self.password: @@ -247,8 +247,8 @@ def __init__(self, """ self.host = host self.port = port - self.username = username | os.getenv("OPENSEARCH_USERNAME") - self.password = password | os.getenv("OPENSEARCH_PASSWORD") + self.username = username or os.getenv("OPENSEARCH_USERNAME") + self.password = password or os.getenv("OPENSEARCH_PASSWORD") self.batch_size = batch_size self.embedded_columns = embedded_columns diff --git a/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py b/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py index 70397550241f..c43d9af181ba 100644 --- a/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py +++ b/examples/notebooks/beam-ml/rag_usecase/opensearch_enrichment.py @@ -77,8 +77,8 @@ def __init__( """ self.opensearch_host = opensearch_host self.opensearch_port = opensearch_port - self.username = username | os.getenv("OPENSEARCH_USERNAME") - self.password = password | os.getenv("OPENSEARCH_PASSWORD") + self.username = username or os.getenv("OPENSEARCH_USERNAME") + self.password = password or os.getenv("OPENSEARCH_PASSWORD") self.index_name = index_name self.vector_field = vector_field self.k = k diff --git a/it/build.gradle b/it/build.gradle index 42a9ad9f4ee8..c550e26d305f 100644 --- a/it/build.gradle +++ b/it/build.gradle @@ -34,4 +34,8 @@ tasks.register('BigTablePerformanceTest') { tasks.register('BigQueryStorageApiStreamingPerformanceTest') { dependsOn(":it:google-cloud-platform:BigQueryStorageApiStreamingPerformanceTest") -} \ No newline at end of file +} + +tasks.register('IcebergPerformanceTest') { + dependsOn(":it:iceberg:IcebergPerformanceTest") +} diff --git a/it/clickhouse/build.gradle b/it/clickhouse/build.gradle index 858d0dea9207..0727f0328e45 100644 --- a/it/clickhouse/build.gradle +++ b/it/clickhouse/build.gradle @@ -29,6 +29,12 @@ ext.summary = "Integration test utilities for ClickHouse." def clickhouse_java_client_version = "0.9.6" +configurations.configureEach { + resolutionStrategy.capabilitiesResolution.withCapability("org.lz4:lz4-java") { + select("at.yawk.lz4:lz4-java:1.10.2") + } +} + dependencies { implementation project(path: ":it:common") implementation project(path: ":it:testcontainers") @@ -52,4 +58,4 @@ dependencies { permitUnusedDeclared "com.clickhouse:client-v2:$clickhouse_java_client_version:all" -} \ No newline at end of file +} diff --git a/it/common/build.gradle b/it/common/build.gradle index c92ec551cb25..62dd45ddacf5 100644 --- a/it/common/build.gradle +++ b/it/common/build.gradle @@ -28,6 +28,19 @@ ext.summary = "Code used by all integration test utilities." dependencies { implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation project(path: ":runners:google-cloud-dataflow-java") + implementation 'com.google.cloud:google-cloud-bigquery' + implementation 'com.google.cloud:google-cloud-monitoring' + implementation 'com.google.cloud:google-cloud-storage' + implementation 'org.apache.hadoop:hadoop-common:3.3.6' + implementation 'org.apache.avro:avro:1.11.1' + implementation 'org.apache.parquet:parquet-avro:1.15.2' + implementation 'org.apache.parquet:parquet-common:1.15.2' + implementation 'org.apache.parquet:parquet-hadoop:1.15.2' + provided 'com.google.api.grpc:proto-google-cloud-monitoring-v3' + implementation library.java.jackson_core + implementation library.java.jackson_databind + implementation library.java.gax implementation library.java.google_api_services_dataflow implementation library.java.google_auth_library_credentials implementation library.java.google_auth_library_oauth2_http @@ -36,14 +49,20 @@ dependencies { implementation library.java.commons_lang3 implementation library.java.failsafe implementation library.java.google_code_gson + implementation library.java.google_cloud_core + implementation library.java.google_api_client implementation library.java.google_http_client implementation library.java.guava - + implementation library.java.protobuf_java_util + implementation library.java.protobuf_java + implementation library.java.junit + testImplementation library.java.mockito_inline + testRuntimeOnly library.java.slf4j_simple // TODO: excluding Guava until Truth updates it to >32.1.x testImplementation(library.java.truth) { - exclude group: 'com.google.guava', module: 'guava' + exclude group: 'com.google.guava', module: 'guava' } - testImplementation library.java.junit - testImplementation library.java.mockito_inline - testRuntimeOnly library.java.slf4j_simple + testImplementation project(path: ":runners:direct-java", configuration: "shadow") + testImplementation project(path: ":sdks:java:testing:test-utils") + testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") } diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java similarity index 97% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java index 830f44598be7..6371d6eedb1e 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/Artifact.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/Artifact.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; /** * Represents a single artifact. diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java index e463baffeaf1..b275d35703e1 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/ArtifactClient.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/ArtifactClient.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; import java.io.IOException; import java.nio.file.Path; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java similarity index 96% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java index c60e3a0fe53a..5f4360ddbb2b 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/GcsArtifact.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/GcsArtifact.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; import com.google.cloud.storage.Blob; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java similarity index 95% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java index de2a2541d558..44b54f82aeba 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/package-info.java @@ -17,4 +17,4 @@ */ /** Package for working with test artifacts. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java index 854651ae1f14..3a3bcb2d7ec4 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtils.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtils.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import static java.util.Arrays.stream; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java index 5cfe4c137e35..58241f9a1db3 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/AvroTestUtil.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/AvroTestUtil.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import java.io.ByteArrayOutputStream; import java.io.IOException; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java index 9a83558f7bfc..08c3cdfb0c5a 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/JsonTestUtil.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/JsonTestUtil.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.JsonNode; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java index d6ebf825e2c1..fb4326e1b4a7 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/ParquetTestUtil.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/ParquetTestUtil.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import java.io.ByteArrayInputStream; import java.io.File; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java similarity index 94% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java index 53cb4ed1d3af..aaf820b25103 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/utils/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/artifacts/utils/package-info.java @@ -17,4 +17,4 @@ */ /** Package for artifact utilities. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java rename to it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java index 5ca6a8c4655c..728a47b224b4 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManager.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import com.google.api.gax.paging.Page; import com.google.auth.Credentials; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java similarity index 96% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java rename to it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java index 61ef27e45260..ce95dc891551 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerException.java +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerException.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; /** Custom exception for {@link BigQueryResourceManager} implementations. */ public class BigQueryResourceManagerException extends RuntimeException { diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java rename to it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java index f2b6849caa61..56cb82013885 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtils.java +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtils.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import static org.apache.beam.it.common.utils.ResourceManagerUtils.generateResourceId; diff --git a/it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java new file mode 100644 index 000000000000..f2d5978fb4eb --- /dev/null +++ b/it/common/src/main/java/org/apache/beam/it/common/bigquery/package-info.java @@ -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. + */ + +/** Package for managing BigQuery common resources within integration tests. */ +package org.apache.beam.it.common.bigquery; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java rename to it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java index 9d00fa67f145..a6efaeca5afc 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncher.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; import static org.apache.beam.it.common.PipelineLauncher.JobState.ACTIVE_STATES; import static org.apache.beam.it.common.PipelineLauncher.JobState.FAILED; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java similarity index 92% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java rename to it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java index 11a09c4ba749..af6438a561aa 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncher.java +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncher.java @@ -15,10 +15,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; -import static org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.fasterxml.jackson.databind.ObjectMapper; @@ -41,7 +40,6 @@ import java.util.stream.StreamSupport; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.utils.PipelineUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowPipelineJob; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.DistributionResult; @@ -73,6 +71,8 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher { private static final String READ_PIPELINE_NAME_OVERWRITE = "readPipelineNameOverride"; private static final String WRITE_PIPELINE_NAME_OVERWRITE = "writePipelineNameOverride"; private static final Pattern JOB_ID_PATTERN = Pattern.compile("Submitted job: (\\S+)"); + /** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */ + public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS"; // For unsupported runners (other than dataflow), implement launcher methods by operating with // PipelineResult. @@ -101,6 +101,18 @@ public class DefaultPipelineLauncher extends AbstractPipelineLauncher { .put(PipelineResult.State.UNRECOGNIZED, JobState.UNKNOWN) .build(); + // To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and + // pipeline defined + // metrics, here we wrap Beam provided metrics as a pre-defined metrics name + // [name_space:metric_type:metric_name + // which will be recognized by getMetric method + public enum PipelineMetricsType { + COUNTER, + STARTTIME, + ENDTIME, + RUNTIME, + } + private DefaultPipelineLauncher(Builder builder) { super( new Dataflow( @@ -169,7 +181,7 @@ private static void checkIfMetricResultIsUnique( resultCount <= 1, "More than one metric result matches name: %s in namespace %s. Metric results count: %s", name, - IOLoadTestBase.BEAM_METRICS_NAMESPACE, + BEAM_METRICS_NAMESPACE, resultCount); } @@ -180,15 +192,13 @@ private static Iterable> getDistributions( .metrics() .queryMetrics( MetricsFilter.builder() - .addNameFilter( - MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName)) + .addNameFilter(MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) .build()); return metrics.getDistributions(); } /** Pull Beam pipeline defined metrics given the jobId. */ - public Long getBeamMetric( - String jobId, IOLoadTestBase.PipelineMetricsType metricType, String metricName) { + public Long getBeamMetric(String jobId, PipelineMetricsType metricType, String metricName) { PipelineResult pipelineResult = MANAGED_JOBS.getOrDefault(jobId, UNMANAGED_JOBS.getOrDefault(jobId, null)); if (pipelineResult != null) { @@ -197,8 +207,7 @@ public Long getBeamMetric( .metrics() .queryMetrics( MetricsFilter.builder() - .addNameFilter( - MetricNameFilter.named(IOLoadTestBase.BEAM_METRICS_NAMESPACE, metricName)) + .addNameFilter(MetricNameFilter.named(BEAM_METRICS_NAMESPACE, metricName)) .build()); switch (metricType) { @@ -210,9 +219,7 @@ public Long getBeamMetric( return metricResult.getAttempted(); } catch (NoSuchElementException e) { LOG.error( - "Failed to get metric {}, from namespace {}", - metricName, - IOLoadTestBase.BEAM_METRICS_NAMESPACE); + "Failed to get metric {}, from namespace {}", metricName, BEAM_METRICS_NAMESPACE); } return UNKNOWN_METRIC_VALUE; case STARTTIME: @@ -230,9 +237,9 @@ public Long getBeamMetric( .map(element -> Objects.requireNonNull(element.getAttempted()).getMax()) .max(Long::compareTo) .orElse(UNKNOWN_METRIC_VALUE); - if (metricType == IOLoadTestBase.PipelineMetricsType.STARTTIME) { + if (metricType == PipelineMetricsType.STARTTIME) { return lowestMin; - } else if (metricType == IOLoadTestBase.PipelineMetricsType.ENDTIME) { + } else if (metricType == PipelineMetricsType.ENDTIME) { return greatestMax; } else { if (lowestMin != UNKNOWN_METRIC_VALUE && greatestMax != UNKNOWN_METRIC_VALUE) { @@ -254,15 +261,14 @@ public Long getBeamMetric( @Override public Double getMetric(String project, String region, String jobId, String metricName) throws IOException { - if (metricName.startsWith(IOLoadTestBase.BEAM_METRICS_NAMESPACE)) { + if (metricName.startsWith(BEAM_METRICS_NAMESPACE)) { String[] nameSpacedMetrics = metricName.split(":", 3); Preconditions.checkState( nameSpacedMetrics.length == 3, String.format( "Invalid Beam metrics name: %s, expected: '%s:metric_type:metric_name'", - metricName, IOLoadTestBase.BEAM_METRICS_NAMESPACE)); - IOLoadTestBase.PipelineMetricsType metricType = - IOLoadTestBase.PipelineMetricsType.valueOf(nameSpacedMetrics[1]); + metricName, BEAM_METRICS_NAMESPACE)); + PipelineMetricsType metricType = PipelineMetricsType.valueOf(nameSpacedMetrics[1]); // Pipeline defined metrics are long values. Have to cast to double that is what the base // class defined. @@ -437,16 +443,18 @@ private List extractOptions(String project, String region, LaunchConfig // add pipeline options from beamTestPipelineOptions system property to preserve the // pipeline options already set in TestPipeline. @Nullable - String beamTestPipelineOptions = System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS); + String beamTestPipelineOptions = + System.getProperty( + org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS); if (!Strings.isNullOrEmpty(beamTestPipelineOptions)) { try { additionalOptions.addAll(MAPPER.readValue(beamTestPipelineOptions, List.class)); } catch (IOException e) { throw new RuntimeException( "Unable to instantiate test options from system property " - + PROPERTY_BEAM_TEST_PIPELINE_OPTIONS + + org.apache.beam.sdk.testing.TestPipeline.PROPERTY_BEAM_TEST_PIPELINE_OPTIONS + ":" - + System.getProperty(PROPERTY_BEAM_TEST_PIPELINE_OPTIONS), + + beamTestPipelineOptions, e); } } diff --git a/it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java new file mode 100644 index 000000000000..eb10897523bc --- /dev/null +++ b/it/common/src/main/java/org/apache/beam/it/common/dataflow/package-info.java @@ -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. + */ + +/** Package for managing Dataflow jobs from integration tests. */ +package org.apache.beam.it.common.dataflow; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java b/it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java rename to it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java index 66550b5a5073..6b447d84512d 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/MonitoringClient.java +++ b/it/common/src/main/java/org/apache/beam/it/common/monitoring/MonitoringClient.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.monitoring; +package org.apache.beam.it.common.monitoring; import com.google.api.gax.core.CredentialsProvider; import com.google.cloud.monitoring.v3.MetricServiceClient; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java similarity index 94% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java index 9cd79ea9cf12..731bbc3a6b25 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/monitoring/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/monitoring/package-info.java @@ -17,4 +17,4 @@ */ /** Package for querying metrics from cloud monitoring. */ -package org.apache.beam.it.gcp.monitoring; +package org.apache.beam.it.common.monitoring; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/GcsResourceManager.java b/it/common/src/main/java/org/apache/beam/it/common/storage/GcsResourceManager.java similarity index 97% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/GcsResourceManager.java rename to it/common/src/main/java/org/apache/beam/it/common/storage/GcsResourceManager.java index 730ca23ee54b..85f2c8b0b018 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/GcsResourceManager.java +++ b/it/common/src/main/java/org/apache/beam/it/common/storage/GcsResourceManager.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.storage; +package org.apache.beam.it.common.storage; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; @@ -39,10 +39,10 @@ import java.util.function.Consumer; import java.util.regex.Pattern; import org.apache.beam.it.common.ResourceManager; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.gcp.artifacts.ArtifactClient; -import org.apache.beam.it.gcp.artifacts.GcsArtifact; -import org.apache.beam.it.gcp.artifacts.utils.ArtifactUtils; +import org.apache.beam.it.common.artifacts.Artifact; +import org.apache.beam.it.common.artifacts.ArtifactClient; +import org.apache.beam.it.common.artifacts.GcsArtifact; +import org.apache.beam.it.common.artifacts.utils.ArtifactUtils; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.junit.rules.TestName; import org.slf4j.Logger; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/package-info.java b/it/common/src/main/java/org/apache/beam/it/common/storage/package-info.java similarity index 95% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/package-info.java rename to it/common/src/main/java/org/apache/beam/it/common/storage/package-info.java index efcea3aba267..1c5b3b3ce21e 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/storage/package-info.java +++ b/it/common/src/main/java/org/apache/beam/it/common/storage/package-info.java @@ -17,4 +17,4 @@ */ /** Package for managing Google Cloud Storage resources within integration tests. */ -package org.apache.beam.it.gcp.storage; +package org.apache.beam.it.common.storage; diff --git a/it/google-cloud-platform/src/main/resources/test-artifact.json b/it/common/src/main/resources/test-artifact.json similarity index 100% rename from it/google-cloud-platform/src/main/resources/test-artifact.json rename to it/common/src/main/resources/test-artifact.json diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/GcsArtifactTest.java b/it/common/src/test/java/org/apache/beam/it/common/artifacts/GcsArtifactTest.java similarity index 94% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/GcsArtifactTest.java rename to it/common/src/test/java/org/apache/beam/it/common/artifacts/GcsArtifactTest.java index 59b7931b603a..470741b399ad 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/GcsArtifactTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/artifacts/GcsArtifactTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts; +package org.apache.beam.it.common.artifacts; import static com.google.common.truth.Truth.assertThat; import static org.mockito.Mockito.when; @@ -30,7 +30,7 @@ import org.mockito.junit.MockitoJUnit; import org.mockito.junit.MockitoRule; -/** Unit tests for {@link GcsArtifact}. */ +/** Unit tests for {@link org.apache.beam.it.common.artifacts.GcsArtifact}. */ @RunWith(JUnit4.class) public class GcsArtifactTest { @Rule public final MockitoRule mockito = MockitoJUnit.rule(); diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtilsTest.java b/it/common/src/test/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtilsTest.java similarity index 94% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtilsTest.java rename to it/common/src/test/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtilsTest.java index 491d32d0ce6a..137c35aaf4bb 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/artifacts/utils/ArtifactUtilsTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/artifacts/utils/ArtifactUtilsTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.utils; +package org.apache.beam.it.common.artifacts.utils; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -24,7 +24,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Artifacts for {@link ArtifactUtils}. */ +/** Artifacts for {@link org.apache.beam.it.common.artifacts.utils.ArtifactUtils}. */ @RunWith(JUnit4.class) public final class ArtifactUtilsTest { diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java similarity index 99% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java rename to it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java index 53987dbe2670..62aac93b4eed 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java similarity index 93% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java rename to it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java index 432b7ef7362f..7fc848b45292 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManagerUtilsTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/bigquery/BigQueryResourceManagerUtilsTest.java @@ -15,9 +15,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.bigquery; +package org.apache.beam.it.common.bigquery; -import static org.apache.beam.it.gcp.bigquery.BigQueryResourceManagerUtils.checkValidTableId; +import static org.apache.beam.it.common.bigquery.BigQueryResourceManagerUtils.checkValidTableId; import static org.junit.Assert.assertThrows; import java.util.Arrays; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java similarity index 99% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java index f2c33d70116b..72d5514a6e3c 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncherTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/AbstractPipelineLauncherTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java similarity index 91% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java index b6c0f8cdc58f..35510a0f1d1c 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/DefaultPipelineLauncherTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/DefaultPipelineLauncherTest.java @@ -15,16 +15,16 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.dataflow; +package org.apache.beam.it.common.dataflow; +import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import java.io.IOException; import java.time.Instant; import org.apache.beam.it.common.PipelineLauncher; -import org.apache.beam.it.gcp.IOLoadTestBase; -import org.apache.beam.it.gcp.IOLoadTestBase.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testutils.metrics.TimeMonitor; @@ -48,7 +48,7 @@ public void testPipelineMetrics() throws IOException { pipeline .apply(GenerateSequence.from(0).to(numElements)) - .apply(ParDo.of(new TimeMonitor<>(IOLoadTestBase.BEAM_METRICS_NAMESPACE, timeMetrics))) + .apply(ParDo.of(new TimeMonitor<>(BEAM_METRICS_NAMESPACE, timeMetrics))) .apply(ParDo.of(new IOLoadTestBase.CountingFn<>(counterMetrics))); PipelineLauncher.LaunchConfig options = diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOLoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java similarity index 87% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOLoadTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java index 14770a429731..c35f402bd8ac 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOLoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOLoadTestBase.java @@ -15,7 +15,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp; +package org.apache.beam.it.common.dataflow; + +import static org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.BEAM_METRICS_NAMESPACE; import com.google.cloud.Timestamp; import java.io.IOException; @@ -25,7 +27,7 @@ import java.util.UUID; import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.dataflow.DefaultPipelineLauncher; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.testutils.NamedTestResult; @@ -90,21 +92,6 @@ public void processElement(ProcessContext ctx) { } } - // To make PipelineLauncher.getMetric work in a unified way for both runner provided metrics and - // pipeline defined - // metrics, here we wrap Beam provided metrics as a pre-defined metrics name - // [name_space:metric_type:metric_name - // which will be recognized by getMetric method - public enum PipelineMetricsType { - COUNTER, - STARTTIME, - ENDTIME, - RUNTIME, - } - - /** Namespace for Beam provided pipeline metrics (set up by Metrics transform). */ - public static final String BEAM_METRICS_NAMESPACE = "BEAM_METRICS"; - /** Given a metrics name, return Beam metrics name. */ public static String getBeamMetricsName(PipelineMetricsType metricstype, String metricsName) { return BEAM_METRICS_NAMESPACE + ":" + metricstype + ":" + metricsName; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOStressTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOStressTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java index 5c2fb74cd2fb..e08eec031eb5 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/IOStressTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/IOStressTestBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp; +package org.apache.beam.it.common.dataflow; import java.io.Serializable; import java.time.Duration; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/LoadTestBase.java b/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java similarity index 98% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/LoadTestBase.java rename to it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java index cd9ef52ed835..cd1b71dc7552 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/LoadTestBase.java +++ b/it/common/src/test/java/org/apache/beam/it/common/dataflow/LoadTestBase.java @@ -15,10 +15,10 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp; +package org.apache.beam.it.common.dataflow; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.RUNNER_V2; import static org.apache.beam.it.common.logging.LogStrings.formatForLogging; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.RUNNER_V2; import com.google.api.gax.core.CredentialsProvider; import com.google.api.gax.core.FixedCredentialsProvider; @@ -44,8 +44,8 @@ import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.bigquery.BigQueryResourceManager; -import org.apache.beam.it.gcp.monitoring.MonitoringClient; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; +import org.apache.beam.it.common.monitoring.MonitoringClient; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.After; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/GcsResourceManagerTest.java b/it/common/src/test/java/org/apache/beam/it/common/storage/GcsResourceManagerTest.java similarity index 95% rename from it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/GcsResourceManagerTest.java rename to it/common/src/test/java/org/apache/beam/it/common/storage/GcsResourceManagerTest.java index 0153573feaed..70dc6dcbc9c4 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/GcsResourceManagerTest.java +++ b/it/common/src/test/java/org/apache/beam/it/common/storage/GcsResourceManagerTest.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.storage; +package org.apache.beam.it.common.storage; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -35,6 +35,7 @@ import com.google.cloud.storage.Storage; import com.google.cloud.storage.Storage.BlobListOption; import com.google.cloud.storage.Storage.BucketListOption; +import com.google.common.truth.Truth; import java.io.IOException; import java.net.URISyntaxException; import java.nio.file.Files; @@ -43,8 +44,8 @@ import java.util.List; import java.util.UUID; import java.util.regex.Pattern; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.gcp.artifacts.GcsArtifact; +import org.apache.beam.it.common.artifacts.Artifact; +import org.apache.beam.it.common.artifacts.GcsArtifact; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Resources; import org.junit.After; @@ -128,7 +129,7 @@ public void testCreateArtifactInRunDir() { verify(client).create(blobInfoCaptor.capture(), contentsCaptor.capture()); BlobInfo actualInfo = blobInfoCaptor.getValue(); - assertThat(actual.blob).isSameInstanceAs(blob); + Truth.assertThat(actual.blob).isSameInstanceAs(blob); assertThat(actualInfo.getBucket()).isEqualTo(BUCKET); assertThat(actualInfo.getName()) .isEqualTo(String.format("%s/%s/%s", TEST_CLASS, gcsClient.runId(), artifactName)); @@ -144,7 +145,7 @@ public void testUploadArtifact() throws IOException { verify(client).create(blobInfoCaptor.capture(), contentsCaptor.capture()); BlobInfo actualInfo = blobInfoCaptor.getValue(); - assertThat(actual.blob).isSameInstanceAs(blob); + Truth.assertThat(actual.blob).isSameInstanceAs(blob); assertThat(actualInfo.getBucket()).isEqualTo(BUCKET); assertThat(actualInfo.getName()) .isEqualTo(String.format("%s/%s/%s", TEST_CLASS, gcsClient.runId(), ARTIFACT_NAME)); @@ -185,8 +186,8 @@ public void testListArtifactsInMethodDirSinglePage() { BlobListOption actualOptions = listOptionsCaptor.getValue(); assertThat(actual).hasSize(2); - assertThat(actual.get(0).name()).isEqualTo(name1); - assertThat(actual.get(1).name()).isEqualTo(name3); + Truth.assertThat(actual.get(0).name()).isEqualTo(name1); + Truth.assertThat(actual.get(1).name()).isEqualTo(name3); assertThat(actualBucket).isEqualTo(BUCKET); assertThat(actualOptions) .isEqualTo( @@ -221,8 +222,8 @@ public void testListArtifactsInMethodDirMultiplePages() { BlobListOption actualOptions = listOptionsCaptor.getValue(); assertThat(actual).hasSize(2); - assertThat(actual.get(0).name()).isEqualTo(name1); - assertThat(actual.get(1).name()).isEqualTo(name3); + Truth.assertThat(actual.get(0).name()).isEqualTo(name1); + Truth.assertThat(actual.get(1).name()).isEqualTo(name3); assertThat(actualBucket).isEqualTo(BUCKET); assertThat(actualOptions) .isEqualTo( diff --git a/it/google-cloud-platform/build.gradle b/it/google-cloud-platform/build.gradle index 3a46f2b94d83..164a75c06ba0 100644 --- a/it/google-cloud-platform/build.gradle +++ b/it/google-cloud-platform/build.gradle @@ -27,49 +27,33 @@ applyJavaNature( description = "Apache Beam :: IT :: Google Cloud Platform" ext.summary = "Integration test utilities for Google Cloud Platform." +evaluationDependsOn(":it:common") + dependencies { implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) - implementation project(path: ":sdks:java:core", configuration: "shadow") - implementation project(path: ":runners:google-cloud-dataflow-java") implementation project(path: ":it:common") implementation project(path: ":it:conditions") implementation project(path: ":it:truthmatchers") - implementation project(path: ":sdks:java:testing:test-utils") implementation library.java.failsafe implementation library.java.google_api_services_dataflow implementation library.java.junit implementation library.java.slf4j_api - implementation(library.java.truth) { - exclude group: 'com.google.guava', module: 'guava' - } implementation library.java.vendored_guava_32_1_2_jre implementation library.java.jackson_core implementation library.java.jackson_databind - implementation 'org.apache.hadoop:hadoop-common:3.3.5' - implementation 'org.apache.avro:avro:1.11.1' - implementation 'org.apache.parquet:parquet-avro:1.15.2' - implementation 'org.apache.parquet:parquet-common:1.15.2' - implementation 'org.apache.parquet:parquet-hadoop:1.15.2' implementation 'org.apache.commons:commons-lang3:3.9' implementation library.java.gax implementation library.java.google_api_common implementation library.java.google_api_client - implementation library.java.google_api_services_bigquery implementation library.java.google_auth_library_credentials implementation library.java.google_auth_library_oauth2_http - implementation library.java.google_http_client implementation library.java.guava - implementation library.java.protobuf_java_util implementation library.java.protobuf_java implementation library.java.threetenbp implementation 'org.awaitility:awaitility:4.2.0' - implementation 'joda-time:joda-time:2.10.10' // Google Cloud Dependencies implementation library.java.google_cloud_core - implementation 'com.google.cloud:google-cloud-storage' implementation 'com.google.cloud:google-cloud-bigquery' - implementation 'com.google.cloud:google-cloud-monitoring' - provided 'com.google.api.grpc:proto-google-cloud-monitoring-v3' implementation 'com.google.cloud:google-cloud-bigtable' implementation 'com.google.cloud:google-cloud-spanner' implementation 'com.google.cloud:google-cloud-pubsub' @@ -84,6 +68,7 @@ dependencies { implementation 'com.google.cloud:google-cloud-secretmanager' provided 'com.google.api.grpc:proto-google-cloud-secretmanager-v1' + testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:google-cloud-platform") testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:synthetic") diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java index 28626f7edd18..da0cf5e55dbe 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/conditions/BigQueryRowsCheck.java @@ -19,8 +19,8 @@ import com.google.auto.value.AutoValue; import com.google.cloud.bigquery.TableId; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; import org.apache.beam.it.conditions.ConditionCheck; -import org.apache.beam.it.gcp.bigquery.BigQueryResourceManager; import org.checkerframework.checker.nullness.qual.Nullable; /** ConditionCheck to validate if BigQuery has received a certain number of rows. */ diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java index 00aa4003817f..7908246dce13 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClient.java @@ -30,6 +30,7 @@ import com.google.auth.http.HttpCredentialsAdapter; import dev.failsafe.Failsafe; import java.io.IOException; +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java index 0a9731b13c32..496c4b0e980e 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClient.java @@ -32,6 +32,7 @@ import com.google.auth.http.HttpCredentialsAdapter; import dev.failsafe.Failsafe; import java.io.IOException; +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java index 5a101e08d375..1778d70cf60a 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/spanner/matchers/SpannerAsserts.java @@ -17,7 +17,7 @@ */ package org.apache.beam.it.gcp.spanner.matchers; -import static org.apache.beam.it.gcp.artifacts.utils.JsonTestUtil.parseJsonString; +import static org.apache.beam.it.common.artifacts.utils.JsonTestUtil.parseJsonString; import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatRecords; import com.google.cloud.spanner.Mutation; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java index b561c0c71fb6..16cf2b8f4d31 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/WordCountIT.java @@ -29,6 +29,7 @@ import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; import org.apache.beam.it.common.PipelineLauncher.Sdk; import org.apache.beam.it.common.PipelineOperator.Result; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowRunner; import org.apache.beam.sdk.io.TextIO; import org.apache.beam.sdk.options.PipelineOptions; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java index f123ffa5954c..9ac03db1c741 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOLT.java @@ -43,8 +43,10 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java index 14c83e48308a..495a4d4485d8 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java @@ -41,8 +41,10 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.bigquery.BigQueryResourceManager; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java index 6e511bd8e5c6..75a7be50a508 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryStreamingLT.java @@ -41,7 +41,7 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.IOLoadTestBase; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.runners.dataflow.DataflowRunner; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java index 410c992fe2d6..27773121e79d 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOLT.java @@ -34,8 +34,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; import org.apache.beam.sdk.testing.TestPipeline; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java index 052a0839adf5..b86f13a0de10 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigtable/BigTableIOST.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java index 88c35589f2be..a50c0951005f 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/ClassicTemplateClientTest.java @@ -18,10 +18,10 @@ package org.apache.beam.it.gcp.dataflow; import static com.google.common.truth.Truth.assertThat; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; @@ -49,6 +49,7 @@ import org.apache.beam.it.common.PipelineLauncher.JobState; import org.apache.beam.it.common.PipelineLauncher.LaunchConfig; import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; +import org.apache.beam.it.common.dataflow.AbstractPipelineLauncher; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java index 06f44437414a..b94d52e45afb 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/dataflow/FlexTemplateClientTest.java @@ -18,10 +18,10 @@ package org.apache.beam.it.gcp.dataflow; import static com.google.common.truth.Truth.assertThat; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; -import static org.apache.beam.it.gcp.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.LEGACY_RUNNER; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_ID; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_JOB_TYPE; +import static org.apache.beam.it.common.dataflow.AbstractPipelineLauncher.PARAM_RUNNER; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java similarity index 99% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java rename to it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java index 99016b5dd3a4..0b838bddeb6d 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java @@ -17,7 +17,7 @@ */ package org.apache.beam.it.gcp.datagenerator; -import static org.apache.beam.it.gcp.LoadTestBase.createConfig; +import static org.apache.beam.it.common.dataflow.LoadTestBase.createConfig; import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatPipeline; import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java index cb96db40b749..47058c75e399 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubSubIOST.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java index 477c79dd3bc7..da43fca2eaa3 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/pubsub/PubsubIOLT.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.extensions.protobuf.Proto3SchemaMessages.Primitive; import org.apache.beam.sdk.io.Read; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java index 8f3de1b58f40..4441b8dcd6b6 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOLT.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.sdk.io.GenerateSequence; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; import org.apache.beam.sdk.io.synthetic.SyntheticSourceOptions; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java index 76312b3e10fb..4eb7030c23c3 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/spanner/SpannerIOST.java @@ -37,8 +37,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.gcp.spanner.SpannerIO; diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java index 58ba88d900e8..ac1a7fc103c6 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/FileBasedIOLT.java @@ -33,8 +33,10 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.storage.GcsResourceManager; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.sdk.io.Compression; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.TextIO; diff --git a/it/iceberg/build.gradle b/it/iceberg/build.gradle new file mode 100644 index 000000000000..f11374a1ca4f --- /dev/null +++ b/it/iceberg/build.gradle @@ -0,0 +1,134 @@ +/* + * 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 org.apache.beam.gradle.IoPerformanceTestUtilities + +plugins { id 'org.apache.beam.module' } +applyJavaNature( + automaticModuleName: 'org.apache.beam.it.iceberg', + exportJavadoc: false, + // iceberg ended support for Java 8 in 1.7.0 + requireJavaVersion: JavaVersion.VERSION_11, +) + +description = "Apache Beam :: IT :: Iceberg" +ext.summary = "Integration test utilities for Iceberg." + +def iceberg_version = "1.10.0" +def parquet_version = "1.16.0" +def hive_version = "3.1.3" + +def hadoopVersions = [ + "336": "3.3.6", +] + +hadoopVersions.each {kv -> configurations.create("hadoopVersion$kv.key")} + +evaluationDependsOn(":it:common") + +dependencies { + provided "org.immutables:value:2.8.8" + permitUnusedDeclared "org.immutables:value:2.8.8" + runtimeOnly "org.apache.iceberg:iceberg-gcp:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-aws:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-aws-bundle:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-azure:$iceberg_version" + runtimeOnly "org.apache.iceberg:iceberg-azure-bundle:$iceberg_version" + runtimeOnly library.java.bigdataoss_gcs_connector + runtimeOnly library.java.bigdataoss_util_hadoop + runtimeOnly library.java.hadoop_client + + testImplementation project(":sdks:java:managed") + testImplementation library.java.bigdataoss_gcsio + testImplementation library.java.bigdataoss_util_hadoop + testImplementation "org.apache.parquet:parquet-avro:$parquet_version" + testImplementation "org.apache.iceberg:iceberg-data:$iceberg_version" + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation library.java.junit + testImplementation library.java.hamcrest + testImplementation project(path: ":sdks:java:extensions:avro") + testImplementation 'org.awaitility:awaitility:4.2.0' + + // Hive catalog test dependencies + testImplementation project(path: ":sdks:java:io:iceberg:hive") + testImplementation "org.apache.iceberg:iceberg-common:$iceberg_version" + testImplementation ("org.apache.iceberg:iceberg-hive-metastore:$iceberg_version") + testImplementation ("org.apache.hive:hive-metastore:$hive_version") + testImplementation "org.assertj:assertj-core:3.11.1" + testRuntimeOnly ("org.apache.hive.hcatalog:hive-hcatalog-core:$hive_version") { + exclude group: "org.apache.hive", module: "hive-exec" + exclude group: "org.apache.parquet", module: "parquet-hadoop-bundle" + } + + // BigQueryMetastore catalog dep + testImplementation project(path: ":sdks:java:io:iceberg:bqms", configuration: "shadow") + testImplementation project(":sdks:java:io:google-cloud-platform") + testImplementation library.java.google_api_services_bigquery + testImplementation 'com.google.cloud:google-cloud-storage' + + testImplementation library.java.google_auth_library_oauth2_http + testRuntimeOnly library.java.slf4j_jdk14 + testImplementation project(path: ":runners:direct-java", configuration: "shadow") + testRuntimeOnly project(path: ":runners:google-cloud-dataflow-java") + testRuntimeOnly project(path: ":sdks:java:harness") + hadoopVersions.each {kv -> + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-client:$kv.value" + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-minicluster:$kv.value" + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-hdfs-client:$kv.value" + "hadoopVersion$kv.key" "org.apache.hadoop:hadoop-mapreduce-client-core:$kv.value" + } + + testImplementation project(path: ":sdks:java:extensions:protobuf", configuration: "testRuntimeMigration") + testImplementation project(path: ":sdks:java:io:synthetic") + testImplementation project(path: ":sdks:java:io:iceberg") + testImplementation project(path: ":it:truthmatchers") + testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") + testImplementation library.java.mockito_inline + testImplementation project(path: ":sdks:java:extensions:google-cloud-platform-core", configuration: "testRuntimeMigration") + testRuntimeOnly library.java.slf4j_simple +} + +configurations.all { + // iceberg-core needs avro:1.12.0 + resolutionStrategy.force 'org.apache.avro:avro:1.12.0' + // TODO(https://github.com/apache/beam/issues/38515): + // Remove below pins when parquet-hadoop upgrades to hadoop-common:3.4.2 + resolutionStrategy.force 'org.apache.hadoop:hadoop-common:3.3.6' + resolutionStrategy.force 'org.apache.hadoop:hadoop-client:3.3.6' + resolutionStrategy.force 'org.apache.hadoop:hadoop-hdfs:3.3.6' + resolutionStrategy.force 'org.apache.hadoop:hadoop-hdfs-client:3.3.6' +} + +hadoopVersions.each {kv -> + configurations."hadoopVersion$kv.key" { + resolutionStrategy { + force "org.apache.hadoop:hadoop-client:$kv.value" + force "org.apache.hadoop:hadoop-common:$kv.value" + force "org.apache.hadoop:hadoop-mapreduce-client-core:$kv.value" + force "org.apache.hadoop:hadoop-minicluster:$kv.value" + force "org.apache.hadoop:hadoop-hdfs:$kv.value" + force "org.apache.hadoop:hadoop-hdfs-client:$kv.value" + } + } +} + +tasks.register( + "IcebergPerformanceTest", IoPerformanceTestUtilities.IoPerformanceTest, project, 'iceberg', 'IcebergIOLT', + ['configuration':'large','project':'apache-beam-testing', 'artifactBucket':'io-performance-temp'] + + System.properties +) diff --git a/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java new file mode 100644 index 000000000000..2a6098a1518e --- /dev/null +++ b/it/iceberg/src/test/java/org/apache/beam/it/iceberg/IcebergIOLT.java @@ -0,0 +1,363 @@ +/* + * 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. + */ +package org.apache.beam.it.iceberg; + +import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; +import static org.junit.Assert.assertEquals; + +import com.google.auto.value.AutoValue; +import java.io.IOException; +import java.text.ParseException; +import java.time.Duration; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.Map; +import java.util.Random; +import java.util.UUID; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.PipelineOperator; +import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.it.common.storage.GcsResourceManager; +import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.sdk.io.GenerateSequence; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergIO; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestPipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.catalog.TableIdentifier; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; + +/** + * IcebergIO performance tests using a Hadoop catalog backed by GCS. + * + *

The test writes generated rows to an Iceberg table, reads the table back, validates the number + * of rows, and exports write and read performance metrics. + * + *

Example trigger command: + * + *

+ * ./gradlew :it:iceberg:integrationTest \
+ *   --tests "org.apache.beam.it.iceberg.IcebergIOLT" \
+ *   -Dconfiguration=local \
+ *   -Dproject=[gcpProject] \
+ *   -DartifactBucket=[artifactBucket]
+ * 
+ */ +public final class IcebergIOLT extends IOLoadTestBase { + + private static final String READ_ELEMENT_METRIC_NAME = "read_count"; + private static final String NAMESPACE = "default"; + + private static final Schema ROW_SCHEMA = + Schema.builder().addInt64Field("id").addByteArrayField("data").build(); + + private static final Map TEST_CONFIGS = + ImmutableMap.of( + "local", + Configuration.of(1_000L, 5, "DirectRunner", 1_000), // approximately 1 MB + "small", + Configuration.of(1_000_000L, 20, "DataflowRunner", 1_000), // approximately 1 GB + "medium", + Configuration.of(10_000_000L, 40, "DataflowRunner", 1_000), // approximately 10 GB + "large", + Configuration.of(100_000_000L, 100, "DataflowRunner", 1_000) // approximately 100 GB + ); + + private static GcsResourceManager resourceManager; + + private Configuration configuration; + private IcebergCatalogConfig catalogConfig; + private TableIdentifier tableIdentifier; + + @Rule public TestPipeline writePipeline = TestPipeline.create(); + + @Rule public TestPipeline readPipeline = TestPipeline.create(); + + @BeforeClass + public static void beforeClass() { + resourceManager = + GcsResourceManager.builder(TestProperties.artifactBucket(), "icebergiolt", CREDENTIALS) + .build(); + } + + @AfterClass + public static void tearDownClass() { + ResourceManagerUtils.cleanResources(resourceManager); + } + + @Before + public void setup() { + configuration = getTestConfiguration(); + + String uniqueSuffix = + DateTimeFormatter.ofPattern("MMddHHmmssSSS") + .withZone(ZoneOffset.UTC) + .format(java.time.Instant.now()) + + "-" + + UUID.randomUUID().toString().substring(0, 10); + + String warehouseDirectory = "icebergiolt-" + uniqueSuffix; + resourceManager.registerTempDir(warehouseDirectory); + + String warehouseLocation = + String.format("gs://%s/%s", TestProperties.artifactBucket(), warehouseDirectory); + + tableIdentifier = TableIdentifier.of(NAMESPACE, "table_" + uniqueSuffix.replace("-", "_")); + + Map catalogProperties = + ImmutableMap.of("type", "hadoop", "warehouse", warehouseLocation); + + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("iceberg-load-test") + .setCatalogProperties(catalogProperties) + .build(); + + setTempLocation(writePipeline); + setTempLocation(readPipeline); + } + + /** Writes generated rows to Iceberg and verifies that all rows can be read back. */ + @Test + public void testIcebergWriteAndRead() throws IOException { + PipelineLauncher.LaunchInfo writeInfo = testWrite(); + + PipelineOperator.Result writeResult = + pipelineOperator.waitUntilDone( + createConfig(writeInfo, Duration.ofMinutes(configuration.getPipelineTimeout()))); + + assertThatResult(writeResult).isLaunchFinished(); + assertEquals( + PipelineLauncher.JobState.DONE, + pipelineLauncher.getJobStatus(project, region, writeInfo.jobId())); + + PipelineLauncher.LaunchInfo readInfo = testRead(); + + PipelineOperator.Result readResult = + pipelineOperator.waitUntilDone( + createConfig(readInfo, Duration.ofMinutes(configuration.getPipelineTimeout()))); + + assertThatResult(readResult).isLaunchFinished(); + assertEquals( + PipelineLauncher.JobState.DONE, + pipelineLauncher.getJobStatus(project, region, readInfo.jobId())); + + double numRecords = + pipelineLauncher.getMetric( + project, + region, + readInfo.jobId(), + getBeamMetricsName(PipelineMetricsType.COUNTER, READ_ELEMENT_METRIC_NAME)); + + assertEquals((double) configuration.getNumRows(), numRecords, 0.5); + + exportMetrics(writeInfo, readInfo); + } + + private PipelineLauncher.LaunchInfo testWrite() throws IOException { + writePipeline + .apply("Generate records", GenerateSequence.from(0).to(configuration.getNumRows())) + .apply("Map records", ParDo.of(new MapToIcebergFormat(configuration.getValueSizeBytes()))) + .setRowSchema(ROW_SCHEMA) + .apply("Write to Iceberg", IcebergIO.writeRows(catalogConfig).to(tableIdentifier)); + + PipelineLauncher.LaunchConfig options = + PipelineLauncher.LaunchConfig.builder("write-iceberg") + .setSdk(PipelineLauncher.Sdk.JAVA) + .setPipeline(writePipeline) + .addParameter("runner", configuration.getRunner()) + .addParameter( + "maxNumWorkers", + TestProperties.getProperty("maxNumWorkers", "10", TestProperties.Type.PROPERTY)) + .build(); + + return pipelineLauncher.launch(project, region, options); + } + + private PipelineLauncher.LaunchInfo testRead() throws IOException { + readPipeline + .apply("Read from Iceberg", IcebergIO.readRows(catalogConfig).from(tableIdentifier)) + .apply("Counting element", ParDo.of(new CountingFn<>(READ_ELEMENT_METRIC_NAME))); + + PipelineLauncher.LaunchConfig options = + PipelineLauncher.LaunchConfig.builder("read-iceberg") + .setSdk(PipelineLauncher.Sdk.JAVA) + .setPipeline(readPipeline) + .addParameter("runner", configuration.getRunner()) + .addParameter( + "maxNumWorkers", + TestProperties.getProperty("maxNumWorkers", "10", TestProperties.Type.PROPERTY)) + .build(); + + return pipelineLauncher.launch(project, region, options); + } + + private void exportMetrics( + PipelineLauncher.LaunchInfo writeInfo, PipelineLauncher.LaunchInfo readInfo) { + + /* + * Only use PCollection names belonging to transforms defined directly in this test. + * Internal IcebergIO transform names may change when the implementation changes. + */ + MetricsConfiguration writeMetricsConfig = + MetricsConfiguration.builder() + .setInputPCollection("Map records.out0") + .setInputPCollectionV2("Map records/ParMultiDo(MapToIcebergFormat).out0") + .build(); + + MetricsConfiguration readMetricsConfig = + MetricsConfiguration.builder() + .setOutputPCollection("Counting element.out0") + .setOutputPCollectionV2("Counting element/ParMultiDo(Counting).out0") + .build(); + + try { + exportMetricsToBigQuery(writeInfo, getMetrics(writeInfo, writeMetricsConfig)); + + exportMetricsToBigQuery(readInfo, getMetrics(readInfo, readMetricsConfig)); + } catch (IOException | ParseException | InterruptedException e) { + throw new RuntimeException("Unable to collect or export IcebergIO load-test metrics.", e); + } + } + + private Configuration getTestConfiguration() { + String testConfig = + TestProperties.getProperty("configuration", "local", TestProperties.Type.PROPERTY); + + Configuration selectedConfiguration = TEST_CONFIGS.get(testConfig); + + if (selectedConfiguration == null) { + throw new IllegalArgumentException( + String.format( + "Unknown test configuration: [%s]. Known configurations: %s", + testConfig, TEST_CONFIGS.keySet())); + } + + return selectedConfiguration; + } + + private static void setTempLocation(TestPipeline pipeline) { + String tempBucketName = TestProperties.artifactBucket(); + + if (Strings.isNullOrEmpty(tempBucketName)) { + return; + } + + String tempLocation = String.format("gs://%s/temp/", tempBucketName); + + pipeline.getOptions().as(TestPipelineOptions.class).setTempRoot(tempLocation); + + pipeline.getOptions().setTempLocation(tempLocation); + } + + /** Maps a generated sequence value to a deterministic Iceberg-compatible Beam row. */ + private static class MapToIcebergFormat extends DoFn { + + private final int valueSizeBytes; + private transient Random random; + + private MapToIcebergFormat(int valueSizeBytes) { + if (valueSizeBytes <= 0) { + throw new IllegalArgumentException("valueSizeBytes must be positive."); + } + + this.valueSizeBytes = valueSizeBytes; + } + + @Setup + public void setup() { + random = new Random(); + } + + @ProcessElement + public void processElement(@Element Long index, OutputReceiver output) { + + byte[] value = new byte[valueSizeBytes]; + + /* + * Using the row index as the seed makes retries deterministic. Random data also prevents + * compression from making the generated payload significantly smaller than configured. + */ + random.setSeed(index); + random.nextBytes(value); + + output.output(Row.withSchema(ROW_SCHEMA).addValues(index, value).build()); + } + } + + /** Options for the IcebergIO load test. */ + @AutoValue + abstract static class Configuration { + + abstract long getNumRows(); + + abstract int getPipelineTimeout(); + + abstract String getRunner(); + + abstract int getValueSizeBytes(); + + static Configuration of(long numRows, int pipelineTimeout, String runner, int valueSizeBytes) { + + if (numRows <= 0) { + throw new IllegalArgumentException("numRows must be positive."); + } + + if (pipelineTimeout <= 0) { + throw new IllegalArgumentException("pipelineTimeout must be positive."); + } + + if (valueSizeBytes <= 0) { + throw new IllegalArgumentException("valueSizeBytes must be positive."); + } + + return new AutoValue_IcebergIOLT_Configuration.Builder() + .setNumRows(numRows) + .setPipelineTimeout(pipelineTimeout) + .setRunner(runner) + .setValueSizeBytes(valueSizeBytes) + .build(); + } + + @AutoValue.Builder + abstract static class Builder { + + abstract Builder setNumRows(long numRows); + + abstract Builder setPipelineTimeout(int pipelineTimeout); + + abstract Builder setRunner(String runner); + + abstract Builder setValueSizeBytes(int valueSizeBytes); + + abstract Configuration build(); + } + } +} diff --git a/it/kafka/build.gradle b/it/kafka/build.gradle index 0496c423bebd..32ef0623fea9 100644 --- a/it/kafka/build.gradle +++ b/it/kafka/build.gradle @@ -42,8 +42,7 @@ dependencies { permitUsedUndeclared library.java.guava testImplementation library.java.mockito_core testRuntimeOnly library.java.slf4j_simple - // TODO: Remove the below dependencies after KafkaIOLT has been moved to the right directory. - testImplementation project(path: ":it:google-cloud-platform") + testImplementation project(path: ":it:common", configuration: "testRuntimeMigration") testImplementation project(path: ":sdks:java:io:kafka") testImplementation project(path: ":sdks:java:io:synthetic") testImplementation project(path: ":runners:direct-java") diff --git a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java index c20699d61a55..1607ca1cd674 100644 --- a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java +++ b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOLT.java @@ -30,8 +30,9 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; import org.apache.beam.it.common.utils.ResourceManagerUtils; -import org.apache.beam.it.gcp.IOLoadTestBase; import org.apache.beam.runners.direct.DirectOptions; import org.apache.beam.sdk.io.kafka.KafkaIO; import org.apache.beam.sdk.io.synthetic.SyntheticOptions; diff --git a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java index bb1cb37a11b6..d6231a69f1e7 100644 --- a/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java +++ b/it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java @@ -34,7 +34,8 @@ import org.apache.beam.it.common.PipelineLauncher; import org.apache.beam.it.common.PipelineOperator; import org.apache.beam.it.common.TestProperties; -import org.apache.beam.it.gcp.IOStressTestBase; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.dataflow.IOStressTestBase; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.io.kafka.KafkaIO; diff --git a/it/truthmatchers/build.gradle b/it/truthmatchers/build.gradle index fd54ad591a0a..2e21c1b5df04 100644 --- a/it/truthmatchers/build.gradle +++ b/it/truthmatchers/build.gradle @@ -27,11 +27,15 @@ ext.summary = "Truth matchers for integration tests." dependencies { implementation project(path: ":it:common") + implementation 'org.apache.avro:avro:1.11.1' // TODO: excluding Guava until Truth updates it to >32.1.x implementation(library.java.truth) { exclude group: 'com.google.guava', module: 'guava' } + implementation library.java.jackson_core + implementation library.java.jackson_databind implementation library.java.guava + implementation library.java.vendored_guava_32_1_2_jre permitUsedUndeclared library.java.guava implementation library.java.google_code_gson -} \ No newline at end of file +} diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactAsserts.java b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactAsserts.java similarity index 79% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactAsserts.java rename to it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactAsserts.java index 3b1251307d0d..92407e5a2fc9 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactAsserts.java +++ b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactAsserts.java @@ -15,15 +15,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.matchers; - -import static com.google.common.truth.Truth.assertAbout; -import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatRecords; +package org.apache.beam.it.truthmatchers; +import com.google.common.truth.Truth; import java.util.List; import org.apache.avro.generic.GenericRecord; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.truthmatchers.RecordsSubject; +import org.apache.beam.it.common.artifacts.Artifact; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; public class ArtifactAsserts { @@ -36,7 +33,7 @@ public class ArtifactAsserts { * @return Truth Subject to chain assertions. */ public static ArtifactsSubject assertThatArtifacts(List artifacts) { - return assertAbout(ArtifactsSubject.records()).that(artifacts); + return Truth.assertAbout(ArtifactsSubject.records()).that(artifacts); } /** @@ -47,7 +44,7 @@ public static ArtifactsSubject assertThatArtifacts(List artifacts) { * @return Truth Subject to chain assertions. */ public static ArtifactsSubject assertThatArtifact(Artifact artifact) { - return assertAbout(ArtifactsSubject.records()).that(ImmutableList.of(artifact)); + return Truth.assertAbout(ArtifactsSubject.records()).that(ImmutableList.of(artifact)); } /** @@ -57,6 +54,6 @@ public static ArtifactsSubject assertThatArtifact(Artifact artifact) { * @return Truth Subject to chain assertions. */ public static RecordsSubject assertThatGenericRecords(List records) { - return assertThatRecords(ArtifactsSubject.genericRecordToRecords(records)); + return PipelineAsserts.assertThatRecords(ArtifactsSubject.genericRecordToRecords(records)); } } diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactsSubject.java b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactsSubject.java similarity index 87% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactsSubject.java rename to it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactsSubject.java index ab5b49699484..e2518280345d 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/ArtifactsSubject.java +++ b/it/truthmatchers/src/main/java/org/apache/beam/it/truthmatchers/ArtifactsSubject.java @@ -15,11 +15,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.it.gcp.artifacts.matchers; +package org.apache.beam.it.truthmatchers; -import static org.apache.beam.it.gcp.artifacts.matchers.ArtifactAsserts.assertThatGenericRecords; -import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatRecords; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing.sha256; +import static org.apache.beam.it.truthmatchers.ArtifactAsserts.assertThatGenericRecords; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; @@ -31,11 +29,11 @@ import java.util.Map; import org.apache.avro.Schema; import org.apache.avro.generic.GenericRecord; -import org.apache.beam.it.gcp.artifacts.Artifact; -import org.apache.beam.it.gcp.artifacts.utils.AvroTestUtil; -import org.apache.beam.it.gcp.artifacts.utils.JsonTestUtil; -import org.apache.beam.it.gcp.artifacts.utils.ParquetTestUtil; -import org.apache.beam.it.truthmatchers.RecordsSubject; +import org.apache.beam.it.common.artifacts.Artifact; +import org.apache.beam.it.common.artifacts.utils.AvroTestUtil; +import org.apache.beam.it.common.artifacts.utils.JsonTestUtil; +import org.apache.beam.it.common.artifacts.utils.ParquetTestUtil; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; /** * Subject that has assertion operations for artifact lists (GCS files), usually coming from the @@ -98,7 +96,8 @@ public void hasContent(String content) { */ public void hasHash(String hash) { if (actual.stream() - .noneMatch(artifact -> sha256().hashBytes(artifact.contents()).toString().equals(hash))) { + .noneMatch( + artifact -> Hashing.sha256().hashBytes(artifact.contents()).toString().equals(hash))) { failWithActual("expected to contain hash", hash); } } @@ -167,6 +166,6 @@ public RecordsSubject asJsonRecords() { throw new RuntimeException("Error reading " + artifact.name() + " as JSON.", e); } } - return assertThatRecords(allRecords); + return PipelineAsserts.assertThatRecords(allRecords); } } diff --git a/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto b/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto index 6351b2472579..faae80261402 100644 --- a/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto +++ b/model/interactive/src/main/proto/org/apache/beam/model/interactive/v1/beam_interactive_api.proto @@ -50,11 +50,15 @@ message TestStreamFileRecord { org.apache.beam.model.pipeline.v1.TestStreamPayload.Event recorded_event = 1; } +// A gRPC service that serves cached TestStream events to a TestStream source. +// Used by Interactive Beam to replay recorded pipeline elements from a streaming cache. service TestStreamService { // A TestStream will request for events using this RPC. rpc Events(EventsRequest) returns (stream org.apache.beam.model.pipeline.v1.TestStreamPayload.Event) {} } +// A request for TestStream events from a streaming cache. +// The client specifies which PCollection output tags to read events for. message EventsRequest { // The set of PCollections to read from. These are the PTransform outputs // local names. These are a subset of the TestStream's outputs. This allows diff --git a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto index 35cd4f0c73d1..a64352c9289e 100644 --- a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto +++ b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_artifact_api.proto @@ -78,18 +78,23 @@ message ResolveArtifactsResponse { // A request to get an artifact. message GetArtifactRequest { + // (Required) The artifact to retrieve, typically a resolved artifact from ResolveArtifacts. org.apache.beam.model.pipeline.v1.ArtifactInformation artifact = 1; } // Part of a response to getting an artifact. message GetArtifactResponse { + // A chunk of the artifact's binary content. bytes data = 1; } // Wraps an ArtifactRetrievalService request for use in ReverseArtifactRetrievalService. +// The server sends these requests to the client during the reverse staging flow. message ArtifactRequestWrapper { oneof request { + // Request to resolve a set of artifacts into concrete artifact references. ResolveArtifactsRequest resolve_artifact = 1000; + // Request to retrieve the bytes of a specific artifact. GetArtifactRequest get_artifact = 1001; } } @@ -151,16 +156,22 @@ message ArtifactMetadata { // A collection of artifacts. message Manifest { + // (Required) The list of artifacts staged for a job. repeated ArtifactMetadata artifact = 1; } // A manifest with location information. message ProxyManifest { + // (Required) The manifest of staged artifacts. Manifest manifest = 1; + // A location (name + URI) at which a staged artifact can be retrieved. message Location { - string name = 1; - string uri = 2; + // (Required) The name of the artifact, matching the name in the Manifest. + string name = 1; + // (Required) The URI at which the artifact content can be retrieved. + string uri = 2; } + // (Required) The list of artifact locations for retrieval. repeated Location location = 2; } @@ -173,6 +184,7 @@ message GetManifestRequest { // A response containing a job manifest. message GetManifestResponse { + // (Required) The manifest of staged artifacts for the job. Manifest manifest = 1; } @@ -187,9 +199,11 @@ message LegacyGetArtifactRequest { // Part of an artifact. message ArtifactChunk { + // A chunk of the artifact's binary content. bytes data = 1; } +// Metadata for an artifact being staged, sent as the first message in a PutArtifact stream. message PutArtifactMetadata { // (Required) A token for artifact staging session. This token can be obtained // from PrepareJob request in JobService @@ -211,6 +225,7 @@ message PutArtifactRequest { } } +// A response to staging an artifact. An empty response indicating the artifact was staged successfully. message PutArtifactResponse { } diff --git a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto index a4736f8b4938..0e3b1adbffd6 100644 --- a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto +++ b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_expansion_api.proto @@ -33,6 +33,8 @@ import "google/protobuf/struct.proto"; import "org/apache/beam/model/pipeline/v1/beam_runner_api.proto"; import "org/apache/beam/model/pipeline/v1/schema.proto"; +// A request to expand a single PTransform in a remote SDK. The expansion service +// resolves the transform into its constituent subtransforms and outputs. message ExpansionRequest { // Set of components needed to interpret the transform, or which // may be useful for its expansion. This includes the input @@ -64,6 +66,8 @@ message ExpansionRequest { google.protobuf.Struct pipeline_options = 6; } +// The result of expanding a PTransform, containing the expanded transform and +// all newly created components. message ExpansionResponse { // Set of components needed to execute the expanded transform, // including the (original) inputs, outputs, and subtransforms. @@ -82,9 +86,13 @@ message ExpansionResponse { string error = 10; } +// A request to discover all SchemaTransformProviders registered with the expansion service. +// This is used by cross-language pipeline construction to enumerate available transforms. message DiscoverSchemaTransformRequest { } +// Configuration metadata for a SchemaTransform, describing its inputs, outputs, +// and configuration schema. Used to enable cross-language transform discovery. message SchemaTransformConfig { // Config schema of the SchemaTransform org.apache.beam.model.pipeline.v1.Schema config_schema = 1; @@ -102,6 +110,7 @@ message SchemaTransformConfig { string description = 4; } +// A response containing all discovered SchemaTransform configurations. message DiscoverSchemaTransformResponse { // A mapping from SchemaTransform ID to schema transform config of discovered // SchemaTransforms diff --git a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto index 8e196d5bedd3..353eaec57994 100644 --- a/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto +++ b/model/job-management/src/main/proto/org/apache/beam/model/job_management/v1/beam_job_api.proto @@ -112,6 +112,7 @@ message RunJobRequest { } +// The response from submitting a job for execution. message RunJobResponse { string job_id = 1; // (required) The ID for the executing job } @@ -144,6 +145,7 @@ message DrainJobResponse { } // A subset of info provided by ProvisionApi.ProvisionInfo +// Represents metadata about a submitted job. message JobInfo { string job_id = 1; // (required) string job_name = 2; // (required) @@ -155,6 +157,7 @@ message JobInfo { // Throws error GRPC_STATUS_UNAVAILABLE if server is down message GetJobsRequest { } +// The response from requesting a list of all invoked jobs. message GetJobsResponse { repeated JobInfo job_info = 1; // (required) } @@ -168,6 +171,7 @@ message GetJobStateRequest { } +// An event representing a job state transition. Emitted by GetState and GetStateStream. message JobStateEvent { JobState.Enum state = 1; // (required) google.protobuf.Timestamp timestamp = 2; // (required) @@ -182,6 +186,7 @@ message GetJobPipelineRequest { } +// The response from requesting a job's pipeline representation. message GetJobPipelineResponse { org.apache.beam.model.pipeline.v1.Pipeline pipeline = 1; // (required) } @@ -195,25 +200,41 @@ message JobMessagesRequest { string job_id = 1; // (required) } +// A single log message or diagnostic event from a running job. message JobMessage { + // (Required) A unique identifier for this message. string message_id = 1; + // (Required) The time at which the message was emitted, as a string representation. string time = 2; + // (Required) The importance level of this message. MessageImportance importance = 3; + // (Required) The text content of the message. string message_text = 4; + // Importance levels for job messages, ordered from least to most severe. enum MessageImportance { + // The importance was not specified. MESSAGE_IMPORTANCE_UNSPECIFIED = 0; + // Debug-level messages, typically very verbose. JOB_MESSAGE_DEBUG = 1; + // Detailed informational messages. JOB_MESSAGE_DETAILED = 2; + // Basic informational messages. JOB_MESSAGE_BASIC = 3; + // Warning messages indicating potential issues. JOB_MESSAGE_WARNING = 4; + // Error messages indicating failures. JOB_MESSAGE_ERROR = 5; } } +// A streaming response from GetMessageStream, containing either a job message +// or a job state change event. message JobMessagesResponse { oneof response { + // A log message or diagnostic event from the job. JobMessage message_response = 1; + // A job state transition event. JobStateEvent state_response = 2; } } @@ -270,17 +291,22 @@ message JobState { } +// A request to fetch metrics for a given job. message GetJobMetricsRequest { string job_id = 1; // (required) } +// A response containing metrics for a given job. message GetJobMetricsResponse { + // (Required) The metrics results containing both attempted and committed values. MetricResults metrics = 1; } // All metrics for a given job. Runners may support one or the other or both. message MetricResults { + // Metrics reflecting the result of attempted (non-committed) computations. repeated org.apache.beam.model.pipeline.v1.MonitoringInfo attempted = 1; + // Metrics reflecting the result of committed computations. repeated org.apache.beam.model.pipeline.v1.MonitoringInfo committed = 2; } @@ -296,12 +322,17 @@ message DescribePipelineOptionsRequest { // Types mirror those of JSON, since that's how pipeline options are serialized. message PipelineOptionType { enum Enum { + // A string value. STRING = 0; + // A boolean (true/false) value. BOOLEAN = 1; // whole numbers, see https://json-schema.org/understanding-json-schema/reference/numeric.html INTEGER = 2; + // A floating-point number. NUMBER = 3; + // An array of values. ARRAY = 4; + // A nested object. OBJECT = 5; }; } @@ -324,6 +355,7 @@ message PipelineOptionDescriptor { string group = 5; } +// The response from describing pipeline options, containing a list of option descriptors. message DescribePipelineOptionsResponse { // List of pipeline option descriptors. repeated PipelineOptionDescriptor options = 1; diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java index b26e865526dd..d7c8912ded01 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkDetachedRunnerResult.java @@ -120,11 +120,11 @@ public synchronized State drain() throws IOException { private State getDrainState(CompletableFuture drainFuture) throws IOException { if (!drainFuture.isDone()) { - return State.RUNNING; + return State.DRAINING; } try { drainFuture.get(); - return State.DONE; + return State.DRAINED; } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IOException("Failed to drain Flink job", e); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java index 908d940f5efb..09ceede9d586 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerResultTest.java @@ -66,18 +66,18 @@ public void testDrainDoneResultDoesNotThrowAnException() throws Exception { } @Test - public void testDetachedDrainReturnsRunningThenDone() throws Exception { + public void testDetachedDrainReturnsDrainingThenDrained() throws Exception { JobClient jobClient = mock(JobClient.class); CompletableFuture drainFuture = new CompletableFuture<>(); when(jobClient.stopWithSavepoint(true, null, SavepointFormatType.DEFAULT)) .thenReturn(drainFuture); FlinkDetachedRunnerResult result = new FlinkDetachedRunnerResult(jobClient, 1); - assertThat(result.drain(), is(PipelineResult.State.RUNNING)); - assertThat(result.getState(), is(PipelineResult.State.RUNNING)); + assertThat(result.drain(), is(PipelineResult.State.DRAINING)); + assertThat(result.getState(), is(PipelineResult.State.DRAINING)); drainFuture.complete("savepoint"); - assertThat(result.getState(), is(PipelineResult.State.DONE)); + assertThat(result.getState(), is(PipelineResult.State.DRAINED)); verify(jobClient).stopWithSavepoint(true, null, SavepointFormatType.DEFAULT); } @@ -132,11 +132,11 @@ public void testDetachedDrainRetriesAfterFailure() throws Exception { result.drain(); fail("Expected IOException"); } catch (IOException expected) { - assertThat(result.drain(), is(PipelineResult.State.RUNNING)); + assertThat(result.drain(), is(PipelineResult.State.DRAINING)); } retryDrainFuture.complete("savepoint"); - assertThat(result.getState(), is(PipelineResult.State.DONE)); + assertThat(result.getState(), is(PipelineResult.State.DRAINED)); verify(jobClient, times(2)).stopWithSavepoint(true, null, SavepointFormatType.DEFAULT); } } diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java index 0d7e5eaf68d9..278dace6ef48 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineJob.java @@ -365,6 +365,7 @@ private void logTerminalState(State state) { switch (state) { case DONE: case CANCELLED: + case DRAINED: LOG.info("Job {} finished with status {}.", getJobId(), state); break; case UPDATED: diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java index 011f60f4fd14..1b03cc10351e 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java @@ -2535,8 +2535,9 @@ private String getJobIdFromName(String jobName) { listResult = dataflowClient.listJobs(token); token = listResult.getNextPageToken(); for (Job job : listResult.getJobs()) { + State state = MonitoringUtil.toState(job.getCurrentState()); if (job.getName().equals(jobName) - && MonitoringUtil.toState(job.getCurrentState()).equals(State.RUNNING)) { + && (state.equals(State.RUNNING) || state.equals(State.DRAINING))) { return job.getId(); } } diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java index d117cf786129..0e25bbd265ae 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/MonitoringUtil.java @@ -221,17 +221,19 @@ public static State toState(@Nullable String stateName) { return State.CANCELLED; case "JOB_STATE_UPDATED": return State.UPDATED; + case "JOB_STATE_DRAINING": + return State.DRAINING; + case "JOB_STATE_DRAINED": + return State.DRAINED; case "JOB_STATE_RUNNING": case "JOB_STATE_PENDING": // Job has not yet started; closest mapping is RUNNING - case "JOB_STATE_DRAINING": // Job is still active; the closest mapping is RUNNING case "JOB_STATE_CANCELLING": // Job is still active; the closest mapping is RUNNING case "JOB_STATE_PAUSING": // Job is still active; the closest mapping is RUNNING case "JOB_STATE_RESOURCE_CLEANING_UP": // Job is still active; the closest mapping is RUNNING return State.RUNNING; case "JOB_STATE_DONE": - case "JOB_STATE_DRAINED": // Job has successfully terminated; closest mapping is DONE return State.DONE; default: LOG.warn( diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java index 4b088eb41a7d..be7ad2e6e110 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java @@ -417,7 +417,7 @@ public void testDrainUnterminatedJobThatSucceeds() throws IOException { DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); - assertEquals(State.RUNNING, job.drain()); + assertEquals(State.DRAINING, job.drain()); Job content = new Job(); content.setProjectId(PROJECT_ID); content.setId(JOB_ID); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java index 073b30f928dc..6c32dd749d0b 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java @@ -285,6 +285,11 @@ private static Pipeline buildDataflowPipelineWithLargeGraph(DataflowPipelineOpti } static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) throws IOException { + return buildMockDataflow(mockJobs, "JOB_STATE_RUNNING"); + } + + static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs, String currentState) + throws IOException { Dataflow mockDataflowClient = mock(Dataflow.class); Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); Dataflow.Projects.Locations mockLocations = mock(Dataflow.Projects.Locations.class); @@ -308,7 +313,7 @@ static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) thr new Job() .setName("oldjobname") .setId("oldJobId") - .setCurrentState("JOB_STATE_RUNNING")))); + .setCurrentState(currentState)))); Job resultJob = new Job(); resultJob.setId("newid"); @@ -375,6 +380,10 @@ static GcsUtil buildMockGcsUtil() throws IOException { } private DataflowPipelineOptions buildPipelineOptions() throws IOException { + return buildPipelineOptions("JOB_STATE_RUNNING"); + } + + private DataflowPipelineOptions buildPipelineOptions(String currentState) throws IOException { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); options.setRunner(DataflowRunner.class); options.setProject(PROJECT_ID); @@ -382,7 +391,7 @@ private DataflowPipelineOptions buildPipelineOptions() throws IOException { options.setRegion(REGION_ID); // Set FILES_PROPERTY to empty to prevent a default value calculated from classpath. options.setFilesToStage(new ArrayList<>()); - options.setDataflowClient(buildMockDataflow(mockJobs)); + options.setDataflowClient(buildMockDataflow(mockJobs, currentState)); options.setGcsUtil(mockGcsUtil); options.setGcpCredential(new TestCredential()); @@ -793,6 +802,19 @@ public void testUpdate() throws IOException { assertValidJob(jobCaptor.getValue()); } + @Test + public void testUpdateDrainingJob() throws IOException { + DataflowPipelineOptions options = buildPipelineOptions("JOB_STATE_DRAINING"); + options.setUpdate(true); + options.setJobName("oldJobName"); + Pipeline p = buildDataflowPipeline(options); + p.run(); + + ArgumentCaptor jobCaptor = ArgumentCaptor.forClass(Job.class); + Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture()); + assertEquals("oldJobId", jobCaptor.getValue().getReplaceJobId()); + } + @Test public void testUploadGraph() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java index 5f76b6750ffa..35af1be4cdfc 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/MonitoringUtilTest.java @@ -100,9 +100,9 @@ public void testToStateNormal() { // Non-trivially mapped cases assertEquals(State.STOPPED, MonitoringUtil.toState("JOB_STATE_PAUSED")); - assertEquals(State.RUNNING, MonitoringUtil.toState("JOB_STATE_DRAINING")); + assertEquals(State.DRAINING, MonitoringUtil.toState("JOB_STATE_DRAINING")); assertEquals(State.RUNNING, MonitoringUtil.toState("JOB_STATE_PAUSING")); - assertEquals(State.DONE, MonitoringUtil.toState("JOB_STATE_DRAINED")); + assertEquals(State.DRAINED, MonitoringUtil.toState("JOB_STATE_DRAINED")); } @Test diff --git a/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java b/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java index 9da3bf38e575..3bd0a1311b04 100644 --- a/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java +++ b/runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvocation.java @@ -116,6 +116,12 @@ public void onSuccess(PortablePipelineResult pipelineResult) { case RUNNING: setState(JobState.Enum.RUNNING); break; + case DRAINING: + setState(JobState.Enum.DRAINING); + break; + case DRAINED: + setState(JobState.Enum.DRAINED); + break; case CANCELLED: setState(JobState.Enum.CANCELLED); break; @@ -169,9 +175,12 @@ public synchronized void cancel() { new FutureCallback() { @Override public void onSuccess(PortablePipelineResult pipelineResult) { - // Do not cancel when we are already done. - if (pipelineResult != null - && pipelineResult.getState() != PipelineResult.State.DONE) { + // Do not cancel when the runner has already successfully finished. + if (pipelineResult != null) { + PipelineResult.State state = pipelineResult.getState(); + if (state == PipelineResult.State.DONE || state == PipelineResult.State.DRAINED) { + return; + } try { pipelineResult.cancel(); setState(JobState.Enum.CANCELLED); diff --git a/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java b/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java index 31380ace43c3..7100e57717be 100644 --- a/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java +++ b/runners/java-job-service/src/test/java/org/apache/beam/runners/jobsubmission/JobInvocationTest.java @@ -80,6 +80,28 @@ public void testStateAfterCompletion() throws Exception { awaitJobState(jobInvocation, JobApi.JobState.Enum.DONE); } + @Test(timeout = 10_000) + public void testStateAfterDrainCompleted() throws Exception { + jobInvocation.start(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.RUNNING)); + + TestPipelineResult pipelineResult = new TestPipelineResult(PipelineResult.State.DRAINED); + runner.setResult(pipelineResult); + + awaitJobState(jobInvocation, JobApi.JobState.Enum.DRAINED); + } + + @Test(timeout = 10_000) + public void testStateAfterDrainStarted() throws Exception { + jobInvocation.start(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.RUNNING)); + + TestPipelineResult pipelineResult = new TestPipelineResult(PipelineResult.State.DRAINING); + runner.setResult(pipelineResult); + + awaitJobState(jobInvocation, JobApi.JobState.Enum.DRAINING); + } + @Test(timeout = 10_000) public void testStateAfterCompletionWithoutResult() throws Exception { jobInvocation.start(); @@ -128,6 +150,20 @@ public void testNoCancellationWhenDone() throws Exception { assertThat(pipelineResult.cancelLatch.getCount(), is(1L)); } + @Test(timeout = 10_000) + public void testNoCancellationWhenDrained() throws Exception { + jobInvocation.start(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.RUNNING)); + + TestPipelineResult pipelineResult = new TestPipelineResult(PipelineResult.State.DRAINED); + runner.setResult(pipelineResult); + awaitJobState(jobInvocation, JobApi.JobState.Enum.DRAINED); + + jobInvocation.cancel(); + assertThat(jobInvocation.getState(), is(JobApi.JobState.Enum.DRAINED)); + assertThat(pipelineResult.cancelLatch.getCount(), is(1L)); + } + @Test(timeout = 10_000) public void testReturnsMetricsFromJobInvocationAfterSuccess() throws Exception { JobApi.MetricResults expectedMonitoringInfos = JobApi.MetricResults.newBuilder().build(); diff --git a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java index 0f3d3b6d1910..7c72a1039c0e 100644 --- a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java +++ b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java @@ -160,7 +160,7 @@ private void waitForTerminalState() { } private void propagateErrors() { - if (terminalState != State.DONE) { + if (terminalState != State.DONE && terminalState != State.DRAINED) { JobMessagesRequest messageStreamRequest = JobMessagesRequest.newBuilder().setJobIdBytes(jobId).build(); Iterator messageStreamIterator = @@ -196,10 +196,9 @@ private static State getJavaState(JobApi.JobState.Enum protoState) { case UPDATED: return State.UPDATED; case DRAINING: - // TODO: Determine the correct mappings for the states below. - return State.UNKNOWN; + return State.DRAINING; case DRAINED: - return State.UNKNOWN; + return State.DRAINED; case STARTING: return State.RUNNING; case CANCELLING: diff --git a/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java b/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java index 68f3f6eae396..068d70b0d250 100644 --- a/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java +++ b/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java @@ -109,6 +109,27 @@ public void stagesAndRunsJob() throws Exception { assertThat(state, is(State.DONE)); } + @Test + public void mapsDrainingJobState() throws Exception { + createJobServer(JobState.Enum.DRAINING, JobApi.MetricResults.getDefaultInstance()); + PortableRunner runner = PortableRunner.create(options, ManagedChannelFactory.createInProcess()); + PipelineResult result = runner.run(p); + + try { + assertThat(result.getState(), is(State.DRAINING)); + } finally { + ((AutoCloseable) result).close(); + } + } + + @Test + public void mapsDrainedJobState() throws Exception { + createJobServer(JobState.Enum.DRAINED, JobApi.MetricResults.getDefaultInstance()); + PortableRunner runner = PortableRunner.create(options, ManagedChannelFactory.createInProcess()); + State state = runner.run(p).waitUntilFinish(); + assertThat(state, is(State.DRAINED)); + } + @Test public void extractsMetrics() throws Exception { JobApi.MetricResults metricResults = generateMetricResults(); diff --git a/sdks/go.mod b/sdks/go.mod index c438f2b351d2..12bd9e9ba2f4 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -31,12 +31,12 @@ require ( cloud.google.com/go/profiler v0.6.0 cloud.google.com/go/pubsub v1.51.0 cloud.google.com/go/spanner v1.93.0 - cloud.google.com/go/storage v1.63.1 - github.com/aws/aws-sdk-go-v2 v1.42.1 - github.com/aws/aws-sdk-go-v2/config v1.32.30 - github.com/aws/aws-sdk-go-v2/credentials v1.19.29 - github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34 - github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2 + cloud.google.com/go/storage v1.64.0 + github.com/aws/aws-sdk-go-v2 v1.43.0 + github.com/aws/aws-sdk-go-v2/config v1.32.31 + github.com/aws/aws-sdk-go-v2/credentials v1.19.30 + github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.35 + github.com/aws/aws-sdk-go-v2/service/s3 v1.106.0 github.com/aws/smithy-go v1.27.4 github.com/docker/go-connections v0.7.0 // indirect github.com/dustin/go-humanize v1.0.1 @@ -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.4.1 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.5.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 @@ -148,17 +148,17 @@ require ( github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40 // indirect github.com/apache/thrift v0.23.0 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect - github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect - 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.32.1 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 // indirect - github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.24 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.32 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.45.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 ba65418653ef..81b86614757b 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -99,8 +99,8 @@ cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RX cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0= cloud.google.com/go/storage v1.12.0/go.mod h1:fFLk2dp2oAhDz8QFKwqrjdJvxSp/W2g7nillojlL5Ho= cloud.google.com/go/storage v1.21.0/go.mod h1:XmRlxkgPjlBONznT2dDUU/5XlpU2OjMnKuqnZI01LAA= -cloud.google.com/go/storage v1.63.1 h1:CYXILV9G4CH0C18IQ9+V0h4XiqD2LhKnMLO0o7uJWNs= -cloud.google.com/go/storage v1.63.1/go.mod h1:lWyAtwvDZHdL3k68WVKbESP6bmWaV23ZJJ/JEVw/ZaQ= +cloud.google.com/go/storage v1.64.0 h1:KLpxI/oX9LxeRsNqn877d2WyeT3ryiEwnGt8pwcSPZg= +cloud.google.com/go/storage v1.64.0/go.mod h1:lWyAtwvDZHdL3k68WVKbESP6bmWaV23ZJJ/JEVw/ZaQ= cloud.google.com/go/trace v1.0.0/go.mod h1:4iErSByzxkyHWzzlAj63/Gmjz0NH1ASqhJguHpGcr6A= cloud.google.com/go/trace v1.2.0/go.mod h1:Wc8y/uYyOhPy12KEnXG9XGrvfMz5F5SrYecQlbW1rwM= cloud.google.com/go/trace v1.16.0 h1:GmQovzFc5F0CNfl0VLgL64aoTtu7xsM0YajW2GlG9+E= @@ -196,79 +196,79 @@ github.com/aws/aws-sdk-go v1.37.0/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zK github.com/aws/aws-sdk-go v1.43.31/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo= github.com/aws/aws-sdk-go-v2 v1.16.2/go.mod h1:ytwTPBG6fXTZLxxeeCCWj2/EMYp/xDUgX+OET6TLNNU= github.com/aws/aws-sdk-go-v2 v1.23.0/go.mod h1:i1XDttT4rnf6vxc9AuskLc6s7XBee8rlLilKlc03uAA= -github.com/aws/aws-sdk-go-v2 v1.42.1 h1:9eOTgu1z/dVtYpNZ3/8/XbbaX0x/BqE3HUzAzs6K0ek= -github.com/aws/aws-sdk-go-v2 v1.42.1/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= +github.com/aws/aws-sdk-go-v2 v1.43.0 h1:fharf/WhbRAVZ1du0QL7roNFxZ6T/sWr+4Ni617bwSI= +github.com/aws/aws-sdk-go-v2 v1.43.0/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.1/go.mod h1:n8Bs1ElDD2wJ9kCRTczA83gYbBmjSwZp3umc6zF4EeM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.5.1/go.mod h1:t8PYl/6LzdAqsU4/9tz28V/kU+asFePvpOMkdul0gEQ= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHzbkHfPzy4ACPcEjVG0x87KOwtpqGY= 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.30 h1:XwsEzpTJfQYJbFicz/QMLwAZdyeNVVoOEkbF7R3gPJk= -github.com/aws/aws-sdk-go-v2/config v1.32.30/go.mod h1:Ud32SuMc+/9BGxfpSVld7HrE2o05JwKmXY4M3jOQNZU= +github.com/aws/aws-sdk-go-v2/config v1.32.31 h1:n4nY9O3QKoHIkL85EX+V8RcMFtOhlpTFhGArg915PXk= +github.com/aws/aws-sdk-go-v2/config v1.32.31/go.mod h1:PN0NYDCCoOpGGsZ2+elDUidmHfQBPyYzN2GCgl8HEBs= 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.29 h1:WHZGssHH887cO0ox07SIQZsFx3MKD4ps6w0xUEmnKYQ= -github.com/aws/aws-sdk-go-v2/credentials v1.19.29/go.mod h1:Mhl0xR6zjguiuj00XRx2wMx22sAltk7oya39sT7fdg8= +github.com/aws/aws-sdk-go-v2/credentials v1.19.30 h1:TTCvvzFU6gXa4iJecNG/0F/B0oYTiazoRECr2XyLHrY= +github.com/aws/aws-sdk-go-v2/credentials v1.19.30/go.mod h1:jKxAp2AEncnliinzpgOSZDFv6+VjvWhjw/AtbfsWT9U= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.3/go.mod h1:uk1vhHHERfSVCUnqSqz8O48LBYDSC+k6brng09jcMOk= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.14.4/go.mod h1:t4i+yGHMCcUNIX1x7YVYa6bH/Do7civ5I6cG/6PMfyA= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQu0EtP23pP553Cf9KgsDV4= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30/go.mod h1:/3AOgy4K17Dm4ucMZVC/MJkzy5kmfKUcINRHZyo0koQ= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 h1:kfVL5wAunCJycL6MOQ6aNh6PlAYEymflcjuKmrWUA0o= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31/go.mod h1:nWfRNDAppujCQgOUd43lKT4yeLv9z3nJ3bw1G3BgQKo= 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.34 h1:Pn7OsMwBLbkZ6OnCxWHAjf0L/22H8cnhxZC0uPwtMtg= -github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.34/go.mod h1:eToXR/Gk1uqpn04eSmdgVXwfS0WvH8aG4eBFr8ygbpU= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.35 h1:TwCjUC1rnFKTtfqEpQY9ClYPFpGpUaouODrdGPB5b3Y= +github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.35/go.mod h1:V0zqtP3iJk9zu86GxuQGN09RYyQB+3mjcPiyfLC6wlg= 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= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30/go.mod h1:WueJeNDZvK1fMYEWJIkcivBfEzUkTpBhzlrUKKY8EuA= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 h1:Z8F3hfCY33IGpJjFAnv0wvtv1FIKj1GHmRDEYqy64tw= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31/go.mod h1:aVyUoytEyOViR6jhq6jula0xkc5NfBE2hgeF6BvOrao= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.3/go.mod h1:ssOhaLpRlh88H3UmEcsBoVKq309quMvm3Ds8e9d4eJM= github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.5.3/go.mod h1:ify42Rb7nKeDDPkFjKn7q1bPscVPu/+gmHH8d2c+anU= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 h1:jn46zC9LdsVR/ZpMIJqMqb8hHv31BlLx3ulVqNspUOk= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30/go.mod h1:1hTMsAgbdS/AtUi4bw8+gUuh1pceo+eXRLfpSuSQj3M= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 h1:hyOxUyXdh3AyjE93gBgsfziJag9ACwcs+ZpDBLzi8mw= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31/go.mod h1:OERqI9k0draSLB8O8woxY3q25ZWTELRK4RRoLMuMZFo= github.com/aws/aws-sdk-go-v2/internal/ini v1.3.10/go.mod h1:8DcYQcz0+ZJaSxANlHIsbbi6S+zMwjwdDqwW3r9AzaE= github.com/aws/aws-sdk-go-v2/internal/ini v1.7.1/go.mod h1:6fQQgfuGmw8Al/3M2IgIllycxV7ZW7WCdVSqfBeUiCY= github.com/aws/aws-sdk-go-v2/internal/v4a v1.2.3/go.mod h1:5yzAuE9i2RkVAttBl8yxZgQr5OCq4D5yDnG7j9x2L0U= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 h1:3GUprIsfmGcC5SACIyB0e7E0BM1O1b3Erl5CePYIAeQ= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31/go.mod h1:7PuV1yl5e2xnUbm+RqvVg5i2iBM8EyijZNoI9wsOoOc= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 h1:0MrUL35H/Y4kdFfItoR5jCgtDQ4Z/8LudAoIHRfA4hE= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32/go.mod h1:2tNZkuWz54arj8mHVf+8Y7cKkcD8Wr/fBpENgEXpjLc= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.1/go.mod h1:GeUru+8VzrTXV/83XyMJ80KpH8xO89VPoUileyNQ+tc= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.10.1/go.mod h1:l9ymW25HOqymeU2m1gbUQ3rUIsTwKs8gYHXkqDQUhiI= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur/BiHK6SKPjoBIXSE/hJ6g6JGRLuxQy1jGjlN4= github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13/go.mod h1:ITg9em2KbJx1s0y4aqRX5OYWG6HBZ5TVR//OdpEZ2CQ= github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.3/go.mod h1:Seb8KNmD6kVTjwRjVEgOT5hPin6sq+v4C2ycJQDwuH8= github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.2.3/go.mod h1:R+/S1O4TYpcktbVwddeOYg+uwUfLhADP2S/x4QwsCTM= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23 h1:9Fjh6fi/U5JEStVZijmaMpUwE/gvBJj7x2B/PjbO9To= -github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.23/go.mod h1:iMoT2f1tClxrWAAnKCXjZQ6LOmfLrMG14wmnWpM+F14= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.24 h1:mdPwDQPqxlw9Sc62Nt15yjEcARaDbPXkjRYtXsUripo= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.24/go.mod h1:ls5ytnwLTcQaUu32fMYXFI3MjpKuTwL840PAm9iqyEg= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.3/go.mod h1:wlY6SVjuwvh3TVRpTqdy4I1JpBFLX4UGeKZdWntaocw= github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.10.3/go.mod h1:Owv1I59vaghv1Ax8zz8ELY8DN7/Y0rGS+WWAmjgi950= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 h1:/Z5jmNrKsSD7EmDjzAPsm/3L9IuOkzaynklJZ1qX7S4= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30/go.mod h1:lEzEZnOosE7zi8Z6royW1cFJTD9fpab4Ul1SBrllewk= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 h1:w2SIhW92DZPFrSL4ksVCr8IYff5OZwIcxg8+95tzvAI= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31/go.mod h1:wAhpCQbkov+IcvjozJbd2xRCoZybUEHNkcFunssNACg= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.3/go.mod h1:Bm/v2IaN6rZ+Op7zX+bOUMdL4fsrYZiD0dsjLhNKwZc= github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.16.3/go.mod h1:KZgs2ny8HsxRIRbDwgvJcHHBZPOzQr/+NtGwnP+w2ec= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31 h1:uao4A3QZ5UmB326V6KF+qRpv9Tjz7IlnlnTbbANntlU= -github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.31/go.mod h1:I/1+z0VwL1GhQyLgkoHDlygpUZ+iTAwOQ/NsftiUL2I= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.32 h1:jWXtZdCnhXa9sGFixRaU2AxT4DIVse9HS4E2f+/KwV0= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.32/go.mod h1:9JS1UpfVvyD/ZPX8GsKb/Pq8scEM+7GP5fqh9SwH7po= 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.105.2 h1:5C00eQYpTrgQXnp6V3P6P7zPElna3AXvlukbANE6nJI= -github.com/aws/aws-sdk-go-v2/service/s3 v1.105.2/go.mod h1:zdmCoFO/dSI7GlrwsPqFJI+WlFnSU4Tc8TJnlXrM1Do= +github.com/aws/aws-sdk-go-v2/service/s3 v1.106.0 h1:7QZWVJZWzHivHWIa+5TELLaBBkbuoj0GPwQtMlJ0sqk= +github.com/aws/aws-sdk-go-v2/service/s3 v1.106.0/go.mod h1:fcvq5L7dK+5cQFicEJwpI6e6Wn8NY2i6yT5wRLYVc7s= 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.4.1 h1:V7ZZ300WPXGjvkyore5DGe0ljVPOxCXie/thWdtSBXE= -github.com/aws/aws-sdk-go-v2/service/signin v1.4.1/go.mod h1:mxC0nT/C8wMMS97DemZPzvUZxvIt+2Iq+eS3JdFZGgg= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.0 h1:OHH5iTQvVGmfHjX/5Q+vFuA/Rf2x6/95aJ/75QCQSm4= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.0/go.mod h1:mCF3AK9PpL49oOrhniUXWAfhVBVQ/XbytoE5eccZUIs= 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.32.1 h1:gYFYh4iLLcAOJRLNPY2aD2g9DIhKn4eof8UkIrr1rTk= -github.com/aws/aws-sdk-go-v2/service/sso v1.32.1/go.mod h1:u8af9Nqkmqnr96f7v9nHqzZT9XBwbXEkTiqT4ROuJSE= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 h1:CaJyYhxBE0M/HJX/YvSaSmQlsI91VHB0lKU8LtLxL3A= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.0/go.mod h1:+e6BMRMPjBQoCw/WovYR9GLy2IU0z4Q77smOB1DraSg= 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.37.1 h1:arjT9Cm3/WYbGmD5TUZHk4UQn4Lle1fUNZs5FC6CtF0= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 h1:tC323YV77QdafeBr6LUhLDTsboyuyHLNRwAyCP44kGU= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0/go.mod h1:SfLK1sgviHmbI+MozR9iDwDjL4cdCVZtahsjoR+z7wg= 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.44.1 h1:RvfHDg+xvAeZ+5741vUEjpOVtYSIm93W2zhx10Xtydw= -github.com/aws/aws-sdk-go-v2/service/sts v1.44.1/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.0 h1:Pd6PNlp4t8PTXxqzstICl52Wsy78vpjFZ7PRUj44mJc= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.0/go.mod h1:rmQ0TnHzuLPmabgjPcsywhsSOmaBDgzR4zvDxSPsGdg= 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.4 h1:JQcphmBN4f0q/sPqXqROIItRNV/hy10cgu7CsFy616M= diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java index 46cca7833e52..eff031157e12 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/PipelineResult.java @@ -107,6 +107,12 @@ enum State { /** The job has been updated. */ UPDATED(true, true), + /** The job is draining: no longer accepting new input while finishing in-flight work. */ + DRAINING(false, false), + + /** The job has finished draining. */ + DRAINED(true, false), + /** The job state reported by a runner cannot be interpreted by the SDK. */ UNRECOGNIZED(false, false); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 449afd6a0243..b744da10cf6f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -767,6 +767,7 @@ public BeamFnApi.InstructionResponse.Builder trySplit(InstructionRequest request /** Shutdown the bundles, running the tearDown() functions. */ public void shutdown() throws Exception { bundleProcessorCache.shutdown(); + beamFnDataClient.close(); } @VisibleForTesting diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java index 1a50f5b448c5..458ed1209f43 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataClient.java @@ -17,6 +17,8 @@ */ package org.apache.beam.fn.harness.data; +import java.io.Closeable; +import java.io.IOException; import java.util.List; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.model.pipeline.v1.Endpoints; @@ -30,7 +32,7 @@ * provide a receiver of outbound elements. Callers can register themselves as receivers for inbound * elements or can get a handle for a receiver of outbound elements. */ -public interface BeamFnDataClient { +public interface BeamFnDataClient extends Closeable { /** * Registers a receiver for the provided instruction id. * @@ -72,4 +74,9 @@ void unregisterReceiver( /** Get the outbound observer for the specified apiServiceDescriptor and dataStreamId. */ StreamObserver getOutboundObserver( Endpoints.ApiServiceDescriptor apiServiceDescriptor, String dataStreamId); + + @Override + default void close() throws IOException { + // Default to no-op + } } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java index 2f2a6b0fc660..79e34fe5765f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java @@ -118,6 +118,18 @@ public void poisonInstructionId(String instructionId) { } } + @Override + public void close() { + for (BeamFnDataGrpcMultiplexer client : multiplexerCache.values()) { + try { + client.close(); + } catch (Exception e) { + LOG.warn("Failed to close multiplexer", e); + } + } + multiplexerCache.clear(); + } + @Override public StreamObserver getOutboundObserver( ApiServiceDescriptor apiServiceDescriptor, String dataStreamId) { diff --git a/sdks/java/io/cdap/build.gradle b/sdks/java/io/cdap/build.gradle index e0af1717e162..cc7e3ba12cde 100644 --- a/sdks/java/io/cdap/build.gradle +++ b/sdks/java/io/cdap/build.gradle @@ -90,3 +90,8 @@ test { // Open java.lang for Gson reflection on StackTraceElement under Java 17+ jvmArgs '--add-opens=java.base/java.lang=ALL-UNNAMED' } + +// spark3_streaming depends on old 'org.lz4:lz4-java'; conflict with kafka-clients:3.9.2 +configurations.all { + resolveCapabilitiesConflict(it, 'org.lz4:lz4-java', 'at.yawk.lz4') +} \ No newline at end of file diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java index ac6650c354d4..281a5b9a47c6 100644 --- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java +++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaCommitOffset.java @@ -25,6 +25,8 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.MapElements; @@ -62,6 +64,9 @@ public class KafkaCommitOffset static class CommitOffsetDoFn extends DoFn, Void> { private static final Logger LOG = LoggerFactory.getLogger(CommitOffsetDoFn.class); + private static final Counter COMMIT_FAILURES = + Metrics.counter(CommitOffsetDoFn.class, "commit-failures"); + private final Map consumerConfig; private final SerializableFunction, Consumer> consumerFactoryFn; @@ -85,6 +90,7 @@ public void processElement(@Element KV element) { element.getKey().getTopicPartition(), new OffsetAndMetadata(element.getValue() + 1))); } catch (Exception e) { + COMMIT_FAILURES.inc(); // TODO: consider retrying. LOG.warn("Getting exception when committing offset: {}", e.getMessage()); } diff --git a/sdks/java/ml/inference/gemini/build.gradle b/sdks/java/ml/inference/gemini/build.gradle new file mode 100644 index 000000000000..f62984eba886 --- /dev/null +++ b/sdks/java/ml/inference/gemini/build.gradle @@ -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. + */ +plugins { + id 'org.apache.beam.module' +} + +applyJavaNature( + automaticModuleName: 'org.apache.beam.sdk.ml.inference.gemini', +) +provideIntegrationTestingDependencies() +enableJavaPerformanceTesting() + +description = "Apache Beam :: SDKs :: Java :: ML :: Inference :: Gemini" +ext.summary = "Gemini model handler for remote inference" + +dependencies { + implementation project(":sdks:java:ml:inference:remote") + implementation "com.google.genai:google-genai:1.59.0" + + + testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") + testImplementation project(path: ":sdks:java:core", configuration: "shadow") + testImplementation library.java.slf4j_api + testRuntimeOnly library.java.slf4j_simple + testImplementation library.java.junit + testImplementation library.java.mockito_core +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java new file mode 100644 index 000000000000..5d74c21347cd --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiImageResponse.java @@ -0,0 +1,33 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; + +/** Wrapper for image (byte array) responses from Gemini. */ +public class GeminiImageResponse implements BaseResponse { + public final byte[] imageBytes; + + public GeminiImageResponse(byte[] imageBytes) { + this.imageBytes = imageBytes; + } + + public byte[] getImageBytes() { + return imageBytes; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java new file mode 100644 index 000000000000..5ba009b7194e --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java @@ -0,0 +1,67 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponse; +import com.google.genai.types.GenerateImagesConfig; +import com.google.genai.types.GenerateImagesResponse; +import java.util.ArrayList; +import java.util.List; + +/** Common inference functions for Gemini. */ +public class GeminiInferenceFunctions { + + /** Generates content from string prompts using the standard generateContent API. */ + public static GeminiRequestFunction + generateFromString() { + return (modelName, batch, client) -> { + List results = new ArrayList<>(); + for (GeminiStringInput input : batch) { + GenerateContentResponse response = + client.models.generateContent( + modelName, input.getText(), GenerateContentConfig.builder().build()); + String text = response.text(); + results.add(new GeminiStringResponse(text != null ? text : "")); + } + return results; + }; + } + + /** Generates images from string prompts using the generateImages API. */ + public static GeminiRequestFunction + generateImageFromString() { + return (modelName, batch, client) -> { + List results = new ArrayList<>(); + for (GeminiStringInput input : batch) { + GenerateImagesResponse response = + client.models.generateImages( + modelName, input.getText(), GenerateImagesConfig.builder().build()); + // Retrieve the base64 string or bytes from the first generated image + List images = response.images(); + if (images != null && !images.isEmpty()) { + byte[] imageBytes = images.get(0).imageBytes().orElse(new byte[0]); + results.add(new GeminiImageResponse(imageBytes)); + } else { + results.add(new GeminiImageResponse(new byte[0])); + } + } + return results; + }; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java new file mode 100644 index 000000000000..ecdbf807816f --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java @@ -0,0 +1,86 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.Client; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; +import org.apache.beam.sdk.ml.inference.remote.BaseModelHandler; +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; + +/** + * Model handler for Google Gemini API inference requests. + * + *

This handler manages communication with Google's Gemini API, including client initialization, + * request formatting, and response parsing. It allows executing a custom {@link + * GeminiRequestFunction} against a batch of inputs. + */ +@SuppressWarnings("nullness") +public class GeminiModelHandler + implements BaseModelHandler, InputT, OutputT> { + + private transient Client client; + private GeminiModelParameters modelParameters; + + @Override + public void createClient(GeminiModelParameters parameters) { + if (parameters == null) { + throw new NullPointerException("GeminiModelParameters must not be null"); + } + this.modelParameters = parameters; + + // Configure client based on vertex or API key + if (parameters.getApiKey() != null) { + if (parameters.getProject() != null || parameters.getLocation() != null) { + throw new IllegalArgumentException("Project and location must be null if API key is set"); + } + this.client = Client.builder().apiKey(parameters.getApiKey()).build(); + } else { + Client.Builder builder = Client.builder(); + if (parameters.getProject() != null && parameters.getLocation() != null) { + builder.vertexAI(true).project(parameters.getProject()).location(parameters.getLocation()); + } else if (parameters.getProject() != null || parameters.getLocation() != null) { + throw new IllegalArgumentException( + "Project and location must both be provided if one is provided"); + } + this.client = builder.build(); + } + } + + @Override + public Iterable> request(List input) { + try { + GeminiRequestFunction requestFn = modelParameters.getRequestFn(); + List responses = requestFn.apply(modelParameters.getModelName(), input, client); + + if (responses.size() != input.size()) { + throw new IllegalStateException("Number of responses must match number of inputs"); + } + + List> results = new ArrayList<>(); + for (int i = 0; i < input.size(); i++) { + results.add(PredictionResult.create(input.get(i), responses.get(i))); + } + return results; + } catch (Exception e) { + throw new RuntimeException("Error during Gemini inference request", e); + } + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java new file mode 100644 index 000000000000..6ac4f578a24e --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java @@ -0,0 +1,60 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.auto.value.AutoValue; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; +import org.apache.beam.sdk.ml.inference.remote.BaseModelParameters; +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; +import org.checkerframework.checker.nullness.qual.Nullable; + +@AutoValue +public abstract class GeminiModelParameters + implements BaseModelParameters { + + public abstract @Nullable String getApiKey(); + + public abstract @Nullable String getProject(); + + public abstract @Nullable String getLocation(); + + public abstract String getModelName(); + + public abstract GeminiRequestFunction getRequestFn(); + + public static + Builder builder() { + return new AutoValue_GeminiModelParameters.Builder(); + } + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setApiKey(String apiKey); + + public abstract Builder setProject(String project); + + public abstract Builder setLocation(String location); + + public abstract Builder setModelName(String modelName); + + public abstract Builder setRequestFn( + GeminiRequestFunction requestFn); + + public abstract GeminiModelParameters build(); + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java new file mode 100644 index 000000000000..15da8b360c96 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java @@ -0,0 +1,31 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.Client; +import java.io.Serializable; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; + +/** Functional interface for custom request functions to the Gemini API. */ +@FunctionalInterface +public interface GeminiRequestFunction + extends Serializable { + List apply(String modelName, List batch, Client client) throws Exception; +} diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java similarity index 69% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java rename to sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java index 9759071b397e..8b0cc7c26892 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/artifacts/matchers/package-info.java +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringInput.java @@ -15,6 +15,19 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +package org.apache.beam.sdk.ml.inference.gemini; -/** Package for Artifact Truth matchers / subjects to have reusable assertions. */ -package org.apache.beam.it.gcp.artifacts.matchers; +import org.apache.beam.sdk.ml.inference.remote.BaseInput; + +/** Wrapper for string inputs to Gemini. */ +public class GeminiStringInput implements BaseInput { + public final String text; + + public GeminiStringInput(String text) { + this.text = text; + } + + public String getText() { + return text; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java new file mode 100644 index 000000000000..94bb4d7168dc --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiStringResponse.java @@ -0,0 +1,33 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import org.apache.beam.sdk.ml.inference.remote.BaseResponse; + +/** Wrapper for string responses from Gemini. */ +public class GeminiStringResponse implements BaseResponse { + public final String text; + + public GeminiStringResponse(String text) { + this.text = text; + } + + public String getText() { + return text; + } +} diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java similarity index 89% rename from it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java rename to sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java index d563025ebeb4..9dee9da9cc74 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/package-info.java +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java @@ -16,5 +16,5 @@ * limitations under the License. */ -/** Data generator for load tests. */ -package org.apache.beam.it.gcp.datagenerator; +/** Gemini model handler for remote inference. */ +package org.apache.beam.sdk.ml.inference.gemini; diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java new file mode 100644 index 000000000000..c1188ad6d04a --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java @@ -0,0 +1,163 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeNotNull; +import static org.junit.Assume.assumeTrue; + +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; +import org.apache.beam.sdk.ml.inference.remote.RemoteInference; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; + +/** + * Execute Gemini model handler integration test. + * + *

+ * ./gradlew :sdks:java:ml:inference:gemini:integrationTest \
+ *   --tests org.apache.beam.sdk.ml.inference.gemini.GeminiModelHandlerIT \
+ *   --info
+ * 
+ */ +public class GeminiModelHandlerIT { + + public static class GeminiStringContentHandler + extends GeminiModelHandler {} + + public static class GeminiStringImageHandler + extends GeminiModelHandler {} + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private String apiKey; + private static final String API_KEY_ENV = "GEMINI_API_KEY"; + private static final String DEFAULT_MODEL = "gemini-3.5-flash"; + + @Before + public void setUp() { + // Get API key + apiKey = System.getenv(API_KEY_ENV); + + // Skip tests if API key is not provided + assumeNotNull( + "Gemini API key not found. Set " + + API_KEY_ENV + + " environment variable to run integration tests.", + apiKey); + assumeTrue( + "Gemini API key is empty. Set " + + API_KEY_ENV + + " environment variable to run integration tests.", + !apiKey.trim().isEmpty()); + } + + @Test + public void testSentimentAnalysisWithSingleInput() { + String input = + "Classify the sentiment of the following text as either positive or negative: This product is absolutely amazing! I love it!"; + + PCollection inputs = + pipeline.apply("CreateSingleInput", Create.of(new GeminiStringInput(input))); + + PCollection>> results = + inputs.apply( + "SentimentInference", + RemoteInference.invoke() + .handler(GeminiStringContentHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName(DEFAULT_MODEL) + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build())); + + // Verify results + PAssert.that(results) + .satisfies( + batches -> { + int count = 0; + for (Iterable> batch : + batches) { + for (PredictionResult result : batch) { + count++; + assertNotNull("Input should not be null", result.getInput()); + assertNotNull("Output should not be null", result.getOutput()); + + String sentiment = result.getOutput().getText().toLowerCase(); + assertTrue( + "Sentiment should be positive or negative, got: " + sentiment, + sentiment.contains("positive") || sentiment.contains("negative")); + } + } + assertEquals("Should have exactly 1 result", 1, count); + return null; + }); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void testGenerateImageFromString() { + String input = "A beautiful sunset over the mountains, digital art style."; + + PCollection inputs = + pipeline.apply("CreateImageInput", Create.of(new GeminiStringInput(input))); + + PCollection>> results = + inputs.apply( + "ImageInference", + RemoteInference.invoke() + .handler(GeminiStringImageHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName("imagen-4.0-generate-001") + .setRequestFn(GeminiInferenceFunctions.generateImageFromString()) + .build())); + + // Verify results + PAssert.that(results) + .satisfies( + batches -> { + int count = 0; + for (Iterable> batch : + batches) { + for (PredictionResult result : batch) { + count++; + assertNotNull("Input should not be null", result.getInput()); + assertNotNull("Output should not be null", result.getOutput()); + assertTrue( + "Output should have generated images", + result.getOutput().getImageBytes().length > 0); + } + } + assertEquals("Should have exactly 1 result", 1, count); + return null; + }); + + pipeline.run().waitUntilFinish(); + } +} diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java new file mode 100644 index 000000000000..125f1e9c3e92 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java @@ -0,0 +1,133 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class GeminiModelHandlerTest { + + @Test + public void testAllParamsSet() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setProject("test-project") + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexLocationParam() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setProject("test-project") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexProjectParam() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testNullParameters() { + GeminiModelHandler handler = + new GeminiModelHandler<>(); + assertThrows(NullPointerException.class, () -> handler.createClient(null)); + } + + @Test + public void testRequest() throws Exception { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setModelName("gemini-model-123") + .setRequestFn( + (modelName, batch, client) -> + Arrays.asList( + new GeminiStringResponse("response1"), + new GeminiStringResponse("response2"))) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + handler.createClient(parameters); + + List input = + Arrays.asList(new GeminiStringInput("input1"), new GeminiStringInput("input2")); + Iterable> results = + handler.request(input); + + int count = 0; + for (PredictionResult result : results) { + if (count == 0) { + assertEquals("input1", result.getInput().getText()); + assertEquals("response1", result.getOutput().getText()); + } else { + assertEquals("input2", result.getInput().getText()); + assertEquals("response2", result.getOutput().getText()); + } + count++; + } + assertEquals(2, count); + } + + @Test + public void testRequestMismatchedResponseSize() throws Exception { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setModelName("gemini-model-123") + .setRequestFn( + (modelName, batch, client) -> Arrays.asList(new GeminiStringResponse("response1"))) + .build(); + GeminiModelHandler handler = + new GeminiModelHandler<>(); + handler.createClient(parameters); + + List input = + Arrays.asList(new GeminiStringInput("input1"), new GeminiStringInput("input2")); + assertThrows(RuntimeException.class, () -> handler.request(input)); + } +} diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java new file mode 100644 index 000000000000..0d43870b880d --- /dev/null +++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java @@ -0,0 +1,74 @@ +/* + * 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. + */ +package org.apache.beam.sdk.ml.inference.remote; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.StructuredCoder; + +/** A {@link Coder} for {@link PredictionResult}. */ +public class PredictionResultCoder + extends StructuredCoder> { + + private final Coder inputCoder; + private final Coder outputCoder; + + private PredictionResultCoder(Coder inputCoder, Coder outputCoder) { + this.inputCoder = inputCoder; + this.outputCoder = outputCoder; + } + + public static PredictionResultCoder of( + Coder inputCoder, Coder outputCoder) { + return new PredictionResultCoder<>(inputCoder, outputCoder); + } + + @Override + public void encode(PredictionResult value, OutputStream outStream) + throws CoderException, IOException { + if (value == null) { + throw new CoderException("cannot encode a null PredictionResult"); + } + inputCoder.encode(value.getInput(), outStream); + outputCoder.encode(value.getOutput(), outStream); + } + + @Override + public PredictionResult decode(InputStream inStream) + throws CoderException, IOException { + InputT input = inputCoder.decode(inStream); + OutputT output = outputCoder.decode(inStream); + return PredictionResult.create(input, output); + } + + @Override + public List> getCoderArguments() { + return Arrays.asList(inputCoder, outputCoder); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + verifyDeterministic(this, "Input coder must be deterministic", inputCoder); + verifyDeterministic(this, "Output coder must be deterministic", outputCoder); + } +} diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java index c571911a484e..a998f4d0f077 100644 --- a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java +++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java @@ -21,6 +21,7 @@ import com.google.auto.value.AutoValue; import java.util.List; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.io.components.throttling.ReactiveThrottler; import org.apache.beam.sdk.transforms.BatchElements; import org.apache.beam.sdk.transforms.DoFn; @@ -88,6 +89,8 @@ public abstract static class Invoke outputCoder(); + abstract Builder builder(); @AutoValue.Builder @@ -108,6 +111,8 @@ abstract Builder setHandler( abstract Builder setOverloadRatio(Double overloadRatio); + abstract Builder setOutputCoder(Coder outputCoder); + abstract Invoke build(); } @@ -163,6 +168,14 @@ public Invoke withOverloadRatio(double overloadRatio) { return builder().setOverloadRatio(overloadRatio).build(); } + /** + * Configures the coder for the output of the model. If not provided, it will fallback to using + * standard Java serialization for the output element. + */ + public Invoke withOutputCoder(Coder outputCoder) { + return builder().setOutputCoder(outputCoder).build(); + } + @Override public PCollection>> expand( PCollection input) { @@ -177,9 +190,24 @@ public PCollection>> expand( batchedInput = input.apply("BatchElements", BatchElements.withDefaults()); } - return batchedInput - // Pass the list to the inference function - .apply("RemoteInference", ParDo.of(new RemoteInferenceFn(this))); + PCollection>> result = + batchedInput + // Pass the list to the inference function + .apply("RemoteInference", ParDo.of(new RemoteInferenceFn(this))); + + Coder outCoder = outputCoder(); + if (outCoder != null) { + result.setCoder( + org.apache.beam.sdk.coders.IterableCoder.of( + PredictionResultCoder.of(input.getCoder(), outCoder))); + } else { + result.setCoder( + (org.apache.beam.sdk.coders.Coder) + org.apache.beam.sdk.coders.IterableCoder.of( + org.apache.beam.sdk.coders.SerializableCoder.of(PredictionResult.class))); + } + + return result; } /** diff --git a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java index 64691aa1fedd..5fe8ea1cce19 100644 --- a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java +++ b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java @@ -45,7 +45,7 @@ public class RemoteInferenceTest { @Rule public final transient TestPipeline pipeline = TestPipeline.create(); // Test input class - public static class TestInput implements BaseInput { + public static class TestInput implements BaseInput, java.io.Serializable { private final String value; private TestInput(String value) { @@ -84,7 +84,7 @@ public String toString() { } // Test output class - public static class TestOutput implements BaseResponse { + public static class TestOutput implements BaseResponse, java.io.Serializable { private final String result; private TestOutput(String result) { diff --git a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java index 7e4e5da0d853..aaf63705ad2a 100644 --- a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java +++ b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java @@ -515,9 +515,11 @@ private void invokeBuilderForPublishOnlyPipeline(PipelineBuilder case UNRECOGNIZED: case STOPPED: case RUNNING: + case DRAINING: // Keep going. break; case DONE: + case DRAINED: // All done. running = false; break; diff --git a/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt b/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt index 5b996fdfc70a..2fc217218387 100644 --- a/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt +++ b/sdks/python/apache_beam/examples/ml-orchestration/kfp/components/train/requirements.txt @@ -15,4 +15,4 @@ torch==2.12.0 numpy==1.22.4 -Pillow==10.2.0 \ No newline at end of file +Pillow==12.3.0 \ No newline at end of file diff --git a/sdks/python/apache_beam/runners/interactive/recording_manager.py b/sdks/python/apache_beam/runners/interactive/recording_manager.py index c19b60b64fd2..cabcca558dca 100644 --- a/sdks/python/apache_beam/runners/interactive/recording_manager.py +++ b/sdks/python/apache_beam/runners/interactive/recording_manager.py @@ -105,7 +105,7 @@ def _cancel_clicked(self, b): self.cancel() def update_display(self, msg: str, progress: Optional[float] = None): - if not IS_IPYTHON: + if not IS_IPYTHON or self._output_widget is None: print(f'AsyncCompute: {msg}') return diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index a79cb0e8de6e..db25a40a405c 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -80,7 +80,7 @@ MAX_KNOWN_NOT_RUNNING_INSTRUCTIONS = 1000 # The number of ProcessBundleRequest instruction ids that BundleProcessorCache # will remember for failed instructions. -MAX_FAILED_INSTRUCTIONS = 10000 +MAX_FAILED_INSTRUCTIONS = 1000 # retry on transient UNAVAILABLE grpc error from state channels. _GRPC_SERVICE_CONFIG = json.dumps({ @@ -559,7 +559,15 @@ def discard(self, instruction_id, exception): """ processor = None with self._lock: - self.failed_instruction_ids[instruction_id] = exception + tb_str = "".join(traceback.format_exception(exception)) + if len(tb_str) > 10240: + tb_str = ( + tb_str[:5000] + "\n... [traceback truncated] ...\n" + + tb_str[-5000:]) + clean_exception = RuntimeError( + f"Original Exception: {type(exception).__name__}: {str(exception)[:2000]}\n{tb_str}" + ) + self.failed_instruction_ids[instruction_id] = clean_exception while len(self.failed_instruction_ids) > MAX_FAILED_INSTRUCTIONS: self.failed_instruction_ids.popitem(last=False) if instruction_id in self.active_bundle_processors: diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py index 76e428f06464..bea313a4d2fd 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -296,6 +296,44 @@ def test_failed_bundle_processor_returns_failed_split_response(self): worker.do_instruction(split_request).error, hc.contains_string('test message')) + def test_failed_instruction_id_cache_size_is_capped(self): + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + with mock.patch( + 'apache_beam.runners.worker.sdk_worker.MAX_FAILED_INSTRUCTIONS', 10): + for i in range(15): + bundle_processor_cache.discard(f'inst_{i}', RuntimeError(f'error {i}')) + + for i in range(5): + self.assertNotIn( + f'inst_{i}', bundle_processor_cache.failed_instruction_ids) + for i in range(5, 15): + self.assertIn( + f'inst_{i}', bundle_processor_cache.failed_instruction_ids) + + def test_failed_instruction_tracebacks_are_truncated_when_too_long(self): + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + long_message = "x" * 15000 + bundle_processor_cache.discard('instruction_id', RuntimeError(long_message)) + + stored_exception = bundle_processor_cache.failed_instruction_ids[ + 'instruction_id'] + tb_str = str(stored_exception) + + self.assertLessEqual(len(tb_str), 13000) + self.assertIn('[traceback truncated]', tb_str) + def test_data_sampling_response(self): # Create a data sampler with some fake sampled data. This data will be seen # in the sample response. diff --git a/sdks/python/container/Dockerfile b/sdks/python/container/Dockerfile index 7f5aa3604fe9..a1e17edb82d1 100644 --- a/sdks/python/container/Dockerfile +++ b/sdks/python/container/Dockerfile @@ -53,7 +53,8 @@ RUN \ # Install required packages for Beam Python SDK and common dependencies used by users. # use --no-deps to ensure the list includes all transitive dependencies. - pip install --no-deps -r /tmp/base_image_requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu && \ + # use --prefer-binary to avoid compiling wheels from source when prebuilt wheels exist. + pip install --prefer-binary --no-deps -r /tmp/base_image_requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu && \ rm -rf /tmp/base_image_requirements.txt && \ python -c "import nltk; nltk.download('stopwords')" && \ rm /root/nltk_data/corpora/stopwords.zip && \ diff --git a/sdks/python/container/run_generate_requirements.sh b/sdks/python/container/run_generate_requirements.sh index 4c7eea0b7398..ed1f57faac4e 100755 --- a/sdks/python/container/run_generate_requirements.sh +++ b/sdks/python/container/run_generate_requirements.sh @@ -93,8 +93,8 @@ fi # Force torch dependencies to be pulled from the PyTorch CPU wheel # repository so that they don't include GPU dependencies with # non-compliant licenses -pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir "$SDK_TARBALL""$EXTRAS" $INDEX_URL_OPTION -pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir -r "$PWD"/sdks/python/container/base_image_requirements_manual.txt +pip install --prefer-binary ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir "$SDK_TARBALL""$EXTRAS" $INDEX_URL_OPTION +pip install --prefer-binary ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir -r "$PWD"/sdks/python/container/base_image_requirements_manual.txt pip uninstall -y apache-beam echo "Checking for broken dependencies:" diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index 52f713466246..5d78db863366 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -213,7 +213,6 @@ invalid-inheritance = "ignore" not-iterable = "ignore" unexpected-keyword = "ignore" bad-specialization = "ignore" -bad-context-manager = "ignore" invalid-yield = "ignore" bad-argument-count = "ignore" bad-typed-dict-key = "ignore" diff --git a/settings.gradle.kts b/settings.gradle.kts index f71f6de9d167..8a116eeb3944 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -122,6 +122,7 @@ include(":it:conditions") include(":it:datadog") include(":it:elasticsearch") include(":it:google-cloud-platform") +include(":it:iceberg") include(":it:jdbc") include(":it:kafka") include(":it:testcontainers") @@ -284,6 +285,7 @@ include(":sdks:java:maven-archetypes:gcp-bom-examples") include(":sdks:java:maven-archetypes:starter") include(":sdks:java:ml:inference:remote") include(":sdks:java:ml:inference:openai") +include(":sdks:java:ml:inference:gemini") include(":sdks:java:testing:nexmark") include(":sdks:java:testing:expansion-service") include(":sdks:java:testing:jpms-tests") diff --git a/website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md b/website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md new file mode 100644 index 000000000000..6d4e649c7247 --- /dev/null +++ b/website/www/site/content/en/blog/beam-summit-2026-interview-with-raj-katakam.md @@ -0,0 +1,88 @@ +--- +title: "Beam Summit 2026 | Interview with Raj Katakam, Intuit Credit Karma" +date: 2026-07-22 15:00:00 -0500 +authors: + - janaom +categories: + - blog +--- + + +# Beam Summit 2026 | Interview with Raj Katakam, Intuit Credit Karma + +Apache Beam and Google Cloud Dataflow are usually known for ETL and batch/streaming processing. [Intuit Credit Karma](https://www.creditkarma.com/) said, why stop there, and built an entire ML platform on top of them. + +Here is my interview with **Raj Katakam**, Staff Machine Learning Engineer at Intuit Credit Karma. We talked about how they built a unified ML platform on Apache Beam and Dataflow to serve 140 million users, why being active in the Apache Beam community matters to him, and what advice he would give anyone just getting started with Beam and ML. + +-------------- + +**Jana: Intuit Credit Karma’s tech stack looks like a dream! You are on Google Cloud and using services like Dataflow, BigQuery, and Managed Airflow (previously, Cloud Composer). What were the main reasons for choosing this stack, and what trade-offs did you have to consider when selecting GCP and Apache Beam/Dataflow?** + +**Raj**: The selection of Google Cloud, utilizing BigQuery, Dataflow, and Managed Airflow, was a strategic decision to standardize the ML lifecycle and enable high-velocity development at scale. By leveraging these technologies, we established a unified architectural standard that allows for elastic scalability and reliability across petabyte-scale workloads. This approach ensures that the entire ML lifecycle, from initial experimentation to global production execution, is governed by a consistent framework, providing the necessary infrastructure to manage complex data movements and distributed computing requirements without sacrificing architectural integrity. + +Our goal was to create a unified abstraction layer that could bridge various data processing tools like Spark and BigQuery with ML libraries such as scikit-learn and TensorFlow. GCP provided the deep, native integrations we required to smoothly transition from local prototypes to massive, distributed cloud execution. + +Talking about trade-offs, the biggest hurdle was simplifying the inherent complexity of a multi-framework environment. Even with powerful connectors for GCS and BigQuery, we still had to develop Vega (Intuit Credit Karma’s internal ML platform) to shield our data scientists from manual orchestration and resource management, allowing them to remain focused on a clean, unified Python API. + +**Jana: Dataflow is widely known for ETL and batch/streaming processing, but you took it a step further by using it for ML. You built Vega, a full ML explainability platform on Apache Beam and Dataflow. What motivated that choice, and what advantages did these technologies bring to the project?** + +**Raj**: We utilized Dataflow as the distributed execution engine to power the broader Vega ecosystem, which serves as our comprehensive ML platform. It is important to emphasize that explainability is a foundational pillar of this ecosystem, integrated directly into the core architecture rather than treated as a standalone task. By building on Apache Beam, we ensure that every model, whether focused on feature engineering, scoring, or interpretability, benefits from a unified processing model. This integration allows us to maintain rigorous standards for model transparency and governance as a native component of our large-scale distributed ML operations. + +The main motivation was to eliminate the experimentation-to-production gap. Traditional platforms forced data scientists to manually rewrite code for production frameworks, but by using Dataflow and Apache Beam as the core engine, we were able to support both batch and streaming processing alongside ML scoring within a single workflow. + +The primary advantage that came out of this was the unified Python API. It allowed data scientists to build complex, end-to-end data and ML pipelines that transition seamlessly from local development on sampled data to cloud-scale production execution, without architectural changes or manual re-engineering. + +**Jana: Looking back at building Vega, what were the biggest challenges you faced, and what would you do differently today?** + +**Raj**: The primary challenges centered on managing the inherent complexities of distributed systems and maintaining architectural coherence across a diverse ML lifecycle. Operating at the scale of over 140 million members requires solving for high-concurrency data access, managing complex dependency graphs in distributed environments, and ensuring that performance remains consistent as workloads transition from sampled data to massive production clusters. Balancing the need for developer flexibility with the strict latency and reliability constraints of a globally distributed platform represents the most significant architectural hurdle in modern ML engineering. + +Reflecting on the evolution of Vega, the biggest challenge was navigating the trade-offs of building critical infrastructure while scaling to support massive business growth with limited resources. We prioritized speed-to-market and immediate utility to enable the business, which was necessary but created technical debt in terms of documentation and onboarding. If I were starting today, I would prioritize earlier investments in developer tooling and standardized templates. We are now shifting our focus from that initial fast-and-lean build phase to long-term sustainability, refining the developer experience, automating maintenance even further, and broadening support for modern ML paradigms. This ensures that the platform remains a scalable, durable asset as it evolves to meet the next generation of business needs. + +**Jana: Are there any features or improvements you would love to see in Apache Beam or Dataflow in the future?** + +**Raj**: While Vega has successfully leveraged Beam for our core workflows, there are a few areas where future improvements could make a real difference. One would be enhanced data engineering primitives. We would value deeper native support for large-scale data manipulation, specifically generalized tiling operations for distributed datasets and built-in capabilities for saving and managing partial processing states. This would significantly reduce the complexity of checkpointing and state recovery in our pipelines. + +The other area is real-time inference. While we have strong batch capabilities, we need more streamlined primitives in Beam to better integrate our data pipelines with real-time serving infrastructure. Reducing the complexity of the hand-off between processing and prediction would be a major win for our personalization engines. + +**Jana: You are a Staff Machine Learning Engineer at Intuit Credit Karma — could you tell us a bit about your background and journey to this role? And for readers who are interested in pursuing a career in ML, what advice would you give them?** + +**Raj**: As a Staff Machine Learning Engineer, my career has been defined by the practice of platform-scale engineering. I have focused on architecting ecosystems that provide high-velocity developer experiences, allowing engineering teams to move from conceptual design to production-scale deployment within a unified framework. By treating the ML lifecycle as a first-class engineering problem, we build systems that automate the complexities of infrastructure management, enabling data scientists to focus on model innovation while the platform handles the rigorous demands of distributed execution and governance. + +My advice to anyone pursuing a career in ML would be to focus on first-principles thinking. Don’t just learn a specific library, understand the underlying data infrastructure, including data movement, latency constraints, and caching. The most impactful engineers I have worked with, including those I have collaborated with at Google and beyond, are the ones who can reason across domains, from low-level feature serving latency to high-level platform architecture. + +**Jana: You are an active member of the Apache Beam community. This is your third Beam Summit. How important do you think community involvement is for engineers, and how has it shaped your career?** + +**Raj**: Community involvement is critical. Participating in initiatives like the Apache Beam Summit has allowed me to stay at the frontier of distributed computing and directly influence the evolution of the tools we depend on. It transforms engineering from a siloed task into a collaborative effort; seeing how other engineers solve (or struggle with) similar problems helps refine your own architectural decisions and often provides the reference architectures that we would otherwise have taken years to develop independently. + +**Jana: If you could give one piece of advice to someone just starting with Apache Beam, what would it be?** + +**Raj**: Focus on the abstraction. Do not treat Beam as just another processing framework; treat it as the “universal language” for your data pipeline. Start by leveraging the high-level transforms and built-in IO connectors rather than trying to optimize low-level custom code immediately. Understanding how to structure your pipeline to take advantage of Beam’s windowing and triggering capabilities is what will eventually separate a good engineer from a great one when you hit the scale of millions of users. + +**Jana: Thank you so much, Raj, for sharing these insights!** + +**Raj**: Thank you for having me, Jana, this was a great conversation! + +------------ + +🌐 To learn more about **Intuit Credit Karma**, visit their [website](https://www.creditkarma.com/) or check out their [LinkedIn](https://www.linkedin.com/company/intuitcreditkarma/). You can also connect with **Raj** directly on [LinkedIn](https://www.linkedin.com/in/rajkiran2190). + +📌 Raj and Pallav Anand presented at Beam Summit 2026. Check out their talk, [‘Beyond the Black Box: How Intuit Credit Karma Runs ML Explainability for 140M Members with Beam’](https://beamsummit.org/sessions/2026/beyond-the-black-box-how-intuit-credit-karma-runs-ml-explainability-for-140m-members-with-beam/), and explore the full program at [beamsummit.org](https://beamsummit.org/sessions/2026/). + +📺 Session recording will be available on the [Apache Beam YouTube channel](https://www.youtube.com/@ApacheBeamYT). + +📰 Missed the Beam Summit 2026? Read my recap: [Beam Summit 2026: Apache Beam Just Got Even More Interesting 🐝](https://medium.com/google-cloud/beam-summit-2026-apache-beam-just-got-even-more-interesting-7c6c9967ffb9) + +– [Jana Polianskaja](https://www.linkedin.com/in/jana-polianskaja/) diff --git a/website/www/site/content/en/performance/_index.md b/website/www/site/content/en/performance/_index.md index db05afd4cbba..08fd368809c8 100644 --- a/website/www/site/content/en/performance/_index.md +++ b/website/www/site/content/en/performance/_index.md @@ -38,6 +38,7 @@ writing to various Beam IOs. - [BigQuery](/performance/bigquery) - [BigTable](/performance/bigtable) - [TextIO](/performance/textio) +- [IcebergIO](/performance/icebergio) # Measured Beam Python ML Pipelines diff --git a/website/www/site/content/en/performance/icebergio/_index.md b/website/www/site/content/en/performance/icebergio/_index.md new file mode 100644 index 000000000000..e618865097b4 --- /dev/null +++ b/website/www/site/content/en/performance/icebergio/_index.md @@ -0,0 +1,50 @@ +--- +title: "IcebergIO Performance" +--- + + + +# IcebergIO Performance + +The following graphs show various metrics when reading from or writing to Apache Iceberg tables using +IcebergIO. See the [glossary](/performance/glossary) for definitions. + +## Read + +### What is the estimated cost of reading from Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="read" section="test_name" >}} + +### How has various metrics changed when reading from Apache Iceberg tables using IcebergIO for different Beam SDK versions? + +{{< performance_looks io="icebergio" read_or_write="read" section="version" >}} + +### How has various metrics changed over time when reading from Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="read" section="date" >}} + +## Write + +### What is the estimated cost of writing to Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="write" section="test_name" >}} + +### How has various metrics changed when writing to Apache Iceberg tables using IcebergIO for different Beam SDK versions? + +{{< performance_looks io="icebergio" read_or_write="write" section="version" >}} + +### How has various metrics changed over time when writing to Apache Iceberg tables using IcebergIO? + +{{< performance_looks io="icebergio" read_or_write="write" section="date" >}} diff --git a/website/www/site/data/authors.yml b/website/www/site/data/authors.yml index b74a6456d609..d98f2281bddc 100644 --- a/website/www/site/data/authors.yml +++ b/website/www/site/data/authors.yml @@ -87,6 +87,9 @@ jamesmalone: name: James Malone email: jamesmalone@apache.org twitter: chimerasaurus +janaom: + name: Jana Polianskaja + email: jana.polianskaja@gmail.com jesseanderson: name: Jesse Anderson twitter: jessetanderson diff --git a/website/www/site/data/performance.yaml b/website/www/site/data/performance.yaml index f0292fab7472..03eaad774b0b 100644 --- a/website/www/site/data/performance.yaml +++ b/website/www/site/data/performance.yaml @@ -107,6 +107,37 @@ looks: title: AvgInputThroughputBytesPerSec by Version - id: fVVHhXCrHNgBG52TJsTjR8VbmWCCQnVN title: AvgInputThroughputElementsPerSec by Version + icebergio: + read: + folder: 113 + test_name: + - id: x5R7w8VnzkccgzrBTWQ82J7JBNtkGxH2 + title: Read IcebergIO RunTime and EstimatedCost + date: + - id: HQsKkmn7cS7m5bCYzPBYqnnFM6MXvJgV + title: AvgOutputThroughputBytesPerSec by Date + - id: MNYqD2ZCwV2XfYjwDg4NDVT7vr4Ytpz2 + title: AvgOutputThroughputElementsPerSec by Date + version: + - id: Cf7wGQfrM4hn8Wdf64ZjwshrHHdZKmKs + title: AvgOutputThroughputBytesPerSec by Version + - id: MGSX85MhCnCFFWYzSCGS36JRWxcBccqd + title: AvgOutputThroughputElementsPerSec by Version + write: + folder: 114 + test_name: + - id: yvvWB2JbbbFQ4R5MTwF2hbVW2X4bDxnY + title: Write IcebergIO RunTime and EstimatedCost + date: + - id: Gcw6jfhvvkkHKDwqqTkhZtYYNNH2Z8DZ + title: AvgInputThroughputBytesPerSec by Date + - id: RYtRGKsxBRkyRgWk6RGjshFnZPk76vhM + title: AvgInputThroughputElementsPerSec by Date + version: + - id: B5q25ktnXNhHxpYrGmqzgcGMgX26zvM3 + title: AvgInputThroughputBytesPerSec by Version + - id: VfMWk6rk26JQNr8wRpsK5Gk6XQPfd7vW + title: AvgInputThroughputElementsPerSec by Version pytorchbertbase: write: folder: 76