Exercise Flink 2.0 Python Validates Runner tests

Author: Abacn

.github/trigger_files/beam_PostCommit_PortableJar_Flink.json

@@ -0,0 +1 @@
+{}

.github/trigger_files/beam_PostCommit_Python_Examples_Flink.json

@@ -0,0 +1 @@
+{}

.github/workflows/beam_PostCommit_Python_Portable_Flink.yml

@@ -63,8 +63,9 @@ jobs:
         job_name: ["beam_PostCommit_Python_Portable_Flink"]
         job_phrase: ["Run Python Portable Flink"]
         # TODO: Enable PROCESS https://github.com/apache/beam/issues/35702
-        # environment_type: ['DOCKER', 'LOOPBACK', 'PROCESS']
-        environment_type: ['DOCKER', 'LOOPBACK']
+        # all environment_type: ['DOCKER', 'LOOPBACK', 'PROCESS']
+        # Run modes not covered by PreCommit_Python_PVR_Flink (i.e. other than 'LOOPBACK')
+        environment_type: ['DOCKER']
     steps:
       - uses: actions/checkout@v4
       - name: Setup repository
@@ -83,7 +84,7 @@ jobs:
           CLOUDSDK_CONFIG: ${{ env.KUBELET_GCLOUD_CONFIG_PATH}}
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
-          gradle-command: :sdks:python:test-suites:portable:py310:flinkCompatibilityMatrix${{ matrix.environment_type }}
+          gradle-command: :sdks:python:test-suites:portable:py310:flink2CompatibilityMatrix${{ matrix.environment_type }}
           arguments: |
             -PpythonVersion=3.10 \
       - name: Archive Python Test Results

.github/workflows/beam_PreCommit_Python_PVR_Flink.yml

@@ -106,7 +106,8 @@ jobs:
         env:
           CLOUDSDK_CONFIG: ${{ env.KUBELET_GCLOUD_CONFIG_PATH}}
         with:
-          gradle-command: :sdks:python:test-suites:portable:py313:flinkValidatesRunner
+          # Run Flink 2 tests. Flink 1.20 is covered by PostCommit_Python_ValidatesRunner_Flink
+          gradle-command: :sdks:python:test-suites:portable:py313:flink2ValidatesRunner
           arguments: |
             -PpythonVersion=3.13 \
       - name: Archive Python Test Results

runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java

@@ -100,6 +100,7 @@
 import org.apache.beam.sdk.values.WindowedValues.WindowedValueCoder;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.BiMap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashMultiset;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
@@ -987,8 +988,11 @@ private <T> void translateTestStream(
       // stage
       String sideInputTag = sideInputId.getLocalName();
       String collectionId =
-          components
-              .getTransformsOrThrow(sideInputId.getTransformId())
+          MoreObjects.firstNonNull(
+                  components.getTransformsOrDefault(sideInputId.getTransformId(), null),
+                  // In the case of optimized pipeline, side input transform may not be found in
+                  // component proto
+                  stagePayload.getComponents().getTransformsOrThrow(sideInputId.getTransformId()))
               .getInputsOrThrow(sideInputId.getLocalName());
       RunnerApi.WindowingStrategy windowingStrategyProto =
           components.getWindowingStrategiesOrThrow(
@@ -1045,8 +1049,11 @@ private TransformedSideInputs transformSideInputs(
       tagToIntMapping.put(tag, count);
       count++;
       String collectionId =
-          components
-              .getTransformsOrThrow(sideInput.getKey().getTransformId())
+          MoreObjects.firstNonNull(
+                  components.getTransformsOrDefault(sideInput.getKey().getTransformId(), null),
+                  stagePayload
+                      .getComponents()
+                      .getTransformsOrThrow(sideInput.getKey().getTransformId()))
               .getInputsOrThrow(sideInput.getKey().getLocalName());
       DataStream<Object> sideInputStream = context.getDataStreamOrThrow(collectionId);
       TypeInformation<Object> tpe = sideInputStream.getType();
@@ -1078,7 +1085,11 @@ private TransformedSideInputs transformSideInputs(
       TupleTag<?> tag = sideInput.getValue().getTagInternal();
       final int intTag = tagToIntMapping.get(tag);
       RunnerApi.PTransform pTransform =
-          components.getTransformsOrThrow(sideInput.getKey().getTransformId());
+          MoreObjects.firstNonNull(
+              components.getTransformsOrDefault(sideInput.getKey().getTransformId(), null),
+              stagePayload
+                  .getComponents()
+                  .getTransformsOrThrow(sideInput.getKey().getTransformId()));
       String collectionId = pTransform.getInputsOrThrow(sideInput.getKey().getLocalName());
       DataStream<WindowedValue<?>> sideInputStream = context.getDataStreamOrThrow(collectionId);
 

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java

@@ -100,6 +100,7 @@
 import org.apache.beam.sdk.values.WindowedValues.WindowedValueCoder;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.BiMap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashMultiset;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
@@ -986,8 +987,11 @@ private <T> void translateTestStream(
       // stage
       String sideInputTag = sideInputId.getLocalName();
       String collectionId =
-          components
-              .getTransformsOrThrow(sideInputId.getTransformId())
+          MoreObjects.firstNonNull(
+                  components.getTransformsOrDefault(sideInputId.getTransformId(), null),
+                  // In the case of optimized pipeline, side input transform may not be found in
+                  // component proto
+                  stagePayload.getComponents().getTransformsOrThrow(sideInputId.getTransformId()))
               .getInputsOrThrow(sideInputId.getLocalName());
       RunnerApi.WindowingStrategy windowingStrategyProto =
           components.getWindowingStrategiesOrThrow(
@@ -1044,8 +1048,11 @@ private TransformedSideInputs transformSideInputs(
       tagToIntMapping.put(tag, count);
       count++;
       String collectionId =
-          components
-              .getTransformsOrThrow(sideInput.getKey().getTransformId())
+          MoreObjects.firstNonNull(
+                  components.getTransformsOrDefault(sideInput.getKey().getTransformId(), null),
+                  stagePayload
+                      .getComponents()
+                      .getTransformsOrThrow(sideInput.getKey().getTransformId()))
               .getInputsOrThrow(sideInput.getKey().getLocalName());
       DataStream<Object> sideInputStream = context.getDataStreamOrThrow(collectionId);
       TypeInformation<Object> tpe = sideInputStream.getType();
@@ -1077,7 +1084,11 @@ private TransformedSideInputs transformSideInputs(
       TupleTag<?> tag = sideInput.getValue().getTagInternal();
       final int intTag = tagToIntMapping.get(tag);
       RunnerApi.PTransform pTransform =
-          components.getTransformsOrThrow(sideInput.getKey().getTransformId());
+          MoreObjects.firstNonNull(
+              components.getTransformsOrDefault(sideInput.getKey().getTransformId(), null),
+              stagePayload
+                  .getComponents()
+                  .getTransformsOrThrow(sideInput.getKey().getTransformId()));
       String collectionId = pTransform.getInputsOrThrow(sideInput.getKey().getLocalName());
       DataStream<WindowedValue<?>> sideInputStream = context.getDataStreamOrThrow(collectionId);
 

sdks/java/extensions/sql/expansion-service/build.gradle

@@ -56,5 +56,4 @@ shadowJar {
   manifest {
     attributes(["Multi-Release": true])
   }
-  outputs.upToDateWhen { false }
 }

sdks/python/apache_beam/runners/portability/flink_runner_test.py

@@ -18,6 +18,7 @@
 
 import argparse
 import logging
+import platform
 import shlex
 import typing
 import unittest
@@ -139,6 +140,7 @@ def _create_conf_dir(cls):
       cls.test_metrics_path = path.join(cls.conf_dir, 'test-metrics.txt')
 
       # path to write Flink configuration to
+      # Flink 1.x conf:
       conf_path = path.join(cls.conf_dir, 'flink-conf.yaml')
       file_reporter = 'org.apache.beam.runners.flink.metrics.FileReporter'
       with open(conf_path, 'w') as f:
@@ -149,6 +151,19 @@ def _create_conf_dir(cls):
                 'metrics.reporter.file.path: %s' % cls.test_metrics_path,
                 'metrics.scope.operator: <operator_name>',
             ]))
+      # Flink 2.x conf:
+      conf_path_2 = path.join(cls.conf_dir, 'config.yaml')
+      with open(conf_path_2, 'w') as f:
+        f.write(
+            '''metrics:
+  reporters: file
+  reporter:
+    file:
+      class: %s
+      path: %s
+  scope:
+    operator: <operator_name>
+''' % (file_reporter, cls.test_metrics_path))
 
   @classmethod
   def _subprocess_command(cls, job_port, expansion_port):
@@ -158,11 +173,12 @@ def _subprocess_command(cls, job_port, expansion_port):
 
     cls._create_conf_dir()
     cls.expansion_port = expansion_port
-
+    platform_specific_opts = []
+    if platform.system() == 'Linux':
+      # UseContainerSupport is supported in Linux and turned on by default
+      platform_specific_opts.append('-XX:-UseContainerSupport')
     try:
-      return [
-          'java',
-          '-XX:-UseContainerSupport',
+      return ['java'] + platform_specific_opts + [
           '--add-opens=java.base/java.lang=ALL-UNNAMED',
           '--add-opens=java.base/java.nio=ALL-UNNAMED',
           '--add-opens=java.base/java.util=ALL-UNNAMED',

sdks/python/test-suites/portable/common.gradle

@@ -22,25 +22,33 @@ import org.apache.tools.ant.taskdefs.condition.Os
 
 def pythonRootDir = "${rootDir}/sdks/python"
 def pythonVersionSuffix = project.ext.pythonVersion.replace('.', '')
+// TODO(https://github.com/apache/beam/issues/36947): Remove when dropping Flink 1.x support
 def latestFlinkVersion = project.ext.latestFlinkVersion
+def latestFlink2Version = '2.0'
 def currentJavaVersion = project.ext.currentJavaVersion
 
 ext {
   pythonContainerTask = ":sdks:python:container:py${pythonVersionSuffix}:docker"
 }
 
-def createFlinkRunnerTestTask(String workerType) {
-  def taskName = "flinkCompatibilityMatrix${workerType}"
-  // project(":runners:flink:${latestFlinkVersion}:job-server").shadowJar.archivePath is not resolvable until runtime, so hard-code it here.
-  def jobServerJar = "${rootDir}/runners/flink/${latestFlinkVersion}/job-server/build/libs/beam-runners-flink-${latestFlinkVersion}-job-server-${version}.jar"
+def createFlinkRunnerTestTask(String workerType, String flinkVersion) {
+  String taskName
+
+  // project(":runners:flink:${flinkVersion}:job-server").shadowJar.archivePath is not resolvable until runtime, so hard-code it here.
+  def jobServerJar = "${rootDir}/runners/flink/${flinkVersion}/job-server/build/libs/beam-runners-flink-${flinkVersion}-job-server-${version}.jar"
   def options = "--flink_job_server_jar=${jobServerJar} --environment_type=${workerType}"
+  if (flinkVersion.startsWith('2')) {
+    taskName = "flink2CompatibilityMatrix${workerType}"
+  } else {
+    taskName = "flinkCompatibilityMatrix${workerType}"
+  }
   if (workerType == 'PROCESS') {
     options += " --environment_options=process_command=${buildDir.absolutePath}/sdk_worker.sh"
   }
   def task = toxTask(taskName, 'flink-runner-test', options)
   // Through the Flink job server, we transitively add dependencies on the expansion services needed in tests.
   task.configure {
-    dependsOn ":runners:flink:${latestFlinkVersion}:job-server:shadowJar"
+    dependsOn ":runners:flink:${flinkVersion}:job-server:shadowJar"
     // The Java SDK worker is required to execute external transforms.
     def suffix = getSupportedJavaVersion()
     dependsOn ":sdks:java:container:${suffix}:docker"
@@ -53,31 +61,19 @@ def createFlinkRunnerTestTask(String workerType) {
   return task
 }
 
-createFlinkRunnerTestTask('DOCKER')
-createFlinkRunnerTestTask('PROCESS')
-createFlinkRunnerTestTask('LOOPBACK')
+createFlinkRunnerTestTask('DOCKER', latestFlinkVersion)
+createFlinkRunnerTestTask('PROCESS', latestFlinkVersion)
+createFlinkRunnerTestTask('LOOPBACK', latestFlinkVersion)
+createFlinkRunnerTestTask('DOCKER', latestFlink2Version)
+createFlinkRunnerTestTask('PROCESS', latestFlink2Version)
+createFlinkRunnerTestTask('LOOPBACK', latestFlink2Version)
 
 task flinkValidatesRunner() {
   dependsOn 'flinkCompatibilityMatrixLOOPBACK'
 }
 
-// TODO(https://github.com/apache/beam/issues/19962): Enable on pre-commit.
-tasks.register("flinkTriggerTranscript") {
-  dependsOn 'setupVirtualenv'
-  dependsOn ":runners:flink:${latestFlinkVersion}:job-server:shadowJar"
-  doLast {
-    exec {
-      executable 'sh'
-      args '-c', """
-          . ${envdir}/bin/activate \\
-          && cd ${pythonRootDir} \\
-          && pip install -e .[test] \\
-          && pytest \\
-              apache_beam/transforms/trigger_test.py::WeakTestStreamTranscriptTest \\
-              --test-pipeline-options='--runner=FlinkRunner --environment_type=LOOPBACK --flink_job_server_jar=${project(":runners:flink:${latestFlinkVersion}:job-server:").shadowJar.archivePath}'
-          """
-    }
-  }
+task flink2ValidatesRunner() {
+  dependsOn 'flink2CompatibilityMatrixLOOPBACK'
 }
 
 // Verifies BEAM-10702.