Dev/partition allow accounts (#404)

* feat: ModifyPartitionAllowedAccounts grpc * feat: Read AllowAccounts config on startup * feat: Add AllowedAccounts entry display in scontrol show partition * feat: ccontrol update partition allow accounts * feat: When submitting a task, check whether the account is in the partition's AllowAccounts. * refactor allow to allowed * feat: config denied * feat: denied account * refactor * refactor * refactor: rebase master, add errcode * refactor * refactor: config.yaml * refactor errcode * refactor
PKUHPC · Mar 5, 2025 · 25e6caa · 25e6caa
1 parent 3621bc2
commit 25e6caa
Show file tree

Hide file tree

Showing 12 changed files with 228 additions and 11 deletions.
diff --git a/etc/config.yaml b/etc/config.yaml
@@ -98,6 +98,10 @@ Partitions:
     DefaultMemPerCpu: 0
     # Optional maximum memory per cpu in MB, 0 indicates no limit
     MaxMemPerCpu: 0
+    # AllowedAccounts and DeniedAccounts are mutually exclusive, only one should be configured.
+    # If AllowedAccounts is set, any DeniedAccounts configuration will be ignored.
+    # AllowedAccounts: ac1, ac2
+    # DeniedAccounts: ac1, ac2
 
 DefaultPartition: CPU
 

diff --git a/protos/Crane.proto b/protos/Crane.proto
@@ -386,6 +386,19 @@ message BlockAccountOrUserReply {
   ErrCode code = 2;
 }
 
+message ModifyPartitionAclRequest {
+  uint32 uid = 1;
+  string partition = 2;
+  bool is_allowed_list = 3;
+  repeated string accounts = 4;
+}
+
+
+message ModifyPartitionAclReply {
+  bool ok = 1;
+  ErrCode 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 ModifyPartitionAcl(ModifyPartitionAclRequest) returns (ModifyPartitionAclReply);
 
   /* RPCs called from cacctmgr */
   rpc AddAccount(AddAccountRequest) returns (AddAccountReply);

diff --git a/protos/PublicDefs.proto b/protos/PublicDefs.proto
@@ -296,6 +296,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 {
@@ -412,6 +415,8 @@ enum ErrCode {
 
   ERR_MAX_JOB_COUNT_PER_USER = 65;
   ERR_USER_NO_PRIVILEGE = 66;
+  ERR_NOT_IN_ALLOWED_LIST = 67;  // The current account is not in the allowed account list for the partition.
+  ERR_IN_DENIED_LIST = 68;  // The current account has been explicitly added to the deny list for the partition.
 }
 
 enum EntityType {

diff --git a/src/CraneCtld/AccountManager.cpp b/src/CraneCtld/AccountManager.cpp
@@ -1052,6 +1052,30 @@ CraneExpected<void> AccountManager::CheckIfUidHasPermOnUser(
   return CheckIfUserHasPermOnUserNoLock_(*op_user, user, read_only_priv);
 }
 
+CraneExpected<void> AccountManager::CheckModifyPartitionAcl(
+    uint32_t uid, const std::string& partition_name,
+    const std::unordered_set<std::string>& accounts) {
+  CraneExpected<void> result{};
+
+  util::read_lock_guard user_guard(m_rw_user_mutex_);
+  util::read_lock_guard account_guard(m_rw_account_mutex_);
+
+  auto user_result = GetUserInfoByUidNoLock_(uid);
+  if (!user_result) return std::unexpected(user_result.error());
+  const User* op_user = user_result.value();
+  result = CheckIfUserHasHigherPrivThan_(*op_user, User::None);
+  if (!result) return result;
+
+  for (const auto& account_name : accounts) {
+    const Account* account = GetAccountInfoNoLock_(account_name);
+    result =
+        CheckPartitionIsAllowedNoLock_(account, partition_name, false, false);
+    if (!result) return std::unexpected(result.error());
+  }
+
+  return result;
+}
+
 CraneExpected<void> AccountManager::CheckAddUserAllowedPartitionNoLock_(
     const User* user, const Account* account, const std::string& partition) {
   auto result = CheckPartitionIsAllowedNoLock_(account, partition, false, true);

diff --git a/src/CraneCtld/AccountManager.h b/src/CraneCtld/AccountManager.h
@@ -40,7 +40,6 @@ class AccountManager {
   using QosMapMutexSharedPtr = util::ScopeConstSharedPtr<
       std::unordered_map<std::string, std::unique_ptr<Qos>>, util::rw_mutex>;
 
-
   AccountManager();
 
   ~AccountManager() = default;
@@ -128,19 +127,23 @@ class AccountManager {
                                       const std::string& account,
                                       const std::string& partition);
 
-  CraneExpected<void> CheckIfUserOfAccountIsEnabled(
-      const std::string& user, const std::string& account);
+  CraneExpected<void> CheckIfUserOfAccountIsEnabled(const std::string& user,
+                                                    const std::string& account);
 
   CraneExpected<void> CheckAndApplyQosLimitOnTask(const std::string& user,
-                                       const std::string& account,
-                                       TaskInCtld* task);
+                                                  const std::string& account,
+                                                  TaskInCtld* task);
 
   CraneExpected<std::string> CheckUidIsAdmin(uint32_t uid);
 
   CraneExpected<void> CheckIfUidHasPermOnUser(uint32_t uid,
                                               const std::string& username,
                                               bool read_only_priv);
 
+  CraneExpected<void> CheckModifyPartitionAcl(
+      uint32_t uid, const std::string& partition_name,
+      const std::unordered_set<std::string>& accounts);
+
  private:
   void InitDataMap_();
 

diff --git a/src/CraneCtld/CraneCtld.cpp b/src/CraneCtld/CraneCtld.cpp
@@ -465,6 +465,29 @@ void ParseConfig(int argc, char** argv) {
             }
           }
 
+          if (partition["AllowedAccounts"] &&
+              !partition["AllowedAccounts"].IsNull()) {
+            auto allowed_accounts_str =
+                partition["AllowedAccounts"].as<std::string>();
+            std::vector<std::string> allowed_accounts =
+                absl::StrSplit(allowed_accounts_str.data(), ",");
+            for (const auto& account_name : allowed_accounts) {
+              part.allowed_accounts.insert(absl::StripAsciiWhitespace(account_name).data());
+            }
+          }
+
+          if (partition["DeniedAccounts"] && !partition["DeniedAccounts"].IsNull()) {
+            auto denied_accounts_str = partition["DeniedAccounts"].as<std::string>();
+            std::vector<std::string> denied_accounts = absl::StrSplit(denied_accounts_str, ",");
+            for (const auto& account_name : denied_accounts) {
+              part.denied_accounts.insert(absl::StripAsciiWhitespace(account_name).data());
+            }
+
+            if (partition["AllowedAccounts"] &&
+              !partition["AllowedAccounts"].IsNull())
+              CRANE_WARN("Hint: When using AllowedAccounts, DeniedAccounts will not take effect.");
+          }
+
           if (partition["DefaultMemPerCpu"] &&
               !partition["DefaultMemPerCpu"].IsNull()) {
             part.default_mem_per_cpu =

diff --git a/src/CraneCtld/CranedMetaContainer.cpp b/src/CraneCtld/CranedMetaContainer.cpp
@@ -284,9 +284,13 @@ void CranedMetaContainer::InitFromConfig(const Config& config) {
     part_meta.partition_global_meta.res_total_inc_dead = part_res;
     part_meta.partition_global_meta.node_cnt = part_meta.craned_ids.size();
     part_meta.partition_global_meta.nodelist_str = partition.nodelist_str;
+    part_meta.partition_global_meta.allowed_accounts =
+        partition.allowed_accounts;
+    part_meta.partition_global_meta.denied_accounts = partition.denied_accounts;
 
     CRANE_DEBUG(
-        "partition [{}]'s Global resource now: (cpu: {}, mem: {}, gres: {}). "
+        "partition [{}]'s Global resource now: (cpu: {}, mem: {}, "
+        "gres: {}). "
         "It has {} craneds.",
         part_name,
         part_meta.partition_global_meta.res_total_inc_dead.CpuCount(),
@@ -347,7 +351,16 @@ CranedMetaContainer::QueryAllPartitionInfo() {
     part_info->set_total_nodes(part_meta->partition_global_meta.node_cnt);
     part_info->set_alive_nodes(
         part_meta->partition_global_meta.alive_craned_cnt);
-
+    auto* allowed_accounts = part_info->mutable_allowed_accounts();
+    for (const auto& account_name :
+         part_meta->partition_global_meta.allowed_accounts) {
+      allowed_accounts->Add()->assign(account_name);
+    }
+    auto* denied_accounts = part_info->mutable_denied_accounts();
+    for (const auto& account_name :
+         part_meta->partition_global_meta.denied_accounts) {
+      denied_accounts->Add()->assign(account_name);
+    }
     *part_info->mutable_res_total() = static_cast<crane::grpc::ResourceView>(
         part_meta->partition_global_meta.res_total);
     *part_info->mutable_res_avail() = static_cast<crane::grpc::ResourceView>(
@@ -379,7 +392,16 @@ crane::grpc::QueryPartitionInfoReply CranedMetaContainer::QueryPartitionInfo(
   part_info->set_name(part_meta->partition_global_meta.name);
   part_info->set_total_nodes(part_meta->partition_global_meta.node_cnt);
   part_info->set_alive_nodes(part_meta->partition_global_meta.alive_craned_cnt);
-
+  auto* allowed_accounts = part_info->mutable_allowed_accounts();
+  for (const auto& account_name :
+       part_meta->partition_global_meta.allowed_accounts) {
+    allowed_accounts->Add()->assign(account_name);
+  }
+  auto* denied_accounts = part_info->mutable_denied_accounts();
+  for (const auto& account_name :
+       part_meta->partition_global_meta.denied_accounts) {
+    denied_accounts->Add()->assign(account_name);
+  }
   if (part_meta->partition_global_meta.alive_craned_cnt > 0)
     part_info->set_state(crane::grpc::PartitionState::PARTITION_UP);
   else
@@ -574,6 +596,73 @@ crane::grpc::ModifyCranedStateReply CranedMetaContainer::ChangeNodeState(
   return reply;
 }
 
+CraneExpected<void> CranedMetaContainer::ModifyPartitionAcl(
+    const std::string& partition_name, bool is_allowed_list,
+    std::unordered_set<std::string>&& accounts) {
+  CraneExpected<void> result{};
+
+  auto part_metas_map = partition_metas_map_.GetMapSharedPtr();
+
+  const auto part_meta_iter = part_metas_map->find(partition_name);
+
+  if (part_meta_iter == part_metas_map->end())
+    return std::unexpected(CraneErrCode::ERR_INVALID_PARTITION);
+
+  auto part_meta = part_meta_iter->second.GetExclusivePtr();
+  auto& allowed_accounts = part_meta->partition_global_meta.allowed_accounts;
+  auto& denied_accounts = part_meta->partition_global_meta.denied_accounts;
+
+  if (is_allowed_list) {
+    allowed_accounts = std::move(accounts);
+  } else {
+    denied_accounts = std::move(accounts);
+  }
+
+  return result;
+}
+
+CraneExpected<void> CranedMetaContainer::CheckIfAccountIsAllowedInPartition(
+    const std::string& partition_name, const std::string& account_name) {
+  auto part_metas_map = partition_metas_map_.GetMapSharedPtr();
+
+  const auto part_meta_iter = part_metas_map->find(partition_name);
+
+  if (part_meta_iter == part_metas_map->end()) {
+    CRANE_DEBUG(
+        "the partition {} does not exist, submission of the task is "
+        "prohibited.",
+        partition_name);
+    return std::unexpected(CraneErrCode::ERR_INVALID_PARTITION);
+  }
+
+  auto part_meta = part_meta_iter->second.GetExclusivePtr();
+  const auto& allowed_accounts =
+      part_meta->partition_global_meta.allowed_accounts;
+
+  const auto& denied_accounts =
+      part_meta->partition_global_meta.denied_accounts;
+
+  if (!allowed_accounts.empty()) {
+    if (!allowed_accounts.contains(account_name)) {
+      CRANE_DEBUG(
+          "The account {} is not in the AllowedAccounts of the partition {}"
+          "specified for the task, submission of the task is prohibited.",
+          account_name, partition_name);
+      return std::unexpected(CraneErrCode::ERR_NOT_IN_ALLOWED_LIST);
+    }
+  } else if (!denied_accounts.empty()) {
+    if (denied_accounts.contains(account_name)) {
+      CRANE_DEBUG(
+          "The account {} is in the DeniedAccounts of the partition {}"
+          "specified for the task, submission of the task is prohibited.",
+          account_name, partition_name);
+      return std::unexpected(CraneErrCode::ERR_IN_DENIED_LIST);
+    }
+  }
+
+  return {};
+}
+
 void CranedMetaContainer::AddDedicatedResource(
     const CranedId& node_id, const DedicatedResourceInNode& resource) {
   if (!craned_meta_map_.Contains(node_id)) {

diff --git a/src/CraneCtld/CranedMetaContainer.h b/src/CraneCtld/CranedMetaContainer.h
@@ -90,6 +90,13 @@ class CranedMetaContainer final {
   crane::grpc::ModifyCranedStateReply ChangeNodeState(
       const crane::grpc::ModifyCranedStateRequest& request);
 
+  CraneExpected<void> ModifyPartitionAcl(
+      const std::string& partition_name, bool is_allowed_list,
+      std::unordered_set<std::string>&& accounts);
+
+  CraneExpected<void> CheckIfAccountIsAllowedInPartition(
+      const std::string& partition_name, const std::string& account_name);
+
   void CranedUp(const CranedId& craned_id);
 
   void CranedDown(const CranedId& craned_id);

diff --git a/src/CraneCtld/CtldGrpcServer.cpp b/src/CraneCtld/CtldGrpcServer.cpp
@@ -254,6 +254,40 @@ grpc::Status CraneCtldServiceImpl::ModifyNode(
   return grpc::Status::OK;
 }
 
+grpc::Status CraneCtldServiceImpl::ModifyPartitionAcl(
+    grpc::ServerContext *context,
+    const crane::grpc::ModifyPartitionAclRequest *request,
+    crane::grpc::ModifyPartitionAclReply *response) {
+  CraneExpected<void> result;
+
+  std::unordered_set<std::string> accounts;
+
+  for (const auto &account_name : request->accounts()) {
+    accounts.insert(account_name);
+  }
+
+  result = g_account_manager->CheckModifyPartitionAcl(
+      request->uid(), request->partition(), accounts);
+
+  if (!result) {
+    response->set_ok(false);
+    response->set_code(result.error());
+    return grpc::Status::OK;
+  }
+
+  result = g_meta_container->ModifyPartitionAcl(
+      request->partition(), request->is_allowed_list(), std::move(accounts));
+
+  if (!result) {
+    response->set_ok(false);
+    response->set_code(result.error());
+  } else {
+    response->set_ok(true);
+  }
+
+  return grpc::Status::OK;
+}
+
 grpc::Status CraneCtldServiceImpl::QueryTasksInfo(
     grpc::ServerContext *context,
     const crane::grpc::QueryTasksInfoRequest *request,
@@ -992,7 +1026,11 @@ CraneExpected<std::future<task_id_t>> CtldServer::SubmitTaskToScheduler(
     return std::unexpected(enable_res.error());
   }
 
-  auto result = g_task_scheduler->AcquireTaskAttributes(task.get());
+  auto result = g_meta_container->CheckIfAccountIsAllowedInPartition(
+      task->partition_id, task->account);
+  if (!result) return std::unexpected(result.error());
+
+  result = g_task_scheduler->AcquireTaskAttributes(task.get());
 
   if (result) result = g_task_scheduler->CheckTaskValidity(task.get());
 

diff --git a/src/CraneCtld/CtldGrpcServer.h b/src/CraneCtld/CtldGrpcServer.h
@@ -227,6 +227,11 @@ class CraneCtldServiceImpl final : public crane::grpc::CraneCtld::Service {
       const crane::grpc::ModifyCranedStateRequest *request,
       crane::grpc::ModifyCranedStateReply *response) override;
 
+  grpc::Status ModifyPartitionAcl(
+      grpc::ServerContext *context,
+      const crane::grpc::ModifyPartitionAclRequest *request,
+      crane::grpc::ModifyPartitionAclReply *response) override;
+
   grpc::Status AddAccount(grpc::ServerContext *context,
                           const crane::grpc::AddAccountRequest *request,
                           crane::grpc::AddAccountReply *response) override;

diff --git a/src/CraneCtld/CtldPublicDefs.h b/src/CraneCtld/CtldPublicDefs.h
@@ -80,7 +80,8 @@ struct Config {
     // optional, 0 indicates no limit
     uint64_t max_mem_per_cpu;
     std::unordered_set<std::string> nodes;
-    std::unordered_set<std::string> AllowAccounts;
+    std::unordered_set<std::string> allowed_accounts;
+    std::unordered_set<std::string> denied_accounts;
   };
 
   struct CraneCtldListenConf {
@@ -215,6 +216,10 @@ struct PartitionGlobalMeta {
 
   std::string name;
   std::string nodelist_str;
+
+  std::unordered_set<std::string> allowed_accounts;
+  std::unordered_set<std::string> denied_accounts;
+
   uint32_t node_cnt;
   uint32_t alive_craned_cnt;
 };

diff --git a/src/Craned/CranedPublicDefs.h b/src/Craned/CranedPublicDefs.h
@@ -45,7 +45,7 @@ struct TaskInfoOfUid {
 
 struct Partition {
   std::unordered_set<std::string> nodes;
-  std::unordered_set<std::string> AllowAccounts;
+  std::unordered_set<std::string> AllowedAccounts;
 };
 
 struct Config {