org.apache.flink.api.common.resources.GPUResource Java Examples

@Test
public void testMerge() throws Exception {
	ResourceProfile rp1 = new ResourceProfile(1.0, 100, 100, 100, 100, 100, Collections.emptyMap());
	ResourceProfile rp2 = new ResourceProfile(2.0, 200, 200, 200, 200, 200,
			Collections.singletonMap("gpu", new GPUResource(2.0)));

	ResourceProfile rp1MergeRp1 = new ResourceProfile(2.0, 200, 200, 200, 200, 200,
			Collections.emptyMap());
	ResourceProfile rp1MergeRp2 = new ResourceProfile(3.0, 300, 300, 300, 300, 300,
			Collections.singletonMap("gpu", new GPUResource(2.0)));
	ResourceProfile rp2MergeRp2 = new ResourceProfile(4.0, 400, 400, 400, 400, 400,
			Collections.singletonMap("gpu", new GPUResource(4.0)));

	assertEquals(rp1MergeRp1, rp1.merge(rp1));
	assertEquals(rp1MergeRp2, rp1.merge(rp2));
	assertEquals(rp1MergeRp2, rp2.merge(rp1));
	assertEquals(rp2MergeRp2, rp2.merge(rp2));

	assertEquals(ResourceProfile.UNKNOWN, rp1.merge(ResourceProfile.UNKNOWN));
	assertEquals(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN.merge(rp1));
	assertEquals(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN.merge(ResourceProfile.UNKNOWN));
	assertEquals(ResourceProfile.ANY, rp1.merge(ResourceProfile.ANY));
	assertEquals(ResourceProfile.ANY, ResourceProfile.ANY.merge(rp1));
	assertEquals(ResourceProfile.ANY, ResourceProfile.ANY.merge(ResourceProfile.ANY));
}
 

@Test(expected = IllegalArgumentException.class)
public void testSubtractWithInfValues() {
	// Does not equals to ANY since it has extended resources.
	final ResourceProfile rp1 = ResourceProfile.newBuilder()
		.setCpuCores(Double.MAX_VALUE)
		.setTaskHeapMemoryMB(Integer.MAX_VALUE)
		.setTaskOffHeapMemoryMB(Integer.MAX_VALUE)
		.setManagedMemoryMB(Integer.MAX_VALUE)
		.setNetworkMemoryMB(Integer.MAX_VALUE)
		.addExtendedResource("gpu", new GPUResource(4.0))
		.build();
	final ResourceProfile rp2 = ResourceProfile.newBuilder()
		.setCpuCores(2.0)
		.setTaskHeapMemoryMB(200)
		.setTaskOffHeapMemoryMB(200)
		.setManagedMemoryMB(200)
		.setNetworkMemoryMB(200)
		.build();

	rp2.subtract(rp1);
}
 

@Test
public void testMultiply() {
	final int by = 3;
	final ResourceProfile rp1 = ResourceProfile.newBuilder()
		.setCpuCores(1.0)
		.setTaskHeapMemoryMB(100)
		.setTaskOffHeapMemoryMB(100)
		.setNetworkMemoryMB(100)
		.setManagedMemoryMB(100)
		.addExtendedResource("gpu", new GPUResource(1.0))
		.build();

	ResourceProfile rp2 = rp1;
	for (int i = 1; i < by; ++i) {
		rp2 = rp2.merge(rp1);
	}

	assertEquals(rp2, rp1.multiply(by));
}
 

@Test
public void testSubtractWithInfValues() {
	// Does not equals to ANY since it has extended resources.
	ResourceProfile rp1 = new ResourceProfile(Double.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE,
			Integer.MAX_VALUE, Collections.singletonMap("gpu", new GPUResource(4.0)));
	ResourceProfile rp2 = new ResourceProfile(2.0, 200, 200, 200, 200, 200,
			Collections.emptyMap());

	assertEquals(rp1, rp1.subtract(rp2));
}
 

@Test
public void testMerge() throws Exception {
	ResourceSpec rs1 = ResourceSpec.newBuilder(1.0, 100).
			setGPUResource(1.1).
			build();
	ResourceSpec rs2 = ResourceSpec.newBuilder(1.0, 100).build();

	ResourceSpec rs3 = rs1.merge(rs2);
	assertEquals(new CPUResource(2.0), rs3.getCpuCores());
	assertEquals(200, rs3.getTaskHeapMemory().getMebiBytes());
	assertEquals(new GPUResource(1.1), rs3.getGPUResource());

	ResourceSpec rs4 = rs1.merge(rs3);
	assertEquals(new GPUResource(2.2), rs4.getGPUResource());
}
 

@Test
public void testSubtract() {
	final ResourceSpec rs1 = ResourceSpec.newBuilder(1.0, 100)
		.setGPUResource(1.1)
		.build();
	final ResourceSpec rs2 = ResourceSpec.newBuilder(0.2, 100)
		.setGPUResource(0.5)
		.build();

	final ResourceSpec subtracted = rs1.subtract(rs2);
	assertEquals(new CPUResource(0.8), subtracted.getCpuCores());
	assertEquals(0, subtracted.getTaskHeapMemory().getMebiBytes());
	assertEquals(new GPUResource(0.6), subtracted.getGPUResource());
}
 

@Test
public void testGet() {
	ResourceSpec rs = ResourceSpec.newBuilder(1.0, 100).
			setGPUResource(1.6).
			build();
	ResourceProfile rp = ResourceProfile.fromResourceSpec(rs, MemorySize.ofMebiBytes(50));

	assertEquals(new CPUResource(1.0), rp.getCpuCores());
	assertEquals(150, rp.getTotalMemory().getMebiBytes());
	assertEquals(100, rp.getOperatorsMemory().getMebiBytes());
	assertEquals(new GPUResource(1.6), rp.getExtendedResources().get(GPUResource.NAME));
}
 

@Test(expected = IllegalArgumentException.class)
public void testMultiplyNegative() {
	final ResourceProfile rp = ResourceProfile.newBuilder()
		.setCpuCores(1.0)
		.setTaskHeapMemoryMB(100)
		.setTaskOffHeapMemoryMB(100)
		.setNetworkMemoryMB(100)
		.setManagedMemoryMB(100)
		.addExtendedResource("gpu", new GPUResource(1.0))
		.build();
	rp.multiply(-2);
}
 

public Builder setGPUResource(double gpus) {
	this.gpuResource = new GPUResource(gpus);
	return this;
}
 

public Builder setGPUResource(double gpus) {
	this.gpuResource = new GPUResource(gpus);
	return this;
}
 

@Test
public void testResourceCalculationOnSlotAllocatingAndReleasing() {
	ResourceProfile rp1 = new ResourceProfile(1.0, 100, 100, 100, 100, 100, Collections.emptyMap());
	ResourceProfile rp2 = new ResourceProfile(2.0, 200, 200, 200, 200, 200, Collections.singletonMap("gpu", new GPUResource(2.0)));
	ResourceProfile rp3 = new ResourceProfile(3.0, 300, 300, 300, 300, 300, Collections.singletonMap("gpu", new GPUResource(3.0)));

	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	SlotSharingManager slotSharingManager = new SlotSharingManager(
			SLOT_SHARING_GROUP_ID,
			allocatedSlotActions,
			SLOT_OWNER);

	SlotSharingManager.MultiTaskSlot unresolvedRootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			new CompletableFuture<>(),
			new SlotRequestId());

	// Allocates the left subtree.
	SlotSharingManager.MultiTaskSlot leftMultiTaskSlot =
			unresolvedRootSlot.allocateMultiTaskSlot(new SlotRequestId(), new SlotSharingGroupId());

	SlotSharingManager.SingleTaskSlot firstChild = leftMultiTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp1,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	SlotSharingManager.SingleTaskSlot secondChild = leftMultiTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp2,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	assertEquals(rp1, firstChild.getReservedResources());
	assertEquals(rp2, secondChild.getReservedResources());
	assertEquals(rp1.merge(rp2), leftMultiTaskSlot.getReservedResources());
	assertEquals(rp1.merge(rp2), unresolvedRootSlot.getReservedResources());

	// Allocates the right subtree.
	SlotSharingManager.SingleTaskSlot thirdChild = unresolvedRootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp3,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	assertEquals(rp3, thirdChild.getReservedResources());
	assertEquals(rp1.merge(rp2).merge(rp3), unresolvedRootSlot.getReservedResources());

	// Releases the second child in the left-side tree.
	secondChild.release(new Throwable("Release for testing"));
	assertEquals(rp1, leftMultiTaskSlot.getReservedResources());
	assertEquals(rp1.merge(rp3), unresolvedRootSlot.getReservedResources());

	// Releases the third child in the right-side tree.
	thirdChild.release(new Throwable("Release for testing"));
	assertEquals(rp1, unresolvedRootSlot.getReservedResources());

	// Releases the first child in the left-side tree.
	firstChild.release(new Throwable("Release for testing"));
	assertEquals(ResourceProfile.ZERO, unresolvedRootSlot.getReservedResources());
}
 

public Resource getGPUResource() {
	throwUnsupportedOperationExceptionIfUnknown();
	return extendedResources.get(GPUResource.NAME);
}
 

public Builder setGPUResource(double gpus) {
	this.gpuResource = new GPUResource(gpus);
	return this;
}
 

@Test
public void testMerge() {
	final ResourceProfile rp1 = ResourceProfile.newBuilder()
		.setCpuCores(1.0)
		.setTaskHeapMemoryMB(100)
		.setTaskOffHeapMemoryMB(100)
		.setManagedMemoryMB(100)
		.setNetworkMemoryMB(100)
		.build();
	final ResourceProfile rp2 = ResourceProfile.newBuilder()
		.setCpuCores(2.0)
		.setTaskHeapMemoryMB(200)
		.setTaskOffHeapMemoryMB(200)
		.setManagedMemoryMB(200)
		.setNetworkMemoryMB(200)
		.addExtendedResource("gpu", new GPUResource(2.0))
		.build();

	final ResourceProfile rp1MergeRp1 = ResourceProfile.newBuilder()
		.setCpuCores(2.0)
		.setTaskHeapMemoryMB(200)
		.setTaskOffHeapMemoryMB(200)
		.setManagedMemoryMB(200)
		.setNetworkMemoryMB(200)
		.build();
	final ResourceProfile rp1MergeRp2 = ResourceProfile.newBuilder()
		.setCpuCores(3.0)
		.setTaskHeapMemoryMB(300)
		.setTaskOffHeapMemoryMB(300)
		.setManagedMemoryMB(300)
		.setNetworkMemoryMB(300)
		.addExtendedResource("gpu", new GPUResource(2.0))
		.build();
	final ResourceProfile rp2MergeRp2 = ResourceProfile.newBuilder()
		.setCpuCores(4.0)
		.setTaskHeapMemoryMB(400)
		.setTaskOffHeapMemoryMB(400)
		.setManagedMemoryMB(400)
		.setNetworkMemoryMB(400)
		.addExtendedResource("gpu", new GPUResource(4.0))
		.build();

	assertEquals(rp1MergeRp1, rp1.merge(rp1));
	assertEquals(rp1MergeRp2, rp1.merge(rp2));
	assertEquals(rp1MergeRp2, rp2.merge(rp1));
	assertEquals(rp2MergeRp2, rp2.merge(rp2));

	assertEquals(ResourceProfile.UNKNOWN, rp1.merge(ResourceProfile.UNKNOWN));
	assertEquals(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN.merge(rp1));
	assertEquals(ResourceProfile.UNKNOWN, ResourceProfile.UNKNOWN.merge(ResourceProfile.UNKNOWN));
	assertEquals(ResourceProfile.ANY, rp1.merge(ResourceProfile.ANY));
	assertEquals(ResourceProfile.ANY, ResourceProfile.ANY.merge(rp1));
	assertEquals(ResourceProfile.ANY, ResourceProfile.ANY.merge(ResourceProfile.ANY));
}
 

@Test
public void testResourceCalculationOnSlotAllocatingAndReleasing() {
	final ResourceProfile rp1 = ResourceProfile.newBuilder()
		.setCpuCores(1.0)
		.setTaskHeapMemoryMB(100)
		.setTaskOffHeapMemoryMB(100)
		.setManagedMemoryMB(100)
		.setNetworkMemoryMB(100)
		.build();
	final ResourceProfile rp2 = ResourceProfile.newBuilder()
		.setCpuCores(2.0)
		.setTaskHeapMemoryMB(200)
		.setTaskOffHeapMemoryMB(200)
		.setManagedMemoryMB(200)
		.setNetworkMemoryMB(200)
		.addExtendedResource("gpu", new GPUResource(2.0))
		.build();
	final ResourceProfile rp3 = ResourceProfile.newBuilder()
		.setCpuCores(3.0)
		.setTaskHeapMemoryMB(300)
		.setTaskOffHeapMemoryMB(300)
		.setManagedMemoryMB(300)
		.setNetworkMemoryMB(300)
		.addExtendedResource("gpu", new GPUResource(3.0))
		.build();

	SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	SlotSharingManager.MultiTaskSlot unresolvedRootSlot = slotSharingManager.createRootSlot(
			new SlotRequestId(),
			new CompletableFuture<>(),
			new SlotRequestId());

	// Allocates the left subtree.
	SlotSharingManager.MultiTaskSlot leftMultiTaskSlot =
			unresolvedRootSlot.allocateMultiTaskSlot(new SlotRequestId(), new SlotSharingGroupId());

	SlotSharingManager.SingleTaskSlot firstChild = leftMultiTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp1,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	SlotSharingManager.SingleTaskSlot secondChild = leftMultiTaskSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp2,
			new SlotSharingGroupId(),
			Locality.LOCAL);

	assertEquals(rp1, firstChild.getReservedResources());
	assertEquals(rp2, secondChild.getReservedResources());
	assertEquals(rp1.merge(rp2), leftMultiTaskSlot.getReservedResources());
	assertEquals(rp1.merge(rp2), unresolvedRootSlot.getReservedResources());

	// Allocates the right subtree.
	SlotSharingManager.SingleTaskSlot thirdChild = unresolvedRootSlot.allocateSingleTaskSlot(
			new SlotRequestId(),
			rp3,
			new SlotSharingGroupId(),
			Locality.LOCAL);
	assertEquals(rp3, thirdChild.getReservedResources());
	assertEquals(rp1.merge(rp2).merge(rp3), unresolvedRootSlot.getReservedResources());

	// Releases the second child in the left-side tree.
	secondChild.release(new Throwable("Release for testing"));
	assertEquals(rp1, leftMultiTaskSlot.getReservedResources());
	assertEquals(rp1.merge(rp3), unresolvedRootSlot.getReservedResources());

	// Releases the third child in the right-side tree.
	thirdChild.release(new Throwable("Release for testing"));
	assertEquals(rp1, unresolvedRootSlot.getReservedResources());

	// Releases the first child in the left-side tree.
	firstChild.release(new Throwable("Release for testing"));
	assertEquals(ResourceProfile.ZERO, unresolvedRootSlot.getReservedResources());
}