add control cmd

Author: huyongqii

internal/cinfo/cinfo.go

@@ -46,24 +46,44 @@ 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 "preparing_sleep":
+			resourceStateList = append(resourceStateList, protos.CranedResourceState_CRANE_PREPARING_SLEEP)
+		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
 		}
 	}
+
+	serviceUnavailableStateList := []protos.CranedResourceState{
+		protos.CranedResourceState_CRANE_SLEEPED,
+		protos.CranedResourceState_CRANE_SHUTDOWN,
+		protos.CranedResourceState_CRANE_WAKING_UP,
+		protos.CranedResourceState_CRANE_PREPARING_SLEEP,
+		protos.CranedResourceState_CRANE_POWERING_UP,
+		protos.CranedResourceState_CRANE_SHUTTING_DOWN,
+		protos.CranedResourceState_CRANE_UNKNOWN,
+	}
+
 	if len(resourceStateList) == 0 {
 		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, serviceUnavailableStateList...)
 		} 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)
+			resourceStateList = append(resourceStateList, serviceUnavailableStateList...)
 		}
 	}
 	if len(controlStateList) == 0 {

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)
 }
 

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)
-}

protos/PublicDefs.proto

@@ -79,10 +79,19 @@ 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 = 3;     
+  CRANE_SHUTDOWN = 4;    
+  CRANE_WAKING_UP = 5;  
+  CRANE_PREPARING_SLEEP = 6; 
+  CRANE_POWERING_UP = 7; 
+  CRANE_SHUTTING_DOWN = 8;
+  CRANE_UNKNOWN = 9;     
 }
 
 enum CranedControlState {