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

Change-Id: I2a9870709b512af1ac6c09c9701d0b3c0791ff32

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/resourceplugin/deviceframework/AssignedDevice.java

@@ -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());
+  }
+
+}

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/resourceplugin/deviceframework/DevicePluginAdapter.java

@@ -20,8 +20,10 @@ package org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugi
 
 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.linux.resource
 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 class DevicePluginAdapter implements ResourcePlugin {
 
   @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() {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/webapp/dao/NMDeviceResourceInfo.java

@@ -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;
+  }
+}

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/resourceplugin/deviceframework/TestDevicePluginAdapter.java

@@ -44,6 +44,7 @@ import org.apache.hadoop.yarn.server.nodemanager.containermanager.resourceplugin
 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 @@ public class TestDevicePluginAdapter {
     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);