YARN-11198. Adding assert and test

YARN-11198. Adding assert and test YARN-11198. Adding assert and test YARN-11198. Fix test cases YARN-11198. remove unused imports YARN-11198. Address to comments and add ResourceHandlerException YARN-11198. Fix documentation error YARN-11198. Fix documentation error
apache · Jul 13, 2022 · 90ba297 · 90ba297
1 parent a99d171
commit 90ba297
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 17 deletions.
diff --git a/.../yarn/server/nodemanager/containermanager/linux/resources/numa/NumaResourceAllocator.java b/.../yarn/server/nodemanager/containermanager/linux/resources/numa/NumaResourceAllocator.java
@@ -73,7 +73,7 @@ public class NumaResourceAllocator {
   private static final String SPACE = "\\s";
   private static final long DEFAULT_NUMA_NODE_MEMORY = 1024;
   private static final int DEFAULT_NUMA_NODE_CPUS = 1;
-  public static final String NUMA_RESOURCE_TYPE = "numa";
+  private static final String NUMA_RESOURCE_TYPE = "numa";
 
   private List<NumaNodeResource> numaNodesList = new ArrayList<>();
   private Map<String, NumaNodeResource> numaNodeIdVsResource = new HashMap<>();
@@ -231,7 +231,7 @@ public synchronized NumaResourceAllocation allocateNumaNodes(
   }
 
   private NumaResourceAllocation allocate(ContainerId containerId,
-      Resource resource) {
+      Resource resource) throws ResourceHandlerException {
     for (int index = 0; index < numaNodesList.size(); index++) {
       NumaNodeResource numaNode = numaNodesList
           .get((currentAssignNode + index) % numaNodesList.size());
@@ -306,14 +306,20 @@ private NumaResourceAllocation allocate(ContainerId containerId,
    * Release assigned NUMA resources for the container.
    *
    * @param containerId the container ID
+   * @throws ResourceHandlerException while releasing numa resource
    */
-  public synchronized void releaseNumaResource(ContainerId containerId) {
+  public synchronized void releaseNumaResource(ContainerId containerId)
+      throws ResourceHandlerException {
     LOG.info("Releasing the assigned NUMA resources for " + containerId);
     for (NumaNodeResource numaNode : numaNodesList) {
       numaNode.releaseResources(containerId);
     }
-    // delete key from NM State store
-    context.getNMStateStore().releaseNumaResource(containerId);
+    // delete from NM State store
+    try {
+      context.getNMStateStore().releaseAssignedResources(containerId, NUMA_RESOURCE_TYPE);
+    } catch (IOException e){
+      throw new ResourceHandlerException(e);
+    }
   }
 
   /**

diff --git a/...n/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMLeveldbStateStoreService.java b/...n/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMLeveldbStateStoreService.java
@@ -50,8 +50,6 @@
 import org.apache.hadoop.yarn.server.nodemanager.NodeStatusUpdater;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.Container;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.ResourceMappings;
-import static org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.resources.
-    numa.NumaResourceAllocator.NUMA_RESOURCE_TYPE;
 import org.apache.hadoop.yarn.server.records.Version;
 import org.apache.hadoop.yarn.server.records.impl.pb.VersionPBImpl;
 import org.apache.hadoop.yarn.server.utils.BuilderUtils;
@@ -1816,20 +1814,20 @@ protected void checkVersion() throws IOException {
     }
   }
   @Override
-  public void releaseNumaResource(ContainerId containerId)  {
-    LOG.info("Release NUMA resource Called for removing from statestore");
+  public void releaseAssignedResources(ContainerId containerId, String resourceType)
+      throws IOException {
+    LOG.debug("releaseAssignedResources: containerId=" + containerId + " resourceType="
+        + resourceType);
     try {
       try (WriteBatch batch = db.createWriteBatch()) {
-        String key = CONTAINERS_KEY_PREFIX + containerId.toString()
-            + CONTAINER_ASSIGNED_RESOURCES_KEY_SUFFIX + NUMA_RESOURCE_TYPE;
+        String key = CONTAINERS_KEY_PREFIX + containerId
+            + CONTAINER_ASSIGNED_RESOURCES_KEY_SUFFIX + resourceType;
         batch.delete(bytes(key));
-      }catch (IOException e) {
-        throw new RuntimeException(e);
+        db.write(batch);
       }
     }catch (DBException e){
       markStoreUnHealthy(e);
-      LOG.error("Failed to delete Key");
-      throw new RuntimeException(e);
+      throw new IOException(e);
     }
   }
 }
diff --git a/...src/main/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMStateStoreService.java b/...src/main/java/org/apache/hadoop/yarn/server/nodemanager/recovery/NMStateStoreService.java
@@ -787,10 +787,13 @@ public abstract void storeAssignedResources(Container container,
       throws IOException;
 
   /**
-   * Store the assigned resources to a container.
+   * Delete the assigned resources of a container of specific resourceType.
    * @param containerId Container Id
+   * @param resourceType resource Type
+   * @throws IOException while releasing resources
    */
-  public void releaseNumaResource(ContainerId containerId){}
+  public void releaseAssignedResources(ContainerId containerId, String resourceType)
+      throws IOException {}
 
   protected abstract void initStorage(Configuration conf) throws IOException;
 

diff --git a/...va/org/apache/hadoop/yarn/server/nodemanager/recovery/TestNMLeveldbStateStoreService.java b/...va/org/apache/hadoop/yarn/server/nodemanager/recovery/TestNMLeveldbStateStoreService.java
@@ -1756,6 +1756,18 @@ public void testStateStoreForResourceMapping() throws IOException {
     resources = rcs.getResourceMappings().getAssignedResources("numa");
     Assert.assertEquals(numaRes, resources);
     Assert.assertEquals(numaRes, resourceMappings.getAssignedResources("numa"));
+    // test removing numa resources from state store
+    stateStore.releaseAssignedResources(containerId, "numa");
+    recoveredContainers = loadContainersState(stateStore.getContainerStateIterator());
+    resourceMappings = recoveredContainers.get(0).getResourceMappings();
+    assertTrue(resourceMappings.getAssignedResources("numa").isEmpty());
+
+    // testing calling deletion of non-existing key doesn't break anything
+    try {
+      stateStore.releaseAssignedResources(containerId, "numa");
+    }catch (RuntimeException e){
+      Assert.fail("Should not throw exception while deleting non existing key from statestore");
+    }
   }
 
   @Test