From d5f88a14e9902ffbafa5c2d08f189b557e96c34b Mon Sep 17 00:00:00 2001
From: Swann Perarnau <swann@anl.gov>
Date: Thu, 14 Dec 2017 10:40:17 -0600
Subject: [PATCH] [refactor] Allow updates in resource tracking

Implement an update allocation function to be able to update resource
tracking when containers are created and deleted.

The commit should make it easier to improve the resource manager later
on.
---
 nrm/containers.py |  1 +
 nrm/resources.py  | 59 ++++++++++++++++++++++-------------------------
 2 files changed, 28 insertions(+), 32 deletions(-)

diff --git a/nrm/containers.py b/nrm/containers.py
index 2bb9551..93ea97d 100644
--- a/nrm/containers.py
+++ b/nrm/containers.py
@@ -89,6 +89,7 @@ class ContainerManager(object):
     def delete(self, uuid):
         """Delete a container and kill all related processes."""
         self.nodeos.delete(uuid, kill=True)
+        self.resourcemanager.update(uuid)
         c = self.containers[uuid]
         del self.containers[uuid]
         del self.pids[c.pid]
diff --git a/nrm/resources.py b/nrm/resources.py
index 953551f..972563d 100644
--- a/nrm/resources.py
+++ b/nrm/resources.py
@@ -17,6 +17,7 @@ class ResourceManager(object):
         self.allresources = self.hwloc.info()
         self.logger.debug("resource info: %r", self.allresources)
         self.available = self.allresources
+        self.allocations = {}
 
     def schedule(self, uuid, request):
         """Schedule a resource request on the available resources.
@@ -27,43 +28,37 @@ class ResourceManager(object):
         #  - memories exclusive if more than one left
         if len(self.available.cpus) >= request.cpus:
             retcpus = self.available.cpus[:request.cpus]
-            availcpus = self.available.cpus[request.cpus:]
         else:
             retcpus = []
-            availcpus = self.available.cpus
         if len(self.available.mems) > 1:
             retmems = self.available.mems[:request.mems]
-            availmems = self.available.mems[request.mems:]
         else:
             retmems = self.available.mems
-            availmems = self.available.mems
-        self.available = resources(availcpus, availmems)
-        return resources(retcpus, retmems)
+        ret = resources(retcpus, retmems)
+        # make sure we don't remember an error
+        if ret.cpus:
+            self.update(uuid, ret)
+        return ret
 
-    def remove(self, uuid):
-        """Free the resources associated with request uuid."""
-        pass
+    def update(self, uuid, allocation=resources([], [])):
+        """Update resource tracking according to new allocation.
 
-#    def oldcode(self):
-#        numcpus = int(manifest.app.isolators.container.cpus.value)
-#
-#        allresources = hwloc.info()
-#        self.logger.debug("resource info: %r", allresources)
-#        ncontainers = len(allresources.cpus) // numcpus
-#        self.logger.debug("will support %s containers", ncontainers)
-#        cur = nodeos.getavailable()
-#        self.logger.debug("%r are available", cur)
-#        sets = hwloc.distrib(ncontainers, restrict=cur, fake=allresources)
-#        self.logger.info("asking for %s cores", numcpus)
-#        self.logger.debug("will search in one of these: %r", sets)
-#        # find a free set
-#        avail = set(cur.cpus)
-#        for s in sets:
-#            cpuset = set(s.cpus)
-#            if cpuset.issubset(avail):
-#                alloc = s
-#                break
-#        else:
-#            self.logger.error("no exclusive cpuset found among %r", avail)
-#            return -2
-#
+        The new allocation is saved, and available resources updated."""
+        added = {}
+        freed = {}
+        prev = self.allocations.get(uuid, resources([], []))
+        for attr, val in prev._asdict().items():
+            added[attr] = set(getattr(allocation, attr)) - set(val)
+            freed[attr] = set(val) - set(getattr(allocation, attr))
+        if allocation != resources([], []):
+            self.allocations[uuid] = allocation
+            self.logger.info("updated allocation for %r: %r", uuid,
+                             self.available)
+        else:
+            del self.allocations[uuid]
+            self.logger.info("deleted allocation for %r", uuid)
+        new = {}
+        for attr, val in self.available._asdict().items():
+            new[attr] = list(set(val) - set(added[attr]) | set(freed[attr]))
+        self.available = resources(**new)
+        self.logger.info("updated available resources: %r", self.available)
-- 
2.22.2