YARN-8885. [DevicePlugin] Support NM APIs to query device resource allocation. (Zhankun Tang via wangda)

Change-Id: I2a9870709b512af1ac6c09c9701d0b3c0791ff32
This commit is contained in:
Wangda Tan 2018-12-12 11:45:47 -08:00
parent 9fba6cc247
commit 37eb919c59
4 changed files with 218 additions and 1 deletions

View File

@ -0,0 +1,78 @@
/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
* <p>
* http://www.apache.org/licenses/LICENSE-2.0
* <p>
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.deviceframework;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.server.nodemanager.api.deviceplugin.Device;
import java.io.Serializable;
import java.util.Objects;
/**
* Device wrapper class used for NM REST API.
* */
public class AssignedDevice implements Serializable, Comparable {
private static final long serialVersionUID = -544285507952217366L;
private Device device;
private String containerId;
public AssignedDevice(ContainerId cId, Device dev) {
this.device = dev;
this.containerId = cId.toString();
}
public Device getDevice() {
return device;
}
public String getContainerId() {
return containerId;
}
@Override
public int compareTo(Object o) {
if (o == null || !(o instanceof AssignedDevice)) {
return -1;
}
AssignedDevice other = (AssignedDevice) o;
int result = getDevice().compareTo(other.getDevice());
if (0 != result) {
return result;
}
return getContainerId().compareTo(other.getContainerId());
}
@Override
public boolean equals(Object o) {
if (o == null || !(o instanceof AssignedDevice)) {
return false;
}
AssignedDevice other = (AssignedDevice) o;
return getDevice().equals(other.getDevice())
&& getContainerId().equals(other.getContainerId());
}
@Override
public int hashCode() {
return Objects.hash(getDevice(), getContainerId());
}
}

View File

@ -20,8 +20,10 @@
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.server.nodemanager.Context;
import org.apache.hadoop.yarn.server.nodemanager.api.deviceplugin.Device;
import org.apache.hadoop.yarn.server.nodemanager.api.deviceplugin.DevicePlugin;
import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.privileged.PrivilegedOperationExecutor;
import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.CGroupsHandler;
@ -29,8 +31,13 @@
import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.DockerCommandPlugin;
import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.NodeResourceUpdaterPlugin;
import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.ResourcePlugin;
import org.apache.hadoop.yarn.server.nodemanager.webapp.dao.NMDeviceResourceInfo;
import org.apache.hadoop.yarn.server.nodemanager.webapp.dao.NMResourceInfo;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
/**
* The {@link DevicePluginAdapter} will adapt existing hooks.
@ -93,7 +100,16 @@ public DockerCommandPlugin getDockerCommandPluginInstance() {
@Override
public NMResourceInfo getNMResourceInfo() throws YarnException {
return null;
List<Device> allowed = new ArrayList<>(
deviceMappingManager.getAllAllowedDevices().get(resourceName));
List<AssignedDevice> assigned = new ArrayList<>();
Map<Device, ContainerId> assignedMap =
deviceMappingManager.getAllUsedDevices().get(resourceName);
for (Map.Entry<Device, ContainerId> entry : assignedMap.entrySet()) {
assigned.add(new AssignedDevice(entry.getValue(),
entry.getKey()));
}
return new NMDeviceResourceInfo(allowed, assigned);
}
public DeviceResourceHandlerImpl getDeviceResourceHandler() {

View File

@ -0,0 +1,56 @@
/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
* <p>
* http://www.apache.org/licenses/LICENSE-2.0
* <p>
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.hadoop.yarn.server.nodemanager.webapp.dao;
import org.apache.hadoop.yarn.server.nodemanager.api.deviceplugin.Device;
import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin.deviceframework.AssignedDevice;
import java.util.List;
/**
* Wrapper class of Device allocation for NMWebServices.
* */
public class NMDeviceResourceInfo extends NMResourceInfo {
private List<Device> totalDevices;
private List<AssignedDevice> assignedDevices;
public NMDeviceResourceInfo(
List<Device> totalDevices, List<AssignedDevice> assignedDevices) {
this.assignedDevices = assignedDevices;
this.totalDevices = totalDevices;
}
public List<Device> getTotalDevices() {
return totalDevices;
}
public void setTotalDevices(List<Device> totalDevices) {
this.totalDevices = totalDevices;
}
public List<AssignedDevice> getAssignedDevices() {
return assignedDevices;
}
public void setAssignedDevices(
List<AssignedDevice> assignedDevices) {
this.assignedDevices = assignedDevices;
}
}

View File

@ -44,6 +44,7 @@
import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeConstants;
import org.apache.hadoop.yarn.server.nodemanager.recovery.NMMemoryStateStoreService;
import org.apache.hadoop.yarn.server.nodemanager.recovery.NMStateStoreService;
import org.apache.hadoop.yarn.server.nodemanager.webapp.dao.NMDeviceResourceInfo;
import org.apache.hadoop.yarn.util.resource.ResourceUtils;
import org.apache.hadoop.yarn.util.resource.TestResourceUtils;
import org.junit.After;
@ -517,6 +518,72 @@ private static Container mockContainerWithDeviceRequest(int id,
return c;
}
/**
* Ensure correct return value generated.
* */
@Test
public void testNMResourceInfoRESTAPI() throws IOException, YarnException {
NodeManager.NMContext context = mock(NodeManager.NMContext.class);
NMStateStoreService storeService = mock(NMStateStoreService.class);
when(context.getNMStateStore()).thenReturn(storeService);
doNothing().when(storeService).storeAssignedResources(isA(Container.class),
isA(String.class),
isA(ArrayList.class));
// Init scheduler manager
DeviceMappingManager dmm = new DeviceMappingManager(context);
ResourcePluginManager rpm = mock(ResourcePluginManager.class);
when(rpm.getDeviceMappingManager()).thenReturn(dmm);
// Init an plugin
MyPlugin plugin = new MyPlugin();
MyPlugin spyPlugin = spy(plugin);
String resourceName = MyPlugin.RESOURCE_NAME;
// Init an adapter for the plugin
DevicePluginAdapter adapter = new DevicePluginAdapter(
resourceName,
spyPlugin, dmm);
// Bootstrap, adding device
adapter.initialize(context);
adapter.createResourceHandler(context,
mockCGroupsHandler, mockPrivilegedExecutor);
adapter.getDeviceResourceHandler().bootstrap(conf);
int size = dmm.getAvailableDevices(resourceName);
Assert.assertEquals(3, size);
// A container c1 requests 1 device
Container c1 = mockContainerWithDeviceRequest(0,
resourceName,
1, false);
// preStart
adapter.getDeviceResourceHandler().preStart(c1);
// check book keeping
Assert.assertEquals(2,
dmm.getAvailableDevices(resourceName));
Assert.assertEquals(1,
dmm.getAllUsedDevices().get(resourceName).size());
Assert.assertEquals(3,
dmm.getAllAllowedDevices().get(resourceName).size());
// get REST return value
NMDeviceResourceInfo response =
(NMDeviceResourceInfo) adapter.getNMResourceInfo();
Assert.assertEquals(1, response.getAssignedDevices().size());
Assert.assertEquals(3, response.getTotalDevices().size());
Device device = response.getAssignedDevices().get(0).getDevice();
String cId = response.getAssignedDevices().get(0).getContainerId();
Assert.assertTrue(dmm.getAllAllowedDevices().get(resourceName)
.contains(device));
Assert.assertTrue(dmm.getAllUsedDevices().get(resourceName)
.containsValue(ContainerId.fromString(cId)));
//finish container
adapter.getDeviceResourceHandler().postComplete(getContainerId(0));
response =
(NMDeviceResourceInfo) adapter.getNMResourceInfo();
Assert.assertEquals(0, response.getAssignedDevices().size());
Assert.assertEquals(3, response.getTotalDevices().size());
}
private static ContainerId getContainerId(int id) {
return ContainerId.newContainerId(ApplicationAttemptId
.newInstance(ApplicationId.newInstance(1234L, 1), 1), id);