feat: add suspend and resume

protos/Crane.proto

@@ -134,6 +134,24 @@ message TerminateStepsReply {
   string reason = 2;
 }
 
+message SuspendStepsRequest {
+  repeated uint32 task_id_list = 1;
+}
+
+message SuspendStepsReply {
+  bool ok = 1;
+  string reason = 2;
+}
+
+message ResumeStepsRequest {
+  repeated uint32 task_id_list = 1;
+}
+
+message ResumeStepsReply {
+  bool ok = 1;
+  string reason = 2;
+}
+
 message TerminateOrphanedStepRequest {
   repeated uint32 task_id_list = 1;
 }
@@ -225,6 +243,8 @@ message ModifyTaskRequest {
     TimeLimit = 0;
     Priority = 1;
     Hold = 2;
+    Suspend = 3;
+    Resume = 4;
   }
 
   uint32 uid = 1;
@@ -1007,6 +1027,8 @@ service Craned {
   If the task is a batch task, just kill it.
   */
   rpc TerminateSteps(TerminateStepsRequest) returns (TerminateStepsReply);
+  rpc SuspendSteps(SuspendStepsRequest) returns (SuspendStepsReply);
+  rpc ResumeSteps(ResumeStepsRequest) returns (ResumeStepsReply);
   rpc TerminateOrphanedStep(TerminateOrphanedStepRequest) returns (TerminateOrphanedStepReply);
   rpc ChangeJobTimeLimit(ChangeJobTimeLimitRequest) returns (ChangeJobTimeLimitReply);
 

protos/PublicDefs.proto

@@ -114,6 +114,7 @@ enum TaskStatus {
   ExceedTimeLimit = 4;
   Cancelled = 5;
   OutOfMemory = 6;
+  Suspended = 7;
 
   Invalid = 15;
 }

protos/Supervisor.proto

@@ -119,6 +119,20 @@ message TerminateTaskReply {
   string reason = 2;
 }
 
+message SuspendTaskRequest {}
+
+message SuspendTaskReply {
+  bool ok = 1;
+  string reason = 2;
+}
+
+message ResumeTaskRequest {}
+
+message ResumeTaskReply {
+  bool ok = 1;
+  string reason = 2;
+}
+
 message ShutdownSupervisorRequest {}
 
 message ShutdownSupervisorReply {}
@@ -129,5 +143,7 @@ service Supervisor {
   rpc CheckStatus(CheckStatusRequest) returns (CheckStatusReply);
   rpc ChangeTaskTimeLimit(ChangeTaskTimeLimitRequest) returns (ChangeTaskTimeLimitReply);
   rpc TerminateTask(TerminateTaskRequest) returns (TerminateTaskReply);
+  rpc SuspendTask(SuspendTaskRequest) returns (SuspendTaskReply);
+  rpc ResumeTask(ResumeTaskRequest) returns (ResumeTaskReply);
   rpc ShutdownSupervisor(ShutdownSupervisorRequest) returns (ShutdownSupervisorReply);
 }

src/CraneCtld/RpcService/CranedKeeper.cpp

@@ -238,6 +238,64 @@ CraneErrCode CranedStub::ReleaseCgroupForJobs(
   return CraneErrCode::SUCCESS;
 }
 
+CraneErrCode CranedStub::SuspendSteps(const std::vector<task_id_t> &task_ids) {
+  using crane::grpc::SuspendStepsReply;
+  using crane::grpc::SuspendStepsRequest;
+
+  ClientContext context;
+  Status status;
+  SuspendStepsRequest request;
+  SuspendStepsReply reply;
+  context.set_deadline(std::chrono::system_clock::now() +
+                       std::chrono::seconds(kCtldRpcTimeoutSeconds));
+
+  request.mutable_task_id_list()->Assign(task_ids.begin(), task_ids.end());
+
+  status = m_stub_->SuspendSteps(&context, request, &reply);
+  if (!status.ok()) {
+    CRANE_DEBUG("SuspendSteps RPC for Node {} failed: {}", m_craned_id_,
+                status.error_message());
+    HandleGrpcErrorCode_(status.error_code());
+    return CraneErrCode::ERR_RPC_FAILURE;
+  }
+  UpdateLastActiveTime();
+
+  if (reply.ok()) return CraneErrCode::SUCCESS;
+
+  CRANE_WARN("SuspendSteps RPC for Node {} declined: {}", m_craned_id_,
+             reply.reason());
+  return CraneErrCode::ERR_GENERIC_FAILURE;
+}
+
+CraneErrCode CranedStub::ResumeSteps(const std::vector<task_id_t> &task_ids) {
+  using crane::grpc::ResumeStepsReply;
+  using crane::grpc::ResumeStepsRequest;
+
+  ClientContext context;
+  Status status;
+  ResumeStepsRequest request;
+  ResumeStepsReply reply;
+  context.set_deadline(std::chrono::system_clock::now() +
+                       std::chrono::seconds(kCtldRpcTimeoutSeconds));
+
+  request.mutable_task_id_list()->Assign(task_ids.begin(), task_ids.end());
+
+  status = m_stub_->ResumeSteps(&context, request, &reply);
+  if (!status.ok()) {
+    CRANE_DEBUG("ResumeSteps RPC for Node {} failed: {}", m_craned_id_,
+                status.error_message());
+    HandleGrpcErrorCode_(status.error_code());
+    return CraneErrCode::ERR_RPC_FAILURE;
+  }
+  UpdateLastActiveTime();
+
+  if (reply.ok()) return CraneErrCode::SUCCESS;
+
+  CRANE_WARN("ResumeSteps RPC for Node {} declined: {}", m_craned_id_,
+             reply.reason());
+  return CraneErrCode::ERR_GENERIC_FAILURE;
+}
+
 CraneErrCode CranedStub::ChangeJobTimeLimit(uint32_t task_id,
                                             uint64_t seconds) {
   using crane::grpc::ChangeJobTimeLimitReply;

src/CraneCtld/RpcService/CranedKeeper.h

@@ -86,6 +86,10 @@ class CranedStub {
   CraneErrCode ReleaseCgroupForJobs(
       const std::vector<std::pair<task_id_t, uid_t>> &task_uid_pairs);
 
+  CraneErrCode SuspendSteps(const std::vector<task_id_t> &task_ids);
+
+  CraneErrCode ResumeSteps(const std::vector<task_id_t> &task_ids);
+
   CraneErrCode TerminateSteps(const std::vector<task_id_t> &task_ids);
 
   CraneErrCode TerminateOrphanedSteps(const std::vector<task_id_t> &task_ids);

src/CraneCtld/RpcService/CtldGrpcServer.cpp

@@ -585,6 +585,44 @@ grpc::Status CraneCtldServiceImpl::ModifyTask(
             fmt::format("Failed to hold/release job: {}.", CraneErrStr(err)));
       }
     }
+  } else if (request->attribute() == ModifyTaskRequest::Suspend) {
+    for (auto task_id : request->task_ids()) {
+      err = g_task_scheduler->SuspendRunningTask(task_id);
+      if (err == CraneErrCode::SUCCESS) {
+        response->add_modified_tasks(task_id);
+      } else {
+        response->add_not_modified_tasks(task_id);
+        if (err == CraneErrCode::ERR_NON_EXISTENT) {
+          response->add_not_modified_reasons(
+              fmt::format("Task #{} was not found in running queue.", task_id));
+        } else if (err == CraneErrCode::ERR_INVALID_PARAM) {
+          response->add_not_modified_reasons(
+              fmt::format("Task #{} is not running.", task_id));
+        } else {
+          response->add_not_modified_reasons(
+              fmt::format("Failed to suspend task: {}.", CraneErrStr(err)));
+        }
+      }
+    }
+  } else if (request->attribute() == ModifyTaskRequest::Resume) {
+    for (auto task_id : request->task_ids()) {
+      err = g_task_scheduler->ResumeSuspendedTask(task_id);
+      if (err == CraneErrCode::SUCCESS) {
+        response->add_modified_tasks(task_id);
+      } else {
+        response->add_not_modified_tasks(task_id);
+        if (err == CraneErrCode::ERR_NON_EXISTENT) {
+          response->add_not_modified_reasons(
+              fmt::format("Task #{} was not found in running queue.", task_id));
+        } else if (err == CraneErrCode::ERR_INVALID_PARAM) {
+          response->add_not_modified_reasons(
+              fmt::format("Task #{} is not suspended.", task_id));
+        } else {
+          response->add_not_modified_reasons(
+              fmt::format("Failed to resume task: {}.", CraneErrStr(err)));
+        }
+      }
+    }
   } else {
     for (auto task_id : request->task_ids()) {
       response->add_not_modified_tasks(task_id);

src/CraneCtld/TaskScheduler.cpp

@@ -996,6 +996,193 @@ std::future<CraneErrCode> TaskScheduler::HoldReleaseTaskAsync(task_id_t task_id,
   return std::move(future);
 }
 
+CraneErrCode TaskScheduler::SuspendRunningTask(task_id_t task_id) {
+  std::vector<CranedId> executing_nodes;
+
+  {
+    LockGuard running_guard(&m_running_task_map_mtx_);
+    auto iter = m_running_task_map_.find(task_id);
+    if (iter == m_running_task_map_.end()) {
+      CRANE_TRACE("Task #{} not in Rn queue for suspend", task_id);
+      return CraneErrCode::ERR_NON_EXISTENT;
+    }
+
+    TaskInCtld* task = iter->second.get();
+    if (task->Status() == crane::grpc::Suspended) {
+      CRANE_TRACE("Task #{} already suspended", task_id);
+      return CraneErrCode::ERR_INVALID_PARAM;
+    }
+    if (task->Status() != crane::grpc::Running) {
+      CRANE_TRACE("Task #{} is not running (status {}) for suspend", task_id,
+                  static_cast<int>(task->Status()));
+      return CraneErrCode::ERR_INVALID_PARAM;
+    }
+
+    executing_nodes = task->executing_craned_ids;
+  }
+
+  if (executing_nodes.empty()) {
+    CRANE_WARN("Task #{} has no executing craned when suspending", task_id);
+    return CraneErrCode::ERR_INVALID_PARAM;
+  }
+
+  std::vector<task_id_t> job_ids{task_id};
+  std::vector<CranedId> suspended_nodes;
+  CraneErrCode failure_code = CraneErrCode::SUCCESS;
+  bool has_failure = false;
+  for (const auto& craned_id : executing_nodes) {
+    auto stub = g_craned_keeper->GetCranedStub(craned_id);
+    if (!stub || stub->Invalid()) {
+      CRANE_WARN("SuspendSteps stub for {} unavailable", craned_id);
+      if (failure_code == CraneErrCode::SUCCESS)
+        failure_code = CraneErrCode::ERR_RPC_FAILURE;
+      has_failure = true;
+      continue;
+    }
+
+    CraneErrCode err = stub->SuspendSteps(job_ids);
+    if (err != CraneErrCode::SUCCESS) {
+      CRANE_ERROR("Failed to suspend task #{} on craned {}: {}", task_id,
+                  craned_id, CraneErrStr(err));
+      if (failure_code == CraneErrCode::SUCCESS) failure_code = err;
+      has_failure = true;
+      continue;
+    }
+
+    suspended_nodes.push_back(craned_id);
+  }
+
+  if (has_failure) {
+    for (const auto& craned_id : suspended_nodes) {
+      auto stub = g_craned_keeper->GetCranedStub(craned_id);
+      if (!stub || stub->Invalid()) {
+        CRANE_ERROR(
+            "Failed to rollback suspended state on craned {}: stub unavailable",
+            craned_id);
+        continue;
+      }
+
+      CraneErrCode err = stub->ResumeSteps(job_ids);
+      if (err != CraneErrCode::SUCCESS) {
+        CRANE_ERROR("Rollback ResumeSteps for task #{} on craned {} failed: {}",
+                    task_id, craned_id, CraneErrStr(err));
+      }
+    }
+
+    return failure_code;
+  }
+
+  {
+    LockGuard running_guard(&m_running_task_map_mtx_);
+    auto iter = m_running_task_map_.find(task_id);
+    if (iter == m_running_task_map_.end())
+      return CraneErrCode::ERR_NON_EXISTENT;
+
+    TaskInCtld* task = iter->second.get();
+    if (task->Status() != crane::grpc::Suspended) {
+      task->SetStatus(crane::grpc::Suspended);
+      if (!g_embedded_db_client->UpdateRuntimeAttrOfTaskIfExists(
+              0, task->TaskDbId(), task->RuntimeAttr())) {
+        CRANE_ERROR("Failed to persist suspended status for task #{}", task_id);
+      }
+    }
+  }
+
+  return CraneErrCode::SUCCESS;
+}
+
+CraneErrCode TaskScheduler::ResumeSuspendedTask(task_id_t task_id) {
+  std::vector<CranedId> executing_nodes;
+
+  {
+    LockGuard running_guard(&m_running_task_map_mtx_);
+    auto iter = m_running_task_map_.find(task_id);
+    if (iter == m_running_task_map_.end()) {
+      CRANE_TRACE("Task #{} not in Rn queue for resume", task_id);
+      return CraneErrCode::ERR_NON_EXISTENT;
+    }
+
+    TaskInCtld* task = iter->second.get();
+    if (task->Status() != crane::grpc::Suspended) {
+      CRANE_TRACE("Task #{} is not suspended (status {}) for resume", task_id,
+                  static_cast<int>(task->Status()));
+      return CraneErrCode::ERR_INVALID_PARAM;
+    }
+
+    executing_nodes = task->executing_craned_ids;
+  }
+
+  if (executing_nodes.empty()) {
+    CRANE_WARN("Task #{} has no executing craned when resuming", task_id);
+    return CraneErrCode::ERR_INVALID_PARAM;
+  }
+
+  std::vector<task_id_t> job_ids{task_id};
+  std::vector<CranedId> resumed_nodes;
+  CraneErrCode failure_code = CraneErrCode::SUCCESS;
+  bool has_failure = false;
+  for (const auto& craned_id : executing_nodes) {
+    auto stub = g_craned_keeper->GetCranedStub(craned_id);
+    if (!stub || stub->Invalid()) {
+      CRANE_WARN("ResumeSteps stub for {} unavailable", craned_id);
+      if (failure_code == CraneErrCode::SUCCESS)
+        failure_code = CraneErrCode::ERR_RPC_FAILURE;
+      has_failure = true;
+      continue;
+    }
+
+    CraneErrCode err = stub->ResumeSteps(job_ids);
+    if (err != CraneErrCode::SUCCESS) {
+      CRANE_ERROR("Failed to resume task #{} on craned {}: {}", task_id,
+                  craned_id, CraneErrStr(err));
+      if (failure_code == CraneErrCode::SUCCESS) failure_code = err;
+      has_failure = true;
+      continue;
+    }
+
+    resumed_nodes.push_back(craned_id);
+  }
+
+  if (has_failure) {
+    for (const auto& craned_id : resumed_nodes) {
+      auto stub = g_craned_keeper->GetCranedStub(craned_id);
+      if (!stub || stub->Invalid()) {
+        CRANE_ERROR(
+            "Failed to rollback resumed state on craned {}: stub unavailable",
+            craned_id);
+        continue;
+      }
+
+      CraneErrCode err = stub->SuspendSteps(job_ids);
+      if (err != CraneErrCode::SUCCESS) {
+        CRANE_ERROR(
+            "Rollback SuspendSteps for task #{} on craned {} failed: {}",
+            task_id, craned_id, CraneErrStr(err));
+      }
+    }
+
+    return failure_code;
+  }
+
+  {
+    LockGuard running_guard(&m_running_task_map_mtx_);
+    auto iter = m_running_task_map_.find(task_id);
+    if (iter == m_running_task_map_.end())
+      return CraneErrCode::ERR_NON_EXISTENT;
+
+    TaskInCtld* task = iter->second.get();
+    if (task->Status() != crane::grpc::Running) {
+      task->SetStatus(crane::grpc::Running);
+      if (!g_embedded_db_client->UpdateRuntimeAttrOfTaskIfExists(
+              0, task->TaskDbId(), task->RuntimeAttr())) {
+        CRANE_ERROR("Failed to persist resumed status for task #{}", task_id);
+      }
+    }
+  }
+
+  return CraneErrCode::SUCCESS;
+}
+
 CraneErrCode TaskScheduler::ChangeTaskTimeLimit(task_id_t task_id,
                                                 int64_t secs) {
   if (!CheckIfTimeLimitSecIsValid(secs)) return CraneErrCode::ERR_INVALID_PARAM;