/** * Construct an endpoint in driver side. * This method will be called by {@link ClientEndpoint}. * @param planStateManager of running plan. * @param clientEndpoint of running plan. */ public DriverEndpoint(final PlanStateManager planStateManager, final ClientEndpoint clientEndpoint) { this.planStateManager = planStateManager; this.clientEndpoint = clientEndpoint; clientEndpoint.connectDriver(this); }
/** * @return the host of the client */ public String getListeningHost() { // Listening host is determined by LocalAddressProvider, in 'run' method. ensureServerState(true); return host; }
@Override public State getState() { return (State) super.getPlanState(); }
builtJobConf = getJobConf(args); driverRPCServer = new DriverRPCServer(); driverRPCServer .registerHandler(ControlMessage.DriverToClientMessageType.DriverStarted, event -> { }) .registerHandler(ControlMessage.DriverToClientMessageType.DriverReady, event -> driverReadyLatch.countDown()) .registerHandler(ControlMessage.DriverToClientMessageType.ExecutionDone, event -> jobDoneLatch.countDown()) .registerHandler(ControlMessage.DriverToClientMessageType.DataCollected, message -> COLLECTED_DATA.addAll( SerializationUtils.deserialize(Base64.getDecoder().decode(message.getDataCollected().getData())))) .run(); final Configuration driverConf = getDriverConf(builtJobConf); final Configuration driverNcsConf = getDriverNcsConf(); final Configuration driverMessageConfig = getDriverMessageConf(); final String defaultExecutorResourceConfig = "[{\"type\":\"Transient\",\"memory_mb\":512,\"capacity\":5}," + "{\"type\":\"Reserved\",\"memory_mb\":512,\"capacity\":5}]"; final Configuration executorResourceConfig = getJSONConf(builtJobConf, JobConf.ExecutorJSONPath.class, JobConf.ExecutorJSONContents.class, defaultExecutorResourceConfig); final Configuration bandwidthConfig = getJSONConf(builtJobConf, JobConf.BandwidthJSONPath.class, JobConf.BandwidthJSONContents.class, ""); final Configuration clientConf = getClientConf(); final Configuration schedulerConf = getSchedulerConf(builtJobConf); executorResourceConfig, bandwidthConfig, driverRPCServer.getListeningConfiguration(), schedulerConf); deployModeConf = Configurations.merge(getDeployModeConf(builtJobConf), clientConf);
/** * Launch application using the application DAG. * Notice that we launch the DAG one at a time, as the result of a DAG has to be immediately returned to the * Java variable before the application can be resumed. * * @param dag the application DAG. */ // When modifying the signature of this method, see CompilerTestUtil#compileDAG and make corresponding changes public static void launchDAG(final DAG dag) { launchDAG(dag, Collections.emptyMap(), ""); }
/** * Wait for this job to be finished and return the final state. * * @return the final state of this job. */ public final Enum waitUntilJobFinish() { if (driverEndpoint.get() != null) { return stateTranslator.translateState(driverEndpoint.get().waitUntilFinish()); } else { // The driver endpoint is not connected yet. final boolean driverIsConnected = waitUntilConnected(); if (driverIsConnected) { return stateTranslator.translateState(driverEndpoint.get().waitUntilFinish()); } else { return PlanState.State.READY; } } } }
setup(new String[]{"-job_id", jobId}); } catch (Exception e) { throw new RuntimeException(e); serializedDAG = Base64.getEncoder().encodeToString(SerializationUtils.serialize(dag)); jobDoneLatch = new CountDownLatch(1); driverRPCServer.send(ControlMessage.ClientToDriverMessage.newBuilder() .setType(ControlMessage.ClientToDriverMessageType.LaunchDAG) .setLaunchDAG(ControlMessage.LaunchDAGMessage.newBuilder() LOG.info("DAG execution done"); shutdown();
/** * Get the current state of the running plan. * * @return the current state of the running plan. */ public final synchronized Enum getPlanState() { if (driverEndpoint.get() != null) { return stateTranslator.translateState(driverEndpoint.get().getState()); } else { return stateTranslator.translateState(PlanState.State.READY); } }
/** * Main JobLauncher method. * * @param args arguments. * @throws Exception exception on the way. */ public static void main(final String[] args) throws Exception { try { setup(args); // Launch client main. The shutdown() method is called inside the launchDAG() method. runUserProgramMain(builtJobConf); } catch (final InjectionException e) { throw new RuntimeException(e); } }
/** * @return the configuration for RPC server listening information */ public Configuration getListeningConfiguration() { return Tang.Factory.getTang().newConfigurationBuilder() .bindNamedParameter(JobConf.ClientSideRPCServerHost.class, getListeningHost()) .bindNamedParameter(JobConf.ClientSideRPCServerPort.class, String.valueOf(getListeningPort())) .build(); }
/** * Runs the RPC server. * Specifically, creates a {@link NettyMessagingTransport} and binds it to a listening port. */ public void run() { // Calling 'run' multiple times is considered invalid, since it will override state variables like // 'transport', and 'host'. ensureServerState(false); try { final Injector injector = Tang.Factory.getTang().newInjector(); final LocalAddressProvider localAddressProvider = injector.getInstance(LocalAddressProvider.class); host = localAddressProvider.getLocalAddress(); injector.bindVolatileParameter(RemoteConfiguration.HostAddress.class, host); injector.bindVolatileParameter(RemoteConfiguration.Port.class, 0); injector.bindVolatileParameter(RemoteConfiguration.RemoteServerStage.class, new SyncStage<>(new ServerEventHandler())); transport = injector.getInstance(NettyMessagingTransport.class); LOG.info("DriverRPCServer running at {}", transport.getListeningPort()); isRunning = true; } catch (final InjectionException e) { throw new RuntimeException(e); } }
private static DAG<IRVertex, IREdge> compileDAG(final String[] args) throws Exception { final String userMainClassName; final String[] userMainMethodArgs; try { final Configuration configuration = JobLauncher.getJobConf(args); final Injector injector = Tang.Factory.getTang().newInjector(configuration); userMainClassName = injector.getNamedInstance(JobConf.UserMainClass.class); userMainMethodArgs = injector.getNamedInstance(JobConf.UserMainArguments.class).split(" "); } catch (final Exception e) { throw new RuntimeException("An exception occurred while processing configuration for invoking user main. " + "Note: Using compileDAG for multiple times will fail, as compileDAG method enables static method mocking " + "on JobLauncher and because of this Tang may misbehave afterwards.", e); } final Class userMainClass = Class.forName(userMainClassName); final Method userMainMethod = userMainClass.getMethod("main", String[].class); final ArgumentCaptor<DAG> captor = ArgumentCaptor.forClass(DAG.class); final ArgumentCaptor<String> stringArg = ArgumentCaptor.forClass(String.class); PowerMockito.mockStatic(JobLauncher.class); PowerMockito.doNothing().when(JobLauncher.class, "launchDAG", captor.capture(), stringArg.capture()); userMainMethod.invoke(null, (Object) userMainMethodArgs); return captor.getValue(); }
/** * @param dag the application DAG. * @param jobId job ID. */ public static void launchDAG(final DAG dag, final String jobId) { launchDAG(dag, Collections.emptyMap(), jobId); }
/** * Wait for this job to be finished (complete or failed) and return the final state. * It wait for at most the given time. * * @param timeout of waiting. * @param unit of the timeout. * @return the final state of this job. */ public final Enum waitUntilJobFinish(final long timeout, final TimeUnit unit) { if (driverEndpoint.get() != null) { return stateTranslator.translateState(driverEndpoint.get().waitUntilFinish(timeout, unit)); } else { // The driver endpoint is not connected yet. final long currentNano = System.nanoTime(); final boolean driverIsConnected; if (DEFAULT_DRIVER_WAIT_IN_MILLIS < unit.toMillis(timeout)) { driverIsConnected = waitUntilConnected(DEFAULT_DRIVER_WAIT_IN_MILLIS, TimeUnit.MILLISECONDS); } else { driverIsConnected = waitUntilConnected(timeout, unit); } if (driverIsConnected) { final long consumedTime = System.nanoTime() - currentNano; return stateTranslator.translateState(driverEndpoint.get(). waitUntilFinish(timeout - unit.convert(consumedTime, TimeUnit.NANOSECONDS), unit)); } else { return PlanState.State.READY; } } }
/** * @return the listening port */ public int getListeningPort() { // We cannot determine listening port if the server is not listening. ensureServerState(true); return transport.getListeningPort(); }
/** * Method to run the Pipeline. * @param pipeline the Pipeline to run. * @return The result of the pipeline. */ public NemoPipelineResult run(final Pipeline pipeline) { final PipelineVisitor pipelineVisitor = new PipelineVisitor(pipeline, nemoPipelineOptions); pipeline.traverseTopologically(pipelineVisitor); final NemoPipelineResult nemoPipelineResult = new NemoPipelineResult(); JobLauncher.launchDAG(pipelineVisitor.getConvertedPipeline(), nemoPipelineOptions.getJobName()); return nemoPipelineResult; } }
/** * Registers handler for the given type of message. * @param type the type of message * @param handler handler implementation * @return {@code this} */ public DriverRPCServer registerHandler(final ControlMessage.DriverToClientMessageType type, final EventHandler<ControlMessage.DriverToClientMessage> handler) { // Registering a handler after running the server is considered not a good practice. ensureServerState(false); if (handlers.putIfAbsent(type, handler) != null) { throw new RuntimeException(String.format("A handler for %s already registered", type)); } return this; }
/** * Shut down the server. */ public void shutdown() { // Shutting down a 'null' transport is invalid. Also, shutting down a server for multiple times is invalid. ensureServerState(true); try { transport.close(); } catch (final Exception e) { throw new RuntimeException(e); } finally { isShutdown = true; } }
/** * Sends a message to driver. * @param message message to send */ public void send(final ControlMessage.ClientToDriverMessage message) { // This needs active 'link' between the driver and client. // For the link to be alive, the driver should connect to DriverRPCServer. // Thus, the server must be running to send a message to the driver. ensureServerState(true); if (link == null) { throw new RuntimeException("The RPC server has not discovered NemoDriver yet"); } link.write(message.toByteArray()); }