[YUNIKORN-1994] fix release containers metrix in queue

Signed-off-by: PoAn Yang <[email protected]>

pkg/metrics/init.go

@@ -54,6 +54,7 @@ type CoreQueueMetrics interface {
 	IncQueueApplicationsCompleted()
 	IncAllocatedContainer()
 	IncReleasedContainer()
+	AddReleasedContainers(value int)
 	SetQueueGuaranteedResourceMetrics(resourceName string, value float64)
 	SetQueueMaxResourceMetrics(resourceName string, value float64)
 	SetQueueAllocatedResourceMetrics(resourceName string, value float64)

pkg/metrics/queue.go

@@ -108,6 +108,10 @@ func (m *QueueMetrics) IncReleasedContainer() {
 	m.appMetrics.With(prometheus.Labels{"state": "released"}).Inc()
 }
 
+func (m *QueueMetrics) AddReleasedContainers(value int) {
+	m.appMetrics.With(prometheus.Labels{"state": "released"}).Add(float64(value))
+}
+
 func (m *QueueMetrics) SetQueueGuaranteedResourceMetrics(resourceName string, value float64) {
 	m.ResourceMetrics.With(prometheus.Labels{"state": "guaranteed", "resource": resourceName}).Set(value)
 }

pkg/metrics/queue_test.go

@@ -86,6 +86,14 @@ func TestReleasedContainers(t *testing.T) {
 	verifyAppMetrics(t, "released")
 }
 
+func TestAddReleasedContainers(t *testing.T) {
+	cqm = getQueueMetrics()
+	defer unregisterQueueMetrics(t)
+
+	cqm.AddReleasedContainers(1)
+	verifyAppMetrics(t, "released")
+}
+
 func TestQueueGuaranteedResourceMetrics(t *testing.T) {
 	cqm = getQueueMetrics()
 	defer unregisterQueueMetrics(t)

pkg/scheduler/partition.go

@@ -672,6 +672,7 @@ func (pc *PartitionContext) removeNode(nodeID string) ([]*objects.Allocation, []
 func (pc *PartitionContext) removeNodeAllocations(node *objects.Node) ([]*objects.Allocation, []*objects.Allocation) {
 	released := make([]*objects.Allocation, 0)
 	confirmed := make([]*objects.Allocation, 0)
+	releasedContainerByQueue := make(map[string]int)
 	// walk over all allocations still registered for this node
 	for _, alloc := range node.GetAllAllocations() {
 		allocID := alloc.GetUUID()
@@ -768,8 +769,6 @@ func (pc *PartitionContext) removeNodeAllocations(node *objects.Node) ([]*object
 			log.Log(log.SchedPartition).Warn("failed to release resources from queue",
 				zap.String("appID", alloc.GetApplicationID()),
 				zap.Error(err))
-		} else {
-			metrics.GetQueueMetrics(queue.GetQueuePath()).IncReleasedContainer()
 		}
 		// remove preempted resources
 		if alloc.IsPreempted() {
@@ -781,12 +780,16 @@ func (pc *PartitionContext) removeNodeAllocations(node *objects.Node) ([]*object
 
 		// the allocation is removed so add it to the list that we return
 		released = append(released, alloc)
+		releasedContainerByQueue[queue.GetQueuePath()]++
 		log.Log(log.SchedPartition).Info("allocation removed from node",
 			zap.String("nodeID", node.NodeID),
 			zap.String("allocationId", allocID))
 	}
 	// track the number of allocations: decrement the released allocation AND increment with the confirmed
 	pc.updateAllocationCount(len(confirmed) - len(released))
+	for queuePath, count := range releasedContainerByQueue {
+		metrics.GetQueueMetrics(queuePath).AddReleasedContainers(count)
+	}
 	return released, confirmed
 }
 
@@ -1356,8 +1359,6 @@ func (pc *PartitionContext) removeAllocation(release *si.AllocationRelease) ([]*
 				zap.String("appID", appID),
 				zap.String("allocationId", uuid),
 				zap.Error(err))
-		} else {
-			metrics.GetQueueMetrics(queue.GetQueuePath()).IncReleasedContainer()
 		}
 	}
 	if resources.StrictlyGreaterThanZero(totalPreempting) {
@@ -1371,6 +1372,7 @@ func (pc *PartitionContext) removeAllocation(release *si.AllocationRelease) ([]*
 	}
 	// track the number of allocations, when we replace the result is no change
 	pc.updateAllocationCount(-len(released))
+	metrics.GetQueueMetrics(queue.GetQueuePath()).AddReleasedContainers(len(released))
 
 	// if the termination type is timeout, we don't notify the shim, because it's
 	// originated from that side