balloons,cache: support memory-type annotations.

Add support for per balloon type memory configuration and per
container overrides using pod annotations. Pass configured or
annotated memory types to libmem for allocation.

TODO(klihub): per balloon configuration still missing (?)

Co-authored-by: Krisztian Litkey <krisztian.litkey@intel.com>
Signed-off-by: Krisztian Litkey <krisztian.litkey@intel.com>

cmd/plugins/balloons/policy/balloons-policy.go

@@ -1470,20 +1470,31 @@ func (p *balloons) pinCpuMem(c cache.Container, cpus cpuset.CPUSet, mems idset.I
 			if err != nil {
 				log.Error("failed to parse CpusetMems: %v", err)
 			} else {
-				zone := p.allocMem(c, preserveMems, true)
+				zone := p.allocMem(c, preserveMems, 0, true)
 				log.Debug("  - allocated preserved memory %s", c.PrettyName, zone)
 				c.SetCpusetMems(zone.MemsetString())
 			}
 		} else {
-			log.Debug("  - requested %s to memory %s", c.PrettyName(), mems)
-			zone := p.allocMem(c, mems, false)
+			memTypeMask, err := c.MemoryTypes()
+			if err != nil {
+				log.Error("%v", err)
+			}
+			if memTypeMask != 0 {
+				// memory-type pod/container-specific
+				// annotation overrides balloon's
+				// memory options that are the default
+				// to all containers in the balloon.
+				log.Debug("  - %s memory-type annotation overrides balloon mems %s", c.PrettyName(), mems)
+			}
+			log.Debug("  - requested %s to memory %s (types %s)", c.PrettyName(), mems, memTypeMask)
+			zone := p.allocMem(c, mems, memTypeMask, false)
 			log.Debug("  - allocated %s to memory %s", c.PrettyName(), zone)
 			c.SetCpusetMems(zone.MemsetString())
 		}
 	}
 }
 
-func (p *balloons) allocMem(c cache.Container, mems idset.IDSet, preserve bool) libmem.NodeMask {
+func (p *balloons) allocMem(c cache.Container, mems idset.IDSet, types libmem.TypeMask, preserve bool) libmem.NodeMask {
 	var (
 		amount  = getMemoryLimit(c)
 		nodes   = libmem.NewNodeMask(mems.Members()...)
@@ -1502,17 +1513,19 @@ func (p *balloons) allocMem(c cache.Container, mems idset.IDSet, preserve bool)
 				nodes,
 			)
 		} else {
-			req = libmem.Container(
+			req = libmem.ContainerWithTypes(
 				c.GetID(),
 				c.PrettyName(),
 				string(c.GetQOSClass()),
 				amount,
 				nodes,
+				types,
 			)
 		}
 		zone, updates, err = p.memAllocator.Allocate(req)
 	} else {
-		zone, updates, err = p.memAllocator.Realloc(c.GetID(), nodes, 0)
+
+		zone, updates, err = p.memAllocator.Realloc(c.GetID(), nodes, types)
 	}
 
 	if err != nil {

pkg/resmgr/cache/cache.go

@@ -32,6 +32,7 @@ import (
 	resmgr "github.com/containers/nri-plugins/pkg/apis/resmgr/v1alpha1"
 	"github.com/containers/nri-plugins/pkg/kubernetes"
 	logger "github.com/containers/nri-plugins/pkg/log"
+	libmem "github.com/containers/nri-plugins/pkg/resmgr/lib/memory"
 	"github.com/containers/nri-plugins/pkg/topology"
 )
 
@@ -65,6 +66,8 @@ const (
 	PreserveCpuKey = "cpu.preserve." + kubernetes.ResmgrKeyNamespace
 	// PreserveMemoryKey means that memory resources should not be touched.
 	PreserveMemoryKey = "memory.preserve." + kubernetes.ResmgrKeyNamespace
+	// MemoryTypeKey defines memory types of containers.
+	MemoryTypeKey = "memory-type." + kubernetes.ResmgrKeyNamespace
 )
 
 // PodState is the pod state in the runtime.
@@ -268,6 +271,8 @@ type Container interface {
 	// PreserveMemoryResources() returns true if memory resources
 	// of the container must not be changed.
 	PreserveMemoryResources() bool
+	// MemoryTypes() returns memory type mask. The default is 0.
+	MemoryTypes() (libmem.TypeMask, error)
 
 	// GetPendingAdjusmentn clears and returns any pending adjustment for the container.
 	GetPendingAdjustment() *nri.ContainerAdjustment

pkg/resmgr/cache/container.go

@@ -27,6 +27,7 @@ import (
 	resmgr "github.com/containers/nri-plugins/pkg/apis/resmgr/v1alpha1"
 	"github.com/containers/nri-plugins/pkg/cgroups"
 	"github.com/containers/nri-plugins/pkg/kubernetes"
+	libmem "github.com/containers/nri-plugins/pkg/resmgr/lib/memory"
 	"github.com/containers/nri-plugins/pkg/topology"
 
 	nri "github.com/containerd/nri/pkg/api"
@@ -777,6 +778,18 @@ func (c *container) PreserveMemoryResources() bool {
 	return ok && value == "true"
 }
 
+func (c *container) MemoryTypes() (libmem.TypeMask, error) {
+	value, ok := c.GetEffectiveAnnotation(MemoryTypeKey)
+	if !ok {
+		return libmem.TypeMask(0), nil
+	}
+	mask, err := libmem.ParseTypeMask(value)
+	if err != nil {
+		return libmem.TypeMask(0), cacheError("container %s has invalid effective %q annotation (%q): %v", c.PrettyName(), MemoryTypeKey, value, err)
+	}
+	return mask, nil
+}
+
 var (
 	// More complext rules, for Kubelet secrets and config maps
 	ignoredTopologyPathRegexps = []*regexp.Regexp{

test/e2e/policies.test-suite/balloons/n6-hbm-cxl/py_consts.var.py

@@ -0,0 +1,6 @@
+dram0 = "node0"
+dram1 = "node1"
+hbm0  = "node2"
+hbm1  = "node3"
+pmem0 = "node4"
+pmem1 = "node5"

test/e2e/policies.test-suite/balloons/n6-hbm-cxl/test01-memory-types/balloons-memory-types.cfg

@@ -0,0 +1,33 @@
+config:
+  # Reserve one of our CPUs (cpu15) for kube-system tasks.
+  reservedResources:
+    cpu: "1"
+  pinCPU: true
+  pinMemory: true
+  balloonTypes:
+    - name: two-cpu
+      minCPUs: 2
+      maxCPUs: 2
+      preferNewBalloons: true
+
+      cpuClass: class4
+
+    - name: five-cpu
+      maxCPUs: 5
+      allocatorPriority: none
+      preferSpreadingPods: true
+      preferNewBalloons: true
+      cpuClass: class5
+
+  instrumentation:
+    httpEndpoint: ":8891"
+    prometheusExport: true
+  log:
+    debug:
+      - cache
+      - policy
+      - nri-plugin
+      - libmem
+    source: true
+    klog:
+      skip_headers: true

test/e2e/policies.test-suite/balloons/n6-hbm-cxl/test01-memory-types/code.var.sh

@@ -0,0 +1,36 @@
+helm-terminate
+helm_config=${TEST_DIR}/balloons-memory-types.cfg helm-launch balloons
+
+cleanup() {
+    vm-command "kubectl delete pods -n kube-system pod0; kubectl delete pods --all --now"
+    return 0
+}
+
+cleanup
+
+# pod0: all memory combinations when there is enough memory.
+# CPUREQ + CONTCOUNT causes ballooon inflation after 5 containers.
+POD_ANNOTATION=()
+POD_ANNOTATION[bln]="balloon.balloons.resource-policy.nri.io/container.pod0c0: two-cpu"
+POD_ANNOTATION[0]="memory-type.resource-policy.nri.io/container.pod0c0: hbm"
+POD_ANNOTATION[1]="memory-type.resource-policy.nri.io/container.pod0c1: dram"
+POD_ANNOTATION[2]="memory-type.resource-policy.nri.io/container.pod0c2: pmem"
+POD_ANNOTATION[3]="memory-type.resource-policy.nri.io/container.pod0c3: hbm,dram"
+POD_ANNOTATION[4]="memory-type.resource-policy.nri.io/container.pod0c4: dram,pmem"
+POD_ANNOTATION[5]="memory-type.resource-policy.nri.io/container.pod0c5: hbm,dram,pmem"
+CPUREQ="200m" MEMREQ="300M" CPULIM="" MEMLIM="300M" CONTCOUNT=7 create balloons-busybox
+report allowed
+verify 'mems["pod0c0"] == {hbm0}             if packages["pod0c0"] == {"package0"} else mems["pod0c0"] == {hbm1}' \
+       'mems["pod0c1"] == {dram0}            if packages["pod0c1"] == {"package0"} else mems["pod0c1"] == {dram1}' \
+       'mems["pod0c2"] == {pmem0}            if packages["pod0c2"] == {"package0"} else mems["pod0c2"] == {pmem1}' \
+       'mems["pod0c3"] == {hbm0,dram0}       if packages["pod0c3"] == {"package0"} else mems["pod0c3"] == {hbm1,dram1}' \
+       'mems["pod0c4"] == {dram0,pmem0}      if packages["pod0c4"] == {"package0"} else mems["pod0c4"] == {dram1,pmem1}' \
+       'mems["pod0c5"] == {hbm0,dram0,pmem0} if packages["pod0c5"] == {"package0"} else mems["pod0c5"] == {hbm1,dram1,pmem1}' \
+       'mems["pod0c6"] == {dram0}            if packages["pod0c6"] == {"package0"} else mems["pod0c6"] == {dram1}'
+
+echo 'TODO: mems["pod0c5"] was missing access to pmem. Suspecting issue in defaultExpand.'
+breakpoint
+
+cleanup
+
+helm-terminate

test/e2e/policies.test-suite/balloons/n6-hbm-cxl/topology.var.json

@@ -0,0 +1,7 @@
+[
+    {"mem": "2G", "threads":2, "cores": 2, "nodes": 1, "packages": 2},
+    {"mem": "1G", "node-dist": {"0": 15, "1": 30, "2": 10, "3": 35}},
+    {"mem": "1G", "node-dist": {"0": 30, "1": 15, "2": 35, "3": 10}},
+    {"mem": "4G", "node-dist": {"0": 60, "1": 70, "2": 62, "3": 72, "4": 10, "5": 75}},
+    {"mem": "4G", "node-dist": {"0": 70, "1": 60, "2": 72, "3": 62, "4": 75, "5": 10}}
+]