Dev/licenses

internal/cbatch/CmdArgParser.go

@@ -39,6 +39,7 @@ var (
 	FlagJob           string
 	FlagAccount       string
 	FlagQos           string
+	FlagLicenses      string
 	FlagCwd           string
 	FlagRepeat        uint32
 	FlagNodelist      string
@@ -144,6 +145,7 @@ func init() {
 	RootCmd.Flags().StringVarP(&FlagAccount, "account", "A", "", "Account used for the job")
 	RootCmd.Flags().StringVarP(&FlagCwd, "chdir", "D", "", "Working directory of the job")
 	RootCmd.Flags().StringVarP(&FlagQos, "qos", "q", "", "QoS used for the job")
+	RootCmd.Flags().StringVarP(&FlagLicenses, "licenses", "L", "", "Licenses used for the job")
 	RootCmd.Flags().Uint32Var(&FlagRepeat, "repeat", 1, "Submit the job multiple times")
 	RootCmd.Flags().StringVarP(&FlagNodelist, "nodelist", "w", "", "Nodes to be allocated to the job (commas separated list)")
 	RootCmd.Flags().StringVarP(&FlagExcludes, "exclude", "x", "", "Exclude specific nodes from allocating (commas separated list)")

internal/cbatch/cbatch.go

@@ -112,6 +112,13 @@ func ProcessCbatchArgs(cmd *cobra.Command, args []CbatchArg) (bool, *protos.Task
 			task.Account = arg.val
 		case "--qos", "Q":
 			task.Qos = arg.val
+		case "--licenses", "-L":
+			lic_count_map, err := util.ParseLicensesString(arg.val)
+			if err != nil {
+				log.Error(err)
+				return false, nil
+			}
+			task.LicensesCount = lic_count_map
 		case "--chdir":
 			task.Cwd = arg.val
 		case "--exclude", "-x":
@@ -206,6 +213,14 @@ func ProcessCbatchArgs(cmd *cobra.Command, args []CbatchArg) (bool, *protos.Task
 	if FlagQos != "" {
 		task.Qos = FlagQos
 	}
+	if FlagLicenses != "" {
+		lic_count_map, err := util.ParseLicensesString(FlagLicenses)
+		if err != nil {
+			log.Error(err)
+			return false, nil
+		}
+		task.LicensesCount = lic_count_map
+	}
 	if FlagCwd != "" {
 		task.Cwd = FlagCwd
 	}

internal/ccontrol/CmdArgParser.go

@@ -28,17 +28,18 @@ import (
 )
 
 var (
-	FlagNodeName       string
-	FlagState          string
-	FlagReason         string
-	FlagPartitionName  string
-	FlagTaskIds        string
-	FlagQueryAll       bool
-	FlagTimeLimit      string
-	FlagPriority       float64
-	FlagHoldTime       string
-	FlagConfigFilePath string
-	FlagJson           bool
+	FlagNodeName        string
+	FlagState           string
+	FlagReason          string
+	FlagPartitionName   string
+	FlagLicenseNameList string
+	FlagTaskIds         string
+	FlagQueryAll        bool
+	FlagTimeLimit       string
+	FlagPriority        float64
+	FlagHoldTime        string
+	FlagConfigFilePath  string
+	FlagJson            bool
 
 	RootCmd = &cobra.Command{
 		Use:     "ccontrol",
@@ -92,6 +93,24 @@ var (
 			}
 		},
 	}
+	showLicenseCmd = &cobra.Command{
+		Use:   "lic [flags] [license_name]",
+		Short: "Display details of the licenses, default is all",
+		Long:  "",
+		Args:  cobra.MaximumNArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			if len(args) == 0 {
+				FlagLicenseNameList = ""
+				FlagQueryAll = true
+			} else {
+				FlagLicenseNameList = args[0]
+				FlagQueryAll = false
+			}
+			if err := ShowLicenses(FlagLicenseNameList, FlagQueryAll); err != util.ErrorSuccess {
+				os.Exit(err)
+			}
+		},
+	}
 	showJobCmd = &cobra.Command{
 		Use:   "job [flags] [job_id,...]",
 		Short: "Display details of the jobs, default is all",
@@ -225,6 +244,7 @@ func init() {
 	{
 		showCmd.AddCommand(showNodeCmd)
 		showCmd.AddCommand(showPartitionCmd)
+		showCmd.AddCommand(showLicenseCmd)
 		showCmd.AddCommand(showJobCmd)
 		showCmd.AddCommand(showConfigCmd)
 	}

internal/ccontrol/ccontrol.go

@@ -228,6 +228,37 @@ func ShowPartitions(partitionName string, queryAll bool) util.CraneCmdError {
 	return util.ErrorSuccess
 }
 
+func ShowLicenses(licenseName string, queryAll bool) util.CraneCmdError {
+	var licenseNameList []string
+
+	if licenseName != "" {
+		licenseNameList = strings.Split(licenseName, ",")
+	}
+
+	req := &protos.QueryLicensesInfoRequest{LicenseNameList: licenseNameList}
+	reply, err := stub.QueryLicensesInfo(context.Background(), req)
+	if err != nil {
+		util.GrpcErrorPrintf(err, "Failed to show license")
+		return util.ErrorNetwork
+	}
+
+	if len(reply.LicenseInfoList) == 0 {
+		if queryAll {
+			fmt.Println("No license is available.")
+		} else {
+			fmt.Printf("license %s not found.\n", licenseName)
+		}
+	}
+
+	for _, licenseInfo := range reply.LicenseInfoList {
+		fmt.Printf("LicenseName=%v \n"+
+			"\tTotal=%v Used=%d Free=%d\n",
+			licenseInfo.Name, licenseInfo.Total, licenseInfo.Used, licenseInfo.Free)
+	}
+
+	return util.ErrorSuccess
+}
+
 func ShowJobs(jobIds string, queryAll bool) util.CraneCmdError {
 	var req *protos.QueryTasksInfoRequest
 	var jobIdList []uint32

internal/cqueue/CmdArgParser.go

@@ -39,6 +39,7 @@ var (
 	FlagFilterStates     string
 	FlagFilterUsers      string
 	FlagFilterAccounts   string
+	FlagFilterLicenses   string
 	FlagFormat           string
 	FlagIterate          uint64
 	FlagNumLimit         uint32
@@ -98,7 +99,8 @@ func init() {
 		"Specify accounts to view (comma separated list), \ndefault is all accounts")
 	RootCmd.Flags().StringVarP(&FlagFilterPartitions, "partition", "p", "",
 		"Specify partitions to view (comma separated list), \ndefault is all partitions")
-
+	RootCmd.Flags().StringVarP(&FlagFilterLicenses, "licenses", "L", "",
+		"Specify licenses to view (comma separated list), default is all licenses")
 	RootCmd.Flags().Uint64VarP(&FlagIterate, "iterate", "i", 0,
 		"Display at specified intervals (seconds), default is 0 (no iteration)")
 	RootCmd.Flags().BoolVarP(&FlagStartTime, "start", "S", false,

internal/cqueue/cqueue.go

@@ -101,7 +101,10 @@ func QueryTasksInfo() (*protos.QueryTasksInfoReply, util.CraneCmdError) {
 		}
 		req.FilterPartitions = filterPartitionList
 	}
-
+	if FlagFilterLicenses != "" {
+		filterLicenseList := strings.Split(FlagFilterLicenses, ",")
+		req.FilterLicenses = filterLicenseList
+	}
 	if FlagFilterJobIDs != "" {
 		filterJobIdList, err := util.ParseJobIdList(FlagFilterJobIDs, ",")
 		if err != nil {
@@ -129,10 +132,10 @@ func QueryTableOutput(reply *protos.QueryTasksInfoReply) util.CraneCmdError {
 	util.SetBorderlessTable(table)
 	var header []string
 	tableData := make([][]string, len(reply.TaskInfoList))
-	if  FlagFull {
-		header = []string{"JobId", "JobName", "UserName", "Partition", 
-		 "Account", "NodeNum", "AllocCPUs", "MemPerNode", "Status", "Time", "TimeLimit", 
-		 "StartTime", "SubmitTime", "Type", "Qos",  "Held", "Priority", "NodeList/Reason"}
+	if FlagFull {
+		header = []string{"JobId", "JobName", "UserName", "Partition",
+			"Account", "NodeNum", "AllocCPUs", "MemPerNode", "Status", "Time", "TimeLimit",
+			"StartTime", "SubmitTime", "Type", "Qos", "Held", "Priority", "NodeList/Reason"}
 		for i := 0; i < len(reply.TaskInfoList); i++ {
 			taskInfo := reply.TaskInfoList[i]
 
@@ -153,7 +156,7 @@ func QueryTableOutput(reply *protos.QueryTasksInfoReply) util.CraneCmdError {
 			startTimeStr := "unknown"
 			startTime := taskInfo.StartTime.AsTime()
 			if !startTime.Before(time.Date(1980, 1, 1, 0, 0, 0, 0, time.UTC)) &&
-			startTime.Before(time.Now()) {
+				startTime.Before(time.Now()) {
 				startTimeStr = startTime.In(time.Local).Format("2006-01-02 15:04:05")
 			}
 
@@ -163,15 +166,14 @@ func QueryTableOutput(reply *protos.QueryTasksInfoReply) util.CraneCmdError {
 				submitTimeStr = submitTime.In(time.Local).Format("2006-01-02 15:04:05")
 			}
 
-
 			var reasonOrListStr string
 			if reply.TaskInfoList[i].Status == protos.TaskStatus_Pending {
 				reasonOrListStr = reply.TaskInfoList[i].GetPendingReason()
 			} else {
 				reasonOrListStr = reply.TaskInfoList[i].GetCranedList()
 			}
 
-			tableData[i] = []string {
+			tableData[i] = []string{
 				strconv.FormatUint(uint64(taskInfo.TaskId), 10),
 				taskInfo.Name,
 				taskInfo.Username,
@@ -259,6 +261,23 @@ func QueryTableOutput(reply *protos.QueryTasksInfoReply) util.CraneCmdError {
 		}
 	}
 
+	if FlagFilterLicenses != "" {
+		header = append(header, "Licenses")
+		for i := 0; i < len(tableData); i++ {
+			// map to lic1:count1,lic2:count2
+			var licenses string
+			first := true
+			for k, v := range reply.TaskInfoList[i].LicensesCount {
+				if !first {
+					licenses += ","
+				}
+				licenses += k + ":" + strconv.FormatUint(uint64(v), 10)
+				first = false
+			}
+			tableData[i] = append(tableData[i], licenses)
+		}
+	}
+
 	if !FlagNoHeader {
 		table.SetHeader(header)
 	}
@@ -400,57 +419,57 @@ func ProcessHeld(task *protos.TaskInfo) string {
 }
 
 type FieldProcessor struct {
-	header string
+	header  string
 	process func(task *protos.TaskInfo) string
 }
 
 var fieldMap = map[string]FieldProcessor{
-	"a":         {"Account", ProcessAccount},
-	"account":   {"Account", ProcessAccount},
-	"c":         {"CpuPerNode", ProcessCpuPerNode},
-	"cpupernode": {"CpuPerNode", ProcessCpuPerNode},
-	"C":         {"AllocCpus", ProcessAllocCpus},
-	"alloccpus": {"AllocCpus", ProcessAllocCpus},
-	"N":         {"NodeNum", ProcessNodeNum},
-	"nodenum":   {"NodeNum", ProcessNodeNum},
-	"e":         {"ElapsedTime", ProcessElapsedTime},
-	"elapsedtime": {"ElapsedTime", ProcessElapsedTime},
-	"j":         {"JobId", ProcessJobId},
-	"jobid":     {"JobId", ProcessJobId},
-	"l":         {"TimeLimit", ProcessTimeLimit},
-	"timelimit": {"TimeLimit", ProcessTimeLimit},
-	"S":         {"StartTime", ProcessStartTime},
-	"starttime": {"StartTime", ProcessStartTime},
-	"s":         {"SubmitTime", ProcessSubmitTime},
-	"submittime": {"SubmitTime", ProcessSubmitTime},
-	"L":         {"NodeList(Reason)", ProcessNodeList},
-	"nodelist":  {"NodeList(Reason)", ProcessNodeList},
-	"m":         {"MemPerNode", ProcessMemPerNode},
-	"mempernode": {"MemPerNode", ProcessMemPerNode},
-	"n":         {"Name", ProcessName},
-	"name":      {"Name", ProcessName},
-	"t":         {"State", ProcessState},
-	"state":     {"State", ProcessState},
-	"p":         {"Priority", ProcessPriority},
-	"priority":  {"Priority", ProcessPriority},
-	"P":         {"Partition", ProcessPartition},
-	"partition": {"Partition", ProcessPartition},
-	"q":         {"QoS", ProcessQoS},
-	"qos":       {"QoS", ProcessQoS},
-	"T":         {"JobType", ProcessJobType},
-	"jobtype":   {"JobType", ProcessJobType},
-	"u":         {"User", ProcessUser},
-	"user":      {"User", ProcessUser},
-	"U":         {"Uid", ProcessUid},
-	"uid":       {"Uid", ProcessUid},
-	"R":         {"Reason", ProcessReason},
-	"reason":    {"Reason", ProcessReason},
-	"r":         {"ReqNodes", ProcessReqNodes},
-	"reqnodes":  {"ReqNodes", ProcessReqNodes},
-	"x":         {"ExcludeNodes", ProcessExcludeNodes},
+	"a":            {"Account", ProcessAccount},
+	"account":      {"Account", ProcessAccount},
+	"c":            {"CpuPerNode", ProcessCpuPerNode},
+	"cpupernode":   {"CpuPerNode", ProcessCpuPerNode},
+	"C":            {"AllocCpus", ProcessAllocCpus},
+	"alloccpus":    {"AllocCpus", ProcessAllocCpus},
+	"N":            {"NodeNum", ProcessNodeNum},
+	"nodenum":      {"NodeNum", ProcessNodeNum},
+	"e":            {"ElapsedTime", ProcessElapsedTime},
+	"elapsedtime":  {"ElapsedTime", ProcessElapsedTime},
+	"j":            {"JobId", ProcessJobId},
+	"jobid":        {"JobId", ProcessJobId},
+	"l":            {"TimeLimit", ProcessTimeLimit},
+	"timelimit":    {"TimeLimit", ProcessTimeLimit},
+	"S":            {"StartTime", ProcessStartTime},
+	"starttime":    {"StartTime", ProcessStartTime},
+	"s":            {"SubmitTime", ProcessSubmitTime},
+	"submittime":   {"SubmitTime", ProcessSubmitTime},
+	"L":            {"NodeList(Reason)", ProcessNodeList},
+	"nodelist":     {"NodeList(Reason)", ProcessNodeList},
+	"m":            {"MemPerNode", ProcessMemPerNode},
+	"mempernode":   {"MemPerNode", ProcessMemPerNode},
+	"n":            {"Name", ProcessName},
+	"name":         {"Name", ProcessName},
+	"t":            {"State", ProcessState},
+	"state":        {"State", ProcessState},
+	"p":            {"Priority", ProcessPriority},
+	"priority":     {"Priority", ProcessPriority},
+	"P":            {"Partition", ProcessPartition},
+	"partition":    {"Partition", ProcessPartition},
+	"q":            {"QoS", ProcessQoS},
+	"qos":          {"QoS", ProcessQoS},
+	"T":            {"JobType", ProcessJobType},
+	"jobtype":      {"JobType", ProcessJobType},
+	"u":            {"User", ProcessUser},
+	"user":         {"User", ProcessUser},
+	"U":            {"Uid", ProcessUid},
+	"uid":          {"Uid", ProcessUid},
+	"R":            {"Reason", ProcessReason},
+	"reason":       {"Reason", ProcessReason},
+	"r":            {"ReqNodes", ProcessReqNodes},
+	"reqnodes":     {"ReqNodes", ProcessReqNodes},
+	"x":            {"ExcludeNodes", ProcessExcludeNodes},
 	"excludenodes": {"ExcludeNodes", ProcessExcludeNodes},
-	"h":         {"Held", ProcessHeld},
-	"held":      {"Held", ProcessHeld},
+	"h":            {"Held", ProcessHeld},
+	"held":         {"Held", ProcessHeld},
 }
 
 // FormatData formats the output data according to the format string.
@@ -515,10 +534,10 @@ func FormatData(reply *protos.QueryTasksInfoReply) (header []string, tableData [
 		//S-StartTime, t-State, T-JobType, u-User, U-Uid, x-ExcludeNodes
 		fieldProcessor, found := fieldMap[field]
 		if !found {
-			log.Errorf("Invalid format specifier or string : %s, string unfold case insensitive, reference:\n" +
-	 		"a/Account, c/CpuPerNode, C/AllocCPUs, e/ElapsedTime, h/Held, j/JobID, l/TimeLimit, L/NodeList,\n" +
-	 		"m/MemPerNode, n/Name, N/NodeNum, p/Priority, P/Partition, q/Qos, R/Reason, r/ReqNodes, s/SubmitTime,\n" +
-	 		"S/StartTime, t/State, T/JobType, u/User, U/Uid, x/ExcludeNodes.", field)
+			log.Errorf("Invalid format specifier or string : %s, string unfold case insensitive, reference:\n"+
+				"a/Account, c/CpuPerNode, C/AllocCPUs, e/ElapsedTime, h/Held, j/JobID, l/TimeLimit, L/NodeList,\n"+
+				"m/MemPerNode, n/Name, N/NodeNum, p/Priority, P/Partition, q/Qos, R/Reason, r/ReqNodes, s/SubmitTime,\n"+
+				"S/StartTime, t/State, T/JobType, u/User, U/Uid, x/ExcludeNodes.", field)
 			os.Exit(util.ErrorInvalidFormat)
 		}
 

internal/util/string.go

@@ -260,6 +260,35 @@ func CheckMailType(mailtype string) bool {
 		mailtype == "ALL"
 }
 
+func ParseLicensesString(val string) (map[string]uint32, error) {
+	tmp_str := strings.ReplaceAll(val, " ", "")
+
+	license_vec := strings.Split(tmp_str, ",")
+	lic_count_map := make(map[string]uint32)
+	pattern := regexp.MustCompile(`(\w+):(\d+)`)
+
+	for _, license_count := range license_vec {
+		if !pattern.MatchString(license_count) {
+			return nil, fmt.Errorf("invalid licenses string format")
+		}
+		matches := pattern.FindStringSubmatch(license_count)
+		if len(matches) != 3 {
+			return nil, fmt.Errorf("invalid licenses string format")
+		}
+		name := matches[1]
+		count, err := strconv.ParseUint(matches[2], 10, 32)
+		if err != nil {
+			return nil, err
+		}
+		if count <= 0 {
+			return nil, fmt.Errorf("license count must > 0")
+		}
+		lic_count_map[name] = uint32(count)
+	}
+
+	return lic_count_map, nil
+}
+
 // CheckNodeList check if the node list is comma separated node names.
 // The node name should contain only letters and numbers, and start with a letter, end with a number.
 func CheckNodeList(nodeStr string) bool {