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"))));
}
Aggregations