add control cmd

PKUHPC · Dec 25, 2024 · ef283d5 · ef283d5
1 parent 053cbd0
commit ef283d5
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 18 deletions.
diff --git a/internal/cinfo/cinfo.go b/internal/cinfo/cinfo.go
@@ -46,12 +46,18 @@ func Query() util.CraneCmdError {
 			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_MIX)
 		case "alloc":
 			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_ALLOC)
-		case "down":
-			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_DOWN)
-		case "none":
-			controlStateList = append(controlStateList, protos.CranedControlState_CRANE_NONE)
-		case "drain":
-			controlStateList = append(controlStateList, protos.CranedControlState_CRANE_DRAIN)
+		case "sleeped":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_SLEEPED)
+		case "shutdown":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_SHUTDOWN)
+		case "waking":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_WAKING_UP)
+		case "poweringup":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_POWERING_UP)
+		case "shuttingdown":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_SHUTTING_DOWN)
+		case "unknown":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_UNKNOWN)
 		default:
 			log.Errorf("Invalid state given: %s.\n", FlagFilterCranedStates[i])
 			return util.ErrorCmdArg
@@ -61,9 +67,9 @@ func Query() util.CraneCmdError {
 		if FlagFilterRespondingOnly {
 			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_IDLE, protos.CranedResourceState_CRANE_MIX, protos.CranedResourceState_CRANE_ALLOC)
 		} else if FlagFilterDownOnly {
-			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_DOWN)
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_SLEEPED, protos.CranedResourceState_CRANE_SHUTDOWN, protos.CranedResourceState_CRANE_WAKING_UP, protos.CranedResourceState_CRANE_POWERING_UP, protos.CranedResourceState_CRANE_SHUTTING_DOWN, protos.CranedResourceState_CRANE_UNKNOWN)
 		} else {
-			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_IDLE, protos.CranedResourceState_CRANE_MIX, protos.CranedResourceState_CRANE_ALLOC, protos.CranedResourceState_CRANE_DOWN)
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_IDLE, protos.CranedResourceState_CRANE_MIX, protos.CranedResourceState_CRANE_ALLOC, protos.CranedResourceState_CRANE_SLEEPED, protos.CranedResourceState_CRANE_SHUTDOWN, protos.CranedResourceState_CRANE_WAKING_UP, protos.CranedResourceState_CRANE_POWERING_UP, protos.CranedResourceState_CRANE_SHUTTING_DOWN, protos.CranedResourceState_CRANE_UNKNOWN)
 		}
 	}
 	if len(controlStateList) == 0 {

diff --git a/plugin/energy/energy.go b/plugin/energy/energy.go
@@ -131,7 +131,7 @@ func (p EnergyPlugin) CreateCgroupHook(ctx *api.PluginContext) {
 	log.Infof("CreateCgroupHook received for cgroup: %s", req.Cgroup)
 
 	boundGPUs := getBoundGPUs(req.Resource, globalMonitor.config.Monitor.GPUType)
-	globalMonitor.monitor.UpdateTaskCount(true)
+	globalMonitor.monitor.NodeMonitor.UpdateTaskCount(true)
 	globalMonitor.monitor.TaskMonitor.Start(req.TaskId, req.Cgroup, boundGPUs)
 }
 
@@ -143,7 +143,7 @@ func (p EnergyPlugin) DestroyCgroupHook(ctx *api.PluginContext) {
 	}
 
 	log.Infof("DestroyCgroupHook received for cgroup: %s", req.Cgroup)
-	globalMonitor.monitor.UpdateTaskCount(false)
+	globalMonitor.monitor.NodeMonitor.UpdateTaskCount(false)
 	globalMonitor.monitor.TaskMonitor.Stop(req.TaskId)
 }
 

diff --git a/plugin/energy/pkg/monitor/monitor.go b/plugin/energy/pkg/monitor/monitor.go
@@ -68,7 +68,3 @@ func (sm *Monitor) Close() {
 	sm.NodeMonitor.Close()
 	sm.TaskMonitor.Close()
 }
-
-func (sm *Monitor) UpdateTaskCount(increment bool) {
-	sm.NodeMonitor.UpdateTaskCount(increment)
-}
diff --git a/protos/PublicDefs.proto b/protos/PublicDefs.proto
@@ -79,10 +79,18 @@ enum PartitionState {
 }
 
 enum CranedResourceState {
-  CRANE_IDLE = 0;
-  CRANE_MIX = 1;
-  CRANE_ALLOC = 2;
-  CRANE_DOWN = 3;
+  // Running
+  CRANE_IDLE = 0;        
+  CRANE_ALLOC = 1;       
+  CRANE_MIX = 2;         
+
+  // Not Running
+  CRANE_SLEEPED = 4;     
+  CRANE_SHUTDOWN = 5;    
+  CRANE_WAKING_UP = 6;   
+  CRANE_POWERING_UP = 7; 
+  CRANE_SHUTTING_DOWN = 8;
+  CRANE_UNKNOWN = 9;     
 }
 
 enum CranedControlState {