org.apache.flink.client.program.ClusterClient Java Examples

private void deployApplication(Configuration configuration) throws Exception {
	try (final YarnClusterDescriptor yarnClusterDescriptor = createYarnClusterDescriptor(configuration)) {

		final int masterMemory = yarnClusterDescriptor.getFlinkConfiguration().get(JobManagerOptions.TOTAL_PROCESS_MEMORY).getMebiBytes();
		final ClusterSpecification clusterSpecification = new ClusterSpecification.ClusterSpecificationBuilder()
			.setMasterMemoryMB(masterMemory)
			.setTaskManagerMemoryMB(1024)
			.setSlotsPerTaskManager(1)
			.createClusterSpecification();

		try (ClusterClient<ApplicationId> clusterClient = yarnClusterDescriptor
				.deployApplicationCluster(
						clusterSpecification,
						ApplicationConfiguration.fromConfiguration(configuration))
				.getClusterClient()) {

			ApplicationId applicationId = clusterClient.getClusterId();

			waitApplicationFinishedElseKillIt(
				applicationId, yarnAppTerminateTimeout, yarnClusterDescriptor, sleepIntervalInMS);
		}
	}
}
 

@Override
protected ClusterClient<ApplicationId> createYarnClusterClient(
    AbstractYarnClusterDescriptor clusterDescriptor,
    int numberTaskManagers,
    int slotPerTaskManager,
    ApplicationReport applicationReport,
    Configuration configuration,
    boolean isNewlyCreatedCluster) throws Exception {
  return new YarnClusterClient(
      clusterDescriptor,
      numberTaskManagers,
      slotPerTaskManager,
      applicationReport,
      configuration,
      isNewlyCreatedCluster);
}
 

@Test
public void testDisposeSavepointSuccess() throws Exception {
	replaceStdOutAndStdErr();

	String savepointPath = "expectedSavepointPath";

	ClusterClient clusterClient = new DisposeSavepointClusterClient(
		(String path) -> CompletableFuture.completedFuture(Acknowledge.get()), getConfiguration());

	try {

		CliFrontend frontend = new MockedCliFrontend(clusterClient);

		String[] parameters = { "-d", savepointPath };
		frontend.savepoint(parameters);

		String outMsg = buffer.toString();
		assertTrue(outMsg.contains(savepointPath));
		assertTrue(outMsg.contains("disposed"));
	}
	finally {
		clusterClient.shutdown();
		restoreStdOutAndStdErr();
	}
}
 

/**
 * Sends a SavepointDisposalRequest to the job manager.
 */
private void disposeSavepoint(ClusterClient<?> clusterClient, String savepointPath) throws FlinkException {
	Preconditions.checkNotNull(savepointPath, "Missing required argument: savepoint path. " +
		"Usage: bin/flink savepoint -d <savepoint-path>");

	logAndSysout("Disposing savepoint '" + savepointPath + "'.");

	final CompletableFuture<Acknowledge> disposeFuture = clusterClient.disposeSavepoint(savepointPath);

	logAndSysout("Waiting for response...");

	try {
		disposeFuture.get(clientTimeout.toMillis(), TimeUnit.MILLISECONDS);
	} catch (Exception e) {
		throw new FlinkException("Disposing the savepoint '" + savepointPath + "' failed.", e);
	}

	logAndSysout("Savepoint '" + savepointPath + "' disposed.");
}
 

@Test
public void testUnknownJobId() throws Exception {
	// test unknown job Id
	JobID jid = new JobID();

	String[] parameters = { "-p", "test-target-dir", jid.toString() };
	String expectedMessage = "Test exception";
	FlinkException testException = new FlinkException(expectedMessage);
	final ClusterClient<String> clusterClient = createClusterClient(testException);
	MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient);

	try {
		testFrontend.stop(parameters);
		fail("Should have failed.");
	} catch (FlinkException e) {
		assertTrue(ExceptionUtils.findThrowableWithMessage(e, expectedMessage).isPresent());
	}
}
 

@Test
public void testTriggerSavepointSuccess() throws Exception {
	replaceStdOutAndStdErr();

	JobID jobId = new JobID();

	String savepointPath = "expectedSavepointPath";

	final ClusterClient<String> clusterClient = createClusterClient(savepointPath);

	try {
		MockedCliFrontend frontend = new MockedCliFrontend(clusterClient);

		String[] parameters = { jobId.toString() };
		frontend.savepoint(parameters);

		verify(clusterClient, times(1))
			.triggerSavepoint(eq(jobId), isNull(String.class));

		assertTrue(buffer.toString().contains(savepointPath));
	}
	finally {
		clusterClient.close();
		restoreStdOutAndStdErr();
	}
}
 

@Test
public void testUnknownJobId() throws Exception {
	// test unknown job Id
	JobID jid = new JobID();

	String[] parameters = { jid.toString() };
	String expectedMessage = "Test exception";
	FlinkException testException = new FlinkException(expectedMessage);
	final ClusterClient<String> clusterClient = createClusterClient(testException);
	MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient);

	try {
		testFrontend.stop(parameters);
		fail("Should have failed.");
	} catch (FlinkException e) {
		assertTrue(ExceptionUtils.findThrowableWithMessage(e, expectedMessage).isPresent());
	}
}
 

private String submitJobAndTakeSavepoint(MiniClusterResourceFactory clusterFactory, int parallelism) throws Exception {
	final JobGraph jobGraph = createJobGraph(parallelism, 0, 1000);
	final JobID jobId = jobGraph.getJobID();
	StatefulCounter.resetForTest(parallelism);

	MiniClusterWithClientResource cluster = clusterFactory.get();
	cluster.before();
	ClusterClient<?> client = cluster.getClusterClient();

	try {
		client.setDetached(true);
		client.submitJob(jobGraph, SavepointITCase.class.getClassLoader());

		StatefulCounter.getProgressLatch().await();

		return client.cancelWithSavepoint(jobId, null);
	} finally {
		cluster.after();
		StatefulCounter.resetForTest(parallelism);
	}
}
 

@Override
public CompletableFuture<JobClient> execute(@Nonnull final Pipeline pipeline, @Nonnull final Configuration configuration) throws Exception {
	final JobGraph jobGraph = PipelineExecutorUtils.getJobGraph(pipeline, configuration);

	try (final ClusterDescriptor<ClusterID> clusterDescriptor = clusterClientFactory.createClusterDescriptor(configuration)) {
		final ClusterID clusterID = clusterClientFactory.getClusterId(configuration);
		checkState(clusterID != null);

		final ClusterClientProvider<ClusterID> clusterClientProvider = clusterDescriptor.retrieve(clusterID);
		ClusterClient<ClusterID> clusterClient = clusterClientProvider.getClusterClient();
		return clusterClient
				.submitJob(jobGraph)
				.thenApplyAsync(jobID -> (JobClient) new ClusterClientJobClientAdapter<>(
						clusterClientProvider,
						jobID))
				.whenComplete((ignored1, ignored2) -> clusterClient.close());
	}
}
 

public <T> String takeSavepoint(Collection<T> data, Function<SourceFunction<T>, StreamExecutionEnvironment> jobGraphFactory) throws Exception {

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableClosureCleaner();

		WaitingSource<T> waitingSource = createSource(data);

		JobGraph jobGraph = jobGraphFactory.apply(waitingSource).getStreamGraph().getJobGraph();
		JobID jobId = jobGraph.getJobID();

		ClusterClient<?> client = miniClusterResource.getClusterClient();

		try {
			JobSubmissionResult result = ClientUtils.submitJob(client, jobGraph);

			return CompletableFuture
				.runAsync(waitingSource::awaitSource)
				.thenCompose(ignore -> triggerSavepoint(client, result.getJobID()))
				.get(5, TimeUnit.MINUTES);
		} catch (Exception e) {
			throw new RuntimeException("Failed to take savepoint", e);
		} finally {
			client.cancel(jobId);
		}
	}
 

@Override
public ClusterClient<ApplicationId> deployJobCluster(
	ClusterSpecification clusterSpecification,
	JobGraph jobGraph,
	boolean detached) throws ClusterDeploymentException {

	// this is required because the slots are allocated lazily
	jobGraph.setAllowQueuedScheduling(true);

	try {
		return deployInternal(
			clusterSpecification,
			"Flink per-job cluster",
			getYarnJobClusterEntrypoint(),
			jobGraph,
			detached);
	} catch (Exception e) {
		throw new ClusterDeploymentException("Could not deploy Yarn job cluster.", e);
	}
}
 

protected void executeProgram(PackagedProgram program, ClusterClient<?> client, int parallelism) throws ProgramMissingJobException, ProgramInvocationException {
	logAndSysout("Starting execution of program");

	final JobSubmissionResult result = client.run(program, parallelism);

	if (null == result) {
		throw new ProgramMissingJobException("No JobSubmissionResult returned, please make sure you called " +
			"ExecutionEnvironment.execute()");
	}

	if (result.isJobExecutionResult()) {
		logAndSysout("Program execution finished");
		JobExecutionResult execResult = result.getJobExecutionResult();
		System.out.println("Job with JobID " + execResult.getJobID() + " has finished.");
		System.out.println("Job Runtime: " + execResult.getNetRuntime() + " ms");
		Map<String, Object> accumulatorsResult = execResult.getAllAccumulatorResults();
		if (accumulatorsResult.size() > 0) {
			System.out.println("Accumulator Results: ");
			System.out.println(AccumulatorHelper.getResultsFormatted(accumulatorsResult));
		}
	} else {
		logAndSysout("Job has been submitted with JobID " + result.getJobID());
	}
}
 

private static <T> CompletableFuture<T> bridgeClientRequest(
		ClusterClientProvider<?> clusterClientProvider,
		Function<ClusterClient<?>, CompletableFuture<T>> resultRetriever) {

	ClusterClient<?> clusterClient = clusterClientProvider.getClusterClient();

	CompletableFuture<T> resultFuture;
	try {
		resultFuture = resultRetriever.apply(clusterClient);
	} catch (Throwable throwable) {
		IOUtils.closeQuietly(clusterClient::close);
		return FutureUtils.completedExceptionally(throwable);
	}

	return resultFuture.whenCompleteAsync(
			(jobResult, throwable) -> IOUtils.closeQuietly(clusterClient::close));
}
 

protected void runAndCancelJob(Plan plan, final int msecsTillCanceling, int maxTimeTillCanceled) throws Exception {
	// submit job
	final JobGraph jobGraph = getJobGraph(plan);

	ClusterClient<?> client = CLUSTER.getClusterClient();
	client.setDetached(true);

	JobSubmissionResult jobSubmissionResult = client.submitJob(jobGraph, CancelingTestBase.class.getClassLoader());

	Deadline submissionDeadLine = new FiniteDuration(2, TimeUnit.MINUTES).fromNow();

	JobStatus jobStatus = client.getJobStatus(jobSubmissionResult.getJobID()).get(GET_FUTURE_TIMEOUT, TimeUnit.MILLISECONDS);
	while (jobStatus != JobStatus.RUNNING && submissionDeadLine.hasTimeLeft()) {
		Thread.sleep(50);
		jobStatus = client.getJobStatus(jobSubmissionResult.getJobID()).get(GET_FUTURE_TIMEOUT, TimeUnit.MILLISECONDS);
	}
	if (jobStatus != JobStatus.RUNNING) {
		Assert.fail("Job not in state RUNNING.");
	}

	Thread.sleep(msecsTillCanceling);

	client.cancel(jobSubmissionResult.getJobID());

	Deadline cancelDeadline = new FiniteDuration(maxTimeTillCanceled, TimeUnit.MILLISECONDS).fromNow();

	JobStatus jobStatusAfterCancel = client.getJobStatus(jobSubmissionResult.getJobID()).get(GET_FUTURE_TIMEOUT, TimeUnit.MILLISECONDS);
	while (jobStatusAfterCancel != JobStatus.CANCELED && cancelDeadline.hasTimeLeft()) {
		Thread.sleep(50);
		jobStatusAfterCancel = client.getJobStatus(jobSubmissionResult.getJobID()).get(GET_FUTURE_TIMEOUT, TimeUnit.MILLISECONDS);
	}
	if (jobStatusAfterCancel != JobStatus.CANCELED) {
		Assert.fail("Failed to cancel job with ID " + jobSubmissionResult.getJobID() + '.');
	}
}
 

private static ClusterClient<String> createFailingClusterClient(Exception expectedException) throws Exception {
	final ClusterClient<String> clusterClient = mock(ClusterClient.class);

	when(clusterClient.triggerSavepoint(any(JobID.class), nullable(String.class)))
		.thenReturn(FutureUtils.completedExceptionally(expectedException));

	return clusterClient;
}
 

/**
 * Tests that a CLI call with a custom savepoint directory target is
 * forwarded correctly to the cluster client.
 */
@Test
public void testTriggerSavepointCustomTarget() throws Exception {
	replaceStdOutAndStdErr();

	JobID jobId = new JobID();

	String savepointDirectory = "customTargetDirectory";

	final ClusterClient<String> clusterClient = createClusterClient(savepointDirectory);

	try {
		MockedCliFrontend frontend = new MockedCliFrontend(clusterClient);

		String[] parameters = { jobId.toString(), savepointDirectory };
		frontend.savepoint(parameters);

		verify(clusterClient, times(1))
			.triggerSavepoint(eq(jobId), eq(savepointDirectory));

		assertTrue(buffer.toString().contains(savepointDirectory));
	}
	finally {
		clusterClient.shutdown();

		restoreStdOutAndStdErr();
	}
}
 

private RestClusterClient<ApplicationId> deploySessionCluster(YarnClusterDescriptor yarnClusterDescriptor) throws ClusterDeploymentException {
	final int masterMemory = yarnClusterDescriptor.getFlinkConfiguration().get(JobManagerOptions.TOTAL_PROCESS_MEMORY).getMebiBytes();
	final int taskManagerMemory = 1024;
	final ClusterClient<ApplicationId> yarnClusterClient = yarnClusterDescriptor
			.deploySessionCluster(new ClusterSpecification.ClusterSpecificationBuilder()
					.setMasterMemoryMB(masterMemory)
					.setTaskManagerMemoryMB(taskManagerMemory)
					.setSlotsPerTaskManager(1)
					.createClusterSpecification())
			.getClusterClient();

	assertThat(yarnClusterClient, is(instanceOf(RestClusterClient.class)));
	return (RestClusterClient<ApplicationId>) yarnClusterClient;
}
 

private String takeSavepoint(JobGraph jobGraph) throws Exception {
	SavepointSource.initializeForTest();

	ClusterClient<?> client = miniClusterResource.getClusterClient();
	JobID jobId = jobGraph.getJobID();

	Deadline deadline = Deadline.fromNow(Duration.ofMinutes(5));

	String dirPath = getTempDirPath(new AbstractID().toHexString());

	try {
		JobSubmissionResult result = ClientUtils.submitJob(client, jobGraph);

		boolean finished = false;
		while (deadline.hasTimeLeft()) {
			if (SavepointSource.isFinished()) {
				finished = true;

				break;
			}

			try {
				Thread.sleep(2L);
			} catch (InterruptedException ignored) {
				Thread.currentThread().interrupt();
			}
		}

		if (!finished) {
			Assert.fail("Failed to initialize state within deadline");
		}

		CompletableFuture<String> path = client.triggerSavepoint(result.getJobID(), dirPath);
		return path.get(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
	} finally {
		client.cancel(jobId).get();
	}
}
 

private <T> LocalExecutor createModifiedExecutor(
		String yamlFile, ClusterClient<T> clusterClient, Map<String, String> replaceVars) throws Exception {
	return new LocalExecutor(
		EnvironmentFileUtil.parseModified(yamlFile, replaceVars),
		Collections.emptyList(),
		clusterClient.getFlinkConfiguration(),
		new DummyCustomCommandLine<T>(clusterClient));
}
 

private static List<JobID> getRunningJobs(ClusterClient<?> client) throws Exception {
	Collection<JobStatusMessage> statusMessages = client.listJobs().get();
	return statusMessages.stream()
		.filter(status -> !status.getJobState().isGloballyTerminalState())
		.map(JobStatusMessage::getJobId)
		.collect(Collectors.toList());
}
 

@Test
public void testCancel() throws Exception {
	// test cancel properly
	JobID jid = new JobID();

	String[] parameters = { jid.toString() };
	final ClusterClient<String> clusterClient = createClusterClient();
	MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient);

	testFrontend.cancel(parameters);

	Mockito.verify(clusterClient, times(1)).cancel(any(JobID.class));
}
 

private <T> LocalExecutor createModifiedExecutor(
		String yamlFile, ClusterClient<T> clusterClient, Map<String, String> replaceVars) throws Exception {
	replaceVars.putIfAbsent("$VAR_RESTART_STRATEGY_TYPE", "failure-rate");
	return new LocalExecutor(
			EnvironmentFileUtil.parseModified(yamlFile, replaceVars),
			Collections.emptyList(),
			clusterClient.getFlinkConfiguration(),
			new DefaultCLI(clusterClient.getFlinkConfiguration()),
			new DefaultClusterClientServiceLoader());
}
 

private static ClusterClient<String> createFailingClusterClient(Exception expectedException) throws Exception {
	final ClusterClient<String> clusterClient = mock(ClusterClient.class);

	when(clusterClient.triggerSavepoint(any(JobID.class), nullable(String.class)))
		.thenReturn(FutureUtils.completedExceptionally(expectedException));

	return clusterClient;
}
 

private void validateModification(String savepointPath) throws ProgramInvocationException {
	StreamExecutionEnvironment sEnv = StreamExecutionEnvironment.getExecutionEnvironment();
	sEnv.setStateBackend(backend);

	CollectSink.accountList.clear();

	DataStream<Account> stream = sEnv.fromCollection(accounts)
		.keyBy(acc -> acc.id)
		.flatMap(new UpdateAndGetAccount())
		.uid(ACCOUNT_UID);

	stream.addSink(new CollectSink());

	stream
		.map(acc -> acc.id)
		.map(new StatefulOperator())
		.uid(MODIFY_UID)
		.addSink(new DiscardingSink<>());

	JobGraph jobGraph = sEnv.getStreamGraph().getJobGraph();
	jobGraph.setSavepointRestoreSettings(SavepointRestoreSettings.forPath(savepointPath, false));

	ClusterClient<?> client = miniClusterResource.getClusterClient();
	client.submitJob(jobGraph, SavepointWriterITCase.class.getClassLoader());

	Assert.assertEquals("Unexpected output", 3, CollectSink.accountList.size());
}
 

@Override
protected ClusterClient<ApplicationId> createYarnClusterClient(
		AbstractYarnClusterDescriptor descriptor,
		int numberTaskManagers,
		int slotsPerTaskManager,
		ApplicationReport report,
		Configuration flinkConfiguration,
		boolean perJobCluster) {
	return clusterClient;
}
 

private void stopAfterJob(ClusterClient client, JobID jobID) {
  Preconditions.checkNotNull(jobID, "The job id must not be null");
  try {
    Future<Object> replyFuture =
        client.getJobManagerGateway().ask(
            new ShutdownClusterAfterJob(jobID),
            AKKA_TIMEOUT);
    Await.ready(replyFuture, AKKA_TIMEOUT);
  } catch (Exception e) {
    throw new RuntimeException("Unable to tell application master to stop"
        + " once the specified job has been finished", e);
  }
}
 

@Test
public void testTriggerSavepointWithCheckpointingDisabled() throws Exception {
	// Config
	final int numTaskManagers = 1;
	final int numSlotsPerTaskManager = 1;

	final Configuration config = new Configuration();

	final MiniClusterWithClientResource cluster = new MiniClusterWithClientResource(
		new MiniClusterResourceConfiguration.Builder()
			.setConfiguration(config)
			.setNumberTaskManagers(numTaskManagers)
			.setNumberSlotsPerTaskManager(numSlotsPerTaskManager)
			.build());
	cluster.before();
	final ClusterClient<?> client = cluster.getClusterClient();

	final JobVertex vertex = new JobVertex("Blocking vertex");
	vertex.setInvokableClass(BlockingNoOpInvokable.class);
	vertex.setParallelism(1);

	final JobGraph graph = new JobGraph(vertex);

	try {
		client.setDetached(true);
		client.submitJob(graph, SavepointITCase.class.getClassLoader());

		client.triggerSavepoint(graph.getJobID(), null).get();

		fail();
	} catch (ExecutionException e) {
		assertTrue(ExceptionUtils.findThrowable(e, IllegalStateException.class).isPresent());
		assertTrue(ExceptionUtils.findThrowableWithMessage(e, graph.getJobID().toString()).isPresent());
		assertTrue(ExceptionUtils.findThrowableWithMessage(e, "is not a streaming job").isPresent());
	} finally {
		cluster.after();
	}
}
 

/**
 * Creates a YarnClusterClient; may be overridden in tests.
 */
protected abstract ClusterClient<ApplicationId> createYarnClusterClient(
		AbstractYarnClusterDescriptor descriptor,
		int numberTaskManagers,
		int slotsPerTaskManager,
		ApplicationReport report,
		org.apache.flink.configuration.Configuration flinkConfiguration,
		boolean perJobCluster) throws Exception;
 

@Override
public ClusterClient<ApplicationId> deploySessionCluster(ClusterSpecification clusterSpecification) throws ClusterDeploymentException {
	try {
		return deployInternal(
			clusterSpecification,
			"Flink session cluster",
			getYarnSessionClusterEntrypoint(),
			null,
			false);
	} catch (Exception e) {
		throw new ClusterDeploymentException("Couldn't deploy Yarn session cluster", e);
	}
}
 

/**
 * Tests that the configuration is properly passed via the DefaultCLI to the
 * created ClusterDescriptor.
 */
@Test
public void testConfigurationPassing() throws Exception {
	final Configuration configuration = getConfiguration();

	final String localhost = "localhost";
	final int port = 1234;

	configuration.setString(JobManagerOptions.ADDRESS, localhost);
	configuration.setInteger(JobManagerOptions.PORT, port);

	@SuppressWarnings("unchecked")
	final AbstractCustomCommandLine<StandaloneClusterId> defaultCLI =
		(AbstractCustomCommandLine<StandaloneClusterId>) getCli(configuration);

	final String[] args = {};

	CommandLine commandLine = defaultCLI.parseCommandLineOptions(args, false);

	final ClusterDescriptor<StandaloneClusterId> clusterDescriptor =
		defaultCLI.createClusterDescriptor(commandLine);

	final ClusterClient<?> clusterClient = clusterDescriptor.retrieve(defaultCLI.getClusterId(commandLine));

	final LeaderConnectionInfo clusterConnectionInfo = clusterClient.getClusterConnectionInfo();

	assertThat(clusterConnectionInfo.getHostname(), Matchers.equalTo(localhost));
	assertThat(clusterConnectionInfo.getPort(), Matchers.equalTo(port));
}