Refactor Session Clone API to include synchronous and asynchronous versions.

LiteRT-LM-PiperOrigin-RevId: 861944789

Author: ai-edge-bot

runtime/core/session_advanced.cc

@@ -329,7 +329,18 @@ absl::StatusOr<BenchmarkInfo*> SessionAdvanced::GetMutableBenchmarkInfo() {
   return execution_manager_lock->GetMutableBenchmarkInfo(session_id_);
 }
 
-absl::StatusOr<std::unique_ptr<Engine::Session>> SessionAdvanced::Clone(
+absl::StatusOr<std::unique_ptr<Engine::Session>> SessionAdvanced::Clone() {
+  absl::Status status = absl::OkStatus();
+  ASSIGN_OR_RETURN(auto session,
+                   CloneAsync([&status](absl::StatusOr<Responses> responses) {
+                     status = responses.status();
+                   }));
+  RETURN_IF_ERROR(WaitUntilDone());
+  RETURN_IF_ERROR(status);
+  return session;
+}
+
+absl::StatusOr<std::unique_ptr<Engine::Session>> SessionAdvanced::CloneAsync(
     absl::AnyInvocable<void(absl::StatusOr<Responses>)> callback) {
   auto execution_manager_lock = execution_manager_.lock();
   if (execution_manager_lock == nullptr) {

runtime/core/session_advanced.h

@@ -91,20 +91,8 @@ class SessionAdvanced : public Engine::Session {
       std::optional<BenchmarkInfo> benchmark_info);
 
   // TODO b/409401231 - Call execution manager's release session instead.
+  // Wait until all tasks are done before destroying the session.
   ~SessionAdvanced() override {
-    CancelProcess();
-    auto execution_manager_lock = execution_manager_.lock();
-    if (execution_manager_lock == nullptr) {
-      ABSL_LOG(ERROR) << "Execution manager should not be deleted before "
-                         "Session is deleted.";
-      return;
-    }
-
-    auto status = execution_manager_lock->CancelAllTasksInSession(session_id_);
-    if (!status.ok()) {
-      ABSL_LOG(ERROR) << "Failed to release session: " << status;
-    }
-
     WaitUntilDone().IgnoreError();
   };
 
@@ -193,7 +181,10 @@ class SessionAdvanced : public Engine::Session {
   }
 
   // TODO b/409401231 - Add unit tests for this function.
-  absl::StatusOr<std::unique_ptr<Session>> Clone(
+  absl::StatusOr<std::unique_ptr<Session>> Clone() override;
+
+  // TODO b/409401231 - Add unit tests for this function.
+  absl::StatusOr<std::unique_ptr<Session>> CloneAsync(
       absl::AnyInvocable<void(absl::StatusOr<Responses>)> callback) override;
 
  private: