Search in sources :

Example 1 with ResourceHintsOptions

use of org.apache.beam.sdk.transforms.resourcehints.ResourceHintsOptions in project beam by apache.

the class DataflowPipelineTranslatorTest method testResourceHintsTranslationsResolvesHintsOnOptionsAndComposites.

@Test
public void testResourceHintsTranslationsResolvesHintsOnOptionsAndComposites() {
    ResourceHintsOptions options = PipelineOptionsFactory.as(ResourceHintsOptions.class);
    options.setResourceHints(Arrays.asList("accelerator=set_via_options", "minRam=1B"));
    Pipeline pipeline = Pipeline.create(options);
    PCollection<byte[]> root = pipeline.apply(Impulse.create());
    root.apply(new Outer().setResourceHints(ResourceHints.create().withAccelerator("set_on_outer_transform").withMinRam(20)));
    root.apply("Leaf", ParDo.of(new IdentityDoFn<byte[]>()));
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, false);
    assertThat(pipelineProto.getComponents().getEnvironmentsMap().get(getLeafTransform(pipelineProto, "Leaf").getEnvironmentId()).getResourceHintsMap(), org.hamcrest.Matchers.allOf(org.hamcrest.Matchers.hasEntry("beam:resources:min_ram_bytes:v1", ByteString.copyFromUtf8("1")), org.hamcrest.Matchers.hasEntry("beam:resources:accelerator:v1", ByteString.copyFromUtf8("set_via_options"))));
    assertThat(pipelineProto.getComponents().getEnvironmentsMap().get(getLeafTransform(pipelineProto, "Innermost").getEnvironmentId()).getResourceHintsMap(), org.hamcrest.Matchers.allOf(org.hamcrest.Matchers.hasEntry("beam:resources:min_ram_bytes:v1", ByteString.copyFromUtf8("20")), org.hamcrest.Matchers.hasEntry("beam:resources:accelerator:v1", ByteString.copyFromUtf8("set_in_inner_transform"))));
}
Also used : ResourceHintsOptions(org.apache.beam.sdk.transforms.resourcehints.ResourceHintsOptions) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Aggregations

RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)1 Pipeline (org.apache.beam.sdk.Pipeline)1 ResourceHintsOptions (org.apache.beam.sdk.transforms.resourcehints.ResourceHintsOptions)1 Test (org.junit.Test)1