use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.
the class DefaultJobBundleFactoryTest method creatingMultipleEnvironmentFromMultipleTypes.
@Test
public void creatingMultipleEnvironmentFromMultipleTypes() throws Exception {
ServerFactory serverFactory = ServerFactory.createDefault();
Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").build();
EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
EnvironmentFactory.Provider environmentProviderFactoryA = mock(EnvironmentFactory.Provider.class);
when(environmentProviderFactoryA.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryA);
when(environmentProviderFactoryA.getServerFactory()).thenReturn(serverFactory);
Environment environmentB = Environment.newBuilder().setUrn("env:urn:b").build();
EnvironmentFactory envFactoryB = mock(EnvironmentFactory.class);
when(envFactoryB.createEnvironment(eq(environmentB), any())).thenReturn(remoteEnvironment);
EnvironmentFactory.Provider environmentProviderFactoryB = mock(EnvironmentFactory.Provider.class);
when(environmentProviderFactoryB.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryB);
when(environmentProviderFactoryB.getServerFactory()).thenReturn(serverFactory);
Map<String, Provider> environmentFactoryProviderMap = ImmutableMap.of(environmentA.getUrn(), environmentProviderFactoryA, environmentB.getUrn(), environmentProviderFactoryB);
try (DefaultJobBundleFactory bundleFactory = DefaultJobBundleFactory.create(JobInfo.create("testJob", "testJob", "token", Struct.getDefaultInstance()), environmentFactoryProviderMap)) {
bundleFactory.forStage(getExecutableStage(environmentB));
bundleFactory.forStage(getExecutableStage(environmentA));
}
verify(envFactoryA).createEnvironment(eq(environmentA), any());
verify(envFactoryB).createEnvironment(eq(environmentB), any());
}
use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.
the class DefaultJobBundleFactoryTest method expiresEnvironment.
@Test
public void expiresEnvironment() throws Exception {
ServerFactory serverFactory = ServerFactory.createDefault();
Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").build();
EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
EnvironmentFactory.Provider environmentProviderFactoryA = mock(EnvironmentFactory.Provider.class);
when(environmentProviderFactoryA.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryA);
when(environmentProviderFactoryA.getServerFactory()).thenReturn(serverFactory);
Map<String, Provider> environmentFactoryProviderMap = ImmutableMap.of(environmentA.getUrn(), environmentProviderFactoryA);
PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
portableOptions.setEnvironmentExpirationMillis(1);
Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
OutputReceiverFactory orf = mock(OutputReceiverFactory.class);
StateRequestHandler srh = mock(StateRequestHandler.class);
when(srh.getCacheTokens()).thenReturn(Collections.emptyList());
StageBundleFactory sbf = bundleFactory.forStage(getExecutableStage(environmentA));
// allow environment to expire
Thread.sleep(10);
sbf.getBundle(orf, srh, BundleProgressHandler.ignored()).close();
// allow environment to expire
Thread.sleep(10);
sbf.getBundle(orf, srh, BundleProgressHandler.ignored()).close();
}
verify(envFactoryA, Mockito.times(3)).createEnvironment(eq(environmentA), any());
verify(remoteEnvironment, Mockito.times(3)).close();
}
use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.
the class DefaultJobBundleFactoryTest method createsMultipleEnvironmentsWithSdkWorkerParallelism.
@Test
public void createsMultipleEnvironmentsWithSdkWorkerParallelism() throws Exception {
ServerFactory serverFactory = ServerFactory.createDefault();
Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[1])).build();
EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
EnvironmentFactory.Provider environmentProviderFactoryA = mock(EnvironmentFactory.Provider.class);
when(environmentProviderFactoryA.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryA);
when(environmentProviderFactoryA.getServerFactory()).thenReturn(serverFactory);
Map<String, Provider> environmentFactoryProviderMap = ImmutableMap.of(environmentA.getUrn(), environmentProviderFactoryA);
PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
portableOptions.setSdkWorkerParallelism(2);
Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
bundleFactory.forStage(getExecutableStage(environmentA));
verify(environmentProviderFactoryA, Mockito.times(1)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
bundleFactory.forStage(getExecutableStage(environmentA));
verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
verify(envFactoryA, Mockito.times(2)).createEnvironment(eq(environmentA), any());
// round robin, no new environment created
bundleFactory.forStage(getExecutableStage(environmentA));
verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
verify(envFactoryA, Mockito.times(2)).createEnvironment(eq(environmentA), any());
}
portableOptions.setSdkWorkerParallelism(0);
pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
Mockito.reset(envFactoryA);
when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
int expectedParallelism = Math.max(1, Runtime.getRuntime().availableProcessors() - 1);
try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
HashSet<StageBundleFactory> stageBundleFactorySet = new HashSet<>();
// more factories than parallelism for round-robin
int numStageBundleFactories = expectedParallelism + 5;
for (int i = 0; i < numStageBundleFactories; i++) {
stageBundleFactorySet.add(bundleFactory.forStage(getExecutableStage(environmentA)));
}
verify(envFactoryA, Mockito.times(expectedParallelism)).createEnvironment(eq(environmentA), any());
Assert.assertEquals(numStageBundleFactories, stageBundleFactorySet.size());
}
}
Aggregations