org.apache.flink.api.common.cache.DistributedCache Java Examples

public PythonDependencyInfoTest() {
	Map<String, Future<Path>> distributeCachedFiles = new HashMap<>();
	distributeCachedFiles.put(
		"python_file_{SHA256_0}",
		CompletableFuture.completedFuture(new Path("/distributed_cache/file0")));
	distributeCachedFiles.put(
		"python_file_{SHA256_1}",
		CompletableFuture.completedFuture(new Path("/distributed_cache/file1")));
	distributeCachedFiles.put(
		"python_requirements_file_{SHA256}",
		CompletableFuture.completedFuture(new Path("/distributed_cache/file2")));
	distributeCachedFiles.put(
		"python_requirements_cache_{SHA256}",
		CompletableFuture.completedFuture(new Path("/distributed_cache/file3")));
	distributeCachedFiles.put(
		"python_archive_{SHA256_0}",
		CompletableFuture.completedFuture(new Path("/distributed_cache/file4")));
	distributeCachedFiles.put(
		"python_archive_{SHA256_1}",
		CompletableFuture.completedFuture(new Path("/distributed_cache/file5")));
	distributedCache = new DistributedCache(distributeCachedFiles);
}
 

@Test
public void testLoadingCachedFilesFromConfiguration() {
	StreamExecutionEnvironment envFromConfiguration = StreamExecutionEnvironment.getExecutionEnvironment();
	envFromConfiguration.registerCachedFile("/tmp3", "file3", true);

	Configuration configuration = new Configuration();
	configuration.setString("pipeline.cached-files", "name:file1,path:/tmp1,executable:true;name:file2,path:/tmp2");

	// mutate config according to configuration
	envFromConfiguration.configure(configuration, Thread.currentThread().getContextClassLoader());

	assertThat(envFromConfiguration.getCachedFiles(), equalTo(Arrays.asList(
		Tuple2.of("file1", new DistributedCache.DistributedCacheEntry("/tmp1", true)),
		Tuple2.of("file2", new DistributedCache.DistributedCacheEntry("/tmp2", false))
	)));
}
 

/**
 * Sets all relevant options contained in the {@link ReadableConfig} such as e.g.
 * {@link StreamPipelineOptions#TIME_CHARACTERISTIC}. It will reconfigure
 * {@link StreamExecutionEnvironment}, {@link ExecutionConfig} and {@link CheckpointConfig}.
 *
 * <p>It will change the value of a setting only if a corresponding option was set in the
 * {@code configuration}. If a key is not present, the current value of a field will remain
 * untouched.
 *
 * @param configuration a configuration to read the values from
 * @param classLoader a class loader to use when loading classes
 */
@PublicEvolving
public void configure(ReadableConfig configuration, ClassLoader classLoader) {
	configuration.getOptional(StreamPipelineOptions.TIME_CHARACTERISTIC)
		.ifPresent(this::setStreamTimeCharacteristic);
	Optional.ofNullable(loadStateBackend(configuration, classLoader))
		.ifPresent(this::setStateBackend);
	configuration.getOptional(PipelineOptions.OPERATOR_CHAINING)
		.ifPresent(c -> this.isChainingEnabled = c);
	configuration.getOptional(ExecutionOptions.BUFFER_TIMEOUT)
		.ifPresent(t -> this.setBufferTimeout(t.toMillis()));
	configuration.getOptional(DeploymentOptions.JOB_LISTENERS)
		.ifPresent(listeners -> registerCustomListeners(classLoader, listeners));
	configuration.getOptional(PipelineOptions.CACHED_FILES)
		.ifPresent(f -> {
			this.cacheFile.clear();
			this.cacheFile.addAll(DistributedCache.parseCachedFilesFromString(f));
		});
	config.configure(configuration, classLoader);
	checkpointCfg.configure(configuration);
}
 

@Test
public void testFileDownloadedFromBlob() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey));
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID);

	final Path dstPath = copyResult.get();
	final String actualContent = Files.toString(new File(dstPath.toUri()), StandardCharsets.UTF_8);
	assertTrue(dstPath.getFileSystem().exists(dstPath));
	assertEquals(testFileContent, actualContent);
}
 

private void testDirectoryDownloaded(DistributedCache.DistributedCacheEntry entry) throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID = new ExecutionAttemptID();

	// copy / create the file
	final String fileName = "test_file";
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID);

	final Path dstPath = copyResult.get();
	final FileSystem fs = dstPath.getFileSystem();
	final FileStatus fileStatus = fs.getFileStatus(dstPath);
	assertTrue(fileStatus.isDir());

	final Path cacheFile = new Path(dstPath, "cacheFile");
	assertTrue(fs.exists(cacheFile));
	final String actualContent = FileUtils.readFileUtf8(new File(cacheFile.getPath()));
	assertEquals(testFileContent, actualContent);
}
 

@Test
public void testFileDownloadedFromBlob() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey));
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID);

	final Path dstPath = copyResult.get();
	final String actualContent = Files.toString(new File(dstPath.toUri()), StandardCharsets.UTF_8);
	assertTrue(dstPath.getFileSystem().exists(dstPath));
	assertEquals(testFileContent, actualContent);
}
 

@Test
public void testDirectoryDownloadedFromBlob() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey),
		true);
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID);

	final Path dstPath = copyResult.get();
	final FileSystem fs = dstPath.getFileSystem();
	final FileStatus fileStatus = fs.getFileStatus(dstPath);
	assertTrue(fileStatus.isDir());

	final Path cacheFile = new Path(dstPath, "cacheFile");
	assertTrue(fs.exists(cacheFile));
	final String actualContent = FileUtils.readFileUtf8(new File(cacheFile.getPath()));
	assertEquals(testFileContent, actualContent);
}
 

@Test
public void testDirectoryDownloadedFromBlob() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey),
		true);
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID);

	final Path dstPath = copyResult.get();
	final FileSystem fs = dstPath.getFileSystem();
	final FileStatus fileStatus = fs.getFileStatus(dstPath);
	assertTrue(fileStatus.isDir());

	final Path cacheFile = new Path(dstPath, "cacheFile");
	assertTrue(fs.exists(cacheFile));
	final String actualContent = FileUtils.readFileUtf8(new File(cacheFile.getPath()));
	assertEquals(testFileContent, actualContent);
}
 

@Test
public void testFileDownloadedFromBlob() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey));
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID);

	final Path dstPath = copyResult.get();
	final String actualContent = Files.toString(new File(dstPath.toUri()), StandardCharsets.UTF_8);
	assertTrue(dstPath.getFileSystem().exists(dstPath));
	assertEquals(testFileContent, actualContent);
}
 

private static void assertState(DistributedCache.DistributedCacheEntry original, DistributedCache.DistributedCacheEntry actual, boolean isBlobKeyNull, JobID jobId) throws Exception {
	assertEquals(original.isZipped, actual.isZipped);
	assertEquals(original.isExecutable, actual.isExecutable);
	assertEquals(original.filePath, actual.filePath);
	assertEquals(isBlobKeyNull, actual.blobKey == null);
	if (!isBlobKeyNull) {
		blobServer.getFile(
			jobId,
			InstantiationUtil.<PermanentBlobKey>deserializeObject(actual.blobKey, ClientUtilsTest.class.getClassLoader()));
	}
}
 

public PlanGenerator(
		List<DataSink<?>> sinks,
		ExecutionConfig config,
		int defaultParallelism,
		List<Tuple2<String, DistributedCache.DistributedCacheEntry>> cacheFile,
		String jobName) {
	this.sinks = checkNotNull(sinks);
	this.config = checkNotNull(config);
	this.cacheFile = checkNotNull(cacheFile);
	this.jobName = checkNotNull(jobName);
	this.defaultParallelism = defaultParallelism;
}
 

private static void assertState(DistributedCache.DistributedCacheEntry entry, boolean isExecutable, boolean isZipped) throws IOException {
	assertNotNull(entry);
	assertEquals(isExecutable, entry.isExecutable);
	assertEquals(isZipped, entry.isZipped);
	org.apache.flink.core.fs.Path filePath = new org.apache.flink.core.fs.Path(entry.filePath);
	assertTrue(filePath.getFileSystem().exists(filePath));
	assertFalse(filePath.getFileSystem().getFileStatus(filePath).isDir());
}
 

@Test
public void testDirectoryCleanUp() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID1 = new ExecutionAttemptID();
	ExecutionAttemptID attemptID2 = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey),
		true);
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID1);
	fileCache.createTmpFile(fileName, entry, jobID, attemptID2);

	final Path dstPath = copyResult.get();
	final FileSystem fs = dstPath.getFileSystem();
	final FileStatus fileStatus = fs.getFileStatus(dstPath);
	final Path cacheFile = new Path(dstPath, "cacheFile");
	assertTrue(fileStatus.isDir());
	assertTrue(fs.exists(cacheFile));

	fileCache.releaseJob(jobID, attemptID1);
	// still should be available
	assertTrue(fileStatus.isDir());
	assertTrue(fs.exists(cacheFile));

	fileCache.releaseJob(jobID, attemptID2);
	// still should be available, file will be deleted after cleanupInterval
	assertTrue(fileStatus.isDir());
	assertTrue(fs.exists(cacheFile));

	// after a while, the file should disappear
	assertEquals(CLEANUP_INTERVAL, executorService.lastDelayMillis);
	executorService.lastDeleteProcess.run();

	assertFalse(fs.exists(dstPath));
	assertFalse(fs.exists(cacheFile));
}
 

public void setUserArtifactBlobKey(String entryName, PermanentBlobKey blobKey) throws IOException {
	byte[] serializedBlobKey;
	serializedBlobKey = InstantiationUtil.serializeObject(blobKey);

	userArtifacts.computeIfPresent(entryName, (key, originalEntry) -> new DistributedCache.DistributedCacheEntry(
		originalEntry.filePath,
		originalEntry.isExecutable,
		serializedBlobKey,
		originalEntry.isZipped
	));
}
 

public void writeUserArtifactEntriesToConfiguration() {
	for (Map.Entry<String, DistributedCache.DistributedCacheEntry> userArtifact : userArtifacts.entrySet()) {
		DistributedCache.writeFileInfoToConfig(
			userArtifact.getKey(),
			userArtifact.getValue(),
			jobConfiguration
		);
	}
}
 

public void setUserArtifactBlobKey(String entryName, PermanentBlobKey blobKey) throws IOException {
	byte[] serializedBlobKey;
	serializedBlobKey = InstantiationUtil.serializeObject(blobKey);

	userArtifacts.computeIfPresent(entryName, (key, originalEntry) -> new DistributedCache.DistributedCacheEntry(
		originalEntry.filePath,
		originalEntry.isExecutable,
		serializedBlobKey,
		originalEntry.isZipped
	));
}
 

public void writeUserArtifactEntriesToConfiguration() {
	for (Map.Entry<String, DistributedCache.DistributedCacheEntry> userArtifact : userArtifacts.entrySet()) {
		DistributedCache.writeFileInfoToConfig(
			userArtifact.getKey(),
			userArtifact.getValue(),
			jobConfiguration
		);
	}
}
 

private static void assertState(DistributedCache.DistributedCacheEntry original, DistributedCache.DistributedCacheEntry actual, boolean isBlobKeyNull, JobID jobId) throws Exception {
	assertEquals(original.isZipped, actual.isZipped);
	assertEquals(original.isExecutable, actual.isExecutable);
	assertEquals(original.filePath, actual.filePath);
	assertEquals(isBlobKeyNull, actual.blobKey == null);
	if (!isBlobKeyNull) {
		blobServer.getFile(
			jobId,
			InstantiationUtil.<PermanentBlobKey>deserializeObject(actual.blobKey, ClientUtilsTest.class.getClassLoader()));
	}
}
 

@Test
public void testGenerate() {

	final String fileA = "fileA";
	final String fileB = "fileB";

	final Map<String, DistributedCache.DistributedCacheEntry> originalArtifacts = Stream.of(
			Tuple2.of(fileA, new DistributedCache.DistributedCacheEntry("test1", true)),
			Tuple2.of(fileB, new DistributedCache.DistributedCacheEntry("test2", false))
	).collect(Collectors.toMap(x -> x.f0, x -> x.f1));

	final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);
	env.registerCachedFile("test1", fileA, true);
	env.registerCachedFile("test2", fileB, false);

	env.fromElements(1, 3, 5)
			.map((MapFunction<Integer, String>) value -> String.valueOf(value + 1))
			.writeAsText("/tmp/csv");

	final Plan generatedPlanUnderTest = env.createProgramPlan("test");

	final Map<String, DistributedCache.DistributedCacheEntry> retrievedArtifacts =
			generatedPlanUnderTest
					.getCachedFiles()
					.stream()
					.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

	assertEquals(1, generatedPlanUnderTest.getDataSinks().size());
	assertEquals(10, generatedPlanUnderTest.getDefaultParallelism());
	assertEquals(env.getConfig(), generatedPlanUnderTest.getExecutionConfig());
	assertEquals("test", generatedPlanUnderTest.getJobName());

	assertEquals(originalArtifacts.size(), retrievedArtifacts.size());
	assertEquals(originalArtifacts.get(fileA), retrievedArtifacts.get(fileA));
	assertEquals(originalArtifacts.get(fileB), retrievedArtifacts.get(fileB));
}
 

public AbstractRuntimeUDFContext(TaskInfo taskInfo,
									ClassLoader userCodeClassLoader,
									ExecutionConfig executionConfig,
									Map<String, Accumulator<?, ?>> accumulators,
									Map<String, Future<Path>> cpTasks,
									MetricGroup metrics) {
	this.taskInfo = checkNotNull(taskInfo);
	this.userCodeClassLoader = userCodeClassLoader;
	this.executionConfig = executionConfig;
	this.distributedCache = new DistributedCache(checkNotNull(cpTasks));
	this.accumulators = checkNotNull(accumulators);
	this.metrics = metrics;
}
 

/**
 * Adds the path of a custom file required to run the job on a task manager.
 *
 * @param name a name under which this artifact will be accessible through {@link DistributedCache}
 * @param file path of a custom file required to run the job on a task manager
 */
public void addUserArtifact(String name, DistributedCache.DistributedCacheEntry file) {
	if (file == null) {
		throw new IllegalArgumentException();
	}

	userArtifacts.putIfAbsent(name, file);
}
 

public AbstractRuntimeUDFContext(TaskInfo taskInfo,
									ClassLoader userCodeClassLoader,
									ExecutionConfig executionConfig,
									Map<String, Accumulator<?, ?>> accumulators,
									Map<String, Future<Path>> cpTasks,
									MetricGroup metrics) {
	this.taskInfo = checkNotNull(taskInfo);
	this.userCodeClassLoader = userCodeClassLoader;
	this.executionConfig = executionConfig;
	this.distributedCache = new DistributedCache(checkNotNull(cpTasks));
	this.accumulators = checkNotNull(accumulators);
	this.metrics = metrics;
}
 

public void setUserArtifactRemotePath(String entryName, String remotePath) {
	userArtifacts.computeIfPresent(entryName, (key, originalEntry) -> new DistributedCache.DistributedCacheEntry(
		remotePath,
		originalEntry.isExecutable,
		null,
		originalEntry.isZipped
	));
}
 

public void writeUserArtifactEntriesToConfiguration() {
	for (Map.Entry<String, DistributedCache.DistributedCacheEntry> userArtifact : userArtifacts.entrySet()) {
		DistributedCache.writeFileInfoToConfig(
			userArtifact.getKey(),
			userArtifact.getValue(),
			jobConfiguration
		);
	}
}
 

private static void assertState(DistributedCache.DistributedCacheEntry original, DistributedCache.DistributedCacheEntry actual, boolean isBlobKeyNull, JobID jobId) throws Exception {
	assertEquals(original.isZipped, actual.isZipped);
	assertEquals(original.isExecutable, actual.isExecutable);
	assertEquals(original.filePath, actual.filePath);
	assertEquals(isBlobKeyNull, actual.blobKey == null);
	if (!isBlobKeyNull) {
		blobServer.getFile(
			jobId,
			InstantiationUtil.<PermanentBlobKey>deserializeObject(actual.blobKey, ClientUtilsTest.class.getClassLoader()));
	}
}
 

@Test
public void testDirectoryDownloadedFromBlob() throws Exception {
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		"test_file",
		false,
		InstantiationUtil.serializeObject(permanentBlobKey),
		true);

	testDirectoryDownloaded(entry);
}
 

@Test
public void testDirectoryDownloadedFromDFS() throws Exception {
	final String zippedFile = blobService.getFile(new JobID(), permanentBlobKey).getAbsolutePath();
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		zippedFile,
		false,
		null,
		true);

	testDirectoryDownloaded(entry);
}
 

@Test
public void testDirectoryCleanUp() throws Exception {
	JobID jobID = new JobID();
	ExecutionAttemptID attemptID1 = new ExecutionAttemptID();
	ExecutionAttemptID attemptID2 = new ExecutionAttemptID();

	final String fileName = "test_file";
	// copy / create the file
	final DistributedCache.DistributedCacheEntry entry = new DistributedCache.DistributedCacheEntry(
		fileName,
		false,
		InstantiationUtil.serializeObject(permanentBlobKey),
		true);
	Future<Path> copyResult = fileCache.createTmpFile(fileName, entry, jobID, attemptID1);
	fileCache.createTmpFile(fileName, entry, jobID, attemptID2);

	final Path dstPath = copyResult.get();
	final FileSystem fs = dstPath.getFileSystem();
	final FileStatus fileStatus = fs.getFileStatus(dstPath);
	final Path cacheFile = new Path(dstPath, "cacheFile");
	assertTrue(fileStatus.isDir());
	assertTrue(fs.exists(cacheFile));

	fileCache.releaseJob(jobID, attemptID1);
	// still should be available
	assertTrue(fileStatus.isDir());
	assertTrue(fs.exists(cacheFile));

	fileCache.releaseJob(jobID, attemptID2);
	// still should be available, file will be deleted after cleanupInterval
	assertTrue(fileStatus.isDir());
	assertTrue(fs.exists(cacheFile));

	// after a while, the file should disappear
	assertEquals(CLEANUP_INTERVAL, executorService.lastDelayMillis);
	executorService.lastDeleteProcess.run();

	assertFalse(fs.exists(dstPath));
	assertFalse(fs.exists(cacheFile));
}
 

@Test
public void testNotOverridingCachedFilesFromConfiguration() {
	StreamExecutionEnvironment envFromConfiguration = StreamExecutionEnvironment.getExecutionEnvironment();
	envFromConfiguration.registerCachedFile("/tmp3", "file3", true);

	Configuration configuration = new Configuration();

	// mutate config according to configuration
	envFromConfiguration.configure(configuration, Thread.currentThread().getContextClassLoader());

	assertThat(envFromConfiguration.getCachedFiles(), equalTo(Arrays.asList(
		Tuple2.of("file3", new DistributedCache.DistributedCacheEntry("/tmp3", true))
	)));
}
 

private static void assertState(DistributedCache.DistributedCacheEntry entry, boolean isExecutable, boolean isZipped) throws IOException {
	assertNotNull(entry);
	assertEquals(isExecutable, entry.isExecutable);
	assertEquals(isZipped, entry.isZipped);
	org.apache.flink.core.fs.Path filePath = new org.apache.flink.core.fs.Path(entry.filePath);
	assertTrue(filePath.getFileSystem().exists(filePath));
	assertFalse(filePath.getFileSystem().getFileStatus(filePath).isDir());
}