Java Code Examples for org.apache.flink.runtime.jobgraph.JobVertex#getSlotSharingGroup()

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);
		}
	}
}
 

/**
 * 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);
}
 

private static void setManagedMemoryFraction(
		final Map<Integer, JobVertex> jobVertices,
		final Map<Integer, StreamConfig> operatorConfigs,
		final Map<Integer, Map<Integer, StreamConfig>> vertexChainedConfigs,
		final java.util.function.Function<Integer, ResourceSpec> operatorResourceRetriever,
		final java.util.function.Function<Integer, Integer> operatorManagedMemoryWeightRetriever) {

	// all slot sharing groups in this job
	final Set<SlotSharingGroup> slotSharingGroups = Collections.newSetFromMap(new IdentityHashMap<>());

	// maps a job vertex ID to its head operator ID
	final Map<JobVertexID, Integer> vertexHeadOperators = new HashMap<>();

	// maps a job vertex ID to IDs of all operators in the vertex
	final Map<JobVertexID, Set<Integer>> vertexOperators = new HashMap<>();

	for (Entry<Integer, JobVertex> entry : jobVertices.entrySet()) {
		final int headOperatorId = entry.getKey();
		final JobVertex jobVertex = entry.getValue();

		final SlotSharingGroup jobVertexSlotSharingGroup = jobVertex.getSlotSharingGroup();

		checkState(jobVertexSlotSharingGroup != null, "JobVertex slot sharing group must not be null");
		slotSharingGroups.add(jobVertexSlotSharingGroup);

		vertexHeadOperators.put(jobVertex.getID(), headOperatorId);

		final Set<Integer> operatorIds = new HashSet<>();
		operatorIds.add(headOperatorId);
		operatorIds.addAll(vertexChainedConfigs.getOrDefault(headOperatorId, Collections.emptyMap()).keySet());
		vertexOperators.put(jobVertex.getID(), operatorIds);
	}

	for (SlotSharingGroup slotSharingGroup : slotSharingGroups) {
		setManagedMemoryFractionForSlotSharingGroup(
			slotSharingGroup,
			vertexHeadOperators,
			vertexOperators,
			operatorConfigs,
			vertexChainedConfigs,
			operatorResourceRetriever,
			operatorManagedMemoryWeightRetriever);
	}
}