org.apache.flink.runtime.clusterframework.ContaineredTaskManagerParameters Java Examples

private KubernetesTaskManagerParameters createKubernetesTaskManagerParameters(WorkerResourceSpec workerResourceSpec) {
	final TaskExecutorProcessSpec taskExecutorProcessSpec =
		TaskExecutorProcessUtils.processSpecFromWorkerResourceSpec(flinkConfig, workerResourceSpec);

	final String podName = String.format(
		TASK_MANAGER_POD_FORMAT,
		clusterId,
		currentMaxAttemptId,
		++currentMaxPodId);

	final ContaineredTaskManagerParameters taskManagerParameters =
		ContaineredTaskManagerParameters.create(flinkConfig, taskExecutorProcessSpec);

	final Configuration taskManagerConfig = new Configuration(flinkConfig);
	taskManagerConfig.set(TaskManagerOptions.TASK_MANAGER_RESOURCE_ID, podName);

	final String dynamicProperties =
		BootstrapTools.getDynamicPropertiesAsString(flinkClientConfig, taskManagerConfig);

	return new KubernetesTaskManagerParameters(
		flinkConfig,
		podName,
		dynamicProperties,
		taskManagerParameters,
		ExternalResourceUtils.getExternalResources(flinkConfig, KubernetesConfigOptions.EXTERNAL_RESOURCE_KUBERNETES_CONFIG_KEY_SUFFIX));
}
 

public MesosTaskManagerParameters(
		double cpus,
		int gpus,
		ContainerType containerType,
		Option<String> containerImageName,
		ContaineredTaskManagerParameters containeredParameters,
		List<Protos.Volume> containerVolumes,
		List<Protos.Parameter> dockerParameters,
		boolean dockerForcePullImage,
		List<ConstraintEvaluator> constraints,
		String command,
		Option<String> bootstrapCommand,
		Option<String> taskManagerHostname,
		List<String> uris) {

	this.cpus = cpus;
	this.gpus = gpus;
	this.containerType = Preconditions.checkNotNull(containerType);
	this.containerImageName = Preconditions.checkNotNull(containerImageName);
	this.containeredParameters = Preconditions.checkNotNull(containeredParameters);
	this.containerVolumes = Preconditions.checkNotNull(containerVolumes);
	this.dockerParameters = Preconditions.checkNotNull(dockerParameters);
	this.dockerForcePullImage = dockerForcePullImage;
	this.constraints = Preconditions.checkNotNull(constraints);
	this.command = Preconditions.checkNotNull(command);
	this.bootstrapCommand = Preconditions.checkNotNull(bootstrapCommand);
	this.taskManagerHostname = Preconditions.checkNotNull(taskManagerHostname);
	this.uris = Preconditions.checkNotNull(uris);
}
 

public KubernetesTaskManagerParameters(
		Configuration flinkConfig,
		String podName,
		String dynamicProperties,
		ContaineredTaskManagerParameters containeredTaskManagerParameters,
		Map<String, Long> taskManagerExternalResources) {
	super(flinkConfig);
	this.podName = checkNotNull(podName);
	this.dynamicProperties = checkNotNull(dynamicProperties);
	this.containeredTaskManagerParameters = checkNotNull(containeredTaskManagerParameters);
	this.taskManagerExternalResources = checkNotNull(taskManagerExternalResources);
}
 

private static String getTaskManagerStartCommand(
		Configuration flinkConfig,
		ContaineredTaskManagerParameters tmParams,
		String configDirectory,
		String logDirectory,
		boolean hasLogback,
		boolean hasLog4j,
		String mainClass,
		String mainArgs) {
	final TaskExecutorProcessSpec taskExecutorProcessSpec = tmParams.getTaskExecutorProcessSpec();
	final String jvmMemOpts = ProcessMemoryUtils.generateJvmParametersStr(taskExecutorProcessSpec);
	String args = TaskExecutorProcessUtils.generateDynamicConfigsStr(taskExecutorProcessSpec);
	if (mainArgs != null) {
		args += " " + mainArgs;
	}

	return KubernetesUtils.getCommonStartCommand(
		flinkConfig,
		KubernetesUtils.ClusterComponent.TASK_MANAGER,
		jvmMemOpts,
		configDirectory,
		logDirectory,
		hasLogback,
		hasLog4j,
		mainClass,
		args);
}
 

@Override
protected void onSetup() throws Exception {
	super.onSetup();

	final TaskExecutorProcessSpec taskExecutorProcessSpec =
		TaskExecutorProcessUtils.processSpecFromConfig(flinkConfig);
	final ContaineredTaskManagerParameters containeredTaskManagerParameters =
		ContaineredTaskManagerParameters.create(flinkConfig, taskExecutorProcessSpec);

	this.kubernetesTaskManagerParameters = new KubernetesTaskManagerParameters(flinkConfig,
		POD_NAME,
		DYNAMIC_PROPERTIES,
		containeredTaskManagerParameters,
		Collections.emptyMap());
}
 

public static Configuration createActiveResourceManagerConfiguration(Configuration originalConfiguration) {
	final int taskManagerMemoryMB = ConfigurationUtils.getTaskManagerHeapMemory(originalConfiguration).getMebiBytes();
	final long cutoffMB = ContaineredTaskManagerParameters.calculateCutoffMB(originalConfiguration, taskManagerMemoryMB);
	final long processMemoryBytes = (taskManagerMemoryMB - cutoffMB) << 20; // megabytes to bytes
	final long managedMemoryBytes = TaskManagerServices.getManagedMemoryFromProcessMemory(originalConfiguration, processMemoryBytes);

	final Configuration resourceManagerConfig = new Configuration(originalConfiguration);
	resourceManagerConfig.setString(TaskManagerOptions.MANAGED_MEMORY_SIZE, managedMemoryBytes + "b");

	return resourceManagerConfig;
}
 

/**
 * Method to validate cluster specification before deploy it, it will throw
 * an {@link FlinkException} if the {@link ClusterSpecification} is invalid.
 *
 * @param clusterSpecification cluster specification to check against the configuration of the
 *                             AbstractYarnClusterDescriptor
 * @throws FlinkException if the cluster cannot be started with the provided {@link ClusterSpecification}
 */
private void validateClusterSpecification(ClusterSpecification clusterSpecification) throws FlinkException {
	try {
		final long taskManagerMemorySize = clusterSpecification.getTaskManagerMemoryMB();
		// We do the validation by calling the calculation methods here
		// Internally these methods will check whether the cluster can be started with the provided
		// ClusterSpecification and the configured memory requirements
		final long cutoff = ContaineredTaskManagerParameters.calculateCutoffMB(flinkConfiguration, taskManagerMemorySize);
		TaskManagerServices.calculateHeapSizeMB(taskManagerMemorySize - cutoff, flinkConfiguration);
	} catch (IllegalArgumentException iae) {
		throw new FlinkException("Cannot fulfill the minimum memory requirements with the provided " +
			"cluster specification. Please increase the memory of the cluster.", iae);
	}
}
 

private ContainerLaunchContext createTaskExecutorLaunchContext(Resource resource, String containerId, String host)
		throws Exception {
	// init the ContainerLaunchContext
	final String currDir = env.get(ApplicationConstants.Environment.PWD.key());

	final ContaineredTaskManagerParameters taskManagerParameters =
			ContaineredTaskManagerParameters.create(flinkConfig, resource.getMemory(), numberOfTaskSlots);

	log.debug("TaskExecutor {} will be started with container size {} MB, JVM heap size {} MB, " +
			"JVM direct memory limit {} MB",
			containerId,
			taskManagerParameters.taskManagerTotalMemoryMB(),
			taskManagerParameters.taskManagerHeapSizeMB(),
			taskManagerParameters.taskManagerDirectMemoryLimitMB());

	Configuration taskManagerConfig = BootstrapTools.cloneConfiguration(flinkConfig);

	log.debug("TaskManager configuration: {}", taskManagerConfig);

	ContainerLaunchContext taskExecutorLaunchContext = Utils.createTaskExecutorContext(
		flinkConfig,
		yarnConfig,
		env,
		taskManagerParameters,
		taskManagerConfig,
		currDir,
		YarnTaskExecutorRunner.class,
		log);

	// set a special environment variable to uniquely identify this container
	taskExecutorLaunchContext.getEnvironment()
			.put(ENV_FLINK_CONTAINER_ID, containerId);
	taskExecutorLaunchContext.getEnvironment()
			.put(ENV_FLINK_NODE_ID, host);
	return taskExecutorLaunchContext;
}
 

@Override
protected void onSetup() throws Exception {
	taskExecutorProcessSpec = TaskExecutorProcessUtils.processSpecFromConfig(flinkConfig);
	containeredTaskManagerParameters = ContaineredTaskManagerParameters.create(flinkConfig, taskExecutorProcessSpec);
	kubernetesTaskManagerParameters = new KubernetesTaskManagerParameters(
			flinkConfig,
			POD_NAME,
			DYNAMIC_PROPERTIES,
			containeredTaskManagerParameters,
			ExternalResourceUtils.getExternalResources(flinkConfig, KubernetesConfigOptions.EXTERNAL_RESOURCE_KUBERNETES_CONFIG_KEY_SUFFIX));
}
 

public MesosTaskManagerParameters(
		int gpus,
		int disk,
		int network,
		ContainerType containerType,
		Option<String> containerImageName,
		ContaineredTaskManagerParameters containeredParameters,
		List<Protos.Volume> containerVolumes,
		List<Protos.Parameter> dockerParameters,
		boolean dockerForcePullImage,
		List<ConstraintEvaluator> constraints,
		String command,
		Option<String> bootstrapCommand,
		Option<String> taskManagerHostname,
		List<String> uris) {

	this.gpus = gpus;
	this.disk = disk;
	this.network = network;
	this.containerType = Preconditions.checkNotNull(containerType);
	this.containerImageName = Preconditions.checkNotNull(containerImageName);
	this.containeredParameters = Preconditions.checkNotNull(containeredParameters);
	this.containerVolumes = Preconditions.checkNotNull(containerVolumes);
	this.dockerParameters = Preconditions.checkNotNull(dockerParameters);
	this.dockerForcePullImage = dockerForcePullImage;
	this.constraints = Preconditions.checkNotNull(constraints);
	this.command = Preconditions.checkNotNull(command);
	this.bootstrapCommand = Preconditions.checkNotNull(bootstrapCommand);
	this.taskManagerHostname = Preconditions.checkNotNull(taskManagerHostname);
	this.uris = Preconditions.checkNotNull(uris);
}
 

public MesosTaskManagerParameters(
		double cpus,
		int gpus,
		int disk,
		ContainerType containerType,
		Option<String> containerImageName,
		ContaineredTaskManagerParameters containeredParameters,
		List<Protos.Volume> containerVolumes,
		List<Protos.Parameter> dockerParameters,
		boolean dockerForcePullImage,
		List<ConstraintEvaluator> constraints,
		String command,
		Option<String> bootstrapCommand,
		Option<String> taskManagerHostname,
		List<String> uris) {

	this.cpus = cpus;
	this.gpus = gpus;
	this.disk = disk;
	this.containerType = Preconditions.checkNotNull(containerType);
	this.containerImageName = Preconditions.checkNotNull(containerImageName);
	this.containeredParameters = Preconditions.checkNotNull(containeredParameters);
	this.containerVolumes = Preconditions.checkNotNull(containerVolumes);
	this.dockerParameters = Preconditions.checkNotNull(dockerParameters);
	this.dockerForcePullImage = dockerForcePullImage;
	this.constraints = Preconditions.checkNotNull(constraints);
	this.command = Preconditions.checkNotNull(command);
	this.bootstrapCommand = Preconditions.checkNotNull(bootstrapCommand);
	this.taskManagerHostname = Preconditions.checkNotNull(taskManagerHostname);
	this.uris = Preconditions.checkNotNull(uris);
}
 

/**
 * Method to validate cluster specification before deploy it, it will throw
 * an {@link FlinkException} if the {@link ClusterSpecification} is invalid.
 *
 * @param clusterSpecification cluster specification to check against the configuration of the
 *                             AbstractYarnClusterDescriptor
 * @throws FlinkException if the cluster cannot be started with the provided {@link ClusterSpecification}
 */
private void validateClusterSpecification(ClusterSpecification clusterSpecification) throws FlinkException {
	try {
		final long taskManagerMemorySize = clusterSpecification.getTaskManagerMemoryMB();
		// We do the validation by calling the calculation methods here
		// Internally these methods will check whether the cluster can be started with the provided
		// ClusterSpecification and the configured memory requirements
		final long cutoff = ContaineredTaskManagerParameters.calculateCutoffMB(flinkConfiguration, taskManagerMemorySize);
		TaskManagerServices.calculateHeapSizeMB(taskManagerMemorySize - cutoff, flinkConfiguration);
	} catch (IllegalArgumentException iae) {
		throw new FlinkException("Cannot fulfill the minimum memory requirements with the provided " +
			"cluster specification. Please increase the memory of the cluster.", iae);
	}
}
 

private ContainerLaunchContext createTaskExecutorLaunchContext(Resource resource, String containerId, String host)
		throws Exception {
	// init the ContainerLaunchContext
	final String currDir = env.get(ApplicationConstants.Environment.PWD.key());

	final ContaineredTaskManagerParameters taskManagerParameters =
			ContaineredTaskManagerParameters.create(flinkConfig, resource.getMemory(), numberOfTaskSlots);

	log.debug("TaskExecutor {} will be started with container size {} MB, JVM heap size {} MB, " +
			"JVM direct memory limit {} MB",
			containerId,
			taskManagerParameters.taskManagerTotalMemoryMB(),
			taskManagerParameters.taskManagerHeapSizeMB(),
			taskManagerParameters.taskManagerDirectMemoryLimitMB());

	Configuration taskManagerConfig = BootstrapTools.cloneConfiguration(flinkConfig);

	log.debug("TaskManager configuration: {}", taskManagerConfig);

	ContainerLaunchContext taskExecutorLaunchContext = Utils.createTaskExecutorContext(
		flinkConfig,
		yarnConfig,
		env,
		taskManagerParameters,
		taskManagerConfig,
		currDir,
		YarnTaskExecutorRunner.class,
		log);

	// set a special environment variable to uniquely identify this container
	taskExecutorLaunchContext.getEnvironment()
			.put(ENV_FLINK_CONTAINER_ID, containerId);
	taskExecutorLaunchContext.getEnvironment()
			.put(ENV_FLINK_NODE_ID, host);
	return taskExecutorLaunchContext;
}
 

private static ContaineredTaskManagerParameters createContaineredTaskManagerParameters(final Configuration flinkConfig) {
	double cpus = getCpuCores(flinkConfig);
	TaskExecutorProcessSpec taskExecutorProcessSpec = TaskExecutorProcessUtils
		.newProcessSpecBuilder(flinkConfig)
		.withCpuCores(cpus)
		.build();

	return ContaineredTaskManagerParameters.create(
		flinkConfig,
		taskExecutorProcessSpec);
}
 

private ContainerLaunchContext createTaskExecutorLaunchContext(
	String containerId,
	String host,
	TaskExecutorProcessSpec taskExecutorProcessSpec) throws Exception {

	// init the ContainerLaunchContext
	final String currDir = env.get(ApplicationConstants.Environment.PWD.key());

	final ContaineredTaskManagerParameters taskManagerParameters =
			ContaineredTaskManagerParameters.create(flinkConfig, taskExecutorProcessSpec);

	log.info("TaskExecutor {} will be started on {} with {}.",
		containerId,
		host,
		taskExecutorProcessSpec);

	final Configuration taskManagerConfig = BootstrapTools.cloneConfiguration(flinkConfig);
	taskManagerConfig.set(TaskManagerOptions.TASK_MANAGER_RESOURCE_ID, containerId);

	final String taskManagerDynamicProperties =
		BootstrapTools.getDynamicPropertiesAsString(flinkClientConfig, taskManagerConfig);

	log.debug("TaskManager configuration: {}", taskManagerConfig);

	ContainerLaunchContext taskExecutorLaunchContext = Utils.createTaskExecutorContext(
		flinkConfig,
		yarnConfig,
		env,
		taskManagerParameters,
		taskManagerDynamicProperties,
		currDir,
		YarnTaskExecutorRunner.class,
		log);

	taskExecutorLaunchContext.getEnvironment()
			.put(ENV_FLINK_NODE_ID, host);
	return taskExecutorLaunchContext;
}
 

/**
 * Get the common containered parameters.
 */
public ContaineredTaskManagerParameters containeredParameters() {
	return containeredParameters;
}
 

/**
 * Create the Mesos TaskManager parameters.
 *
 * @param flinkConfig the TM configuration.
 */
public static MesosTaskManagerParameters create(Configuration flinkConfig) {

	List<ConstraintEvaluator> constraints = parseConstraints(flinkConfig.getString(MESOS_CONSTRAINTS_HARD_HOSTATTR));

	// parse the common parameters
	ContaineredTaskManagerParameters containeredParameters = createContaineredTaskManagerParameters(flinkConfig);

	int gpus = flinkConfig.getInteger(MESOS_RM_TASKS_GPUS);

	if (gpus < 0) {
		throw new IllegalConfigurationException(MESOS_RM_TASKS_GPUS.key() +
			" cannot be negative");
	}

	int disk = flinkConfig.getInteger(MESOS_RM_TASKS_DISK_MB);

	int network = flinkConfig.getInteger(MESOS_RM_TASKS_NETWORK_MB_PER_SEC);

	// parse the containerization parameters
	String imageName = flinkConfig.getString(MESOS_RM_CONTAINER_IMAGE_NAME);

	ContainerType containerType;
	String containerTypeString = flinkConfig.getString(MESOS_RM_CONTAINER_TYPE);
	switch (containerTypeString) {
		case MESOS_RESOURCEMANAGER_TASKS_CONTAINER_TYPE_MESOS:
			containerType = ContainerType.MESOS;
			break;
		case MESOS_RESOURCEMANAGER_TASKS_CONTAINER_TYPE_DOCKER:
			containerType = ContainerType.DOCKER;
			if (imageName == null || imageName.length() == 0) {
				throw new IllegalConfigurationException(MESOS_RM_CONTAINER_IMAGE_NAME.key() +
					" must be specified for docker container type");
			}
			break;
		default:
			throw new IllegalConfigurationException("invalid container type: " + containerTypeString);
	}

	Option<String> containerVolOpt = Option.<String>apply(flinkConfig.getString(MESOS_RM_CONTAINER_VOLUMES));

	Option<String> dockerParamsOpt = Option.<String>apply(flinkConfig.getString(MESOS_RM_CONTAINER_DOCKER_PARAMETERS));

	Option<String> uriParamsOpt = Option.<String>apply(flinkConfig.getString(MESOS_TM_URIS));

	boolean dockerForcePullImage = flinkConfig.getBoolean(MESOS_RM_CONTAINER_DOCKER_FORCE_PULL_IMAGE);

	List<Protos.Volume> containerVolumes = buildVolumes(containerVolOpt);

	List<Protos.Parameter> dockerParameters = buildDockerParameters(dockerParamsOpt);

	List<String> uris = buildUris(uriParamsOpt);

	//obtain Task Manager Host Name from the configuration
	Option<String> taskManagerHostname = Option.apply(flinkConfig.getString(MESOS_TM_HOSTNAME));

	//obtain command-line from the configuration
	String tmCommand = flinkConfig.getString(MESOS_TM_CMD);
	Option<String> tmBootstrapCommand = Option.apply(flinkConfig.getString(MESOS_TM_BOOTSTRAP_CMD));

	return new MesosTaskManagerParameters(
		gpus,
		disk,
		network,
		containerType,
		Option.apply(imageName),
		containeredParameters,
		containerVolumes,
		dockerParameters,
		dockerForcePullImage,
		constraints,
		tmCommand,
		tmBootstrapCommand,
		taskManagerHostname,
		uris);
}
 

/**
 * Create mock RM dependencies.
 */
Context() throws Exception {
	rpcService = new TestingRpcService();
	fatalErrorHandler = new TestingFatalErrorHandler();
	rmServices = new MockMesosResourceManagerRuntimeServices();
	mesosServices = new MockMesosServices();

	// TaskExecutor templating
	ContainerSpecification containerSpecification = new ContainerSpecification();

	MemorySize totalProcessMemory = MemorySize.parse("2g");
	TaskExecutorProcessSpec spec = TaskExecutorProcessUtils
		.newProcessSpecBuilder(flinkConfig)
		.withCpuCores(1.0)
		.withTotalProcessMemory(totalProcessMemory)
		.build();
	ContaineredTaskManagerParameters containeredParams =
		new ContaineredTaskManagerParameters(spec, new HashMap<String, String>());
	MesosTaskManagerParameters tmParams = new MesosTaskManagerParameters(
		1, 0, 0, MesosTaskManagerParameters.ContainerType.MESOS, Option.<String>empty(), containeredParams,
		Collections.<Protos.Volume>emptyList(), Collections.<Protos.Parameter>emptyList(), false,
		Collections.<ConstraintEvaluator>emptyList(), "", Option.<String>empty(),
		Option.<String>empty(), Collections.<String>emptyList());

	// resource manager
	rmResourceID = ResourceID.generate();
	resourceManager =
		new TestingMesosResourceManager(
			rpcService,
			RM_ADDRESS,
			rmResourceID,
			rmServices.highAvailabilityServices,
			rmServices.heartbeatServices,
			rmServices.slotManager,
			rmServices.jobLeaderIdService,
			fatalErrorHandler,
			// Mesos specifics
			flinkConfig,
			mesosServices,
			rmServices.mesosConfig,
			tmParams,
			containerSpecification,
			UnregisteredMetricGroups.createUnregisteredResourceManagerMetricGroup());
	workerResourceSpec = WorkerResourceSpec.fromTaskExecutorProcessSpec(spec);

	// TaskExecutors
	task1Executor = mockTaskExecutor(task1);
	task2Executor = mockTaskExecutor(task2);
	task3Executor = mockTaskExecutor(task3);

	// JobMaster
	jobMaster1 = mockJobMaster(rmServices, new JobID(1, 0));
}
 

@Test
public void testCreateTaskExecutorCredentials() throws Exception {
	File root = temporaryFolder.getRoot();
	File home = new File(root, "home");
	boolean created = home.mkdir();
	assertTrue(created);

	Configuration flinkConf = new Configuration();
	YarnConfiguration yarnConf = new YarnConfiguration();

	Map<String, String> env = new HashMap<>();
	env.put(YarnConfigKeys.ENV_APP_ID, "foo");
	env.put(YarnConfigKeys.ENV_CLIENT_HOME_DIR, home.getAbsolutePath());
	env.put(YarnConfigKeys.ENV_CLIENT_SHIP_FILES, "");
	env.put(YarnConfigKeys.ENV_FLINK_CLASSPATH, "");
	env.put(YarnConfigKeys.ENV_HADOOP_USER_NAME, "foo");
	env.put(YarnConfigKeys.FLINK_DIST_JAR, new YarnLocalResourceDescriptor(
		"flink.jar",
		new Path(root.toURI()),
		0,
		System.currentTimeMillis(),
		LocalResourceVisibility.APPLICATION).toString());
	env = Collections.unmodifiableMap(env);

	File credentialFile = temporaryFolder.newFile("container_tokens");
	final Text amRmTokenKind = AMRMTokenIdentifier.KIND_NAME;
	final Text hdfsDelegationTokenKind = new Text("HDFS_DELEGATION_TOKEN");
	final Text service = new Text("test-service");
	Credentials amCredentials = new Credentials();
	amCredentials.addToken(amRmTokenKind, new Token<>(new byte[4], new byte[4], amRmTokenKind, service));
	amCredentials.addToken(hdfsDelegationTokenKind, new Token<>(new byte[4], new byte[4],
		hdfsDelegationTokenKind, service));
	amCredentials.writeTokenStorageFile(new org.apache.hadoop.fs.Path(credentialFile.getAbsolutePath()), yarnConf);

	TaskExecutorProcessSpec spec = TaskExecutorProcessUtils
		.newProcessSpecBuilder(flinkConf)
		.withTotalProcessMemory(MemorySize.parse("1g"))
		.build();
	ContaineredTaskManagerParameters tmParams = new ContaineredTaskManagerParameters(spec, new HashMap<>(1));
	Configuration taskManagerConf = new Configuration();

	String workingDirectory = root.getAbsolutePath();
	Class<?> taskManagerMainClass = YarnTaskExecutorRunner.class;
	ContainerLaunchContext ctx;

	final Map<String, String> originalEnv = System.getenv();
	try {
		Map<String, String> systemEnv = new HashMap<>(originalEnv);
		systemEnv.put("HADOOP_TOKEN_FILE_LOCATION", credentialFile.getAbsolutePath());
		CommonTestUtils.setEnv(systemEnv);
		ctx = Utils.createTaskExecutorContext(flinkConf, yarnConf, env, tmParams,
			"", workingDirectory, taskManagerMainClass, LOG);
	} finally {
		CommonTestUtils.setEnv(originalEnv);
	}

	Credentials credentials = new Credentials();
	try (DataInputStream dis = new DataInputStream(new ByteArrayInputStream(ctx.getTokens().array()))) {
		credentials.readTokenStorageStream(dis);
	}
	Collection<Token<? extends TokenIdentifier>> tokens = credentials.getAllTokens();
	boolean hasHdfsDelegationToken = false;
	boolean hasAmRmToken = false;
	for (Token<? extends TokenIdentifier> token : tokens) {
		if (token.getKind().equals(amRmTokenKind)) {
			hasAmRmToken = true;
		} else if (token.getKind().equals(hdfsDelegationTokenKind)) {
			hasHdfsDelegationToken = true;
		}
	}
	assertTrue(hasHdfsDelegationToken);
	assertFalse(hasAmRmToken);
}
 

public ContaineredTaskManagerParameters getContaineredTaskManagerParameters() {
	return containeredTaskManagerParameters;
}
 

@Test
public void testCreateTaskExecutorCredentials() throws Exception {
	File root = temporaryFolder.getRoot();
	File home = new File(root, "home");
	boolean created = home.mkdir();
	assertTrue(created);

	Configuration flinkConf = new Configuration();
	YarnConfiguration yarnConf = new YarnConfiguration();

	Map<String, String> env = new HashMap<>();
	env.put(YarnConfigKeys.ENV_APP_ID, "foo");
	env.put(YarnConfigKeys.ENV_CLIENT_HOME_DIR, home.getAbsolutePath());
	env.put(YarnConfigKeys.ENV_CLIENT_SHIP_FILES, "");
	env.put(YarnConfigKeys.ENV_FLINK_CLASSPATH, "");
	env.put(YarnConfigKeys.ENV_HADOOP_USER_NAME, "foo");
	env.put(YarnConfigKeys.FLINK_JAR_PATH, root.toURI().toString());
	env = Collections.unmodifiableMap(env);

	File credentialFile = temporaryFolder.newFile("container_tokens");
	final Text amRmTokenKind = AMRMTokenIdentifier.KIND_NAME;
	final Text hdfsDelegationTokenKind = new Text("HDFS_DELEGATION_TOKEN");
	final Text service = new Text("test-service");
	Credentials amCredentials = new Credentials();
	amCredentials.addToken(amRmTokenKind, new Token<>(new byte[4], new byte[4], amRmTokenKind, service));
	amCredentials.addToken(hdfsDelegationTokenKind, new Token<>(new byte[4], new byte[4],
		hdfsDelegationTokenKind, service));
	amCredentials.writeTokenStorageFile(new org.apache.hadoop.fs.Path(credentialFile.getAbsolutePath()), yarnConf);

	ContaineredTaskManagerParameters tmParams = new ContaineredTaskManagerParameters(64,
		64, 16, 1, new HashMap<>(1));
	Configuration taskManagerConf = new Configuration();

	String workingDirectory = root.getAbsolutePath();
	Class<?> taskManagerMainClass = YarnTaskExecutorRunner.class;
	ContainerLaunchContext ctx;

	final Map<String, String> originalEnv = System.getenv();
	try {
		Map<String, String> systemEnv = new HashMap<>(originalEnv);
		systemEnv.put("HADOOP_TOKEN_FILE_LOCATION", credentialFile.getAbsolutePath());
		CommonTestUtils.setEnv(systemEnv);
		ctx = Utils.createTaskExecutorContext(flinkConf, yarnConf, env, tmParams,
			taskManagerConf, workingDirectory, taskManagerMainClass, LOG);
	} finally {
		CommonTestUtils.setEnv(originalEnv);
	}

	Credentials credentials = new Credentials();
	try (DataInputStream dis = new DataInputStream(new ByteArrayInputStream(ctx.getTokens().array()))) {
		credentials.readTokenStorageStream(dis);
	}
	Collection<Token<? extends TokenIdentifier>> tokens = credentials.getAllTokens();
	boolean hasHdfsDelegationToken = false;
	boolean hasAmRmToken = false;
	for (Token<? extends TokenIdentifier> token : tokens) {
		if (token.getKind().equals(amRmTokenKind)) {
			hasAmRmToken = true;
		} else if (token.getKind().equals(hdfsDelegationTokenKind)) {
			hasHdfsDelegationToken = true;
		}
	}
	assertTrue(hasHdfsDelegationToken);
	assertFalse(hasAmRmToken);
}
 

/**
 * Create mock RM dependencies.
 */
Context() throws Exception {
	rpcService = new TestingRpcService();
	fatalErrorHandler = new TestingFatalErrorHandler();
	rmServices = new MockMesosResourceManagerRuntimeServices();
	mesosServices = new MockMesosServices();

	// TaskExecutor templating
	ContainerSpecification containerSpecification = new ContainerSpecification();
	ContaineredTaskManagerParameters containeredParams =
		new ContaineredTaskManagerParameters(1024, 768, 256, 4, new HashMap<String, String>());
	MesosTaskManagerParameters tmParams = new MesosTaskManagerParameters(
		1.0, 1, 0, MesosTaskManagerParameters.ContainerType.MESOS, Option.<String>empty(), containeredParams,
		Collections.<Protos.Volume>emptyList(), Collections.<Protos.Parameter>emptyList(), false,
		Collections.<ConstraintEvaluator>emptyList(), "", Option.<String>empty(),
		Option.<String>empty(), Collections.<String>emptyList());

	// resource manager
	rmResourceID = ResourceID.generate();
	resourceManager =
		new TestingMesosResourceManager(
			rpcService,
			RM_ADDRESS,
			rmResourceID,
			rmServices.highAvailabilityServices,
			rmServices.heartbeatServices,
			rmServices.slotManager,
			rmServices.metricRegistry,
			rmServices.jobLeaderIdService,
			fatalErrorHandler,
			// Mesos specifics
			flinkConfig,
			mesosServices,
			rmServices.mesosConfig,
			tmParams,
			containerSpecification,
			UnregisteredMetricGroups.createUnregisteredJobManagerMetricGroup());

	// TaskExecutors
	task1Executor = mockTaskExecutor(task1);
	task2Executor = mockTaskExecutor(task2);
	task3Executor = mockTaskExecutor(task3);

	// JobMaster
	jobMaster1 = mockJobMaster(rmServices, new JobID(1, 0));
}
 

/**
 * Create the Mesos TaskManager parameters.
 *
 * @param flinkConfig the TM configuration.
 */
public static MesosTaskManagerParameters create(Configuration flinkConfig) {

	List<ConstraintEvaluator> constraints = parseConstraints(flinkConfig.getString(MESOS_CONSTRAINTS_HARD_HOSTATTR));
	// parse the common parameters
	ContaineredTaskManagerParameters containeredParameters = ContaineredTaskManagerParameters.create(
		flinkConfig,
		flinkConfig.getInteger(MESOS_RM_TASKS_MEMORY_MB),
		flinkConfig.getInteger(MESOS_RM_TASKS_SLOTS));

	double cpus = flinkConfig.getDouble(MESOS_RM_TASKS_CPUS);
	if (cpus <= 0.0) {
		cpus = Math.max(containeredParameters.numSlots(), 1.0);
	}

	int gpus = flinkConfig.getInteger(MESOS_RM_TASKS_GPUS);

	if (gpus < 0) {
		throw new IllegalConfigurationException(MESOS_RM_TASKS_GPUS.key() +
			" cannot be negative");
	}

	int disk = flinkConfig.getInteger(MESOS_RM_TASKS_DISK_MB);

	// parse the containerization parameters
	String imageName = flinkConfig.getString(MESOS_RM_CONTAINER_IMAGE_NAME);

	ContainerType containerType;
	String containerTypeString = flinkConfig.getString(MESOS_RM_CONTAINER_TYPE);
	switch (containerTypeString) {
		case MESOS_RESOURCEMANAGER_TASKS_CONTAINER_TYPE_MESOS:
			containerType = ContainerType.MESOS;
			break;
		case MESOS_RESOURCEMANAGER_TASKS_CONTAINER_TYPE_DOCKER:
			containerType = ContainerType.DOCKER;
			if (imageName == null || imageName.length() == 0) {
				throw new IllegalConfigurationException(MESOS_RM_CONTAINER_IMAGE_NAME.key() +
					" must be specified for docker container type");
			}
			break;
		default:
			throw new IllegalConfigurationException("invalid container type: " + containerTypeString);
	}

	Option<String> containerVolOpt = Option.<String>apply(flinkConfig.getString(MESOS_RM_CONTAINER_VOLUMES));

	Option<String> dockerParamsOpt = Option.<String>apply(flinkConfig.getString(MESOS_RM_CONTAINER_DOCKER_PARAMETERS));

	Option<String> uriParamsOpt = Option.<String>apply(flinkConfig.getString(MESOS_TM_URIS));

	boolean dockerForcePullImage = flinkConfig.getBoolean(MESOS_RM_CONTAINER_DOCKER_FORCE_PULL_IMAGE);

	List<Protos.Volume> containerVolumes = buildVolumes(containerVolOpt);

	List<Protos.Parameter> dockerParameters = buildDockerParameters(dockerParamsOpt);

	List<String> uris = buildUris(uriParamsOpt);

	//obtain Task Manager Host Name from the configuration
	Option<String> taskManagerHostname = Option.apply(flinkConfig.getString(MESOS_TM_HOSTNAME));

	//obtain command-line from the configuration
	String tmCommand = flinkConfig.getString(MESOS_TM_CMD);
	Option<String> tmBootstrapCommand = Option.apply(flinkConfig.getString(MESOS_TM_BOOTSTRAP_CMD));

	return new MesosTaskManagerParameters(
		cpus,
		gpus,
		disk,
		containerType,
		Option.apply(imageName),
		containeredParameters,
		containerVolumes,
		dockerParameters,
		dockerForcePullImage,
		constraints,
		tmCommand,
		tmBootstrapCommand,
		taskManagerHostname,
		uris);
}
 

/**
 * Get the common containered parameters.
 */
public ContaineredTaskManagerParameters containeredParameters() {
	return containeredParameters;
}
 

@Test
public void testCreateTaskExecutorCredentials() throws Exception {
	File root = temporaryFolder.getRoot();
	File home = new File(root, "home");
	boolean created = home.mkdir();
	assertTrue(created);

	Configuration flinkConf = new Configuration();
	YarnConfiguration yarnConf = new YarnConfiguration();

	Map<String, String> env = new HashMap<>();
	env.put(YarnConfigKeys.ENV_APP_ID, "foo");
	env.put(YarnConfigKeys.ENV_CLIENT_HOME_DIR, home.getAbsolutePath());
	env.put(YarnConfigKeys.ENV_CLIENT_SHIP_FILES, "");
	env.put(YarnConfigKeys.ENV_FLINK_CLASSPATH, "");
	env.put(YarnConfigKeys.ENV_HADOOP_USER_NAME, "foo");
	env.put(YarnConfigKeys.FLINK_JAR_PATH, root.toURI().toString());
	env = Collections.unmodifiableMap(env);

	File credentialFile = temporaryFolder.newFile("container_tokens");
	final Text amRmTokenKind = AMRMTokenIdentifier.KIND_NAME;
	final Text hdfsDelegationTokenKind = new Text("HDFS_DELEGATION_TOKEN");
	final Text service = new Text("test-service");
	Credentials amCredentials = new Credentials();
	amCredentials.addToken(amRmTokenKind, new Token<>(new byte[4], new byte[4], amRmTokenKind, service));
	amCredentials.addToken(hdfsDelegationTokenKind, new Token<>(new byte[4], new byte[4],
		hdfsDelegationTokenKind, service));
	amCredentials.writeTokenStorageFile(new org.apache.hadoop.fs.Path(credentialFile.getAbsolutePath()), yarnConf);

	ContaineredTaskManagerParameters tmParams = new ContaineredTaskManagerParameters(64,
		64, 16, 1, new HashMap<>(1));
	Configuration taskManagerConf = new Configuration();

	String workingDirectory = root.getAbsolutePath();
	Class<?> taskManagerMainClass = YarnTaskExecutorRunner.class;
	ContainerLaunchContext ctx;

	final Map<String, String> originalEnv = System.getenv();
	try {
		Map<String, String> systemEnv = new HashMap<>(originalEnv);
		systemEnv.put("HADOOP_TOKEN_FILE_LOCATION", credentialFile.getAbsolutePath());
		CommonTestUtils.setEnv(systemEnv);
		ctx = Utils.createTaskExecutorContext(flinkConf, yarnConf, env, tmParams,
			taskManagerConf, workingDirectory, taskManagerMainClass, LOG);
	} finally {
		CommonTestUtils.setEnv(originalEnv);
	}

	Credentials credentials = new Credentials();
	try (DataInputStream dis = new DataInputStream(new ByteArrayInputStream(ctx.getTokens().array()))) {
		credentials.readTokenStorageStream(dis);
	}
	Collection<Token<? extends TokenIdentifier>> tokens = credentials.getAllTokens();
	boolean hasHdfsDelegationToken = false;
	boolean hasAmRmToken = false;
	for (Token<? extends TokenIdentifier> token : tokens) {
		if (token.getKind().equals(amRmTokenKind)) {
			hasAmRmToken = true;
		} else if (token.getKind().equals(hdfsDelegationTokenKind)) {
			hasHdfsDelegationToken = true;
		}
	}
	assertTrue(hasHdfsDelegationToken);
	assertFalse(hasAmRmToken);
}
 

/**
 * Create mock RM dependencies.
 */
Context() throws Exception {
	rpcService = new TestingRpcService();
	fatalErrorHandler = new TestingFatalErrorHandler();
	rmServices = new MockMesosResourceManagerRuntimeServices();
	mesosServices = new MockMesosServices();

	// TaskExecutor templating
	ContainerSpecification containerSpecification = new ContainerSpecification();
	ContaineredTaskManagerParameters containeredParams =
		new ContaineredTaskManagerParameters(1024, 768, 256, 4, new HashMap<String, String>());
	MesosTaskManagerParameters tmParams = new MesosTaskManagerParameters(
		1.0, 1, MesosTaskManagerParameters.ContainerType.MESOS, Option.<String>empty(), containeredParams,
		Collections.<Protos.Volume>emptyList(), Collections.<Protos.Parameter>emptyList(), false,
		Collections.<ConstraintEvaluator>emptyList(), "", Option.<String>empty(),
		Option.<String>empty(), Collections.<String>emptyList());

	// resource manager
	rmResourceID = ResourceID.generate();
	resourceManager =
		new TestingMesosResourceManager(
			rpcService,
			RM_ADDRESS,
			rmResourceID,
			rmServices.highAvailabilityServices,
			rmServices.heartbeatServices,
			rmServices.slotManager,
			rmServices.metricRegistry,
			rmServices.jobLeaderIdService,
			fatalErrorHandler,
			// Mesos specifics
			flinkConfig,
			mesosServices,
			rmServices.mesosConfig,
			tmParams,
			containerSpecification,
			UnregisteredMetricGroups.createUnregisteredJobManagerMetricGroup());

	// TaskExecutors
	task1Executor = mockTaskExecutor(task1);
	task2Executor = mockTaskExecutor(task2);
	task3Executor = mockTaskExecutor(task3);

	// JobMaster
	jobMaster1 = mockJobMaster(rmServices, new JobID(1, 0));
}
 

public MesosResourceManager(
		// base class
		RpcService rpcService,
		String resourceManagerEndpointId,
		ResourceID resourceId,
		HighAvailabilityServices highAvailabilityServices,
		HeartbeatServices heartbeatServices,
		SlotManager slotManager,
		MetricRegistry metricRegistry,
		JobLeaderIdService jobLeaderIdService,
		ClusterInformation clusterInformation,
		FatalErrorHandler fatalErrorHandler,
		// Mesos specifics
		Configuration flinkConfig,
		MesosServices mesosServices,
		MesosConfiguration mesosConfig,
		MesosTaskManagerParameters taskManagerParameters,
		ContainerSpecification taskManagerContainerSpec,
		@Nullable String webUiUrl,
		JobManagerMetricGroup jobManagerMetricGroup) {
	super(
		rpcService,
		resourceManagerEndpointId,
		resourceId,
		highAvailabilityServices,
		heartbeatServices,
		slotManager,
		metricRegistry,
		jobLeaderIdService,
		clusterInformation,
		fatalErrorHandler,
		jobManagerMetricGroup);

	this.mesosServices = Preconditions.checkNotNull(mesosServices);
	this.actorSystem = Preconditions.checkNotNull(mesosServices.getLocalActorSystem());

	this.flinkConfig = Preconditions.checkNotNull(flinkConfig);
	this.mesosConfig = Preconditions.checkNotNull(mesosConfig);

	this.artifactServer = Preconditions.checkNotNull(mesosServices.getArtifactServer());

	this.taskManagerParameters = Preconditions.checkNotNull(taskManagerParameters);
	this.taskManagerContainerSpec = Preconditions.checkNotNull(taskManagerContainerSpec);
	this.webUiUrl = webUiUrl;

	this.workersInNew = new HashMap<>(8);
	this.workersInLaunch = new HashMap<>(8);
	this.workersBeingReturned = new HashMap<>(8);

	final ContaineredTaskManagerParameters containeredTaskManagerParameters = taskManagerParameters.containeredParameters();
	this.slotsPerWorker = createSlotsPerWorker(containeredTaskManagerParameters.numSlots());
}
 

/**
 * Create the Mesos TaskManager parameters.
 *
 * @param flinkConfig the TM configuration.
 */
public static MesosTaskManagerParameters create(Configuration flinkConfig) {

	List<ConstraintEvaluator> constraints = parseConstraints(flinkConfig.getString(MESOS_CONSTRAINTS_HARD_HOSTATTR));
	// parse the common parameters
	ContaineredTaskManagerParameters containeredParameters = ContaineredTaskManagerParameters.create(
		flinkConfig,
		flinkConfig.getInteger(MESOS_RM_TASKS_MEMORY_MB),
		flinkConfig.getInteger(MESOS_RM_TASKS_SLOTS));

	double cpus = flinkConfig.getDouble(MESOS_RM_TASKS_CPUS);
	if (cpus <= 0.0) {
		cpus = Math.max(containeredParameters.numSlots(), 1.0);
	}

	int gpus = flinkConfig.getInteger(MESOS_RM_TASKS_GPUS);

	if (gpus < 0) {
		throw new IllegalConfigurationException(MESOS_RM_TASKS_GPUS.key() +
			" cannot be negative");
	}

	// parse the containerization parameters
	String imageName = flinkConfig.getString(MESOS_RM_CONTAINER_IMAGE_NAME);

	ContainerType containerType;
	String containerTypeString = flinkConfig.getString(MESOS_RM_CONTAINER_TYPE);
	switch (containerTypeString) {
		case MESOS_RESOURCEMANAGER_TASKS_CONTAINER_TYPE_MESOS:
			containerType = ContainerType.MESOS;
			break;
		case MESOS_RESOURCEMANAGER_TASKS_CONTAINER_TYPE_DOCKER:
			containerType = ContainerType.DOCKER;
			if (imageName == null || imageName.length() == 0) {
				throw new IllegalConfigurationException(MESOS_RM_CONTAINER_IMAGE_NAME.key() +
					" must be specified for docker container type");
			}
			break;
		default:
			throw new IllegalConfigurationException("invalid container type: " + containerTypeString);
	}

	Option<String> containerVolOpt = Option.<String>apply(flinkConfig.getString(MESOS_RM_CONTAINER_VOLUMES));

	Option<String> dockerParamsOpt = Option.<String>apply(flinkConfig.getString(MESOS_RM_CONTAINER_DOCKER_PARAMETERS));

	Option<String> uriParamsOpt = Option.<String>apply(flinkConfig.getString(MESOS_TM_URIS));

	boolean dockerForcePullImage = flinkConfig.getBoolean(MESOS_RM_CONTAINER_DOCKER_FORCE_PULL_IMAGE);

	List<Protos.Volume> containerVolumes = buildVolumes(containerVolOpt);

	List<Protos.Parameter> dockerParameters = buildDockerParameters(dockerParamsOpt);

	List<String> uris = buildUris(uriParamsOpt);

	//obtain Task Manager Host Name from the configuration
	Option<String> taskManagerHostname = Option.apply(flinkConfig.getString(MESOS_TM_HOSTNAME));

	//obtain command-line from the configuration
	String tmCommand = flinkConfig.getString(MESOS_TM_CMD);
	Option<String> tmBootstrapCommand = Option.apply(flinkConfig.getString(MESOS_TM_BOOTSTRAP_CMD));

	return new MesosTaskManagerParameters(
		cpus,
		gpus,
		containerType,
		Option.apply(imageName),
		containeredParameters,
		containerVolumes,
		dockerParameters,
		dockerForcePullImage,
		constraints,
		tmCommand,
		tmBootstrapCommand,
		taskManagerHostname,
		uris);
}
 

/**
 * Get the common containered parameters.
 */
public ContaineredTaskManagerParameters containeredParameters() {
	return containeredParameters;
}