YARN-9291. Backport YARN-7637 to branch-2
This commit is contained in:
parent
e102c3d74e
commit
22ba6939eb
|
@ -272,6 +272,7 @@ public class NMNullStateStoreService extends NMStateStoreService {
|
|||
public void storeAssignedResources(Container container,
|
||||
String resourceType, List<Serializable> assignedResources)
|
||||
throws IOException {
|
||||
updateContainerResourceMapping(container, resourceType, assignedResources);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
|
@ -38,6 +38,7 @@ import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resource
|
|||
import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.gpu.GpuDevice;
|
||||
import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.gpu.GpuDiscoverer;
|
||||
import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeConstants;
|
||||
import org.apache.hadoop.yarn.server.nodemanager.recovery.NMNullStateStoreService;
|
||||
import org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService;
|
||||
import org.apache.hadoop.yarn.util.resource.TestResourceUtils;
|
||||
import org.junit.Assert;
|
||||
|
@ -348,6 +349,35 @@ public class TestGpuResourceHandler {
|
|||
anyListOf(Serializable.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testAllocationStoredWithNULLStateStore() throws Exception {
|
||||
NMNullStateStoreService mockNMNULLStateStore = mock(NMNullStateStoreService.class);
|
||||
|
||||
Context nmnctx = mock(Context.class);
|
||||
when(nmnctx.getNMStateStore()).thenReturn(mockNMNULLStateStore);
|
||||
|
||||
GpuResourceHandlerImpl gpuNULLStateResourceHandler =
|
||||
new GpuResourceHandlerImpl(nmnctx, mockCGroupsHandler,
|
||||
mockPrivilegedExecutor);
|
||||
|
||||
Configuration conf = new YarnConfiguration();
|
||||
conf.set(YarnConfiguration.NM_GPU_ALLOWED_DEVICES, "0:0,1:1,2:3,3:4");
|
||||
GpuDiscoverer.getInstance().initialize(conf);
|
||||
|
||||
gpuNULLStateResourceHandler.bootstrap(conf);
|
||||
Assert.assertEquals(4,
|
||||
gpuNULLStateResourceHandler.getGpuAllocator().getAvailableGpus());
|
||||
|
||||
/* Start container 1, asks 3 containers */
|
||||
Container container = mockContainerWithGpuRequest(1, 3);
|
||||
gpuNULLStateResourceHandler.preStart(container);
|
||||
|
||||
verify(nmnctx.getNMStateStore()).storeAssignedResources(container,
|
||||
ResourceInformation.GPU_URI, Arrays
|
||||
.<Serializable>asList(new GpuDevice(0, 0), new GpuDevice(1, 1),
|
||||
new GpuDevice(2, 3)));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testRecoverResourceAllocation() throws Exception {
|
||||
Configuration conf = new YarnConfiguration();
|
||||
|
|
Loading…
Reference in New Issue