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

Example 1 with EnvironmentFactory

use of org.apache.beam.runners.fnexecution.environment.EnvironmentFactory 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 2 with EnvironmentFactory

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

the class DefaultJobBundleFactory method createEnvironmentCaches.

private ImmutableList<EnvironmentCacheAndLock> createEnvironmentCaches(ThrowingFunction<ServerFactory, ServerInfo> serverInfoCreator, int count) {
    ImmutableList.Builder<EnvironmentCacheAndLock> caches = ImmutableList.builder();
    for (int i = 0; i < count; i++) {
        final Lock refLock;
        if (environmentExpirationMillis > 0) {
            // The lock ensures there is no race condition between expiring an environment and a client
            // still attempting to use it, hence referencing it.
            refLock = new ReentrantLock(true);
        } else {
            refLock = NoopLock.get();
        }
        CacheBuilder<Environment, WrappedSdkHarnessClient> cacheBuilder = CacheBuilder.newBuilder().removalListener(notification -> {
            WrappedSdkHarnessClient client = notification.getValue();
            final int refCount;
            // We need to use a lock here to ensure we are not causing the environment to
            // be removed if beforehand a StageBundleFactory has retrieved it but not yet
            // issued ref() on it.
            refLock.lock();
            try {
                refCount = client.unref();
            } finally {
                refLock.unlock();
            }
            if (refCount > 0) {
                LOG.warn("Expiring environment {} with {} remaining bundle references. Taking note to clean it up during shutdown if the references are not removed by then.", notification.getKey(), refCount);
                evictedActiveClients.add(client);
            }
        });
        if (environmentExpirationMillis > 0) {
            cacheBuilder.expireAfterWrite(environmentExpirationMillis, TimeUnit.MILLISECONDS);
        }
        LoadingCache<Environment, WrappedSdkHarnessClient> cache = cacheBuilder.build(new CacheLoader<Environment, WrappedSdkHarnessClient>() {

            @Override
            public WrappedSdkHarnessClient load(Environment environment) throws Exception {
                EnvironmentFactory.Provider environmentFactoryProvider = environmentFactoryProviderMap.get(environment.getUrn());
                ServerFactory serverFactory = environmentFactoryProvider.getServerFactory();
                ServerInfo serverInfo = serverInfoCreator.apply(serverFactory);
                String workerId = stageIdGenerator.getId();
                serverInfo.getProvisioningServer().getService().registerEnvironment(workerId, environment);
                EnvironmentFactory environmentFactory = environmentFactoryProvider.createEnvironmentFactory(serverInfo.getControlServer(), serverInfo.getLoggingServer(), serverInfo.getRetrievalServer(), serverInfo.getProvisioningServer(), clientPool, stageIdGenerator);
                return WrappedSdkHarnessClient.wrapping(environmentFactory.createEnvironment(environment, workerId), serverInfo);
            }
        });
        caches.add(new EnvironmentCacheAndLock(cache, refLock));
    }
    return caches.build();
}

Example 3 with EnvironmentFactory

use of org.apache.beam.runners.fnexecution.environment.EnvironmentFactory 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 4 with EnvironmentFactory

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

Example 5 with EnvironmentFactory

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