Examples with RemoteEnvironment org.apache.beam.runners.fnexecution.environment.RemoteEnvironment used on opensource projects

Example 1 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment in project beam by apache.

the class DefaultJobBundleFactoryTest method doesNotCacheDifferentEnvironments.

@Test
public void doesNotCacheDifferentEnvironments() throws Exception {
    Environment envFoo = Environment.newBuilder().setUrn("dummy:urn:another").build();
    RemoteEnvironment remoteEnvFoo = mock(RemoteEnvironment.class);
    InstructionRequestHandler fooInstructionHandler = mock(InstructionRequestHandler.class);
    Map<String, EnvironmentFactory.Provider> envFactoryProviderMapFoo = ImmutableMap.of(environment.getUrn(), envFactoryProvider, envFoo.getUrn(), envFactoryProvider);
    when(envFactory.createEnvironment(eq(envFoo), any())).thenReturn(remoteEnvFoo);
    when(remoteEnvFoo.getInstructionRequestHandler()).thenReturn(fooInstructionHandler);
    // Don't bother creating a distinct instruction response because we don't use it here.
    when(fooInstructionHandler.handle(any())).thenReturn(CompletableFuture.completedFuture(instructionResponse));
    try (DefaultJobBundleFactory bundleFactory = createDefaultJobBundleFactory(envFactoryProviderMapFoo)) {
        bundleFactory.forStage(getExecutableStage(environment));
        bundleFactory.forStage(getExecutableStage(envFoo));
        verify(envFactory).createEnvironment(eq(environment), any());
        verify(envFactory).createEnvironment(eq(envFoo), any());
        verifyNoMoreInteractions(envFactory);
    }
}

Example 2 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment in project beam by apache.

the class DefaultJobBundleFactoryTest method createsMultipleEnvironmentOfSingleType.

@Test
public void createsMultipleEnvironmentOfSingleType() throws Exception {
    ServerFactory serverFactory = ServerFactory.createDefault();
    Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[1])).build();
    Environment environmentAA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[2])).build();
    EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
    when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
    when(envFactoryA.createEnvironment(eq(environmentAA), 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 = createDefaultJobBundleFactory(environmentFactoryProviderMap)) {
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(1)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(environmentProviderFactoryB, Mockito.times(0)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
        verify(envFactoryA, Mockito.times(0)).createEnvironment(eq(environmentAA), any());
        bundleFactory.forStage(getExecutableStage(environmentAA));
        verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(environmentProviderFactoryB, Mockito.times(0)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentAA), any());
    }
}

Example 3 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment in project beam by apache.

the class SingleEnvironmentInstanceJobBundleFactoryTest method closeShutsDownEnvironmentsWhenSomeFail.

@Test
public void closeShutsDownEnvironmentsWhenSomeFail() throws Exception {
    Pipeline p = Pipeline.create();
    ExperimentalOptions.addExperiment(p.getOptions().as(ExperimentalOptions.class), "beam_fn_api");
    p.apply("Create", Create.of(1, 2, 3));
    ExecutableStage firstEnvStage = GreedyPipelineFuser.fuse(PipelineTranslation.toProto(p)).getFusedStages().stream().findFirst().get();
    ExecutableStagePayload basePayload = ExecutableStagePayload.parseFrom(firstEnvStage.toPTransform("foo").getSpec().getPayload());
    Environment secondEnv = Environments.createDockerEnvironment("second_env");
    ExecutableStage secondEnvStage = ExecutableStage.fromPayload(basePayload.toBuilder().setEnvironment(secondEnv).build());
    Environment thirdEnv = Environments.createDockerEnvironment("third_env");
    ExecutableStage thirdEnvStage = ExecutableStage.fromPayload(basePayload.toBuilder().setEnvironment(thirdEnv).build());
    RemoteEnvironment firstRemoteEnv = mock(RemoteEnvironment.class, "First Remote Env");
    RemoteEnvironment secondRemoteEnv = mock(RemoteEnvironment.class, "Second Remote Env");
    RemoteEnvironment thirdRemoteEnv = mock(RemoteEnvironment.class, "Third Remote Env");
    when(environmentFactory.createEnvironment(firstEnvStage.getEnvironment(), GENERATED_ID)).thenReturn(firstRemoteEnv);
    when(environmentFactory.createEnvironment(secondEnvStage.getEnvironment(), GENERATED_ID)).thenReturn(secondRemoteEnv);
    when(environmentFactory.createEnvironment(thirdEnvStage.getEnvironment(), GENERATED_ID)).thenReturn(thirdRemoteEnv);
    when(firstRemoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    when(secondRemoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    when(thirdRemoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    factory.forStage(firstEnvStage);
    factory.forStage(secondEnvStage);
    factory.forStage(thirdEnvStage);
    IllegalStateException firstException = new IllegalStateException("first stage");
    doThrow(firstException).when(firstRemoteEnv).close();
    IllegalStateException thirdException = new IllegalStateException("third stage");
    doThrow(thirdException).when(thirdRemoteEnv).close();
    try {
        factory.close();
        fail("Factory close should have thrown");
    } catch (IllegalStateException expected) {
        if (expected.equals(firstException)) {
            assertThat(ImmutableList.copyOf(expected.getSuppressed()), contains(thirdException));
        } else if (expected.equals(thirdException)) {
            assertThat(ImmutableList.copyOf(expected.getSuppressed()), contains(firstException));
        } else {
            throw expected;
        }
        verify(firstRemoteEnv).close();
        verify(secondRemoteEnv).close();
        verify(thirdRemoteEnv).close();
    }
}

Example 4 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment 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());
}

Example 5 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment 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();
}