org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup Java Examples

private ExecutionGraph createExecutionGraphWithCoLocationConstraint() throws Exception {
	JobVertex[] jobVertices = new JobVertex[2];
	int parallelism = 3;
	jobVertices[0] = createNoOpVertex("v1", parallelism);
	jobVertices[1] = createNoOpVertex("v2", parallelism);
	jobVertices[1].connectNewDataSetAsInput(jobVertices[0], ALL_TO_ALL, PIPELINED);

	SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
	jobVertices[0].setSlotSharingGroup(slotSharingGroup);
	jobVertices[1].setSlotSharingGroup(slotSharingGroup);

	CoLocationGroup coLocationGroup = new CoLocationGroup();
	coLocationGroup.addVertex(jobVertices[0]);
	coLocationGroup.addVertex(jobVertices[1]);
	jobVertices[0].updateCoLocationGroup(coLocationGroup);
	jobVertices[1].updateCoLocationGroup(coLocationGroup);

	return createSimpleTestGraph(
		taskManagerGateway,
		triggeredRestartStrategy,
		jobVertices);
}
 

private JobGraph createJobGraph(final boolean containsCoLocationConstraint) {
	final JobVertex[] jobVertices = new JobVertex[3];
	final int parallelism = 3;
	jobVertices[0] = createNoOpVertex("v1", parallelism);
	jobVertices[1] = createNoOpVertex("v2", parallelism);
	jobVertices[2] = createNoOpVertex("v3", parallelism);
	jobVertices[1].connectNewDataSetAsInput(jobVertices[0], ALL_TO_ALL, PIPELINED);
	jobVertices[2].connectNewDataSetAsInput(jobVertices[1], ALL_TO_ALL, BLOCKING);

	if (containsCoLocationConstraint) {
		final SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
		jobVertices[0].setSlotSharingGroup(slotSharingGroup);
		jobVertices[1].setSlotSharingGroup(slotSharingGroup);

		final CoLocationGroup coLocationGroup = new CoLocationGroup();
		coLocationGroup.addVertex(jobVertices[0]);
		coLocationGroup.addVertex(jobVertices[1]);
		jobVertices[0].updateCoLocationGroup(coLocationGroup);
		jobVertices[1].updateCoLocationGroup(coLocationGroup);
	}

	return new JobGraph(jobVertices);
}
 

private void setSlotSharingAndCoLocation() {
	final HashMap<String, SlotSharingGroup> slotSharingGroups = new HashMap<>();
	final HashMap<String, Tuple2<SlotSharingGroup, CoLocationGroup>> coLocationGroups = new HashMap<>();

	for (Entry<Integer, JobVertex> entry : jobVertices.entrySet()) {

		final StreamNode node = streamGraph.getStreamNode(entry.getKey());
		final JobVertex vertex = entry.getValue();

		// configure slot sharing group
		final String slotSharingGroupKey = node.getSlotSharingGroup();
		final SlotSharingGroup sharingGroup;

		if (slotSharingGroupKey != null) {
			sharingGroup = slotSharingGroups.computeIfAbsent(
					slotSharingGroupKey, (k) -> new SlotSharingGroup());
			vertex.setSlotSharingGroup(sharingGroup);
		} else {
			sharingGroup = null;
		}

		// configure co-location constraint
		final String coLocationGroupKey = node.getCoLocationGroup();
		if (coLocationGroupKey != null) {
			if (sharingGroup == null) {
				throw new IllegalStateException("Cannot use a co-location constraint without a slot sharing group");
			}

			Tuple2<SlotSharingGroup, CoLocationGroup> constraint = coLocationGroups.computeIfAbsent(
					coLocationGroupKey, (k) -> new Tuple2<>(sharingGroup, new CoLocationGroup()));

			if (constraint.f0 != sharingGroup) {
				throw new IllegalStateException("Cannot co-locate operators from different slot sharing groups");
			}

			vertex.updateCoLocationGroup(constraint.f1);
			constraint.f1.addVertex(vertex);
		}
	}
}
 

private void setCoLocation() {
	final Map<String, Tuple2<SlotSharingGroup, CoLocationGroup>> coLocationGroups = new HashMap<>();

	for (Entry<Integer, JobVertex> entry : jobVertices.entrySet()) {

		final StreamNode node = streamGraph.getStreamNode(entry.getKey());
		final JobVertex vertex = entry.getValue();
		final SlotSharingGroup sharingGroup = vertex.getSlotSharingGroup();

		// configure co-location constraint
		final String coLocationGroupKey = node.getCoLocationGroup();
		if (coLocationGroupKey != null) {
			if (sharingGroup == null) {
				throw new IllegalStateException("Cannot use a co-location constraint without a slot sharing group");
			}

			Tuple2<SlotSharingGroup, CoLocationGroup> constraint = coLocationGroups.computeIfAbsent(
					coLocationGroupKey, k -> new Tuple2<>(sharingGroup, new CoLocationGroup()));

			if (constraint.f0 != sharingGroup) {
				throw new IllegalStateException("Cannot co-locate operators from different slot sharing groups");
			}

			vertex.updateCoLocationGroup(constraint.f1);
			constraint.f1.addVertex(vertex);
		}
	}
}
 

private ExecutionGraph createExecutionGraph(boolean addCollocationConstraints) throws Exception {
	JobVertex[] jobVertices = new JobVertex[3];
	int parallelism = 3;
	jobVertices[0] = createNoOpVertex("v1", parallelism);
	jobVertices[1] = createNoOpVertex("v2", parallelism);
	jobVertices[2] = createNoOpVertex("v3", parallelism);
	jobVertices[1].connectNewDataSetAsInput(jobVertices[0], ALL_TO_ALL, BLOCKING);
	jobVertices[2].connectNewDataSetAsInput(jobVertices[1], POINTWISE, PIPELINED);

	if (addCollocationConstraints) {
		SlotSharingGroup slotSharingGroup = new SlotSharingGroup();
		jobVertices[1].setSlotSharingGroup(slotSharingGroup);
		jobVertices[2].setSlotSharingGroup(slotSharingGroup);

		CoLocationGroup coLocationGroup = new CoLocationGroup();
		coLocationGroup.addVertex(jobVertices[1]);
		coLocationGroup.addVertex(jobVertices[2]);
		jobVertices[1].updateCoLocationGroup(coLocationGroup);
		jobVertices[2].updateCoLocationGroup(coLocationGroup);
	}

	return createSimpleTestGraph(
		new JobID(),
		taskManagerGateway,
		triggeredRestartStrategy,
		jobVertices);
}
 

@Test(expected = IllegalStateException.class)
public void testCoLocationConstraintThrowsException() {
	final ExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();

	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Collections.singletonList(
		new ExecutionVertexSchedulingRequirements.Builder()
			.withExecutionVertexId(new ExecutionVertexID(new JobVertexID(), 0))
			.withCoLocationConstraint(coLocationConstraint)
			.build()
	);

	executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
}
 

/**
 * Tests that validate the parameters when calling allocateSlot in SlotProvider.
 */
@Test
public void testAllocateSlotsParameters() {
	final ExecutionVertexID executionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
	final AllocationID allocationId = new AllocationID();
	final SlotSharingGroupId sharingGroupId = new SlotSharingGroupId();
	final ResourceProfile resourceProfile = new ResourceProfile(0.5, 250);
	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final Collection<TaskManagerLocation> taskManagerLocations = Collections.singleton(new LocalTaskManagerLocation());

	final DefaultExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator();

	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Arrays.asList(
			new ExecutionVertexSchedulingRequirements.Builder()
					.withExecutionVertexId(executionVertexId)
					.withPreviousAllocationId(allocationId)
					.withSlotSharingGroupId(sharingGroupId)
					.withPreferredLocations(taskManagerLocations)
					.withResourceProfile(resourceProfile)
					.withCoLocationConstraint(coLocationConstraint)
					.build()
	);

	executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
	assertThat(slotProvider.getSlotAllocationRequests(), hasSize(1));

	ScheduledUnit expectedTask = slotProvider.getSlotAllocationRequests().get(0).f1;
	SlotProfile expectedSlotProfile = slotProvider.getSlotAllocationRequests().get(0).f2;

	assertEquals(sharingGroupId, expectedTask.getSlotSharingGroupId());
	assertEquals(coLocationConstraint, expectedTask.getCoLocationConstraint());
	assertThat(expectedSlotProfile.getPreferredAllocations(), contains(allocationId));
	assertThat(expectedSlotProfile.getPreviousExecutionGraphAllocations(), contains(allocationId));
	assertEquals(resourceProfile, expectedSlotProfile.getResourceProfile());
	assertThat(expectedSlotProfile.getPreferredLocations(), contains(taskManagerLocations.toArray()));
}
 

/**
 * Tells this vertex to strictly co locate its subtasks with the subtasks of the given vertex.
 * Strict co-location implies that the n'th subtask of this vertex will run on the same parallel computing
 * instance (TaskManager) as the n'th subtask of the given vertex.
 * 
 * NOTE: Co-location is only possible between vertices in a slot sharing group.
 * 
 * NOTE: This vertex must (transitively) depend on the vertex to be co-located with. That means that the
 * respective vertex must be a (transitive) input of this vertex.
 * 
 * @param strictlyCoLocatedWith The vertex whose subtasks to co-locate this vertex's subtasks with.
 * 
 * @throws IllegalArgumentException Thrown, if this vertex and the vertex to co-locate with are not in a common
 *                                  slot sharing group.
 * 
 * @see #setSlotSharingGroup(SlotSharingGroup)
 */
public void setStrictlyCoLocatedWith(JobVertex strictlyCoLocatedWith) {
	if (this.slotSharingGroup == null || this.slotSharingGroup != strictlyCoLocatedWith.slotSharingGroup) {
		throw new IllegalArgumentException("Strict co-location requires that both vertices are in the same slot sharing group.");
	}

	CoLocationGroup thisGroup = this.coLocationGroup;
	CoLocationGroup otherGroup = strictlyCoLocatedWith.coLocationGroup;

	if (otherGroup == null) {
		if (thisGroup == null) {
			CoLocationGroup group = new CoLocationGroup(this, strictlyCoLocatedWith);
			this.coLocationGroup = group;
			strictlyCoLocatedWith.coLocationGroup = group;
		}
		else {
			thisGroup.addVertex(strictlyCoLocatedWith);
			strictlyCoLocatedWith.coLocationGroup = thisGroup;
		}
	}
	else {
		if (thisGroup == null) {
			otherGroup.addVertex(this);
			this.coLocationGroup = otherGroup;
		}
		else {
			// both had yet distinct groups, we need to merge them
			thisGroup.mergeInto(otherGroup);
		}
	}
}
 

@Test
public void testImmediateReleaseTwoLevel() {
	try {
		JobVertexID vid = new JobVertexID();
		JobVertex vertex = new JobVertex("vertex", vid);
		
		SlotSharingGroup sharingGroup = new SlotSharingGroup(vid);
		SlotSharingGroupAssignment assignment = sharingGroup.getTaskAssignment();

		CoLocationGroup coLocationGroup = new CoLocationGroup(vertex);
		CoLocationConstraint constraint = coLocationGroup.getLocationConstraint(0);
		
		Instance instance = SchedulerTestUtils.getRandomInstance(1);
		
		SharedSlot sharedSlot = instance.allocateSharedSlot(assignment);

		SimpleSlot sub = assignment.addSharedSlotAndAllocateSubSlot(sharedSlot, Locality.UNCONSTRAINED, constraint);
		
		assertNull(sub.getGroupID());
		assertEquals(constraint.getSharedSlot(), sub.getParent());
		
		sub.releaseSlot();

		assertTrue(sub.isReleased());
		assertTrue(sharedSlot.isReleased());

		assertEquals(1, instance.getNumberOfAvailableSlots());
		assertEquals(0, instance.getNumberOfAllocatedSlots());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tells this vertex to strictly co locate its subtasks with the subtasks of the given vertex.
 * Strict co-location implies that the n'th subtask of this vertex will run on the same parallel computing
 * instance (TaskManager) as the n'th subtask of the given vertex.
 * 
 * NOTE: Co-location is only possible between vertices in a slot sharing group.
 * 
 * NOTE: This vertex must (transitively) depend on the vertex to be co-located with. That means that the
 * respective vertex must be a (transitive) input of this vertex.
 * 
 * @param strictlyCoLocatedWith The vertex whose subtasks to co-locate this vertex's subtasks with.
 * 
 * @throws IllegalArgumentException Thrown, if this vertex and the vertex to co-locate with are not in a common
 *                                  slot sharing group.
 * 
 * @see #setSlotSharingGroup(SlotSharingGroup)
 */
public void setStrictlyCoLocatedWith(JobVertex strictlyCoLocatedWith) {
	if (this.slotSharingGroup == null || this.slotSharingGroup != strictlyCoLocatedWith.slotSharingGroup) {
		throw new IllegalArgumentException("Strict co-location requires that both vertices are in the same slot sharing group.");
	}

	CoLocationGroup thisGroup = this.coLocationGroup;
	CoLocationGroup otherGroup = strictlyCoLocatedWith.coLocationGroup;

	if (otherGroup == null) {
		if (thisGroup == null) {
			CoLocationGroup group = new CoLocationGroup(this, strictlyCoLocatedWith);
			this.coLocationGroup = group;
			strictlyCoLocatedWith.coLocationGroup = group;
		}
		else {
			thisGroup.addVertex(strictlyCoLocatedWith);
			strictlyCoLocatedWith.coLocationGroup = thisGroup;
		}
	}
	else {
		if (thisGroup == null) {
			otherGroup.addVertex(this);
			this.coLocationGroup = otherGroup;
		}
		else {
			// both had yet distinct groups, we need to merge them
			thisGroup.mergeInto(otherGroup);
		}
	}
}
 

private void resetTasks(final Set<ExecutionVertex> vertices, final long globalModVersion) throws Exception {
	final Set<CoLocationGroup> colGroups = new HashSet<>();
	final long restartTimestamp = System.currentTimeMillis();

	for (ExecutionVertex ev : vertices) {
		CoLocationGroup cgroup = ev.getJobVertex().getCoLocationGroup();
		if (cgroup != null && !colGroups.contains(cgroup)){
			cgroup.resetConstraints();
			colGroups.add(cgroup);
		}

		ev.resetForNewExecution(restartTimestamp, globalModVersion);
	}

	// if there is checkpointed state, reload it into the executions
	if (executionGraph.getCheckpointCoordinator() != null) {
		// abort pending checkpoints to
		// i) enable new checkpoint triggering without waiting for last checkpoint expired.
		// ii) ensure the EXACTLY_ONCE semantics if needed.
		executionGraph.getCheckpointCoordinator().abortPendingCheckpoints(
			new CheckpointException(CheckpointFailureReason.JOB_FAILOVER_REGION));

		final Map<JobVertexID, ExecutionJobVertex> involvedExecutionJobVertices =
			getInvolvedExecutionJobVertices(vertices);
		executionGraph.getCheckpointCoordinator().restoreLatestCheckpointedState(
			involvedExecutionJobVertices, false, true);
	}
}
 

protected void resetForNewExecutions(final Collection<ExecutionVertexID> vertices) {
	final Set<CoLocationGroup> colGroups = new HashSet<>();
	vertices.forEach(executionVertexId -> {
		final ExecutionVertex ev = getExecutionVertex(executionVertexId);

		final CoLocationGroup cgroup = ev.getJobVertex().getCoLocationGroup();
		if (cgroup != null && !colGroups.contains(cgroup)){
			cgroup.resetConstraints();
			colGroups.add(cgroup);
		}

		ev.resetForNewExecution();
	});

}
 

/**
 * Tells this vertex to strictly co locate its subtasks with the subtasks of the given vertex.
 * Strict co-location implies that the n'th subtask of this vertex will run on the same parallel computing
 * instance (TaskManager) as the n'th subtask of the given vertex.
 *
 * <p>NOTE: Co-location is only possible between vertices in a slot sharing group.
 *
 * <p>NOTE: This vertex must (transitively) depend on the vertex to be co-located with. That means that the
 * respective vertex must be a (transitive) input of this vertex.
 *
 * @param strictlyCoLocatedWith The vertex whose subtasks to co-locate this vertex's subtasks with.
 *
 * @throws IllegalArgumentException Thrown, if this vertex and the vertex to co-locate with are not in a common
 *                                  slot sharing group.
 *
 * @see #setSlotSharingGroup(SlotSharingGroup)
 */
public void setStrictlyCoLocatedWith(JobVertex strictlyCoLocatedWith) {
	if (this.slotSharingGroup == null || this.slotSharingGroup != strictlyCoLocatedWith.slotSharingGroup) {
		throw new IllegalArgumentException("Strict co-location requires that both vertices are in the same slot sharing group.");
	}

	CoLocationGroup thisGroup = this.coLocationGroup;
	CoLocationGroup otherGroup = strictlyCoLocatedWith.coLocationGroup;

	if (otherGroup == null) {
		if (thisGroup == null) {
			CoLocationGroup group = new CoLocationGroup(this, strictlyCoLocatedWith);
			this.coLocationGroup = group;
			strictlyCoLocatedWith.coLocationGroup = group;
		}
		else {
			thisGroup.addVertex(strictlyCoLocatedWith);
			strictlyCoLocatedWith.coLocationGroup = thisGroup;
		}
	}
	else {
		if (thisGroup == null) {
			otherGroup.addVertex(this);
			this.coLocationGroup = otherGroup;
		}
		else {
			// both had yet distinct groups, we need to merge them
			thisGroup.mergeInto(otherGroup);
		}
	}
}
 

/**
 * Creates an ExecutionVertex.
 *
 * @param timeout
 *            The RPC timeout to use for deploy / cancel calls
 * @param initialGlobalModVersion
 *            The global modification version to initialize the first Execution with.
 * @param createTimestamp
 *            The timestamp for the vertex creation, used to initialize the first Execution with.
 * @param maxPriorExecutionHistoryLength
 *            The number of prior Executions (= execution attempts) to keep.
 */
public ExecutionVertex(
		ExecutionJobVertex jobVertex,
		int subTaskIndex,
		IntermediateResult[] producedDataSets,
		Time timeout,
		long initialGlobalModVersion,
		long createTimestamp,
		int maxPriorExecutionHistoryLength) {

	this.jobVertex = jobVertex;
	this.subTaskIndex = subTaskIndex;
	this.executionVertexId = new ExecutionVertexID(jobVertex.getJobVertexId(), subTaskIndex);
	this.taskNameWithSubtask = String.format("%s (%d/%d)",
			jobVertex.getJobVertex().getName(), subTaskIndex + 1, jobVertex.getParallelism());

	this.resultPartitions = new LinkedHashMap<>(producedDataSets.length, 1);

	for (IntermediateResult result : producedDataSets) {
		IntermediateResultPartition irp = new IntermediateResultPartition(result, this, subTaskIndex);
		result.setPartition(subTaskIndex, irp);

		resultPartitions.put(irp.getPartitionId(), irp);
	}

	this.inputEdges = new ExecutionEdge[jobVertex.getJobVertex().getInputs().size()][];

	this.priorExecutions = new EvictingBoundedList<>(maxPriorExecutionHistoryLength);

	this.currentExecution = new Execution(
		getExecutionGraph().getFutureExecutor(),
		this,
		0,
		initialGlobalModVersion,
		createTimestamp,
		timeout);

	// create a co-location scheduling hint, if necessary
	CoLocationGroup clg = jobVertex.getCoLocationGroup();
	if (clg != null) {
		this.locationConstraint = clg.getLocationConstraint(subTaskIndex);
	}
	else {
		this.locationConstraint = null;
	}

	getExecutionGraph().registerExecution(currentExecution);

	this.timeout = timeout;
	this.inputSplits = new ArrayList<>();
}
 

/**
 * Test iteration job, check slot sharing group and co-location group.
 */
@Test
public void testIteration() {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	DataStream<Integer> source = env.fromElements(1, 2, 3).name("source");
	IterativeStream<Integer> iteration = source.iterate(3000);
	iteration.name("iteration").setParallelism(2);
	DataStream<Integer> map = iteration.map(x -> x + 1).name("map").setParallelism(2);
	DataStream<Integer> filter = map.filter((x) -> false).name("filter").setParallelism(2);
	iteration.closeWith(filter).print();

	JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());

	SlotSharingGroup slotSharingGroup = jobGraph.getVerticesAsArray()[0].getSlotSharingGroup();
	assertNotNull(slotSharingGroup);

	CoLocationGroup iterationSourceCoLocationGroup = null;
	CoLocationGroup iterationSinkCoLocationGroup = null;

	for (JobVertex jobVertex : jobGraph.getVertices()) {
		// all vertices have same slot sharing group by default
		assertEquals(slotSharingGroup, jobVertex.getSlotSharingGroup());

		// all iteration vertices have same co-location group,
		// others have no co-location group by default
		if (jobVertex.getName().startsWith(StreamGraph.ITERATION_SOURCE_NAME_PREFIX)) {
			iterationSourceCoLocationGroup = jobVertex.getCoLocationGroup();
			assertTrue(iterationSourceCoLocationGroup.getVertices().contains(jobVertex));
		} else if (jobVertex.getName().startsWith(StreamGraph.ITERATION_SINK_NAME_PREFIX)) {
			iterationSinkCoLocationGroup = jobVertex.getCoLocationGroup();
			assertTrue(iterationSinkCoLocationGroup.getVertices().contains(jobVertex));
		} else {
			assertNull(jobVertex.getCoLocationGroup());
		}
	}

	assertNotNull(iterationSourceCoLocationGroup);
	assertNotNull(iterationSinkCoLocationGroup);
	assertEquals(iterationSourceCoLocationGroup, iterationSinkCoLocationGroup);
}
 

/**
 * Test slot sharing group is enabled or disabled for iteration.
 */
@Test
public void testDisableSlotSharingForIteration() {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	DataStream<Integer> source = env.fromElements(1, 2, 3).name("source");
	IterativeStream<Integer> iteration = source.iterate(3000);
	iteration.name("iteration").setParallelism(2);
	DataStream<Integer> map = iteration.map(x -> x + 1).name("map").setParallelism(2);
	DataStream<Integer> filter = map.filter((x) -> false).name("filter").setParallelism(2);
	iteration.closeWith(filter).print();

	List<Transformation<?>> transformations = new ArrayList<>();
	transformations.add(source.getTransformation());
	transformations.add(iteration.getTransformation());
	transformations.add(map.getTransformation());
	transformations.add(filter.getTransformation());
	// when slot sharing group is disabled
	// all job vertices except iteration vertex would have no slot sharing group
	// iteration vertices would be set slot sharing group automatically
	StreamGraphGenerator generator = new StreamGraphGenerator(transformations, env.getConfig(), env.getCheckpointConfig());
	generator.setSlotSharingEnabled(false);

	JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(generator.generate());

	SlotSharingGroup iterationSourceSlotSharingGroup = null;
	SlotSharingGroup iterationSinkSlotSharingGroup = null;

	CoLocationGroup iterationSourceCoLocationGroup = null;
	CoLocationGroup iterationSinkCoLocationGroup = null;

	for (JobVertex jobVertex : jobGraph.getVertices()) {
		if (jobVertex.getName().startsWith(StreamGraph.ITERATION_SOURCE_NAME_PREFIX)) {
			iterationSourceSlotSharingGroup = jobVertex.getSlotSharingGroup();
			iterationSourceCoLocationGroup = jobVertex.getCoLocationGroup();
		} else if (jobVertex.getName().startsWith(StreamGraph.ITERATION_SINK_NAME_PREFIX)) {
			iterationSinkSlotSharingGroup = jobVertex.getSlotSharingGroup();
			iterationSinkCoLocationGroup = jobVertex.getCoLocationGroup();
		} else {
			assertNull(jobVertex.getSlotSharingGroup());
		}
	}

	assertNotNull(iterationSourceSlotSharingGroup);
	assertNotNull(iterationSinkSlotSharingGroup);
	assertEquals(iterationSourceSlotSharingGroup, iterationSinkSlotSharingGroup);

	assertNotNull(iterationSourceCoLocationGroup);
	assertNotNull(iterationSinkCoLocationGroup);
	assertEquals(iterationSourceCoLocationGroup, iterationSinkCoLocationGroup);
}
 

public CoLocationGroup getCoLocationGroup() {
	return coLocationGroup;
}
 

public void updateCoLocationGroup(CoLocationGroup group) {
	this.coLocationGroup = group;
}
 

public void restart(long expectedGlobalVersion) {

		assertRunningInJobMasterMainThread();

		try {
			// check the global version to see whether this recovery attempt is still valid
			if (globalModVersion != expectedGlobalVersion) {
				LOG.info("Concurrent full restart subsumed this restart.");
				return;
			}

			final JobStatus current = state;

			if (current == JobStatus.CANCELED) {
				LOG.info("Canceled job during restart. Aborting restart.");
				return;
			} else if (current == JobStatus.FAILED) {
				LOG.info("Failed job during restart. Aborting restart.");
				return;
			} else if (current == JobStatus.SUSPENDED) {
				LOG.info("Suspended job during restart. Aborting restart.");
				return;
			} else if (current != JobStatus.RESTARTING) {
				throw new IllegalStateException("Can only restart job from state restarting.");
			}

			this.currentExecutions.clear();

			final Collection<CoLocationGroup> colGroups = new HashSet<>();
			final long resetTimestamp = System.currentTimeMillis();

			for (ExecutionJobVertex jv : this.verticesInCreationOrder) {

				CoLocationGroup cgroup = jv.getCoLocationGroup();
				if (cgroup != null && !colGroups.contains(cgroup)){
					cgroup.resetConstraints();
					colGroups.add(cgroup);
				}

				jv.resetForNewExecution(resetTimestamp, expectedGlobalVersion);
			}

			for (int i = 0; i < stateTimestamps.length; i++) {
				if (i != JobStatus.RESTARTING.ordinal()) {
					// Only clear the non restarting state in order to preserve when the job was
					// restarted. This is needed for the restarting time gauge
					stateTimestamps[i] = 0;
				}
			}

			transitionState(JobStatus.RESTARTING, JobStatus.CREATED);

			// if we have checkpointed state, reload it into the executions
			if (checkpointCoordinator != null) {
				checkpointCoordinator.restoreLatestCheckpointedState(getAllVertices(), false, false);
			}

			scheduleForExecution();
		}
		// TODO remove the catch block if we align the schematics to not fail global within the restarter.
		catch (Throwable t) {
			LOG.warn("Failed to restart the job.", t);
			failGlobal(t);
		}
	}
 

@Test
public void testCoLocatedSlotRequestsFailAfterResolved() throws ExecutionException, InterruptedException {
	final ResourceProfile rp1 = new ResourceProfile(1.0, 100);
	final ResourceProfile rp2 = new ResourceProfile(2.0, 200);
	final ResourceProfile rp3 = new ResourceProfile(5.0, 500);

	final ResourceProfile allocatedSlotRp = new ResourceProfile(3.0, 300);

	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(1);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
			(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();
	JobVertexID jobVertexId3 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture1 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId1,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp1),
			TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture2 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId2,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp2),
			TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
			taskManagerLocation,
			new SimpleAckingTaskManagerGateway(),
			Collections.singletonList(new SlotOffer(
					allocationId1,
					0,
					allocatedSlotRp)));

	assertFalse(slotOfferFuture1.isEmpty());

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = slotProvider.allocateSlot(
			new ScheduledUnit(
					jobVertexId3,
					slotSharingGroupId,
					coLocationConstraint1),
			true,
			SlotProfile.noLocality(rp3),
			TestingUtils.infiniteTime());

	LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(allocationId1, logicalSlot1.getAllocationId());
	assertEquals(allocationId1, logicalSlot2.getAllocationId());

	assertTrue(logicalSlotFuture3.isDone() && logicalSlotFuture3.isCompletedExceptionally());
	logicalSlotFuture3.whenComplete((LogicalSlot ignored, Throwable throwable) -> {
		assertTrue(throwable instanceof SharedSlotOversubscribedException);
		assertTrue(((SharedSlotOversubscribedException) throwable).canRetry());
	});
}
 

private Execution resetForNewExecutionInternal(final long timestamp, final long originatingGlobalModVersion) {
	final Execution oldExecution = currentExecution;
	final ExecutionState oldState = oldExecution.getState();

	if (oldState.isTerminal()) {
		if (oldState == FINISHED) {
			// pipelined partitions are released in Execution#cancel(), covering both job failures and vertex resets
			// do not release pipelined partitions here to save RPC calls
			oldExecution.handlePartitionCleanup(false, true);
			getExecutionGraph().getPartitionReleaseStrategy().vertexUnfinished(executionVertexId);
		}

		priorExecutions.add(oldExecution.archive());

		final Execution newExecution = new Execution(
			getExecutionGraph().getFutureExecutor(),
			this,
			oldExecution.getAttemptNumber() + 1,
			originatingGlobalModVersion,
			timestamp,
			timeout);

		currentExecution = newExecution;

		synchronized (inputSplits) {
			InputSplitAssigner assigner = jobVertex.getSplitAssigner();
			if (assigner != null) {
				assigner.returnInputSplit(inputSplits, getParallelSubtaskIndex());
				inputSplits.clear();
			}
		}

		CoLocationGroup grp = jobVertex.getCoLocationGroup();
		if (grp != null) {
			locationConstraint = grp.getLocationConstraint(subTaskIndex);
		}

		// register this execution at the execution graph, to receive call backs
		getExecutionGraph().registerExecution(newExecution);

		// if the execution was 'FINISHED' before, tell the ExecutionGraph that
		// we take one step back on the road to reaching global FINISHED
		if (oldState == FINISHED) {
			getExecutionGraph().vertexUnFinished();
		}

		// reset the intermediate results
		for (IntermediateResultPartition resultPartition : resultPartitions.values()) {
			resultPartition.resetForNewExecution();
		}

		return newExecution;
	}
	else {
		throw new IllegalStateException("Cannot reset a vertex that is in non-terminal state " + oldState);
	}
}
 

public CoLocationGroup getCoLocationGroup() {
	return coLocationGroup;
}
 

/**
 * Tests the scheduling of two tasks with a parallelism of 2 and a co-location constraint.
 */
@Test
public void testSimpleCoLocatedSlotScheduling() throws ExecutionException, InterruptedException {
	final BlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(2);

	final TestingResourceManagerGateway testingResourceManagerGateway = slotPoolResource.getTestingResourceManagerGateway();

	testingResourceManagerGateway.setRequestSlotConsumer(
		(SlotRequest slotRequest) -> allocationIds.offer(slotRequest.getAllocationId()));

	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final SlotPool slotPoolGateway = slotPoolResource.getSlotPool();
	slotPoolGateway.registerTaskManager(taskManagerLocation.getResourceID());

	CoLocationGroup group = new CoLocationGroup();
	CoLocationConstraint coLocationConstraint1 = group.getLocationConstraint(0);
	CoLocationConstraint coLocationConstraint2 = group.getLocationConstraint(1);

	final SlotSharingGroupId slotSharingGroupId = new SlotSharingGroupId();

	JobVertexID jobVertexId1 = new JobVertexID();
	JobVertexID jobVertexId2 = new JobVertexID();

	final SlotProvider slotProvider = slotPoolResource.getSlotProvider();
	CompletableFuture<LogicalSlot> logicalSlotFuture11 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint1),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture22 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint2),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture12 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId2,
			slotSharingGroupId,
			coLocationConstraint1),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	CompletableFuture<LogicalSlot> logicalSlotFuture21 = slotProvider.allocateSlot(
		new ScheduledUnit(
			jobVertexId1,
			slotSharingGroupId,
			coLocationConstraint2),
		SlotProfile.noRequirements(),
		TestingUtils.infiniteTime());

	final AllocationID allocationId1 = allocationIds.take();
	final AllocationID allocationId2 = allocationIds.take();

	Collection<SlotOffer> slotOfferFuture1 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId1,
			0,
			ResourceProfile.ANY)));

	Collection<SlotOffer> slotOfferFuture2 = slotPoolGateway.offerSlots(
		taskManagerLocation,
		new SimpleAckingTaskManagerGateway(),
		Collections.singletonList(new SlotOffer(
			allocationId2,
			0,
			ResourceProfile.ANY)));

	assertFalse(slotOfferFuture1.isEmpty());
	assertFalse(slotOfferFuture2.isEmpty());

	LogicalSlot logicalSlot11 = logicalSlotFuture11.get();
	LogicalSlot logicalSlot12 = logicalSlotFuture12.get();
	LogicalSlot logicalSlot21 = logicalSlotFuture21.get();
	LogicalSlot logicalSlot22 = logicalSlotFuture22.get();

	assertEquals(logicalSlot11.getAllocationId(), logicalSlot12.getAllocationId());
	assertEquals(logicalSlot21.getAllocationId(), logicalSlot22.getAllocationId());
	assertNotEquals(logicalSlot11.getAllocationId(), logicalSlot21.getAllocationId());
}
 

/**
 * Tests that validate the parameters when calling allocateSlot in SlotProvider.
 */
@Test
public void testAllocateSlotsParameters() {
	final ExecutionVertexID executionVertexId = new ExecutionVertexID(new JobVertexID(), 0);
	final AllocationID allocationId = new AllocationID();
	final SlotSharingGroupId sharingGroupId = new SlotSharingGroupId();
	final ResourceProfile taskResourceProfile = ResourceProfile.fromResources(0.5, 250);
	final ResourceProfile physicalSlotResourceProfile = ResourceProfile.fromResources(1.0, 300);
	final CoLocationConstraint coLocationConstraint = new CoLocationGroup().getLocationConstraint(0);
	final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();

	final TestingStateLocationRetriever stateLocationRetriever = new TestingStateLocationRetriever();
	stateLocationRetriever.setStateLocation(executionVertexId, taskManagerLocation);

	final ExecutionSlotAllocator executionSlotAllocator = createExecutionSlotAllocator(
		stateLocationRetriever,
		new TestingInputsLocationsRetriever.Builder().build());

	final List<ExecutionVertexSchedulingRequirements> schedulingRequirements = Arrays.asList(
		new ExecutionVertexSchedulingRequirements.Builder()
			.withExecutionVertexId(executionVertexId)
			.withPreviousAllocationId(allocationId)
			.withSlotSharingGroupId(sharingGroupId)
			.withPhysicalSlotResourceProfile(physicalSlotResourceProfile)
			.withTaskResourceProfile(taskResourceProfile)
			.withCoLocationConstraint(coLocationConstraint)
			.build()
	);

	executionSlotAllocator.allocateSlotsFor(schedulingRequirements);
	assertThat(slotProvider.getSlotAllocationRequests(), hasSize(1));

	ScheduledUnit expectedTask = slotProvider.getSlotAllocationRequests().get(0).f1;
	SlotProfile expectedSlotProfile = slotProvider.getSlotAllocationRequests().get(0).f2;

	assertEquals(sharingGroupId, expectedTask.getSlotSharingGroupId());
	assertEquals(coLocationConstraint, expectedTask.getCoLocationConstraint());
	assertThat(expectedSlotProfile.getPreferredAllocations(), contains(allocationId));
	assertThat(expectedSlotProfile.getPreviousExecutionGraphAllocations(), contains(allocationId));
	assertEquals(taskResourceProfile, expectedSlotProfile.getTaskResourceProfile());
	assertEquals(physicalSlotResourceProfile, expectedSlotProfile.getPhysicalSlotResourceProfile());
	assertThat(expectedSlotProfile.getPreferredLocations(), contains(taskManagerLocation));
}
 

/**
 * Test iteration job, check slot sharing group and co-location group.
 */
@Test
public void testIteration() {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	DataStream<Integer> source = env.fromElements(1, 2, 3).name("source");
	IterativeStream<Integer> iteration = source.iterate(3000);
	iteration.name("iteration").setParallelism(2);
	DataStream<Integer> map = iteration.map(x -> x + 1).name("map").setParallelism(2);
	DataStream<Integer> filter = map.filter((x) -> false).name("filter").setParallelism(2);
	iteration.closeWith(filter).print();

	JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());

	SlotSharingGroup slotSharingGroup = jobGraph.getVerticesAsArray()[0].getSlotSharingGroup();
	assertNotNull(slotSharingGroup);

	CoLocationGroup iterationSourceCoLocationGroup = null;
	CoLocationGroup iterationSinkCoLocationGroup = null;

	for (JobVertex jobVertex : jobGraph.getVertices()) {
		// all vertices have same slot sharing group by default
		assertEquals(slotSharingGroup, jobVertex.getSlotSharingGroup());

		// all iteration vertices have same co-location group,
		// others have no co-location group by default
		if (jobVertex.getName().startsWith(StreamGraph.ITERATION_SOURCE_NAME_PREFIX)) {
			iterationSourceCoLocationGroup = jobVertex.getCoLocationGroup();
			assertTrue(iterationSourceCoLocationGroup.getVertices().contains(jobVertex));
		} else if (jobVertex.getName().startsWith(StreamGraph.ITERATION_SINK_NAME_PREFIX)) {
			iterationSinkCoLocationGroup = jobVertex.getCoLocationGroup();
			assertTrue(iterationSinkCoLocationGroup.getVertices().contains(jobVertex));
		} else {
			assertNull(jobVertex.getCoLocationGroup());
		}
	}

	assertNotNull(iterationSourceCoLocationGroup);
	assertNotNull(iterationSinkCoLocationGroup);
	assertEquals(iterationSourceCoLocationGroup, iterationSinkCoLocationGroup);
}
 

private void setSlotSharingAndCoLocation() {
	final HashMap<String, SlotSharingGroup> slotSharingGroups = new HashMap<>();
	final HashMap<String, Tuple2<SlotSharingGroup, CoLocationGroup>> coLocationGroups = new HashMap<>();

	for (Entry<Integer, JobVertex> entry : jobVertices.entrySet()) {

		final StreamNode node = streamGraph.getStreamNode(entry.getKey());
		final JobVertex vertex = entry.getValue();

		// configure slot sharing group
		final String slotSharingGroupKey = node.getSlotSharingGroup();
		final SlotSharingGroup sharingGroup;

		if (slotSharingGroupKey != null) {
			sharingGroup = slotSharingGroups.computeIfAbsent(
					slotSharingGroupKey, (k) -> new SlotSharingGroup());
			vertex.setSlotSharingGroup(sharingGroup);
		} else {
			sharingGroup = null;
		}

		// configure co-location constraint
		final String coLocationGroupKey = node.getCoLocationGroup();
		if (coLocationGroupKey != null) {
			if (sharingGroup == null) {
				throw new IllegalStateException("Cannot use a co-location constraint without a slot sharing group");
			}

			Tuple2<SlotSharingGroup, CoLocationGroup> constraint = coLocationGroups.computeIfAbsent(
					coLocationGroupKey, (k) -> new Tuple2<>(sharingGroup, new CoLocationGroup()));

			if (constraint.f0 != sharingGroup) {
				throw new IllegalStateException("Cannot co-locate operators from different slot sharing groups");
			}

			vertex.updateCoLocationGroup(constraint.f1);
		}
	}

	for (Tuple2<StreamNode, StreamNode> pair : streamGraph.getIterationSourceSinkPairs()) {

		CoLocationGroup ccg = new CoLocationGroup();

		JobVertex source = jobVertices.get(pair.f0.getId());
		JobVertex sink = jobVertices.get(pair.f1.getId());

		ccg.addVertex(source);
		ccg.addVertex(sink);
		source.updateCoLocationGroup(ccg);
		sink.updateCoLocationGroup(ccg);
	}

}
 

public CoLocationGroup getCoLocationGroup() {
	return coLocationGroup;
}
 

public void updateCoLocationGroup(CoLocationGroup group) {
	this.coLocationGroup = group;
}
 

public void restart(long expectedGlobalVersion) {

		assertRunningInJobMasterMainThread();

		try {
			synchronized (progressLock) {
				// check the global version to see whether this recovery attempt is still valid
				if (globalModVersion != expectedGlobalVersion) {
					LOG.info("Concurrent full restart subsumed this restart.");
					return;
				}

				final JobStatus current = state;

				if (current == JobStatus.CANCELED) {
					LOG.info("Canceled job during restart. Aborting restart.");
					return;
				} else if (current == JobStatus.FAILED) {
					LOG.info("Failed job during restart. Aborting restart.");
					return;
				} else if (current == JobStatus.SUSPENDED) {
					LOG.info("Suspended job during restart. Aborting restart.");
					return;
				} else if (current != JobStatus.RESTARTING) {
					throw new IllegalStateException("Can only restart job from state restarting.");
				}

				this.currentExecutions.clear();

				final Collection<CoLocationGroup> colGroups = new HashSet<>();
				final long resetTimestamp = System.currentTimeMillis();

				for (ExecutionJobVertex jv : this.verticesInCreationOrder) {

					CoLocationGroup cgroup = jv.getCoLocationGroup();
					if (cgroup != null && !colGroups.contains(cgroup)){
						cgroup.resetConstraints();
						colGroups.add(cgroup);
					}

					jv.resetForNewExecution(resetTimestamp, expectedGlobalVersion);
				}

				for (int i = 0; i < stateTimestamps.length; i++) {
					if (i != JobStatus.RESTARTING.ordinal()) {
						// Only clear the non restarting state in order to preserve when the job was
						// restarted. This is needed for the restarting time gauge
						stateTimestamps[i] = 0;
					}
				}

				transitionState(JobStatus.RESTARTING, JobStatus.CREATED);

				// if we have checkpointed state, reload it into the executions
				if (checkpointCoordinator != null) {
					checkpointCoordinator.restoreLatestCheckpointedState(getAllVertices(), false, false);
				}
			}

			scheduleForExecution();
		}
		catch (Throwable t) {
			LOG.warn("Failed to restart the job.", t);
			failGlobal(t);
		}
	}
 

/**
 * Creates an ExecutionVertex.
 *
 * @param timeout
 *            The RPC timeout to use for deploy / cancel calls
 * @param initialGlobalModVersion
 *            The global modification version to initialize the first Execution with.
 * @param createTimestamp
 *            The timestamp for the vertex creation, used to initialize the first Execution with.
 * @param maxPriorExecutionHistoryLength
 *            The number of prior Executions (= execution attempts) to keep.
 */
public ExecutionVertex(
		ExecutionJobVertex jobVertex,
		int subTaskIndex,
		IntermediateResult[] producedDataSets,
		Time timeout,
		long initialGlobalModVersion,
		long createTimestamp,
		int maxPriorExecutionHistoryLength) {

	this.jobVertex = jobVertex;
	this.subTaskIndex = subTaskIndex;
	this.taskNameWithSubtask = String.format("%s (%d/%d)",
			jobVertex.getJobVertex().getName(), subTaskIndex + 1, jobVertex.getParallelism());

	this.resultPartitions = new LinkedHashMap<>(producedDataSets.length, 1);

	for (IntermediateResult result : producedDataSets) {
		IntermediateResultPartition irp = new IntermediateResultPartition(result, this, subTaskIndex);
		result.setPartition(subTaskIndex, irp);

		resultPartitions.put(irp.getPartitionId(), irp);
	}

	this.inputEdges = new ExecutionEdge[jobVertex.getJobVertex().getInputs().size()][];

	this.priorExecutions = new EvictingBoundedList<>(maxPriorExecutionHistoryLength);

	this.currentExecution = new Execution(
		getExecutionGraph().getFutureExecutor(),
		this,
		0,
		initialGlobalModVersion,
		createTimestamp,
		timeout);

	// create a co-location scheduling hint, if necessary
	CoLocationGroup clg = jobVertex.getCoLocationGroup();
	if (clg != null) {
		this.locationConstraint = clg.getLocationConstraint(subTaskIndex);
	}
	else {
		this.locationConstraint = null;
	}

	getExecutionGraph().registerExecution(currentExecution);

	this.timeout = timeout;
}