YARN-8885. [DevicePlugin] Support NM APIs to query device resource allocation. (Zhankun Tang via wangda)
Change-Id: I2a9870709b512af1ac6c09c9701d0b3c0791ff32
This commit is contained in:
parent
9fba6cc247
commit
37eb919c59
|
@ -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());
|
||||
}
|
||||
|
||||
}
|
|
@ -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() {
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
}
|
|
@ -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);
|
||||
|
|
Loading…
Reference in New Issue