org.apache.flink.client.program.rest.RestClusterClient Java Examples

private static Optional<Metric> getJobMetric(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final String metricName) throws Exception {

	final JobMetricsMessageParameters messageParameters = new JobMetricsMessageParameters();
	messageParameters.jobPathParameter.resolve(jobId);
	messageParameters.metricsFilterParameter.resolveFromString(metricName);

	final Collection<Metric> metrics = restClusterClient.sendRequest(
		JobMetricsHeaders.getInstance(),
		messageParameters,
		EmptyRequestBody.getInstance()).get().getMetrics();

	final Metric metric = Iterables.getOnlyElement(metrics, null);
	checkState(metric == null || metric.getId().equals(metricName));
	return Optional.ofNullable(metric);
}
 

@Test
public void testRemoteExecutionWithSavepoint() throws Exception {
	SavepointRestoreSettings restoreSettings = SavepointRestoreSettings.forPath("fakePath");
	RemoteStreamEnvironment env = new RemoteStreamEnvironment("fakeHost", 1,
		null, new String[]{}, null, restoreSettings);
	env.fromElements(1).map(x -> x * 2);

	RestClusterClient mockedClient = Mockito.mock(RestClusterClient.class);
	JobExecutionResult expectedResult = new JobExecutionResult(null, 0, null);

	PowerMockito.whenNew(RestClusterClient.class).withAnyArguments().thenReturn(mockedClient);
	when(mockedClient.run(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.eq(restoreSettings)))
		.thenReturn(expectedResult);

	JobExecutionResult actualResult = env.execute("fakeJobName");
	Assert.assertEquals(expectedResult, actualResult);
}
 

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

	try {
		CliFrontend frontend = new MockedCliFrontend(new RestClusterClient<>(getConfiguration(), StandaloneClusterId.getInstance()));

		String[] parameters = { "invalid job id" };
		try {
			frontend.savepoint(parameters);
			fail("Should have failed.");
		} catch (CliArgsException e) {
			assertThat(e.getMessage(), Matchers.containsString("Cannot parse JobID"));
		}
	}
	finally {
		restoreStdOutAndStdErr();
	}
}
 

@Test
public void testJobRecoversAfterKillingTaskManager() throws Exception {
	runTest(() -> {
		final YarnClusterDescriptor yarnClusterDescriptor = setupYarnClusterDescriptor();
		yarnClusterDescriptor.addShipFiles(Arrays.asList(flinkShadedHadoopDir.listFiles()));

		final RestClusterClient<ApplicationId> restClusterClient = deploySessionCluster(yarnClusterDescriptor);
		try {
			final JobID jobId = submitJob(restClusterClient);
			waitUntilJobIsRunning(restClusterClient, jobId);

			stopTaskManagerContainer();
			waitUntilJobIsRestarted(restClusterClient, jobId, 1);

			waitForJobTermination(restClusterClient, jobId);

			killApplicationAndWait(restClusterClient.getClusterId());
		} finally {
			restClusterClient.shutdown();
		}
	});
}
 

private static Optional<Metric> getJobMetric(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final String metricName) throws Exception {

	final JobMetricsMessageParameters messageParameters = new JobMetricsMessageParameters();
	messageParameters.jobPathParameter.resolve(jobId);
	messageParameters.metricsFilterParameter.resolveFromString(metricName);

	final Collection<Metric> metrics = restClusterClient.sendRequest(
		JobMetricsHeaders.getInstance(),
		messageParameters,
		EmptyRequestBody.getInstance()).get().getMetrics();

	final Metric metric = Iterables.getOnlyElement(metrics, null);
	checkState(metric == null || metric.getId().equals(metricName));
	return Optional.ofNullable(metric);
}
 

private static Optional<Metric> getJobMetric(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final String metricName) throws Exception {

	final JobMetricsMessageParameters messageParameters = new JobMetricsMessageParameters();
	messageParameters.jobPathParameter.resolve(jobId);
	messageParameters.metricsFilterParameter.resolveFromString(metricName);

	final Collection<Metric> metrics = restClusterClient.sendRequest(
		JobMetricsHeaders.getInstance(),
		messageParameters,
		EmptyRequestBody.getInstance()).get().getMetrics();

	final Metric metric = Iterables.getOnlyElement(metrics, null);
	checkState(metric == null || metric.getId().equals(metricName));
	return Optional.ofNullable(metric);
}
 

public ClusterClient<ApplicationId> deploy(JobGraph jobGraph, boolean detached)
        throws Exception
{
    // this is required because the slots are allocated lazily
    jobGraph.setAllowQueuedScheduling(true);
    //
    ApplicationReport report = startAppMaster(jobGraph);

    Configuration flinkConfiguration = getFlinkConfiguration();

    ClusterEntrypoint.ExecutionMode executionMode = detached ? ClusterEntrypoint.ExecutionMode.DETACHED : ClusterEntrypoint.ExecutionMode.NORMAL;
    flinkConfiguration.setString(ClusterEntrypoint.EXECUTION_MODE, executionMode.toString());
    String host = report.getHost();
    int port = report.getRpcPort();
    flinkConfiguration.setString(JobManagerOptions.ADDRESS, host);
    flinkConfiguration.setInteger(JobManagerOptions.PORT, port);
    flinkConfiguration.setString(RestOptions.ADDRESS, host);
    flinkConfiguration.setInteger(RestOptions.PORT, port);
    return new RestClusterClient<>(flinkConfiguration, report.getApplicationId());
}
 

/**
 * All the Flink Standalone, Yarn, Mesos, Kubernetes sessions are using {@link RestClusterClient#submitJob(JobGraph)}
 * to submit a job to an existing session. This test will cover this cases.
 */
@Test(timeout = 30000)
public void testSubmittingJobViaRestClusterClient() throws Exception {
	RestClusterClient<String> restClusterClient = new RestClusterClient<>(
		MINI_CLUSTER_RESOURCE.getClientConfiguration(),
		"testSubmittingJobViaRestClusterClient");

	final JobGraph jobGraph = createJobWithRegisteredCachedFiles()
			.getStreamGraph()
			.getJobGraph();

	final JobResult jobResult = restClusterClient
		.submitJob(jobGraph)
		.thenCompose(restClusterClient::requestJobResult)
		.get();

	final String messageInCaseOfFailure = jobResult.getSerializedThrowable().isPresent() ?
			jobResult.getSerializedThrowable().get().getFullStringifiedStackTrace()
			: "Job failed.";
	assertTrue(messageInCaseOfFailure, jobResult.isSuccess());
}
 

@Test
public void testJobRecoversAfterKillingTaskManager() throws Exception {
	final YarnClusterDescriptor yarnClusterDescriptor = setupYarnClusterDescriptor();
	yarnClusterDescriptor.addShipFiles(Arrays.asList(flinkShadedHadoopDir.listFiles()));

	final RestClusterClient<ApplicationId> restClusterClient = deploySessionCluster(yarnClusterDescriptor);
	try {
		final JobID jobId = submitJob(restClusterClient);
		waitUntilJobIsRunning(restClusterClient, jobId);

		stopTaskManagerContainer();
		waitUntilJobIsRestarted(restClusterClient, jobId, 1);

		waitForJobTermination(restClusterClient, jobId);

		killApplicationAndWait(restClusterClient.getClusterId());
	} finally {
		restClusterClient.shutdown();
	}
}
 

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

	try {
		CliFrontend frontend = new MockedCliFrontend(new RestClusterClient<>(getConfiguration(), StandaloneClusterId.getInstance()));

		String[] parameters = { "invalid job id" };
		try {
			frontend.savepoint(parameters);
			fail("Should have failed.");
		} catch (CliArgsException e) {
			assertThat(e.getMessage(), Matchers.containsString("Cannot parse JobID"));
		}
	}
	finally {
		restoreStdOutAndStdErr();
	}
}
 

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

	try {
		CliFrontend frontend = new MockedCliFrontend(new RestClusterClient<>(getConfiguration(), StandaloneClusterId.getInstance()));

		String[] parameters = { "invalid job id" };
		try {
			frontend.savepoint(parameters);
			fail("Should have failed.");
		} catch (CliArgsException e) {
			assertThat(e.getMessage(), Matchers.containsString("Cannot parse JobID"));
		}
	}
	finally {
		restoreStdOutAndStdErr();
	}
}
 

@Test
public void testJobRecoversAfterKillingTaskManager() throws Exception {
	runTest(() -> {
		final YarnClusterDescriptor yarnClusterDescriptor = setupYarnClusterDescriptor();
		final RestClusterClient<ApplicationId> restClusterClient = deploySessionCluster(yarnClusterDescriptor);
		try {
			final JobID jobId = submitJob(restClusterClient);
			waitUntilJobIsRunning(restClusterClient, jobId);

			stopTaskManagerContainer();
			waitUntilJobIsRestarted(restClusterClient, jobId, 1);

			waitForJobTermination(restClusterClient, jobId);

			killApplicationAndWait(restClusterClient.getClusterId());
		} finally {
			restClusterClient.close();
		}
	});
}
 

@Test
public void testRemoteExecutionWithSavepoint() throws Exception {
	SavepointRestoreSettings restoreSettings = SavepointRestoreSettings.forPath("fakePath");
	RemoteStreamEnvironment env = new RemoteStreamEnvironment("fakeHost", 1,
		null, new String[]{}, null, restoreSettings);
	env.fromElements(1).map(x -> x * 2);

	RestClusterClient mockedClient = Mockito.mock(RestClusterClient.class);
	JobExecutionResult expectedResult = new JobExecutionResult(null, 0, null);

	PowerMockito.whenNew(RestClusterClient.class).withAnyArguments().thenReturn(mockedClient);
	when(mockedClient.run(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.eq(restoreSettings)))
		.thenReturn(expectedResult);

	JobExecutionResult actualResult = env.execute("fakeJobName");
	Assert.assertEquals(expectedResult, actualResult);
}
 

private static void waitUntilJobIsRunning(RestClusterClient<ApplicationId> restClusterClient, JobID jobId) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> {
			final JobDetailsInfo jobDetails = restClusterClient.getJobDetails(jobId).get();
			return jobDetails.getJobVertexInfos()
				.stream()
				.map(toExecutionState())
				.allMatch(isRunning());
		},
		Deadline.fromNow(TIMEOUT));
}
 

/**
 * Verifies that the port passed to the RemoteStreamEnvironment is used for connecting to the cluster.
 */
@Test
public void testPortForwarding() throws Exception {

	String host = "fakeHost";
	int port = 99;
	JobExecutionResult expectedResult = new JobExecutionResult(null, 0, null);

	RestClusterClient mockedClient = Mockito.mock(RestClusterClient.class);
	when(mockedClient.run(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any(), Mockito.any()))
		.thenReturn(expectedResult);

	PowerMockito.whenNew(RestClusterClient.class).withAnyArguments().thenAnswer((invocation) -> {
			Object[] args = invocation.getArguments();
			Configuration config = (Configuration) args[0];

			Assert.assertEquals(host, config.getString(RestOptions.ADDRESS));
			Assert.assertEquals(port, config.getInteger(RestOptions.PORT));
			return mockedClient;
		}
	);

	final Configuration clientConfiguration = new Configuration();
	final StreamExecutionEnvironment env = StreamExecutionEnvironment.createRemoteEnvironment(
		host, port, clientConfiguration);
	env.fromElements(1).map(x -> x * 2);
	JobExecutionResult actualResult = env.execute("fakeJobName");
	Assert.assertEquals(expectedResult, actualResult);
}
 

/**
 * Tests that Yarn will restart a killed {@link YarnSessionClusterEntrypoint} which will then resume
 * a persisted {@link JobGraph}.
 */
@Test
public void testKillYarnSessionClusterEntrypoint() throws Exception {
	runTest(() -> {
		assumeTrue(
			"This test kills processes via the pkill command. Thus, it only runs on Linux, Mac OS, Free BSD and Solaris.",
			OperatingSystem.isLinux() || OperatingSystem.isMac() || OperatingSystem.isFreeBSD() || OperatingSystem.isSolaris());

		final YarnClusterDescriptor yarnClusterDescriptor = setupYarnClusterDescriptor();
		yarnClusterDescriptor.addShipFiles(Arrays.asList(flinkShadedHadoopDir.listFiles()));

		final RestClusterClient<ApplicationId> restClusterClient = deploySessionCluster(yarnClusterDescriptor);

		try {
			final JobID jobId = submitJob(restClusterClient);
			final ApplicationId id = restClusterClient.getClusterId();

			waitUntilJobIsRunning(restClusterClient, jobId);

			killApplicationMaster(yarnClusterDescriptor.getYarnSessionClusterEntrypoint());
			waitForApplicationAttempt(id, 2);

			waitForJobTermination(restClusterClient, jobId);

			killApplicationAndWait(id);
		} finally {
			restClusterClient.shutdown();
		}
	});
}
 

private static int getJobFullRestarts(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId) throws Exception {

	return getJobMetric(restClusterClient, jobId, "fullRestarts")
		.map(Metric::getValue)
		.map(Integer::parseInt)
		.orElse(0);
}
 

private static void waitUntilJobIsRestarted(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final int expectedFullRestarts) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> getJobFullRestarts(restClusterClient, jobId) >= expectedFullRestarts,
		Deadline.fromNow(TIMEOUT));
}
 

private void waitForJobTermination(
		final RestClusterClient<ApplicationId> restClusterClient,
		final JobID jobId) throws Exception {
	stopJobSignal.signal();
	final CompletableFuture<JobResult> jobResult = restClusterClient.requestJobResult(jobId);
	jobResult.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
}
 

private RestClusterClient<ApplicationId> deploySessionCluster(YarnClusterDescriptor yarnClusterDescriptor) throws ClusterDeploymentException {
	final int containerMemory = 256;
	final ClusterClient<ApplicationId> yarnClusterClient = yarnClusterDescriptor.deploySessionCluster(
		new ClusterSpecification.ClusterSpecificationBuilder()
			.setMasterMemoryMB(containerMemory)
			.setTaskManagerMemoryMB(containerMemory)
			.setSlotsPerTaskManager(1)
			.createClusterSpecification());

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

private JobID submitJob(RestClusterClient<ApplicationId> restClusterClient) throws InterruptedException, java.util.concurrent.ExecutionException {
	final CompletableFuture<JobSubmissionResult> jobSubmissionResultCompletableFuture =
		restClusterClient.submitJob(job);

	final JobSubmissionResult jobSubmissionResult = jobSubmissionResultCompletableFuture.get();
	return jobSubmissionResult.getJobID();
}
 

private static void waitUntilJobIsRunning(RestClusterClient<ApplicationId> restClusterClient, JobID jobId) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> {
			final JobDetailsInfo jobDetails = restClusterClient.getJobDetails(jobId).get();
			return jobDetails.getJobVertexInfos()
				.stream()
				.map(toExecutionState())
				.allMatch(isRunning());
		},
		Deadline.fromNow(TIMEOUT));
}
 

private static void waitUntilJobIsRestarted(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final int expectedFullRestarts) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> getJobFullRestarts(restClusterClient, jobId) >= expectedFullRestarts,
		Deadline.fromNow(TIMEOUT));
}
 

private static int getJobFullRestarts(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId) throws Exception {

	return getJobMetric(restClusterClient, jobId, "fullRestarts")
		.map(Metric::getValue)
		.map(Integer::parseInt)
		.orElse(0);
}
 

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 static FlinkClient createClient(Configuration configuration, JarLoader jarLoader, List<String> dependencies, String webUrl) {
    try {
        RestClusterClient restClient = new RestClusterClient<>(configuration, "RemoteExecutor");
        restClient.setPrintStatusDuringExecution(true);
        restClient.setDetached(true);
        return new StandaloneClusterFlinkClient(restClient, jarLoader, dependencies, webUrl);
    } catch (Exception e) {
        throw new RuntimeException("Cannot establish connection to JobManager: " + e.getMessage(), e);
    }
}
 

@Override
protected ClusterClient<ApplicationId> createYarnClusterClient(
        AbstractYarnClusterDescriptor descriptor,
        int numberTaskManagers,
        int slotsPerTaskManager,
        ApplicationReport report,
        Configuration flinkConfiguration,
        boolean perJobCluster)
        throws Exception
{
    return new RestClusterClient<>(flinkConfiguration, report.getApplicationId());
}
 

private void waitForJobTermination(
		final RestClusterClient<ApplicationId> restClusterClient,
		final JobID jobId) throws Exception {
	stopJobSignal.signal();
	final CompletableFuture<JobResult> jobResult = restClusterClient.requestJobResult(jobId);
	jobResult.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
}
 

@Before
public void setUp() throws Exception {
	final Configuration clientConfig = new Configuration();
	clientConfig.setInteger(RestOptions.RETRY_MAX_ATTEMPTS, 0);
	clientConfig.setLong(RestOptions.RETRY_DELAY, 0);
	clientConfig.addAll(CLUSTER.getClientConfiguration());

	client = new RestClusterClient<>(
		clientConfig,
		StandaloneClusterId.getInstance()
	);
}
 

private ClusterClientProvider<String> createClusterClientProvider(String clusterId) {
	return () -> {
		final Configuration configuration = new Configuration(flinkConfig);

		final Optional<Endpoint> restEndpoint = client.getRestEndpoint(clusterId);

		if (restEndpoint.isPresent()) {
			configuration.setString(RestOptions.ADDRESS, restEndpoint.get().getAddress());
			configuration.setInteger(RestOptions.PORT, restEndpoint.get().getPort());
		} else {
			throw new RuntimeException(
					new ClusterRetrieveException(
							"Could not get the rest endpoint of " + clusterId));
		}

		try {
			// Flink client will always use Kubernetes service to contact with jobmanager. So we have a pre-configured web
			// monitor address. Using StandaloneClientHAServices to create RestClusterClient is reasonable.
			return new RestClusterClient<>(
				configuration,
				clusterId,
				new StandaloneClientHAServices(HighAvailabilityServicesUtils.getWebMonitorAddress(
					configuration, HighAvailabilityServicesUtils.AddressResolution.TRY_ADDRESS_RESOLUTION)));
		} catch (Exception e) {
			client.handleException(e);
			throw new RuntimeException(new ClusterRetrieveException("Could not create the RestClusterClient.", e));
		}
	};
}