Examples with Environment org.apache.beam.model.pipeline.v1.RunnerApi.Environment used on opensource projects

Example 11 with Environment

use of org.apache.beam.model.pipeline.v1.RunnerApi.Environment 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 12 with Environment

use of org.apache.beam.model.pipeline.v1.RunnerApi.Environment in project beam by apache.

the class DockerEnvironmentFactory method createEnvironment.

/**
 * Creates a new, active {@link RemoteEnvironment} backed by a local Docker container.
 */
@Override
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
    Preconditions.checkState(environment.getUrn().equals(BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.DOCKER)), "The passed environment does not contain a DockerPayload.");
    final RunnerApi.DockerPayload dockerPayload = RunnerApi.DockerPayload.parseFrom(environment.getPayload());
    // Prepare docker invocation.
    String containerImage = dockerPayload.getContainerImage();
    // TODO: https://issues.apache.org/jira/browse/BEAM-4148 The default service address will not
    // work for Docker for Mac.
    String provisionEndpoint = provisioningServiceServer.getApiServiceDescriptor().getUrl();
    ImmutableList.Builder<String> dockerOptsBuilder = ImmutableList.<String>builder().addAll(gcsCredentialArgs()).add("--network=host").add("--env=DOCKER_MAC_CONTAINER=" + System.getenv("DOCKER_MAC_CONTAINER"));
    final boolean retainDockerContainer = pipelineOptions.as(ManualDockerEnvironmentOptions.class).getRetainDockerContainers();
    String semiPersistDir = pipelineOptions.as(RemoteEnvironmentOptions.class).getSemiPersistDir();
    ImmutableList.Builder<String> argsBuilder = ImmutableList.<String>builder().add(String.format("--id=%s", workerId)).add(String.format("--provision_endpoint=%s", provisionEndpoint));
    if (semiPersistDir != null) {
        argsBuilder.add(String.format("--semi_persist_dir=%s", semiPersistDir));
    }
    LOG.debug("Creating Docker Container with ID {}", workerId);
    // Wrap the blocking call to clientSource.get in case an exception is thrown.
    String containerId = null;
    InstructionRequestHandler instructionHandler = null;
    try {
        containerId = docker.runImage(containerImage, dockerOptsBuilder.build(), argsBuilder.build());
        LOG.debug("Created Docker Container with Container ID {}", containerId);
        // Wait on a client from the gRPC server.
        while (instructionHandler == null) {
            try {
                // If the docker is not alive anymore, we abort.
                if (!docker.isContainerRunning(containerId)) {
                    IllegalStateException illegalStateException = new IllegalStateException(String.format("No container running for id %s", containerId));
                    try {
                        String containerLogs = docker.getContainerLogs(containerId);
                        LOG.error("Docker container {} logs:\n{}", containerId, containerLogs);
                    } catch (Exception getLogsException) {
                        illegalStateException.addSuppressed(getLogsException);
                    }
                    throw illegalStateException;
                }
                instructionHandler = clientSource.take(workerId, Duration.ofSeconds(5));
            } catch (TimeoutException timeoutEx) {
                LOG.info("Still waiting for startup of environment {} for worker id {}", dockerPayload.getContainerImage(), workerId);
            } catch (InterruptedException interruptEx) {
                Thread.currentThread().interrupt();
                throw new RuntimeException(interruptEx);
            }
        }
    } catch (Exception e) {
        if (containerId != null) {
            // Kill the launched docker container if we can't retrieve a client for it.
            try {
                docker.killContainer(containerId);
                if (!retainDockerContainer) {
                    docker.removeContainer(containerId);
                }
            } catch (Exception dockerException) {
                e.addSuppressed(dockerException);
            }
        }
        throw e;
    }
    return DockerContainerEnvironment.create(docker, environment, containerId, instructionHandler, retainDockerContainer);
}

Example 13 with Environment

use of org.apache.beam.model.pipeline.v1.RunnerApi.Environment in project beam by apache.

the class ProcessEnvironmentFactory method createEnvironment.

/**
 * Creates a new, active {@link RemoteEnvironment} backed by a forked process.
 */
@Override
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
    Preconditions.checkState(environment.getUrn().equals(BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.PROCESS)), "The passed environment does not contain a ProcessPayload.");
    final RunnerApi.ProcessPayload processPayload = RunnerApi.ProcessPayload.parseFrom(environment.getPayload());
    String executable = processPayload.getCommand();
    String provisionEndpoint = provisioningServiceServer.getApiServiceDescriptor().getUrl();
    String semiPersistDir = pipelineOptions.as(RemoteEnvironmentOptions.class).getSemiPersistDir();
    ImmutableList.Builder<String> argsBuilder = ImmutableList.<String>builder().add(String.format("--id=%s", workerId)).add(String.format("--provision_endpoint=%s", provisionEndpoint));
    if (semiPersistDir != null) {
        argsBuilder.add(String.format("--semi_persist_dir=%s", semiPersistDir));
    }
    LOG.debug("Creating Process for worker ID {}", workerId);
    // Wrap the blocking call to clientSource.get in case an exception is thrown.
    InstructionRequestHandler instructionHandler = null;
    try {
        ProcessManager.RunningProcess process = processManager.startProcess(workerId, executable, argsBuilder.build(), processPayload.getEnvMap());
        // Wait on a client from the gRPC server.
        while (instructionHandler == null) {
            try {
                // If the process is not alive anymore, we abort.
                process.isAliveOrThrow();
                instructionHandler = clientSource.take(workerId, Duration.ofSeconds(5));
            } catch (TimeoutException timeoutEx) {
                LOG.info("Still waiting for startup of environment '{}' for worker id {}", processPayload.getCommand(), workerId);
            } catch (InterruptedException interruptEx) {
                Thread.currentThread().interrupt();
                throw new RuntimeException(interruptEx);
            }
        }
    } catch (Exception e) {
        try {
            processManager.stopProcess(workerId);
        } catch (Exception processKillException) {
            e.addSuppressed(processKillException);
        }
        throw e;
    }
    return ProcessEnvironment.create(processManager, environment, workerId, instructionHandler);
}

Example 14 with Environment

use of org.apache.beam.model.pipeline.v1.RunnerApi.Environment in project beam by apache.

the class FnHarnessTest method testLaunchFnHarnessAndTeardownCleanly.

@Test
// failure will cause test to timeout.
@SuppressWarnings("FutureReturnValueIgnored")
public void testLaunchFnHarnessAndTeardownCleanly() throws Exception {
    Function<String, String> environmentVariableMock = mock(Function.class);
    PipelineOptions options = PipelineOptionsFactory.create();
    when(environmentVariableMock.apply("HARNESS_ID")).thenReturn("id");
    when(environmentVariableMock.apply("PIPELINE_OPTIONS")).thenReturn(PipelineOptionsTranslation.toJson(options));
    List<BeamFnApi.LogEntry> logEntries = new ArrayList<>();
    List<BeamFnApi.InstructionResponse> instructionResponses = mock(List.class);
    BeamFnLoggingGrpc.BeamFnLoggingImplBase loggingService = new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<LogControl> responseObserver) {
            return TestStreams.withOnNext((BeamFnApi.LogEntry.List entries) -> logEntries.addAll(entries.getLogEntriesList())).withOnCompleted(responseObserver::onCompleted).build();
        }
    };
    BeamFnControlGrpc.BeamFnControlImplBase controlService = new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<InstructionResponse> control(StreamObserver<InstructionRequest> responseObserver) {
            CountDownLatch waitForResponses = new CountDownLatch(1);
            options.as(GcsOptions.class).getExecutorService().submit(() -> {
                responseObserver.onNext(INSTRUCTION_REQUEST);
                Uninterruptibles.awaitUninterruptibly(waitForResponses);
                responseObserver.onCompleted();
            });
            return TestStreams.withOnNext((InstructionResponse t) -> {
                instructionResponses.add(t);
                waitForResponses.countDown();
            }).withOnCompleted(waitForResponses::countDown).build();
        }
    };
    Server loggingServer = ServerBuilder.forPort(0).addService(loggingService).build();
    loggingServer.start();
    try {
        Server controlServer = ServerBuilder.forPort(0).addService(controlService).build();
        controlServer.start();
        try {
            Endpoints.ApiServiceDescriptor loggingDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + loggingServer.getPort()).build();
            Endpoints.ApiServiceDescriptor controlDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + controlServer.getPort()).build();
            when(environmentVariableMock.apply("LOGGING_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(loggingDescriptor));
            when(environmentVariableMock.apply("CONTROL_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(controlDescriptor));
            FnHarness.main(environmentVariableMock);
        } finally {
            controlServer.shutdownNow();
        }
    } finally {
        loggingServer.shutdownNow();
    }
    // Verify that we first run onStartup functions before even reading the environment, and that
    // we then call beforeProcessing functions before executing instructions.
    InOrder inOrder = inOrder(onStartupMock, beforeProcessingMock, environmentVariableMock, instructionResponses);
    inOrder.verify(onStartupMock).run();
    inOrder.verify(environmentVariableMock, atLeastOnce()).apply(any());
    inOrder.verify(beforeProcessingMock).accept(any());
    inOrder.verify(instructionResponses).add(INSTRUCTION_RESPONSE);
}

Example 15 with Environment

use of org.apache.beam.model.pipeline.v1.RunnerApi.Environment in project beam by apache.

the class DataflowWorkerHarnessHelper method getPipelineFromEnv.

// TODO: make env logic private to main() so it is never done outside of initializing the process
public static RunnerApi.@Nullable Pipeline getPipelineFromEnv() throws IOException {
    String pipelinePath = System.getenv(PIPELINE_PATH);
    if (pipelinePath == null) {
        LOG.warn("Missing pipeline environment variable '{}'", PIPELINE_PATH);
        return null;
    }
    File pipelineFile = new File(System.getenv(PIPELINE_PATH));
    if (!pipelineFile.exists()) {
        LOG.warn("Pipeline path '{}' does not exist", pipelineFile);
        return null;
    }
    try (FileInputStream inputStream = new FileInputStream(pipelineFile)) {
        RunnerApi.Pipeline pipelineProto = RunnerApi.Pipeline.parseFrom(inputStream);
        LOG.info("Found portable pipeline:\n{}", TextFormat.printToString(pipelineProto));
        return pipelineProto;
    }
}