use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.
the class EmbeddedEnvironmentFactory method createEnvironment.
@Override
// no need to monitor shutdown thread
@SuppressWarnings("FutureReturnValueIgnored")
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
ExecutorService executor = Executors.newSingleThreadExecutor();
Future<?> fnHarness = executor.submit(() -> {
try {
FnHarness.main(workerId, options, // Runner capabilities.
Collections.emptySet(), loggingServer.getApiServiceDescriptor(), controlServer.getApiServiceDescriptor(), null, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.clientDirect(), Caches.fromOptions(options));
} catch (NoClassDefFoundError e) {
// TODO: https://issues.apache.org/jira/browse/BEAM-4384 load the FnHarness in a
// Restricted classpath that we control for any user.
LOG.error("{} while executing an in-process FnHarness. " + "To use the {}, " + "the 'org.apache.beam:beam-sdks-java-harness' artifact " + "and its dependencies must be on the classpath", NoClassDefFoundError.class.getSimpleName(), EmbeddedEnvironmentFactory.class.getSimpleName(), e);
throw e;
}
return null;
});
executor.submit(() -> {
try {
fnHarness.get();
} catch (Throwable t) {
executor.shutdownNow();
}
});
InstructionRequestHandler handler = null;
// Wait on a client from the gRPC server.
while (handler == null) {
try {
// If the thread is not alive anymore, we abort.
if (executor.isShutdown()) {
throw new IllegalStateException("FnHarness startup failed");
}
handler = clientSource.take(workerId, Duration.ofSeconds(5L));
} catch (TimeoutException timeoutEx) {
LOG.info("Still waiting for startup of FnHarness");
} catch (InterruptedException interruptEx) {
Thread.currentThread().interrupt();
throw new RuntimeException(interruptEx);
}
}
return RemoteEnvironment.forHandler(environment, handler);
}
use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.
the class ExternalEnvironmentFactory method createEnvironment.
/**
* Creates a new, active {@link RemoteEnvironment} backed by an unmanaged worker.
*/
@Override
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
Preconditions.checkState(environment.getUrn().equals(BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.EXTERNAL)), "The passed environment does not contain an ExternalPayload.");
final RunnerApi.ExternalPayload externalPayload = RunnerApi.ExternalPayload.parseFrom(environment.getPayload());
BeamFnApi.StartWorkerRequest startWorkerRequest = BeamFnApi.StartWorkerRequest.newBuilder().setWorkerId(workerId).setControlEndpoint(controlServiceServer.getApiServiceDescriptor()).setLoggingEndpoint(loggingServiceServer.getApiServiceDescriptor()).setArtifactEndpoint(retrievalServiceServer.getApiServiceDescriptor()).setProvisionEndpoint(provisioningServiceServer.getApiServiceDescriptor()).putAllParams(externalPayload.getParamsMap()).build();
LOG.debug("Requesting worker ID {}", workerId);
final ManagedChannel managedChannel = ManagedChannelFactory.createDefault().forDescriptor(externalPayload.getEndpoint());
BeamFnApi.StartWorkerResponse startWorkerResponse = BeamFnExternalWorkerPoolGrpc.newBlockingStub(managedChannel).startWorker(startWorkerRequest);
if (!startWorkerResponse.getError().isEmpty()) {
throw new RuntimeException(startWorkerResponse.getError());
}
// Wait on a client from the gRPC server.
InstructionRequestHandler instructionHandler = null;
while (instructionHandler == null) {
try {
instructionHandler = clientSource.take(workerId, Duration.ofMinutes(2));
} catch (TimeoutException timeoutEx) {
LOG.info("Still waiting for startup of environment from {} for worker id {}", externalPayload.getEndpoint().getUrl(), workerId);
} catch (InterruptedException interruptEx) {
Thread.currentThread().interrupt();
managedChannel.shutdownNow();
throw new RuntimeException(interruptEx);
}
}
final InstructionRequestHandler finalInstructionHandler = instructionHandler;
return new RemoteEnvironment() {
@Override
public Environment getEnvironment() {
return environment;
}
@Override
public InstructionRequestHandler getInstructionRequestHandler() {
return finalInstructionHandler;
}
@Override
public void close() throws Exception {
try {
finalInstructionHandler.close();
BeamFnApi.StopWorkerRequest stopWorkerRequest = BeamFnApi.StopWorkerRequest.newBuilder().setWorkerId(workerId).build();
LOG.debug("Closing worker ID {}", workerId);
BeamFnApi.StopWorkerResponse stopWorkerResponse = BeamFnExternalWorkerPoolGrpc.newBlockingStub(managedChannel).stopWorker(stopWorkerRequest);
if (!stopWorkerResponse.getError().isEmpty()) {
throw new RuntimeException(stopWorkerResponse.getError());
}
} finally {
managedChannel.shutdown();
managedChannel.awaitTermination(10, TimeUnit.SECONDS);
if (!managedChannel.isTerminated()) {
managedChannel.shutdownNow();
}
}
}
};
}
use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.
the class ProcessEnvironmentTest method closeClosesInstructionRequestHandler.
@Test
public void closeClosesInstructionRequestHandler() throws Exception {
InstructionRequestHandler handler = mock(InstructionRequestHandler.class);
RemoteEnvironment env = ProcessEnvironment.create(mock(ProcessManager.class), Environment.getDefaultInstance(), "1", handler);
env.close();
verify(handler).close();
}
use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.
the class RemoteEnvironmentTest method closeClosesInstructionRequestHandler.
@Test
public void closeClosesInstructionRequestHandler() throws Exception {
InstructionRequestHandler handler = mock(InstructionRequestHandler.class);
RemoteEnvironment env = new RemoteEnvironment() {
@Override
public Environment getEnvironment() {
throw new UnsupportedOperationException();
}
@Override
public InstructionRequestHandler getInstructionRequestHandler() {
return handler;
}
};
env.close();
verify(handler).close();
}
Aggregations