Search in sources :

Example 1 with DataflowRunner.getContainerImageForJob

use of org.apache.beam.runners.dataflow.DataflowRunner.getContainerImageForJob in project beam by apache.

the class DataflowRunnerTest method testSdkHarnessConfiguration.

@Test
public void testSdkHarnessConfiguration() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    ExperimentalOptions.addExperiment(options, "use_runner_v2");
    Pipeline p = Pipeline.create(options);
    p.apply(Create.of(Arrays.asList(1, 2, 3)));
    String defaultSdkContainerImage = DataflowRunner.getContainerImageForJob(options);
    SdkComponents sdkComponents = SdkComponents.create();
    RunnerApi.Environment defaultEnvironmentForDataflow = Environments.createDockerEnvironment(defaultSdkContainerImage);
    sdkComponents.registerEnvironment(defaultEnvironmentForDataflow.toBuilder().build());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
    Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList()).getJob();
    DataflowRunner.configureSdkHarnessContainerImages(options, pipelineProto, job);
    List<SdkHarnessContainerImage> sdks = job.getEnvironment().getWorkerPools().get(0).getSdkHarnessContainerImages();
    Map<String, String> expectedEnvIdsAndContainerImages = pipelineProto.getComponents().getEnvironmentsMap().entrySet().stream().filter(x -> BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.DOCKER).equals(x.getValue().getUrn())).collect(Collectors.toMap(x -> x.getKey(), x -> {
        RunnerApi.DockerPayload payload;
        try {
            payload = RunnerApi.DockerPayload.parseFrom(x.getValue().getPayload());
        } catch (InvalidProtocolBufferException e) {
            throw new RuntimeException(e);
        }
        return payload.getContainerImage();
    }));
    assertEquals(1, expectedEnvIdsAndContainerImages.size());
    assertEquals(1, sdks.size());
    assertEquals(expectedEnvIdsAndContainerImages, sdks.stream().collect(Collectors.toMap(SdkHarnessContainerImage::getEnvironmentId, SdkHarnessContainerImage::getContainerImage)));
}
Also used : ExpectedLogs(org.apache.beam.sdk.testing.ExpectedLogs) Arrays(java.util.Arrays) Matchers.not(org.hamcrest.Matchers.not) ValueState(org.apache.beam.sdk.state.ValueState) SimpleFunction(org.apache.beam.sdk.transforms.SimpleFunction) DefaultGcpRegionFactory(org.apache.beam.runners.dataflow.options.DefaultGcpRegionFactory) Create(org.apache.beam.sdk.transforms.Create) DataflowPipelineDebugOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions) Map(java.util.Map) Node(org.apache.beam.sdk.runners.TransformHierarchy.Node) Window(org.apache.beam.sdk.transforms.windowing.Window) JsonSerializer(com.fasterxml.jackson.databind.JsonSerializer) JsonNode(com.fasterxml.jackson.databind.JsonNode) Dataflow(com.google.api.services.dataflow.Dataflow) TimestampedValue(org.apache.beam.sdk.values.TimestampedValue) ValueProvider(org.apache.beam.sdk.options.ValueProvider) Files.getFileExtension(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Files.getFileExtension) GcpOptions(org.apache.beam.sdk.extensions.gcp.options.GcpOptions) StreamingOptions(org.apache.beam.sdk.options.StreamingOptions) GcsUtil(org.apache.beam.sdk.extensions.gcp.util.GcsUtil) ShardedKey(org.apache.beam.sdk.util.ShardedKey) NoopCredentialFactory(org.apache.beam.sdk.extensions.gcp.auth.NoopCredentialFactory) Category(org.junit.experimental.categories.Category) Serializable(java.io.Serializable) Matchers.any(org.mockito.Matchers.any) Assert.assertFalse(org.junit.Assert.assertFalse) StreamingShardedWriteFactory(org.apache.beam.runners.dataflow.DataflowRunner.StreamingShardedWriteFactory) NoopPathValidator(org.apache.beam.sdk.extensions.gcp.storage.NoopPathValidator) PipelineVisitor(org.apache.beam.sdk.Pipeline.PipelineVisitor) Matchers.is(org.hamcrest.Matchers.is) Matchers.containsString(org.hamcrest.Matchers.containsString) Matchers.endsWith(org.hamcrest.Matchers.endsWith) JsonDeserialize(com.fasterxml.jackson.databind.annotation.JsonDeserialize) Mockito.mock(org.mockito.Mockito.mock) KV(org.apache.beam.sdk.values.KV) DataflowPipelineWorkerPoolOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions) SetState(org.apache.beam.sdk.state.SetState) ExperimentalOptions(org.apache.beam.sdk.options.ExperimentalOptions) JsonGenerator(com.fasterxml.jackson.core.JsonGenerator) Assume.assumeFalse(org.junit.Assume.assumeFalse) Duration(org.joda.time.Duration) RunWith(org.junit.runner.RunWith) Environments(org.apache.beam.runners.core.construction.Environments) ArrayList(java.util.ArrayList) Matchers.hasProperty(org.hamcrest.Matchers.hasProperty) FileBasedSink(org.apache.beam.sdk.io.FileBasedSink) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) ValidatesRunner(org.apache.beam.sdk.testing.ValidatesRunner) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) Pipeline(org.apache.beam.sdk.Pipeline) PowerMockito(org.powermock.api.mockito.PowerMockito) AppliedPTransform(org.apache.beam.sdk.runners.AppliedPTransform) InvalidProtocolBufferException(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.InvalidProtocolBufferException) Before(org.junit.Before) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) DoFn(org.apache.beam.sdk.transforms.DoFn) DeserializationContext(com.fasterxml.jackson.databind.DeserializationContext) ResourceId(org.apache.beam.sdk.io.fs.ResourceId) Files(java.nio.file.Files) PAssert(org.apache.beam.sdk.testing.PAssert) Assert.assertTrue(org.junit.Assert.assertTrue) IOException(java.io.IOException) Test(org.junit.Test) DataflowPipelineOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineOptions) File(java.io.File) BeamUrns(org.apache.beam.runners.core.construction.BeamUrns) ResourceHints(org.apache.beam.sdk.transforms.resourcehints.ResourceHints) ReplacementOutput(org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput) Matchers.hasItem(org.hamcrest.Matchers.hasItem) Assert.assertNull(org.junit.Assert.assertNull) AutoService(com.google.auto.service.AutoService) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) VoidCoder(org.apache.beam.sdk.coders.VoidCoder) FileSystems(org.apache.beam.sdk.io.FileSystems) Assert.assertEquals(org.junit.Assert.assertEquals) TextIO(org.apache.beam.sdk.io.TextIO) Module(com.fasterxml.jackson.databind.Module) DynamicFileDestinations(org.apache.beam.sdk.io.DynamicFileDestinations) StateSpec(org.apache.beam.sdk.state.StateSpec) UsesStatefulParDo(org.apache.beam.sdk.testing.UsesStatefulParDo) WriteFilesResult(org.apache.beam.sdk.io.WriteFilesResult) SdkHarnessContainerImage(com.google.api.services.dataflow.model.SdkHarnessContainerImage) WriteFiles(org.apache.beam.sdk.io.WriteFiles) Matchers.hasKey(org.hamcrest.Matchers.hasKey) Job(com.google.api.services.dataflow.model.Job) Sessions(org.apache.beam.sdk.transforms.windowing.Sessions) SimpleModule(com.fasterxml.jackson.databind.module.SimpleModule) Matchers.eq(org.mockito.Matchers.eq) ListJobsResponse(com.google.api.services.dataflow.model.ListJobsResponse) Iterables(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables) Assert.fail(org.junit.Assert.fail) JsonDeserializer(com.fasterxml.jackson.databind.JsonDeserializer) StorageObject(com.google.api.services.storage.model.StorageObject) MapElements(org.apache.beam.sdk.transforms.MapElements) Matchers.isA(org.mockito.Matchers.isA) Matchers.lessThanOrEqualTo(org.hamcrest.Matchers.lessThanOrEqualTo) PaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo) BigEndianIntegerCoder(org.apache.beam.sdk.coders.BigEndianIntegerCoder) StandardOpenOption(java.nio.file.StandardOpenOption) PValues(org.apache.beam.sdk.values.PValues) Collectors(java.util.stream.Collectors) CheckEnabled(org.apache.beam.sdk.options.PipelineOptions.CheckEnabled) TransformHierarchy(org.apache.beam.sdk.runners.TransformHierarchy) TypeSafeMatcher(org.hamcrest.TypeSafeMatcher) FileNotFoundException(java.io.FileNotFoundException) List(java.util.List) ParDo(org.apache.beam.sdk.transforms.ParDo) Matchers.containsInAnyOrder(org.hamcrest.Matchers.containsInAnyOrder) Matchers.equalTo(org.hamcrest.Matchers.equalTo) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Pattern(java.util.regex.Pattern) WindowingStrategy(org.apache.beam.sdk.values.WindowingStrategy) DataflowPackage(com.google.api.services.dataflow.model.DataflowPackage) PowerMockito.mockStatic(org.powermock.api.mockito.PowerMockito.mockStatic) Assert.assertThrows(org.junit.Assert.assertThrows) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) HashMap(java.util.HashMap) PipelineTranslation(org.apache.beam.runners.core.construction.PipelineTranslation) PipelineOptionsFactory(org.apache.beam.sdk.options.PipelineOptionsFactory) SerializableFunctions(org.apache.beam.sdk.transforms.SerializableFunctions) DataflowRunner.getContainerImageForJob(org.apache.beam.runners.dataflow.DataflowRunner.getContainerImageForJob) PTransform(org.apache.beam.sdk.transforms.PTransform) Matchers.matchesRegex(org.hamcrest.Matchers.matchesRegex) ArgumentCaptor(org.mockito.ArgumentCaptor) MapState(org.apache.beam.sdk.state.MapState) JsonSerialize(com.fasterxml.jackson.databind.annotation.JsonSerialize) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) SerializerProvider(com.fasterxml.jackson.databind.SerializerProvider) PowerMockRunner(org.powermock.modules.junit4.PowerMockRunner) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) ExpectedException(org.junit.rules.ExpectedException) Nullable(org.checkerframework.checker.nullness.qual.Nullable) GroupIntoBatches(org.apache.beam.sdk.transforms.GroupIntoBatches) Matchers.hasEntry(org.hamcrest.Matchers.hasEntry) GcsPath(org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath) Description(org.hamcrest.Description) SdkComponents(org.apache.beam.runners.core.construction.SdkComponents) JsonParser(com.fasterxml.jackson.core.JsonParser) Assert.assertNotNull(org.junit.Assert.assertNotNull) Matchers.anyListOf(org.mockito.Matchers.anyListOf) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) Matchers(org.hamcrest.Matchers) JsonProcessingException(com.fasterxml.jackson.core.JsonProcessingException) Mockito.when(org.mockito.Mockito.when) JUnit4(org.junit.runners.JUnit4) PCollection(org.apache.beam.sdk.values.PCollection) TestCredential(org.apache.beam.sdk.extensions.gcp.auth.TestCredential) Mockito(org.mockito.Mockito) StateSpecs(org.apache.beam.sdk.state.StateSpecs) Rule(org.junit.Rule) Instant(org.joda.time.Instant) StaticValueProvider(org.apache.beam.sdk.options.ValueProvider.StaticValueProvider) FileChannel(java.nio.channels.FileChannel) Collections(java.util.Collections) TemporaryFolder(org.junit.rules.TemporaryFolder) PropertyNames(org.apache.beam.runners.dataflow.util.PropertyNames) DataflowPipelineOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineOptions) SdkHarnessContainerImage(com.google.api.services.dataflow.model.SdkHarnessContainerImage) InvalidProtocolBufferException(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.InvalidProtocolBufferException) Matchers.containsString(org.hamcrest.Matchers.containsString) SdkComponents(org.apache.beam.runners.core.construction.SdkComponents) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) Pipeline(org.apache.beam.sdk.Pipeline) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) Job(com.google.api.services.dataflow.model.Job) DataflowRunner.getContainerImageForJob(org.apache.beam.runners.dataflow.DataflowRunner.getContainerImageForJob) Test(org.junit.Test) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest)

Aggregations

JsonGenerator (com.fasterxml.jackson.core.JsonGenerator)1 JsonParser (com.fasterxml.jackson.core.JsonParser)1 JsonProcessingException (com.fasterxml.jackson.core.JsonProcessingException)1 DeserializationContext (com.fasterxml.jackson.databind.DeserializationContext)1 JsonDeserializer (com.fasterxml.jackson.databind.JsonDeserializer)1 JsonNode (com.fasterxml.jackson.databind.JsonNode)1 JsonSerializer (com.fasterxml.jackson.databind.JsonSerializer)1 Module (com.fasterxml.jackson.databind.Module)1 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)1 SerializerProvider (com.fasterxml.jackson.databind.SerializerProvider)1 JsonDeserialize (com.fasterxml.jackson.databind.annotation.JsonDeserialize)1 JsonSerialize (com.fasterxml.jackson.databind.annotation.JsonSerialize)1 SimpleModule (com.fasterxml.jackson.databind.module.SimpleModule)1 Dataflow (com.google.api.services.dataflow.Dataflow)1 DataflowPackage (com.google.api.services.dataflow.model.DataflowPackage)1 Job (com.google.api.services.dataflow.model.Job)1 ListJobsResponse (com.google.api.services.dataflow.model.ListJobsResponse)1 SdkHarnessContainerImage (com.google.api.services.dataflow.model.SdkHarnessContainerImage)1 StorageObject (com.google.api.services.storage.model.StorageObject)1 AutoService (com.google.auto.service.AutoService)1