Search in sources :

Example 1 with JobServiceBlockingStub

use of org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub in project beam by apache.

the class JobServicePipelineResult method cancel.

@Override
public State cancel() {
    JobServiceBlockingStub stub = jobService.get();
    CancelJobResponse response = stub.cancel(CancelJobRequest.newBuilder().setJobIdBytes(jobId).build());
    return getJavaState(response.getState());
}
Also used : JobServiceBlockingStub(org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub) CancelJobResponse(org.apache.beam.model.jobmanagement.v1.JobApi.CancelJobResponse)

Example 2 with JobServiceBlockingStub

use of org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub in project beam by apache.

the class JobServicePipelineResult method waitForTerminalState.

private void waitForTerminalState() {
    JobServiceBlockingStub stub = jobService.get().withDeadlineAfter(jobServerTimeout, TimeUnit.SECONDS);
    GetJobStateRequest request = GetJobStateRequest.newBuilder().setJobIdBytes(jobId).build();
    JobStateEvent response = stub.getState(request);
    State lastState = getJavaState(response.getState());
    while (!lastState.isTerminal()) {
        try {
            Thread.sleep(POLL_INTERVAL_MS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(e);
        }
        response = stub.withDeadlineAfter(jobServerTimeout, TimeUnit.SECONDS).getState(request);
        lastState = getJavaState(response.getState());
    }
    terminalState = lastState;
}
Also used : JobServiceBlockingStub(org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub) JobStateEvent(org.apache.beam.model.jobmanagement.v1.JobApi.JobStateEvent) GetJobStateRequest(org.apache.beam.model.jobmanagement.v1.JobApi.GetJobStateRequest)

Example 3 with JobServiceBlockingStub

use of org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub in project beam by apache.

the class JobServicePipelineResult method getState.

@Override
public State getState() {
    if (terminalState != null) {
        return terminalState;
    }
    JobServiceBlockingStub stub = jobService.get().withDeadlineAfter(jobServerTimeout, TimeUnit.SECONDS);
    JobStateEvent response = stub.getState(GetJobStateRequest.newBuilder().setJobIdBytes(jobId).build());
    return getJavaState(response.getState());
}
Also used : JobServiceBlockingStub(org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub) JobStateEvent(org.apache.beam.model.jobmanagement.v1.JobApi.JobStateEvent)

Example 4 with JobServiceBlockingStub

use of org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub in project beam by apache.

the class PortableRunner method run.

@Override
public PipelineResult run(Pipeline pipeline) {
    Runnable cleanup;
    if (Environments.ENVIRONMENT_LOOPBACK.equals(options.as(PortablePipelineOptions.class).getDefaultEnvironmentType())) {
        GrpcFnServer<ExternalWorkerService> workerService;
        try {
            workerService = new ExternalWorkerService(options).start();
        } catch (Exception exn) {
            throw new RuntimeException("Failed to start GrpcFnServer for ExternalWorkerService", exn);
        }
        LOG.info("Starting worker service at {}", workerService.getApiServiceDescriptor().getUrl());
        options.as(PortablePipelineOptions.class).setDefaultEnvironmentConfig(workerService.getApiServiceDescriptor().getUrl());
        cleanup = () -> {
            try {
                LOG.warn("closing worker service {}", workerService);
                workerService.close();
            } catch (Exception exn) {
                throw new RuntimeException(exn);
            }
        };
    } else {
        cleanup = null;
    }
    ImmutableList.Builder<String> filesToStageBuilder = ImmutableList.builder();
    List<String> stagingFiles = options.as(PortablePipelineOptions.class).getFilesToStage();
    if (stagingFiles == null) {
        List<String> classpathResources = detectClassPathResourcesToStage(Environments.class.getClassLoader(), options);
        if (classpathResources.isEmpty()) {
            throw new IllegalArgumentException("No classpath elements found.");
        }
        LOG.debug("PortablePipelineOptions.filesToStage was not specified. " + "Defaulting to files from the classpath: {}", classpathResources.size());
        filesToStageBuilder.addAll(classpathResources);
    } else {
        filesToStageBuilder.addAll(stagingFiles);
    }
    // TODO(heejong): remove jar_packages experimental flag when cross-language dependency
    // management is implemented for all runners.
    List<String> experiments = options.as(ExperimentalOptions.class).getExperiments();
    if (experiments != null) {
        Optional<String> jarPackages = experiments.stream().filter((String flag) -> flag.startsWith("jar_packages=")).findFirst();
        jarPackages.ifPresent(s -> filesToStageBuilder.addAll(Arrays.asList(s.replaceFirst("jar_packages=", "").split(","))));
    }
    options.as(PortablePipelineOptions.class).setFilesToStage(filesToStageBuilder.build());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, SdkComponents.create(options));
    pipelineProto = DefaultArtifactResolver.INSTANCE.resolveArtifacts(pipelineProto);
    PrepareJobRequest prepareJobRequest = PrepareJobRequest.newBuilder().setJobName(options.getJobName()).setPipeline(pipelineProto).setPipelineOptions(PipelineOptionsTranslation.toProto(options)).build();
    LOG.info("Using job server endpoint: {}", endpoint);
    ManagedChannel jobServiceChannel = channelFactory.forDescriptor(ApiServiceDescriptor.newBuilder().setUrl(endpoint).build());
    JobServiceBlockingStub jobService = JobServiceGrpc.newBlockingStub(jobServiceChannel);
    try (CloseableResource<JobServiceBlockingStub> wrappedJobService = CloseableResource.of(jobService, unused -> jobServiceChannel.shutdown())) {
        final int jobServerTimeout = options.as(PortablePipelineOptions.class).getJobServerTimeout();
        PrepareJobResponse prepareJobResponse = jobService.withDeadlineAfter(jobServerTimeout, TimeUnit.SECONDS).withWaitForReady().prepare(prepareJobRequest);
        LOG.info("PrepareJobResponse: {}", prepareJobResponse);
        ApiServiceDescriptor artifactStagingEndpoint = prepareJobResponse.getArtifactStagingEndpoint();
        String stagingSessionToken = prepareJobResponse.getStagingSessionToken();
        try (CloseableResource<ManagedChannel> artifactChannel = CloseableResource.of(channelFactory.forDescriptor(artifactStagingEndpoint), ManagedChannel::shutdown)) {
            ArtifactStagingService.offer(new ArtifactRetrievalService(), ArtifactStagingServiceGrpc.newStub(artifactChannel.get()), stagingSessionToken);
        } catch (CloseableResource.CloseException e) {
            LOG.warn("Error closing artifact staging channel", e);
        // CloseExceptions should only be thrown while closing the channel.
        } catch (Exception e) {
            throw new RuntimeException("Error staging files.", e);
        }
        RunJobRequest runJobRequest = RunJobRequest.newBuilder().setPreparationId(prepareJobResponse.getPreparationId()).build();
        // Run the job and wait for a result, we don't set a timeout here because
        // it may take a long time for a job to complete and streaming
        // jobs never return a response.
        RunJobResponse runJobResponse = jobService.run(runJobRequest);
        LOG.info("RunJobResponse: {}", runJobResponse);
        ByteString jobId = runJobResponse.getJobIdBytes();
        return new JobServicePipelineResult(jobId, jobServerTimeout, wrappedJobService.transfer(), cleanup);
    } catch (CloseException e) {
        throw new RuntimeException(e);
    }
}
Also used : JobServiceBlockingStub(org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub) PrepareJobResponse(org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobResponse) ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) CloseException(org.apache.beam.runners.portability.CloseableResource.CloseException) ExperimentalOptions(org.apache.beam.sdk.options.ExperimentalOptions) ArtifactRetrievalService(org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService) RunJobResponse(org.apache.beam.model.jobmanagement.v1.JobApi.RunJobResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) RunJobRequest(org.apache.beam.model.jobmanagement.v1.JobApi.RunJobRequest) ExternalWorkerService(org.apache.beam.fn.harness.ExternalWorkerService) PrepareJobRequest(org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobRequest) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) CloseException(org.apache.beam.runners.portability.CloseableResource.CloseException) CloseException(org.apache.beam.runners.portability.CloseableResource.CloseException) Environments(org.apache.beam.runners.core.construction.Environments) PortablePipelineOptions(org.apache.beam.sdk.options.PortablePipelineOptions)

Aggregations

JobServiceBlockingStub (org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub)4 JobStateEvent (org.apache.beam.model.jobmanagement.v1.JobApi.JobStateEvent)2 ExternalWorkerService (org.apache.beam.fn.harness.ExternalWorkerService)1 CancelJobResponse (org.apache.beam.model.jobmanagement.v1.JobApi.CancelJobResponse)1 GetJobStateRequest (org.apache.beam.model.jobmanagement.v1.JobApi.GetJobStateRequest)1 PrepareJobRequest (org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobRequest)1 PrepareJobResponse (org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobResponse)1 RunJobRequest (org.apache.beam.model.jobmanagement.v1.JobApi.RunJobRequest)1 RunJobResponse (org.apache.beam.model.jobmanagement.v1.JobApi.RunJobResponse)1 ApiServiceDescriptor (org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor)1 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)1 Environments (org.apache.beam.runners.core.construction.Environments)1 ArtifactRetrievalService (org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService)1 CloseException (org.apache.beam.runners.portability.CloseableResource.CloseException)1 ExperimentalOptions (org.apache.beam.sdk.options.ExperimentalOptions)1 PortablePipelineOptions (org.apache.beam.sdk.options.PortablePipelineOptions)1 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)1 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)1 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)1