Dev/partition allow accounts

internal/ccontrol/CmdArgParser.go

@@ -28,17 +28,20 @@ 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
+	FlagAllowedAccounts string
+	FlagDeniedAccounts  string
+	FlagTaskId          uint32
+	FlagTaskIds         string
+	FlagQueryAll        bool
+	FlagTimeLimit       string
+	FlagPriority        float64
+	FlagHoldTime        string
+	FlagConfigFilePath  string
+	FlagJson            bool
 
 	RootCmd = &cobra.Command{
 		Use:     "ccontrol",
@@ -49,6 +52,7 @@ var (
 			util.DetectNetworkProxy()
 			config := util.ParseConfig(FlagConfigFilePath)
 			stub = util.GetStubToCtldByConfig(config)
+			userUid = uint32(os.Getuid())
 		},
 	}
 	showCmd = &cobra.Command{
@@ -160,6 +164,31 @@ var (
 			}
 		},
 	}
+	updatePartitionCmd = &cobra.Command{
+		Use:   "partition [flags] partition_name",
+		Short: "Modify partition partition attributes",
+		Long:  "",
+		Args: func(cmd *cobra.Command, args []string) error {
+			err := cobra.ExactArgs(1)(cmd, args)
+			if err != nil {
+				return err
+			}
+
+			return nil
+		},
+		Run: func(cmd *cobra.Command, args []string) {
+			if cmd.Flags().Changed("allowed-accounts") {
+				if err := ModifyPartitionAllowedOrDeniedAccounts(args[0], true, FlagAllowedAccounts); err != util.ErrorSuccess {
+					os.Exit(err)
+				}
+			} else if cmd.Flags().Changed("denied-accounts") {
+				if err := ModifyPartitionAllowedOrDeniedAccounts(args[0], false, FlagDeniedAccounts); err != util.ErrorSuccess {
+					os.Exit(err)
+				}
+				log.Warning("Hint: When using AllowedAccounts, DeniedAccounts will not take effect.")
+			}
+		},
+	}
 	holdCmd = &cobra.Command{
 		Use:   "hold [flags] job_id[,job_id...]",
 		Short: "prevent specified job from starting. ",
@@ -249,6 +278,15 @@ func init() {
 				return
 			}
 		}
+
+		updateCmd.AddCommand(updatePartitionCmd)
+		{
+			updatePartitionCmd.Flags().StringVarP(&FlagAllowedAccounts, "allowed-accounts", "A", "", "Set the allow account list for the partition")
+			updatePartitionCmd.Flags().StringVarP(&FlagDeniedAccounts, "denied-accounts", "D", "", "Set the denied account list for the partition")
+
+			updatePartitionCmd.MarkFlagsMutuallyExclusive("allowed-accounts", "denied-accounts")
+			updatePartitionCmd.MarkFlagsOneRequired("allowed-accounts", "denied-accounts")
+		}
 	}
 	RootCmd.AddCommand(holdCmd)
 	{

internal/ccontrol/ccontrol.go

@@ -37,7 +37,8 @@ import (
 )
 
 var (
-	stub protos.CraneCtldClient
+	userUid uint32
+	stub    protos.CraneCtldClient
 )
 
 func formatDeviceMap(data *protos.DeviceMap) string {
@@ -103,6 +104,21 @@ func formatMemToMB(data uint64) string {
 	}
 }
 
+func formatAllowedAccounts(allowedAccounts []string) string {
+	if len(allowedAccounts) == 0 {
+		return "ALL"
+	}
+	return strings.Join(allowedAccounts, ",")
+}
+
+func formatDeniedAccounts(allowedAccounts []string, deniedAccounts []string) string {
+	if len(deniedAccounts) == 0 {
+		return "None"
+	}
+
+	return strings.Join(deniedAccounts, ",")
+}
+
 func ShowNodes(nodeName string, queryAll bool) util.CraneCmdError {
 	req := &protos.QueryCranedInfoRequest{CranedName: nodeName}
 	reply, err := stub.QueryCranedInfo(context.Background(), req)
@@ -206,12 +222,15 @@ func ShowPartitions(partitionName string, queryAll bool) util.CraneCmdError {
 	} else {
 		for _, partitionInfo := range reply.PartitionInfo {
 			fmt.Printf("PartitionName=%v State=%v\n"+
+				"\tAllowedAccounts=%s DeniedAccounts=%s\n"+
 				"\tTotalNodes=%d AliveNodes=%d\n"+
 				"\tTotalCPU=%.2f AvailCPU=%.2f AllocCPU=%.2f\n"+
 				"\tTotalMem=%s AvailMem=%s AllocMem=%s\n"+
 				"\tTotalGres=%s AvailGres=%s AllocGres=%s\n"+
 				"\tHostList=%v\n\n",
 				partitionInfo.Name, partitionInfo.State.String()[10:],
+				formatAllowedAccounts(partitionInfo.AllowedAccounts),
+				formatDeniedAccounts(partitionInfo.AllowedAccounts, partitionInfo.DeniedAccounts),
 				partitionInfo.TotalNodes, partitionInfo.AliveNodes,
 				math.Abs(partitionInfo.ResTotal.AllocatableRes.CpuCoreLimit),
 				math.Abs(partitionInfo.ResAvail.AllocatableRes.CpuCoreLimit),
@@ -627,3 +646,22 @@ func ChangeNodeState(nodeRegex string, state string, reason string) util.CraneCm
 
 	return SummarizeReply(reply)
 }
+
+func ModifyPartitionAllowedOrDeniedAccounts(partition string, isModifyAllowed bool, accounts string) util.CraneCmdError {
+	accountList, _ := util.ParseStringParamList(accounts, ",")
+
+	req := protos.ModifyPartitionAllowedOrDeniedAccountsRequest{Uid: userUid, PartitionName: partition, IsModifyAllowed: isModifyAllowed, Accounts: accountList}
+	reply, err := stub.ModifyPartitionAllowedOrDeniedAccounts(context.Background(), &req)
+	if err != nil {
+		util.GrpcErrorPrintf(err, "Faild to modify partition %s", partition)
+		return util.ErrorNetwork
+	}
+
+	if !reply.GetOk() {
+		fmt.Printf("Modify partition %s failed: %s.\n", partition, util.ErrMsg(reply.GetErrCode()))
+		return util.ErrorBackend
+	}
+
+	fmt.Printf("Modify partition %s succeeded.\n", partition)
+	return util.ErrorSuccess
+}

protos/Crane.proto

@@ -386,6 +386,19 @@ message BlockAccountOrUserReply {
   ErrCode reason = 2;
 }
 
+message ModifyPartitionAllowedOrDeniedAccountsRequest {
+  uint32 uid = 1;
+  string partition_name = 2;
+  bool is_modify_allowed = 3;
+  repeated string accounts = 4;
+}
+
+
+message ModifyPartitionAllowedOrDeniedAccountsReply {
+  bool ok = 1;
+  ErrCode err_code = 2;
+}
+
 message MigrateSshProcToCgroupRequest {
   int32 pid = 1;
   uint32 task_id = 2;
@@ -768,6 +781,7 @@ service CraneCtld {
   rpc QueryPartitionInfo(QueryPartitionInfoRequest) returns (QueryPartitionInfoReply);
   rpc ModifyTask(ModifyTaskRequest) returns (ModifyTaskReply);
   rpc ModifyNode(ModifyCranedStateRequest) returns (ModifyCranedStateReply);
+  rpc ModifyPartitionAllowedOrDeniedAccounts(ModifyPartitionAllowedOrDeniedAccountsRequest) returns (ModifyPartitionAllowedOrDeniedAccountsReply);
 
   /* RPCs called from cacctmgr */
   rpc AddAccount(AddAccountRequest) returns (AddAccountReply);

protos/PublicDefs.proto

@@ -285,6 +285,9 @@ message PartitionInfo {
   ResourceView res_total = 6;
   ResourceView res_avail = 7;
   ResourceView res_alloc = 8;
+
+  repeated string allowed_accounts = 9;
+  repeated string denied_accounts = 10;
 }
 
 message CranedInfo {