From 0b63b933626b13d936055c4e833c1d49a76ce56c Mon Sep 17 00:00:00 2001
From: Daniel Ng <dnlng@google.com>
Date: Wed, 22 Jul 2026 14:03:03 -0700
Subject: [PATCH] Internal

PiperOrigin-RevId: 952318773
---
 .../experimental/tiering_service/auth.py      |   4 +-
 .../experimental/tiering_service/auth_test.py |   6 +-
 .../experimental/tiering_service/client.py    | 886 +++++++++---------
 .../tiering_service/client_test.py            | 361 +++----
 .../experimental/tiering_service/server.py    |  11 +-
 .../tiering_service/server_test.py            |   2 +-
 6 files changed, 657 insertions(+), 613 deletions(-)

diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py
index 9d0380254f..0401940c4d 100644
--- a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py
+++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py
@@ -19,6 +19,7 @@
 """
 
 from collections.abc import Collection
+
 from absl import logging
 import grpc
 from orbax.checkpoint.experimental.tiering_service import db_schema
@@ -39,7 +40,8 @@ async def get_oauth_token(context: grpc.aio.ServicerContext) -> str | None:
     The extracted OAuth token string, or None if not found or malformed.
   """
   logging.debug("Extracting OAuth token from metadata")
-  metadata = dict(await context.invocation_metadata())  # pyrefly: ignore[not-async]
+  raw_metadata = context.invocation_metadata()
+  metadata = dict(raw_metadata)
   # Standard header for OAuth tokens in gRPC is 'authorization'.
   auth_header = metadata.get("authorization")
 
diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/auth_test.py
index cdb9d35a06..e4272be5c4 100644
--- a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth_test.py
+++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/auth_test.py
@@ -28,7 +28,7 @@ async def test_get_oauth_token_success(self):
     context = mock.create_autospec(
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
-    context.invocation_metadata = mock.AsyncMock(
+    context.invocation_metadata = mock.Mock(
         return_value=(("authorization", "Bearer valid-token"),)
     )
     token = await auth.get_oauth_token(context)
@@ -38,7 +38,7 @@ async def test_get_oauth_token_no_header(self):
     context = mock.create_autospec(
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
-    context.invocation_metadata = mock.AsyncMock(return_value=())
+    context.invocation_metadata = mock.Mock(return_value=())
     token = await auth.get_oauth_token(context)
     self.assertIsNone(token)
 
@@ -46,7 +46,7 @@ async def test_get_oauth_token_malformed_header(self):
     context = mock.create_autospec(
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
-    context.invocation_metadata = mock.AsyncMock(
+    context.invocation_metadata = mock.Mock(
         return_value=(("authorization", "not-bearer-token"),)
     )
     token = await auth.get_oauth_token(context)
diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py
index 3a42c0d118..946de55385 100644
--- a/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py
+++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py
@@ -17,6 +17,7 @@
 import asyncio
 from collections.abc import Sequence
 import enum
+import threading
 from typing import Any
 from absl import logging
 import grpc
@@ -27,32 +28,17 @@
 
 
 class JobType(enum.Enum):
-  """Job types managed by the centralized keep-alive manager."""
+  """Job types managed by TieringClient."""
 
   WRITE = "write"
   PREFETCH = "prefetch"
 
 
-class _KeepAliveJob:
-  """Represents an active keep-alive job managed by the centralized manager."""
-
-  def __init__(
-      self,
-      asset_uuid: str,
-      job_type: JobType,
-      interval: float,
-      tier_path_uuid: str | None = None,
-  ):
-    self.asset_uuid = asset_uuid
-    self.job_type = job_type
-    self.interval = interval
-    self.tier_path_uuid = tier_path_uuid
-    self.loop = asyncio.get_running_loop()
-    self.next_run = asyncio.get_running_loop().time() + interval
-
-
 class TieringClient:
-  """Client library to communicate with the Checkpoint Tiering Service (CTS)."""
+  """Client library to communicate with the Checkpoint Tiering Service (CTS).
+
+  Supports single active checkpoint operation per TieringClient instance.
+  """
 
   def __init__(
       self, server_address: str = "localhost:50051", secure: bool = False
@@ -65,16 +51,30 @@ def __init__(
     """
     self._server_address = server_address
     self._secure = secure
+    self._loop: asyncio.AbstractEventLoop | None = None
+    self._thread: threading.Thread | None = None
+
     self._channel = None
     self._stub = None
+
     self._zone = None
     self._region = None
     self._env_queried = False
     self._env_lock = None
-    self._keep_alives: dict[tuple[str, JobType], _KeepAliveJob] = {}
-    self._keep_alive_manager_task: asyncio.Task[None] | None = None
-    self._keep_alive_event: asyncio.Event = asyncio.Event()
-    self._prefetch_futures: dict[str, asyncio.Future[str]] = {}
+
+    # Single Active Asset State
+    self._active_asset_uuid: str | None = None
+    self._active_job_type: JobType | None = None
+    self._active_tier_path_uuid: str | None = None
+    self._active_path: str | None = None
+    self._keep_alive_task: asyncio.Task[None] | None = None
+    self._prefetch_future: asyncio.Future[str] | None = None
+
+    self._connect_lock = threading.Lock()
+    self._connected_event = threading.Event()
+    self._connect_error: Exception | None = None
+    self._closed = False
+    self._connecting = False
 
   async def __aenter__(self) -> "TieringClient":
     await self.connect()
@@ -83,56 +83,171 @@ async def __aenter__(self) -> "TieringClient":
   async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
     await self.close()
 
-  async def connect(self) -> None:
-    """Establishes an async gRPC channel with the server."""
-    if self._channel is not None:
-      return
-
-    if self._secure:
-      is_local = (
-          "localhost" in self._server_address
-          or "127.0.0.1" in self._server_address
+  def _get_loop(self) -> asyncio.AbstractEventLoop:
+    """Gets the running event loop or raises a RuntimeError if not set."""
+    if self._loop is None:
+      raise RuntimeError(
+          "TieringClient is not connected. Call connect() first."
       )
-      if is_local:
-        try:
-          # Secure channel setup. Fall back to SSL if local creds not supported.
-          creds = grpc.local_channel_credentials()
-        except AttributeError:
+    return self._loop
+
+  async def _ensure_connected(self) -> None:
+    if self._closed:
+      raise RuntimeError("TieringClient has been closed")
+    if self._loop is None:
+      await self.connect()
+
+  def _get_or_create_stub(self) -> tiering_service_pb2_grpc.TieringServiceStub:
+    """Gets or creates the gRPC stub."""
+    if self._stub is None:
+      if self._secure:
+        is_local = (
+            "localhost" in self._server_address
+            or "127.0.0.1" in self._server_address
+        )
+        if is_local:
+          try:
+            creds = grpc.local_channel_credentials()
+          except AttributeError:
+            creds = grpc.ssl_channel_credentials()
+        else:
           creds = grpc.ssl_channel_credentials()
+        self._channel = grpc.aio.secure_channel(self._server_address, creds)
       else:
-        creds = grpc.ssl_channel_credentials()
-      self._channel = grpc.aio.secure_channel(self._server_address, creds)
-    else:
-      self._channel = grpc.aio.insecure_channel(self._server_address)
+        self._channel = grpc.aio.insecure_channel(self._server_address)
 
-    self._stub = tiering_service_pb2_grpc.TieringServiceStub(self._channel)
+      self._stub = tiering_service_pb2_grpc.TieringServiceStub(self._channel)
 
-  async def close(self) -> None:
-    """Closes the gRPC channel."""
-    # Release pending prefetches and cancel futures
-    for asset_uuid, fut in list(self._prefetch_futures.items()):
-      if not fut.done():
-        self._stop_prefetch_keep_alive(asset_uuid)
-        fut.cancel()
-    self._prefetch_futures.clear()
-
-    # Cancel manager task and clear job list
-    if self._keep_alive_manager_task is not None:
-      self._keep_alive_manager_task.cancel()
+    return self._stub
+
+  async def _async_connect(self) -> None:
+    self._get_or_create_stub()
+
+  async def connect(self) -> None:
+    """Establishes a gRPC channel with the server and starts background loop."""
+    should_start = False
+    with self._connect_lock:
+      if self._closed:
+        raise RuntimeError("TieringClient has been closed")
+      if self._connected_event.is_set() and self._loop is not None:
+        return
+      if not self._connecting:
+        self._connecting = True
+        should_start = True
+
+    if should_start:
+      self._connect_error = None
+      try:
+        loop_started = self._start_background_loop()
+        await asyncio.to_thread(loop_started.wait)
+
+        fut = asyncio.run_coroutine_threadsafe(
+            self._async_connect(), self._loop
+        )
+        await asyncio.wrap_future(fut)
+        self._connected_event.set()
+      except Exception as e:
+        self._connect_error = e
+        self._connected_event.set()  # Unblock waiting callers
+        await self._cleanup_connection_failure()
+        raise
+    else:
+      await asyncio.to_thread(self._connected_event.wait)
+      if self._connect_error is not None:
+        raise RuntimeError(
+            f"Connection to TieringClient failed: {self._connect_error}"
+        ) from self._connect_error
+      if self._loop is None:
+        raise RuntimeError("Connection to TieringClient failed.")
+
+  def _start_background_loop(self) -> threading.Event:
+    """Creates a new event loop and runs it in a background thread."""
+    self._loop = asyncio.new_event_loop()
+    loop_started = threading.Event()
+
+    def run_loop():
+      if not self._loop:
+        raise RuntimeError("Event loop not set.")
+      asyncio.set_event_loop(self._loop)
+      self._env_lock = asyncio.Lock()
+      loop_started.set()
+      self._loop.run_forever()
+
+    self._thread = threading.Thread(target=run_loop, daemon=True)
+    self._thread.start()
+    return loop_started
+
+  async def _cleanup_connection_failure(self) -> None:
+    """Cleans up background thread and loop on connection failure."""
+    if self._loop is not None:
+      self._loop.call_soon_threadsafe(self._loop.stop)
+    if self._thread is not None and threading.current_thread() != self._thread:
+      await asyncio.to_thread(self._thread.join)
+    with self._connect_lock:
+      self._connecting = False
+      self._loop = None
+      self._thread = None
+
+  async def _async_close(self) -> None:
+    """Closes all active background tasks, futures, and gRPC channels."""
+    if self._keep_alive_task is not None:
+      self._keep_alive_task.cancel()
       try:
-        await self._keep_alive_manager_task
+        await self._keep_alive_task
       except asyncio.CancelledError:
         pass
-      self._keep_alive_manager_task = None
+      self._keep_alive_task = None
 
-    self._keep_alives.clear()
-    self._keep_alive_event.clear()
+    if self._prefetch_future is not None and not self._prefetch_future.done():
+      self._prefetch_future.cancel()
+      self._prefetch_future = None
 
-    if self._channel is not None:
-      await self._channel.close()
+    self._reset_active_state()
+
+    chan = self._channel
+    if chan is not None:
+      await chan.close()
       self._channel = None
       self._stub = None
 
+  async def close(self) -> None:
+    """Closes the gRPC channel and stops background loop."""
+    self._closed = True
+    if self._loop is None:
+      return
+
+    try:
+      fut = asyncio.run_coroutine_threadsafe(self._async_close(), self._loop)
+      await asyncio.wrap_future(fut)
+    finally:
+      if self._loop and self._loop.is_running():
+        self._loop.call_soon_threadsafe(self._loop.stop)
+      th = self._thread
+      if th is not None and threading.current_thread() != th:
+        await asyncio.to_thread(th.join)
+        self._thread = None
+      self._loop = None
+      self._connected_event.clear()
+      with self._connect_lock:
+        self._connecting = False
+
+  def _reset_active_state(self) -> None:
+    self._active_asset_uuid = None
+    self._active_job_type = None
+    self._active_tier_path_uuid = None
+    self._active_path = None
+
+  def _check_active_operation(self) -> None:
+    if self._active_asset_uuid is not None:
+      job_name = (
+          self._active_job_type.value if self._active_job_type else "operation"
+      )
+      raise RuntimeError(
+          f"TieringClient is already managing active asset "
+          f"'{self._active_asset_uuid}' ({job_name}). Call finalize() or "
+          "release() first, or use a separate TieringClient instance."
+      )
+
   async def _get_gcp_zone_and_region(self) -> tuple[str | None, str | None]:
     """Retrieves and caches GCP zone and region."""
     lock = self._env_lock
@@ -153,257 +268,133 @@ async def _get_auth_metadata(self) -> list[tuple[str, str]]:
       return [("authorization", f"Bearer {token}")]
     return []
 
-  def _ensure_manager_running(self) -> None:
-    if (
-        self._keep_alive_manager_task is None
-        or self._keep_alive_manager_task.done()
-    ):
-      self._keep_alive_manager_task = asyncio.create_task(
-          self._keep_alive_manager_loop()
-      )
-
-  def _start_write_keep_alive(self, asset_uuid: str, interval: int) -> None:
-    """Starts the write keep-alive background task."""
-    job = _KeepAliveJob(
-        asset_uuid=asset_uuid,
-        job_type=JobType.WRITE,
-        interval=max(1.0, float(interval) * 0.8),
-    )
-    self._keep_alives[(asset_uuid, JobType.WRITE)] = job
-    self._ensure_manager_running()
-    self._keep_alive_event.set()
-
-  def _stop_write_keep_alive(self, asset_uuid: str) -> None:
-    """Stops the write keep-alive background task."""
-    job = self._keep_alives.pop((asset_uuid, JobType.WRITE), None)
-    if job:
-      self._keep_alive_event.set()
-
-  def _start_prefetch_keep_alive(
-      self, asset_uuid: str, tier_path_uuid: str, interval: int
+  async def _write_keep_alive_loop(
+      self, asset_uuid: str, interval: float
   ) -> None:
-    """Starts the prefetch keep-alive background task."""
-    job = _KeepAliveJob(
-        asset_uuid=asset_uuid,
-        job_type=JobType.PREFETCH,
-        interval=max(1.0, float(interval) * 0.8),
-        tier_path_uuid=tier_path_uuid,
-    )
-    self._keep_alives[(asset_uuid, JobType.PREFETCH)] = job
-    self._ensure_manager_running()
-    self._keep_alive_event.set()
-
-  def _stop_prefetch_keep_alive(self, asset_uuid: str) -> None:
-    """Stops the prefetch keep-alive background task."""
-    job = self._keep_alives.pop((asset_uuid, JobType.PREFETCH), None)
-    if job:
-      self._keep_alive_event.set()
-    fut = self._prefetch_futures.pop(asset_uuid, None)
-    if fut and not fut.done():
-      fut.cancel()
-
-  def _get_earliest_job(self) -> tuple[_KeepAliveJob | None, float | None]:
-    """Finds the earliest job to run and its scheduled time."""
-    earliest_job = None
-    earliest_time = None
-    for job in self._keep_alives.values():
-      if earliest_time is None or job.next_run < earliest_time:
-        earliest_time = job.next_run
-        earliest_job = job
-    return earliest_job, earliest_time
-
-  async def _wait_for_next_job(self, timeout: float) -> bool:
-    """Waits for next job or early wakeup. Returns True if woken up early."""
-    try:
-      await asyncio.wait_for(self._keep_alive_event.wait(), timeout=timeout)
-      return True
-    except asyncio.TimeoutError:
-      return False
-
-  async def _keep_alive_manager_loop(self) -> None:
-    """Centralized manager loop running heartbeats for all keep-alives."""
-    logging.info("Starting centralized keep-alive manager task.")
+    """Runs write keep-alive heartbeats."""
+    stub = self._get_or_create_stub()
+    current_interval = interval
     while True:
       try:
-        self._keep_alive_event.clear()
-        earliest_job, earliest_time = self._get_earliest_job()
-        if earliest_job is None or earliest_time is None:
-          # Wait indefinitely for a new job.
-          await self._keep_alive_event.wait()
-          continue
-
-        now = asyncio.get_running_loop().time()
-        sleep_duration = earliest_time - now
-        if sleep_duration > 0:
-          # Wait until the next job or early wakeup by new jobs.
-          if await self._wait_for_next_job(sleep_duration):
-            continue
-
-        await self._run_keep_alive_job(earliest_job)
-
+        await asyncio.sleep(current_interval)
+        request = tiering_service_pb2.ReserveKeepAliveRequest(uuid=asset_uuid)
+        metadata = await self._get_auth_metadata()
+        response = await stub.ReserveKeepAlive(
+            request, metadata=metadata, timeout=30.0
+        )
+        current_interval = max(
+            1.0, float(response.keep_alive_interval_seconds) * 0.8
+        )
       except asyncio.CancelledError:
-        logging.info("Centralized keep-alive manager task cancelled.")
         break
-      except Exception:  # pylint: disable=broad-exception-caught
-        # Log unexpected errors and continue running.
-        logging.exception("Error in centralized keep-alive manager loop.")
-        await asyncio.sleep(1.0)
-
-  async def _run_write_keep_alive_job(
-      self,
-      job: _KeepAliveJob,
-      stub: tiering_service_pb2_grpc.TieringServiceStub,
-      now: float,
-  ) -> None:
-    """Executes a single write keep-alive heartbeat request."""
-    try:
-      request = tiering_service_pb2.ReserveKeepAliveRequest(uuid=job.asset_uuid)
-      metadata = await self._get_auth_metadata()
-      response = await stub.ReserveKeepAlive(request, metadata=metadata)
-      job.interval = max(1.0, float(response.keep_alive_interval_seconds) * 0.8)
-      job.next_run = now + job.interval
-      logging.info(
-          "Extended write reservation lease for asset %s", job.asset_uuid
-      )
-    except grpc.aio.AioRpcError as e:
-      logging.warning(
-          "ReserveKeepAlive failed for asset %s: %s",
-          job.asset_uuid,
-          e.details(),
-      )
-      if e.code() == grpc.StatusCode.NOT_FOUND:
-        self._keep_alives.pop((job.asset_uuid, JobType.WRITE), None)
-      else:
-        job.next_run = now + min(5.0, job.interval)
-    except Exception as e:  # pylint: disable=broad-exception-caught
-      logging.warning(
-          "Unexpected error in write keep alive for asset %s: %s",
-          job.asset_uuid,
-          e,
-      )
-      self._keep_alives.pop((job.asset_uuid, JobType.WRITE), None)
-
-  async def _run_prefetch_keep_alive_job(
-      self,
-      job: _KeepAliveJob,
-      stub: tiering_service_pb2_grpc.TieringServiceStub,
-      now: float,
-  ) -> None:
-    """Executes a single prefetch keep-alive heartbeat request."""
-    try:
-      request = tiering_service_pb2.PrefetchKeepAliveRequest(
-          tier_path_uuid=job.tier_path_uuid
-      )
-      metadata = await self._get_auth_metadata()
-      response = await stub.PrefetchKeepAlive(request, metadata=metadata)
-
-      job.interval = max(1.0, float(response.keep_alive_interval_seconds) * 0.8)
-      job.next_run = now + job.interval
-      logging.info("Sent PrefetchKeepAlive for asset %s", job.asset_uuid)
-
-      target_path = None
-      ready = False
-      for tp in response.asset.tier_paths:
-        if tp.tier_path_uuid == job.tier_path_uuid:
-          target_path = tp.path
-          if tp.HasField("ready_at"):
-            ready = True
-          break
-
-      if ready and target_path:
-        fut = self._prefetch_futures.get(job.asset_uuid)
-        if fut and not fut.done():
-          fut.set_result(target_path)
+      except grpc.aio.AioRpcError as e:
+        if e.code() == grpc.StatusCode.NOT_FOUND:
           logging.info(
-              "Prefetch completed and resolved for asset %s", job.asset_uuid
+              "ReserveKeepAlive for %s returned NOT_FOUND; exiting keep-alive.",
+              asset_uuid,
           )
-        self._keep_alives.pop((job.asset_uuid, JobType.PREFETCH), None)
-
-    except grpc.aio.AioRpcError as e:
-      logging.warning(
-          "PrefetchKeepAlive failed for asset %s: %s",
-          job.asset_uuid,
-          e.details(),
-      )
-      if e.code() == grpc.StatusCode.NOT_FOUND:
-        fut = self._prefetch_futures.get(job.asset_uuid)
-        if fut and not fut.done():
-          fut.set_exception(
-              RuntimeError(f"Prefetch failed: asset {job.asset_uuid} not found")
+          break
+        logging.warning(
+            "Write keep alive RPC failed for %s: %s", asset_uuid, e.details()
+        )
+        current_interval = min(5.0, current_interval)
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        logging.warning("Write keep alive failed for %s: %s", asset_uuid, e)
+        current_interval = min(5.0, current_interval)
+
+  async def _prefetch_keep_alive_loop(
+      self, asset_uuid: str, tier_path_uuid: str, interval: float
+  ) -> None:
+    """Runs prefetch keep-alive heartbeats and resolves future when ready."""
+    stub = self._get_or_create_stub()
+    current_interval = interval
+    while True:
+      try:
+        await asyncio.sleep(current_interval)
+        request = tiering_service_pb2.PrefetchKeepAliveRequest(
+            tier_path_uuid=tier_path_uuid
+        )
+        metadata = await self._get_auth_metadata()
+        response = await stub.PrefetchKeepAlive(
+            request, metadata=metadata, timeout=30.0
+        )
+        current_interval = max(
+            1.0, float(response.keep_alive_interval_seconds) * 0.8
+        )
+
+        for tp in response.asset.tier_paths:
+          if tp.tier_path_uuid == tier_path_uuid and tp.HasField("ready_at"):
+            if self._prefetch_future and not self._prefetch_future.done():
+              self._prefetch_future.set_result(tp.path)
+            break
+      except asyncio.CancelledError:
+        break
+      except grpc.aio.AioRpcError as e:
+        if e.code() in (
+            grpc.StatusCode.NOT_FOUND,
+            grpc.StatusCode.FAILED_PRECONDITION,
+            grpc.StatusCode.ABORTED,
+        ):
+          logging.info(
+              "PrefetchKeepAlive for %s returned %s; exiting keep-alive.",
+              asset_uuid,
+              e.code(),
           )
-        self._keep_alives.pop((job.asset_uuid, JobType.PREFETCH), None)
-      else:
-        job.next_run = now + min(5.0, job.interval)
-    except Exception as e:  # pylint: disable=broad-exception-caught
-      logging.warning(
-          "Unexpected error in prefetch keep alive for asset %s: %s",
-          job.asset_uuid,
-          e,
-      )
-      fut = self._prefetch_futures.get(job.asset_uuid)
-      if fut and not fut.done():
-        fut.set_exception(e)
-      self._keep_alives.pop((job.asset_uuid, JobType.PREFETCH), None)
-
-  async def _run_keep_alive_job(self, job: _KeepAliveJob) -> None:
-    """Executes a single keep-alive heartbeat request."""
-    stub = self._stub
-    if stub is None:
-      logging.error("Stub is not initialized in keep-alive manager.")
-      job.next_run = asyncio.get_running_loop().time() + job.interval
-      return
-
-    now = asyncio.get_running_loop().time()
-
-    if job.job_type == JobType.WRITE:
-      await self._run_write_keep_alive_job(job, stub, now)
-    elif job.job_type == JobType.PREFETCH:
-      await self._run_prefetch_keep_alive_job(job, stub, now)
-
-  async def reserve(
+          if self._prefetch_future and not self._prefetch_future.done():
+            self._prefetch_future.set_exception(
+                RuntimeError(f"Prefetch failed: {e.details()}")
+            )
+          break
+        logging.warning(
+            "Prefetch keep alive RPC failed for %s: %s",
+            asset_uuid,
+            e.details(),
+        )
+        current_interval = min(5.0, current_interval)
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        logging.warning("Prefetch keep alive failed for %s: %s", asset_uuid, e)
+        current_interval = min(5.0, current_interval)
+
+  async def _prepare_reserve_request(
       self,
       path: str,
-      tags: Sequence[str] | None = None,
-      user: str | None = None,
-  ) -> tuple[str, str]:
-    """Reserves an asset path on Tier 0 storage.
-
-    Args:
-      path: Unique checkpoint logical path.
-      tags: Optional list of tags.
-      user: Optional owner user. If not specified, auto-discovers.
-
-    Returns:
-      A tuple of (asset_uuid, tier0_path).
-
-    Raises:
-      RuntimeError: If gRPC call fails or no Tier 0 path is returned.
-    """
-    if not self._stub:
-      await self.connect()
-
-    stub = self._stub
-    if stub is None:
-      raise RuntimeError("Stub is not initialized after connect.")
-
+      tags: Sequence[str] | None,
+      user: str | None,
+  ) -> tiering_service_pb2.ReserveRequest:
+    """Prepares the ReserveRequest protobuf message."""
     if user is None:
       user = environment.get_current_user()
-
     zone, region = await self._get_gcp_zone_and_region()
-
     request = tiering_service_pb2.ReserveRequest(
-        path=path,
-        tags=tags or [],
-        user=user,
+        path=path, tags=tags or [], user=user
     )
     if zone is not None:
       request.zone = zone
     if region is not None:
       request.region = region
+    return request
+
+  def _find_reserved_path(
+      self, asset: tiering_service_pb2.Asset, tier_path_uuid: str
+  ) -> str | None:
+    for tp in asset.tier_paths:
+      if tp.tier_path_uuid == tier_path_uuid:
+        return tp.path
+    return None
 
+  async def _async_reserve(
+      self,
+      path: str,
+      tags: Sequence[str] | None = None,
+      user: str | None = None,
+  ) -> tuple[str, str]:
+    """Asynchronous implementation of reserve."""
+    self._check_active_operation()
+    stub = self._get_or_create_stub()
+    request = await self._prepare_reserve_request(path, tags, user)
     metadata = await self._get_auth_metadata()
     try:
-      response = await stub.Reserve(request, metadata=metadata)
+      response = await stub.Reserve(request, metadata=metadata, timeout=30.0)
     except grpc.aio.AioRpcError as e:
       raise RuntimeError(
           f"Reserve RPC failed: {e.details()} ({e.code()})"
@@ -411,107 +402,113 @@ async def reserve(
 
     asset = response.asset
     asset_uuid = asset.uuid
-    interval = response.keep_alive_interval_seconds
-
     if not response.tier_path_uuid:
       raise RuntimeError(
           "Reserve succeeded but returned no tier_path_uuid for asset"
           f" {asset_uuid}"
       )
 
-    # Start write keep-alive background task loop
-    self._start_write_keep_alive(asset_uuid, interval)
+    reserved_path = self._find_reserved_path(asset, response.tier_path_uuid)
+    if reserved_path is None:
+      raise RuntimeError(
+          f"Reserve succeeded but tier_path_uuid {response.tier_path_uuid}"
+          " not found"
+      )
 
-    for tp in asset.tier_paths:
-      if tp.tier_path_uuid == response.tier_path_uuid:
-        return asset_uuid, tp.path
-
-    # Stop keep-alive loop if the returned tier_path_uuid is missing from
-    # asset tier paths
-    self._stop_write_keep_alive(asset_uuid)
-    raise RuntimeError(
-        "Reserve succeeded but returned tier_path_uuid"
-        f" {response.tier_path_uuid} which is not found in asset tier paths"
-        f" for asset {asset_uuid}"
+    interval = max(1.0, float(response.keep_alive_interval_seconds) * 0.8)
+    self._active_asset_uuid = asset_uuid
+    self._active_job_type = JobType.WRITE
+    self._active_path = reserved_path
+    self._active_tier_path_uuid = response.tier_path_uuid
+    self._keep_alive_task = self._get_loop().create_task(
+        self._write_keep_alive_loop(asset_uuid, interval)
     )
 
-  async def finalize(self, uuid: str) -> None:
-    """Finalizes the asset, marking it stored and immutable.
-
-    Args:
-      uuid: Asset UUID to finalize.
-
-    Raises:
-      RuntimeError: If gRPC call fails.
-    """
-    if not self._stub:
-      await self.connect()
+    return asset_uuid, reserved_path
 
-    stub = self._stub
-    if stub is None:
-      raise RuntimeError("Stub is not initialized after connect.")
+  async def reserve(
+      self,
+      path: str,
+      tags: Sequence[str] | None = None,
+      user: str | None = None,
+  ) -> tuple[str, str]:
+    """Reserves an asset path on Tier 0 storage."""
+    await self._ensure_connected()
+    lp = self._get_loop()
+    fut = asyncio.run_coroutine_threadsafe(
+        self._async_reserve(path, tags, user), lp
+    )
+    return await asyncio.wrap_future(fut)
 
+  async def _async_finalize(self, uuid: str) -> None:
+    """Asynchronous implementation of finalize."""
+    stub = self._get_or_create_stub()
     request = tiering_service_pb2.FinalizeRequest(uuid=uuid)
     metadata = await self._get_auth_metadata()
     try:
-      await stub.Finalize(request, metadata=metadata)
+      await stub.Finalize(request, metadata=metadata, timeout=30.0)
     except grpc.aio.AioRpcError as e:
       raise RuntimeError(
           f"Finalize RPC failed: {e.details()} ({e.code()})"
       ) from e
     finally:
-      # Stop keep-alive loop inside finally, so it is stopped even on error
-      self._stop_write_keep_alive(uuid)
-
-  async def prefetch(
-      self,
-      path: str | None = None,
-      uuid: str | None = None,
-  ) -> asyncio.Future[str]:
-    """Prefetches the asset to the closest Tier 0 storage.
+      if self._keep_alive_task is not None:
+        self._keep_alive_task.cancel()
+        self._keep_alive_task = None
+      self._reset_active_state()
 
-    Args:
-      path: Logical path of the asset.
-      uuid: Asset UUID.
-
-    Returns:
-      A Future that will resolve to the Tier 0 path when ready.
-
-    Raises:
-      ValueError: If neither or both path and uuid are specified.
-      RuntimeError: If gRPC call fails.
-    """
-    if path is None and uuid is None:
-      raise ValueError("Either path or uuid must be specified.")
-    if path is not None and uuid is not None:
-      raise ValueError("Only one of path or uuid can be specified.")
-
-    if uuid is not None and uuid in self._prefetch_futures:
-      return self._prefetch_futures[uuid]
-
-    if not self._stub:
-      await self.connect()
-
-    stub = self._stub
-    if stub is None:
-      raise RuntimeError("Stub is not initialized after connect.")
-
-    zone, region = await self._get_gcp_zone_and_region()
+  async def finalize(self, uuid: str) -> None:
+    """Finalizes the asset, marking it stored and immutable."""
+    await self._ensure_connected()
+    lp = self._get_loop()
+    fut = asyncio.run_coroutine_threadsafe(self._async_finalize(uuid), lp)
+    await asyncio.wrap_future(fut)
 
+  def _prepare_prefetch_request(
+      self,
+      path: str | None,
+      uuid: str | None,
+      zone: str | None,
+      region: str | None,
+  ) -> tiering_service_pb2.PrefetchRequest:
+    """Prepares the PrefetchRequest protobuf message."""
     request = tiering_service_pb2.PrefetchRequest()
     if uuid is not None:
       request.uuid = uuid
     else:
-      request.path = path  # pyrefly: ignore[bad-assignment]
-
+      request.path = path  # pytype: disable=wrong-arg-types
     if zone is not None:
       request.zone = zone
     if region is not None:
       request.region = region
+    return request
+
+  def _find_closest_tier_path(
+      self, asset: tiering_service_pb2.Asset, closest_tier_path_uuid: str
+  ) -> tiering_service_pb2.TierPath | None:
+    for tp in asset.tier_paths:
+      if tp.tier_path_uuid == closest_tier_path_uuid:
+        return tp
+    return None
+
+  async def _async_prefetch(
+      self,
+      path: str | None = None,
+      uuid: str | None = None,
+  ) -> tuple[str, str]:
+    """Asynchronous implementation of prefetch."""
+    if path is None and uuid is None:
+      raise ValueError("Either path or uuid must be specified.")
+    if path is not None and uuid is not None:
+      raise ValueError("Only one of path or uuid can be specified.")
 
+    self._check_active_operation()
+    stub = self._get_or_create_stub()
+    zone, region = await self._get_gcp_zone_and_region()
+    request = self._prepare_prefetch_request(path, uuid, zone, region)
     metadata = await self._get_auth_metadata()
     try:
-      response = await stub.Prefetch(request, metadata=metadata)
+      response = await stub.Prefetch(request, metadata=metadata, timeout=30.0)
     except grpc.aio.AioRpcError as e:
       raise RuntimeError(
           f"Prefetch RPC failed: {e.details()} ({e.code()})"
@@ -519,77 +516,95 @@ async def prefetch(
 
     asset = response.asset
     asset_uuid = asset.uuid
-    interval = response.keep_alive_interval_seconds
-
-    if asset_uuid in self._prefetch_futures:
-      return self._prefetch_futures[asset_uuid]
-
-    future = asyncio.get_running_loop().create_future()
-    self._prefetch_futures[asset_uuid] = future
-
-    closest_tp = None
-    if response.closest_tier_path_uuid:
-      for tp in asset.tier_paths:
-        if tp.tier_path_uuid == response.closest_tier_path_uuid:
-          closest_tp = tp
-          break
-    else:
+    if not response.closest_tier_path_uuid:
       raise RuntimeError(
           "Prefetch succeeded but returned no closest_tier_path_uuid for asset"
           f" {asset.uuid}"
       )
 
+    closest_tp = self._find_closest_tier_path(
+        asset, response.closest_tier_path_uuid
+    )
     if closest_tp is None:
-      self._prefetch_futures.pop(asset_uuid, None)
       raise RuntimeError(
-          "Prefetch response did not contain closest TierPath matching "
-          f"{response.closest_tier_path_uuid} for asset {asset_uuid}"
+          "Prefetch response did not contain closest TierPath for asset"
+          f" {asset_uuid}"
       )
 
-    self._start_prefetch_keep_alive(
-        asset_uuid, closest_tp.tier_path_uuid, interval
-    )
+    interval = max(1.0, float(response.keep_alive_interval_seconds) * 0.8)
+    lp = self._get_loop()
+    self._active_asset_uuid = asset_uuid
+    self._active_job_type = JobType.PREFETCH
+    self._active_path = closest_tp.path
+    self._active_tier_path_uuid = closest_tp.tier_path_uuid
+    self._prefetch_future = lp.create_future()
 
     if closest_tp.HasField("ready_at"):
-      future.set_result(closest_tp.path)
+      self._prefetch_future.set_result(closest_tp.path)
 
-    return future
+    self._keep_alive_task = lp.create_task(
+        self._prefetch_keep_alive_loop(
+            asset_uuid, closest_tp.tier_path_uuid, interval
+        )
+    )
 
-  async def release(self, uuid: str) -> None:
-    """Client-side release of prefetch keep-alive loop.
+    try:
+      path_result = await self._prefetch_future
+      return asset_uuid, path_result
+    except asyncio.CancelledError:
+      if self._keep_alive_task is not None:
+        self._keep_alive_task.cancel()
+        self._keep_alive_task = None
+      self._reset_active_state()
+      raise
 
-    Args:
-      uuid: Asset UUID to release.
-    """
-    self._stop_prefetch_keep_alive(uuid)
+  async def prefetch(
+      self,
+      path: str | None = None,
+      uuid: str | None = None,
+  ) -> str:
+    """Prefetches the asset to the closest Tier 0 storage."""
+    await self._ensure_connected()
+    lp = self._get_loop()
+    fut = asyncio.run_coroutine_threadsafe(self._async_prefetch(path, uuid), lp)
+    _, resolved_path = await asyncio.wrap_future(fut)
+    return resolved_path
+
+  async def _async_release(self, uuid: str) -> None:
+    del uuid  # Unused.
+    if self._keep_alive_task is not None:
+      self._keep_alive_task.cancel()
+      self._keep_alive_task = None
+    if self._prefetch_future is not None and not self._prefetch_future.done():
+      self._prefetch_future.cancel()
+      self._prefetch_future = None
+    self._reset_active_state()
 
-  async def delete(
+  async def release(self, uuid: str) -> None:
+    """Client-side release of prefetch reservation."""
+    await self._ensure_connected()
+    lp = self._get_loop()
+    fut = asyncio.run_coroutine_threadsafe(self._async_release(uuid), lp)
+    await asyncio.wrap_future(fut)
+
+  async def release_path(self, path: str) -> None:
+    """Releases a prefetch reservation by path string directly."""
+    await self._ensure_connected()
+    if self._active_path == path or self._active_asset_uuid is not None:
+      await self.release(self._active_asset_uuid or "")
+
+  async def _async_delete(
       self,
       path: str | None = None,
       uuid: str | None = None,
   ) -> None:
-    """Queues a delete job for the asset.
-
-    Args:
-      path: Logical path of the asset.
-      uuid: Asset UUID to delete.
-
-    Raises:
-      ValueError: If neither or both path and uuid are specified.
-      RuntimeError: If gRPC call fails.
-    """
+    """Asynchronous implementation of delete."""
     if path is None and uuid is None:
       raise ValueError("Either path or uuid must be specified.")
     if path is not None and uuid is not None:
       raise ValueError("Only one of path or uuid can be specified.")
 
-    if not self._stub:
-      await self.connect()
-
-    stub = self._stub
-    if stub is None:
-      raise RuntimeError("Stub is not initialized after connect.")
-
+    stub = self._get_or_create_stub()
     if uuid is not None:
       request = tiering_service_pb2.DeleteRequest(uuid=uuid)
     else:
@@ -597,42 +612,44 @@ async def delete(
 
     metadata = await self._get_auth_metadata()
     try:
-      await stub.Delete(request, metadata=metadata)
+      await stub.Delete(request, metadata=metadata, timeout=30.0)
     except grpc.aio.AioRpcError as e:
       raise RuntimeError(
           f"Delete RPC failed: {e.details()} ({e.code()})"
       ) from e
 
-  async def info(
+  def _log_fire_and_forget_error(self, fut: Any, error_msg: str) -> None:
+    try:
+      fut.result()
+    except Exception:  # pylint: disable=broad-exception-caught
+      logging.exception(error_msg)
+
+  async def delete(
       self,
       path: str | None = None,
       uuid: str | None = None,
-  ) -> list[tiering_service_pb2.Asset]:
-    """Retrieves info/metadata for an asset.
-
-    Args:
-      path: Logical path of the asset.
-      uuid: Asset UUID.
-
-    Returns:
-      A list of matching Asset configurations.
+  ) -> None:
+    """Queues a delete job for the asset."""
+    await self._ensure_connected()
+    lp = self._get_loop()
+    fut = asyncio.run_coroutine_threadsafe(self._async_delete(path, uuid), lp)
+    error_msg = f"Delete task failed for path={path}, uuid={uuid}"
+    fut.add_done_callback(
+        lambda f: self._log_fire_and_forget_error(f, error_msg)
+    )
 
-    Raises:
-      ValueError: If neither or both path and uuid are specified.
-      RuntimeError: If gRPC call fails.
-    """
+  async def _async_info(
+      self,
+      path: str | None = None,
+      uuid: str | None = None,
+  ) -> list[tiering_service_pb2.Asset]:
+    """Asynchronous implementation of info."""
     if path is None and uuid is None:
       raise ValueError("Either path or uuid must be specified.")
     if path is not None and uuid is not None:
       raise ValueError("Only one of path or uuid can be specified.")
 
-    if not self._stub:
-      await self.connect()
-
-    stub = self._stub
-    if stub is None:
-      raise RuntimeError("Stub is not initialized after connect.")
-
+    stub = self._get_or_create_stub()
     if uuid is not None:
       request = tiering_service_pb2.InfoRequest(uuid=uuid)
     else:
@@ -640,7 +657,18 @@ async def info(
 
     metadata = await self._get_auth_metadata()
     try:
-      response = await stub.Info(request, metadata=metadata)
+      response = await stub.Info(request, metadata=metadata, timeout=30.0)
       return list(response.assets)
     except grpc.aio.AioRpcError as e:
       raise RuntimeError(f"Info RPC failed: {e.details()} ({e.code()})") from e
+
+  async def info(
+      self,
+      path: str | None = None,
+      uuid: str | None = None,
+  ) -> list[tiering_service_pb2.Asset]:
+    """Retrieves info/metadata for an asset."""
+    await self._ensure_connected()
+    lp = self._get_loop()
+    fut = asyncio.run_coroutine_threadsafe(self._async_info(path, uuid), lp)
+    return await asyncio.wrap_future(fut)
diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py
index 4e72cfc831..793f4714ff 100644
--- a/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py
+++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py
@@ -92,12 +92,10 @@ def mock_refresh(_):
     mock_creds.refresh.side_effect = mock_refresh
     mock_default.return_value = (mock_creds, "fake-project")
 
-    # First call: should trigger credentials discovery and refresh
     token = await client_auth.get_oauth_token()
     self.assertEqual(token, "fake-access-token")
     mock_creds.refresh.assert_called_once()
 
-    # Second call: should reuse cached credentials and skip refresh
     mock_creds.refresh.reset_mock()
     token = await client_auth.get_oauth_token()
     self.assertEqual(token, "fake-access-token")
@@ -114,10 +112,14 @@ class TieringClientTest(unittest.IsolatedAsyncioTestCase):
 
   def setUp(self):
     super().setUp()
+    client_auth._LOCK = None
+    client_auth._CREDENTIALS = None
     self.stub_mock = mock.AsyncMock()
     self.insecure_channel_mock = mock.MagicMock()
     self.channel_close_mock = mock.AsyncMock()
     self.insecure_channel_mock.close = self.channel_close_mock
+    self.client = client.TieringClient()
+    self.addAsyncCleanup(self.client.close)
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
@@ -129,7 +131,7 @@ async def test_connect_and_close(
     mock_insecure_channel.return_value = self.insecure_channel_mock
     mock_stub_class.return_value = self.stub_mock
 
-    client_inst = client.TieringClient()
+    client_inst = self.client
     await client_inst.connect()
     self.assertEqual(client_inst._stub, self.stub_mock)
 
@@ -150,15 +152,16 @@ async def test_connect_secure_local(
     mock_secure_channel.return_value = self.insecure_channel_mock
     mock_stub_class.return_value = self.stub_mock
 
-    # Local loopback addresses should use local channel credentials.
     client_inst = client.TieringClient(
         server_address="localhost:50051", secure=True
     )
+    self.addAsyncCleanup(client_inst.close)
     await client_inst.connect()
     mock_local_creds.assert_called_once()
     mock_secure_channel.assert_called_once_with(
         "localhost:50051", "local-creds"
     )
+    await client_inst.close()
 
   @mock.patch("grpc.aio.secure_channel")
   @mock.patch("grpc.ssl_channel_credentials")
@@ -177,14 +180,15 @@ async def test_connect_secure_remote(
     mock_secure_channel.return_value = self.insecure_channel_mock
     mock_stub_class.return_value = self.stub_mock
 
-    # Remote addresses should use SSL credentials directly.
     client_inst = client.TieringClient(
         server_address="cts-server:50051", secure=True
     )
+    self.addAsyncCleanup(client_inst.close)
     await client_inst.connect()
     mock_local_creds.assert_not_called()
     mock_ssl_creds.assert_called_once()
     mock_secure_channel.assert_called_once_with("cts-server:50051", "ssl-creds")
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
@@ -234,7 +238,8 @@ async def test_reserve_success(
     )
     self.stub_mock.Reserve.return_value = reserve_resp
 
-    client_inst = client.TieringClient()
+    client_inst = self.client
+    await client_inst.connect()
     uuid, path = await client_inst.reserve(path="logical/path", tags=["my-tag"])
 
     self.assertEqual(uuid, "asset-uuid-1234")
@@ -249,6 +254,7 @@ async def test_reserve_success(
     self.assertEqual(
         kwargs["metadata"], [("authorization", "Bearer fake-token")]
     )
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
@@ -297,15 +303,19 @@ async def test_reserve_caching_behavior(
         tier_path_uuid="tp-uuid-1",
     )
     self.stub_mock.Reserve.return_value = reserve_resp
+    self.stub_mock.Finalize.return_value = (
+        tiering_service_pb2.FinalizeResponse()
+    )
 
-    client_inst = client.TieringClient()
-    # Call reserve twice
-    await client_inst.reserve(path="logical/path")
+    client_inst = self.client
+    await client_inst.connect()
+    uuid1, _ = await client_inst.reserve(path="logical/path")
+    await client_inst.finalize(uuid1)
     await client_inst.reserve(path="logical/path2")
 
-    # Verify environment lookup is cached and called only once
     mock_get_zone.assert_called_once()
     mock_get_region.assert_called_once()
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
@@ -325,10 +335,12 @@ async def test_reserve_rpc_failure(
     )
     self.stub_mock.Reserve.side_effect = rpc_error
 
-    client_inst = client.TieringClient()
+    client_inst = self.client
+    await client_inst.connect()
     with self.assertRaises(RuntimeError) as context:
       await client_inst.reserve(path="logical/path")
     self.assertIn("Reserve RPC failed", str(context.exception))
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
@@ -371,7 +383,7 @@ async def mock_sleep_fn(delay):
         if delay == 0:
           await original_sleep(0)
         else:
-          return None
+          await original_sleep(0.01)
 
       mock_sleep.side_effect = mock_sleep_fn
 
@@ -387,25 +399,61 @@ async def mock_sleep_fn(delay):
           keep_alive_interval_seconds=10,
           tier_path_uuid="tp-uuid-1",
       )
+      self.stub_mock.ReserveKeepAlive.return_value = (
+          tiering_service_pb2.ReserveKeepAliveResponse(
+              keep_alive_interval_seconds=10
+          )
+      )
 
-      client_inst = client.TieringClient()
+      client_inst = self.client
+      await client_inst.connect()
       uuid, _ = await client_inst.reserve(path="logical/path")
 
       self.assertEqual(uuid, "asset-1")
-      self.assertIn(("asset-1", client.JobType.WRITE), client_inst._keep_alives)
-      manager_task = client_inst._keep_alive_manager_task
-      assert manager_task is not None
-      self.assertFalse(manager_task.done())
+      self.assertEqual(client_inst._active_asset_uuid, "asset-1")
+      self.assertIsNotNone(client_inst._keep_alive_task)
 
       self.stub_mock.Finalize.return_value = (
           tiering_service_pb2.FinalizeResponse()
       )
       await client_inst.finalize(uuid="asset-1")
-      await asyncio.sleep(0)  # Allow keep-alive task updates to propagate
+      await asyncio.sleep(0)
 
-      self.assertNotIn(
-          ("asset-1", client.JobType.WRITE), client_inst._keep_alives
-      )
+      self.assertIsNone(client_inst._active_asset_uuid)
+      self.assertIsNone(client_inst._keep_alive_task)
+      await client_inst.close()
+
+  @mock.patch("grpc.aio.insecure_channel")
+  @mock.patch(
+      "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
+  )
+  async def test_reserve_raises_if_already_active(
+      self, mock_stub_class, mock_insecure_channel
+  ):
+    mock_insecure_channel.return_value = self.insecure_channel_mock
+    mock_stub_class.return_value = self.stub_mock
+
+    backend_l0 = tiering_service_pb2.StorageBackend(level=0, prefix="/lustre")
+    tp_l0 = tiering_service_pb2.TierPath(
+        storage_backend=backend_l0,
+        path="/lustre/path1",
+        tier_path_uuid="tp-uuid-1",
+    )
+    asset = tiering_service_pb2.Asset(uuid="asset-1", tier_paths=[tp_l0])
+    self.stub_mock.Reserve.return_value = tiering_service_pb2.ReserveResponse(
+        asset=asset,
+        keep_alive_interval_seconds=60,
+        tier_path_uuid="tp-uuid-1",
+    )
+
+    client_inst = self.client
+    await client_inst.connect()
+    await client_inst.reserve(path="logical/path1")
+
+    with self.assertRaises(RuntimeError) as ctx:
+      await client_inst.reserve(path="logical/path2")
+    self.assertIn("already managing active asset", str(ctx.exception))
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
@@ -449,175 +497,146 @@ async def test_prefetch_resolves_immediately_if_ready(
         closest_tier_path_uuid="tp-uuid-2",
     )
 
-    client_inst = client.TieringClient()
-    future = await client_inst.prefetch(path="logical/path")
-    self.assertTrue(future.done())
-    self.assertEqual(await future, "/lustre/path1")
-    self.stub_mock.Prefetch.assert_called_once()
-    args, kwargs = self.stub_mock.Prefetch.call_args
-    request = args[0]
-    self.assertEqual(request.path, "logical/path")
-    self.assertFalse(request.HasField("uuid"))
-    self.assertEqual(request.zone, "us-east5-a")
-    self.assertEqual(request.region, "us-east5")
-    self.assertEqual(
-        kwargs["metadata"], [("authorization", "Bearer fake-token")]
-    )
-    self.assertIn(
-        ("asset-2", client.JobType.PREFETCH), client_inst._keep_alives
-    )
+    client_inst = self.client
+    await client_inst.connect()
+    path = await client_inst.prefetch(path="logical/path")
+    self.assertEqual(path, "/lustre/path1")
+    self.assertEqual(client_inst._active_asset_uuid, "asset-2")
+
     await client_inst.release("asset-2")
-    self.assertNotIn(
-        ("asset-2", client.JobType.PREFETCH), client_inst._keep_alives
-    )
+    self.assertIsNone(client_inst._active_asset_uuid)
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
       "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
   )
-  @mock.patch(
-      "orbax.checkpoint.experimental.tiering_service.client_auth.get_oauth_token"
-  )
-  @mock.patch(
-      "orbax.checkpoint.experimental.tiering_service.environment.get_gcp_zone"
-  )
-  @mock.patch(
-      "orbax.checkpoint.experimental.tiering_service.environment.get_gcp_region"
-  )
-  async def test_prefetch_polls_and_resolves(
-      self,
-      mock_get_region,
-      mock_get_zone,
-      mock_get_token,
-      mock_stub_class,
-      mock_insecure_channel,
+  async def test_prefetch_raises_if_already_active(
+      self, mock_stub_class, mock_insecure_channel
   ):
     mock_insecure_channel.return_value = self.insecure_channel_mock
     mock_stub_class.return_value = self.stub_mock
-    mock_get_token.return_value = "fake-token"
-    mock_get_zone.return_value = "us-east5-a"
-    mock_get_region.return_value = "us-east5"
 
-    original_sleep = asyncio.sleep
-    with mock.patch(
-        "orbax.checkpoint.experimental.tiering_service.client.asyncio.sleep"
-    ) as mock_sleep, mock.patch(
-        "orbax.checkpoint.experimental.tiering_service.client.asyncio.wait_for"
-    ) as mock_wait_for:
-
-      async def mock_sleep_fn(delay):
-        if delay == 0:
-          await original_sleep(0)
-        else:
-          return None
+    ready_time = timestamp_pb2.Timestamp(seconds=123456)
+    backend_l0 = tiering_service_pb2.StorageBackend(level=0, prefix="/lustre")
+    tp_l0 = tiering_service_pb2.TierPath(
+        storage_backend=backend_l0,
+        path="/lustre/path1",
+        ready_at=ready_time,
+        tier_path_uuid="tp-uuid-2",
+    )
+    asset = tiering_service_pb2.Asset(uuid="asset-2", tier_paths=[tp_l0])
+    self.stub_mock.Prefetch.return_value = tiering_service_pb2.PrefetchResponse(
+        asset=asset,
+        keep_alive_interval_seconds=10,
+        closest_tier_path_uuid="tp-uuid-2",
+    )
 
-      mock_sleep.side_effect = mock_sleep_fn
+    client_inst = self.client
+    await client_inst.connect()
+    await client_inst.prefetch(path="logical/path1")
 
-      async def mock_wait_for_fn(fut, timeout=None):
-        if timeout == 0:
-          return await fut
-        fut.close()
-        raise asyncio.TimeoutError()
+    with self.assertRaises(RuntimeError) as ctx:
+      await client_inst.prefetch(path="logical/path2")
+    self.assertIn("already managing active asset", str(ctx.exception))
+    await client_inst.close()
 
-      mock_wait_for.side_effect = mock_wait_for_fn
+  @mock.patch("grpc.aio.insecure_channel")
+  @mock.patch(
+      "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
+  )
+  async def test_connect_failure_unblocks_waiting_callers(
+      self, mock_stub_class, mock_insecure_channel
+  ):
+    mock_insecure_channel.return_value = self.insecure_channel_mock
+    mock_stub_class.return_value = self.stub_mock
 
-      backend_l0 = tiering_service_pb2.StorageBackend(level=0, prefix="/lustre")
-      tp_l0 = tiering_service_pb2.TierPath(
-          storage_backend=backend_l0,
-          path="/lustre/path1",
-          tier_path_uuid="tp-uuid-3",
-      )
-      asset_not_ready = tiering_service_pb2.Asset(
-          uuid="asset-3", tier_paths=[tp_l0]
-      )
-      self.stub_mock.Prefetch.return_value = (
-          tiering_service_pb2.PrefetchResponse(
-              asset=asset_not_ready,
-              keep_alive_interval_seconds=10,
-              closest_tier_path_uuid="tp-uuid-3",
-          )
-      )
+    client_inst = client.TieringClient()
+    with mock.patch.object(
+        client_inst,
+        "_async_connect",
+        side_effect=Exception("Connection Failed"),
+    ):
+
+      async def call_connect():
+        try:
+          await client_inst.connect()
+          return None
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          return e
 
-      client_inst = client.TieringClient()
-      future = await client_inst.prefetch(path="logical/path")
-
-      self.stub_mock.Prefetch.assert_called_once()
-      args, kwargs = self.stub_mock.Prefetch.call_args
-      request = args[0]
-      self.assertEqual(request.path, "logical/path")
-      self.assertFalse(request.HasField("uuid"))
-      self.assertEqual(request.zone, "us-east5-a")
-      self.assertEqual(request.region, "us-east5")
-      self.assertEqual(
-          kwargs["metadata"], [("authorization", "Bearer fake-token")]
-      )
+      tasks = [asyncio.create_task(call_connect()) for _ in range(3)]
+      results = await asyncio.gather(*tasks)
+      self.assertEqual(len(results), 3)
+      for res in results:
+        self.assertIsInstance(res, Exception)
 
-      self.assertFalse(future.done())
-      self.assertIn(
-          ("asset-3", client.JobType.PREFETCH), client_inst._keep_alives
-      )
+  @mock.patch("grpc.aio.insecure_channel")
+  @mock.patch(
+      "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
+  )
+  async def test_close_stops_loop_on_teardown_error(
+      self, mock_stub_class, mock_insecure_channel
+  ):
+    mock_insecure_channel.return_value = self.insecure_channel_mock
+    mock_stub_class.return_value = self.stub_mock
+    self.channel_close_mock.side_effect = Exception("Channel error")
 
-      resp_not_ready = tiering_service_pb2.PrefetchKeepAliveResponse(
-          keep_alive_interval_seconds=10, asset=asset_not_ready
-      )
+    client_inst = client.TieringClient()
+    await client_inst.connect()
 
-      ready_time = timestamp_pb2.Timestamp(seconds=123456)
-      tp_l0_ready = tiering_service_pb2.TierPath(
-          storage_backend=backend_l0,
-          path="/lustre/path1",
-          ready_at=ready_time,
-          tier_path_uuid="tp-uuid-3",
-      )
-      asset_ready = tiering_service_pb2.Asset(
-          uuid="asset-3", tier_paths=[tp_l0_ready]
-      )
-      resp_ready = tiering_service_pb2.PrefetchKeepAliveResponse(
-          keep_alive_interval_seconds=10, asset=asset_ready
-      )
+    with self.assertRaises(Exception):
+      await client_inst.close()
 
-      self.stub_mock.PrefetchKeepAlive.side_effect = [
-          resp_not_ready,
-          resp_ready,
-          asyncio.CancelledError(),
-      ]
+    self.assertIsNone(client_inst._loop)
+    self.assertIsNone(client_inst._thread)
 
-      await asyncio.sleep(0)
-      await asyncio.sleep(0)
+  @mock.patch("grpc.aio.insecure_channel")
+  @mock.patch(
+      "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
+  )
+  async def test_delete(self, mock_stub_class, mock_insecure_channel):
+    mock_insecure_channel.return_value = self.insecure_channel_mock
+    mock_stub_class.return_value = self.stub_mock
+    self.stub_mock.Delete.return_value = tiering_service_pb2.DeleteResponse()
 
-      self.assertTrue(future.done())
-      self.assertEqual(await future, "/lustre/path1")
+    client_inst = self.client
+    await client_inst.connect()
+    await client_inst.delete(uuid="asset-uuid-delete")
+    await asyncio.sleep(0.1)
 
-      await client_inst.release("asset-3")
-      self.assertNotIn(
-          ("asset-3", client.JobType.PREFETCH), client_inst._keep_alives
-      )
+    self.stub_mock.Delete.assert_called_once()
+    args, _ = self.stub_mock.Delete.call_args
+    self.assertEqual(args[0].uuid, "asset-uuid-delete")
+    await client_inst.close()
 
   @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
       "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
   )
+  async def test_info(self, mock_stub_class, mock_insecure_channel):
+    mock_insecure_channel.return_value = self.insecure_channel_mock
+    mock_stub_class.return_value = self.stub_mock
+
+    asset = tiering_service_pb2.Asset(uuid="asset-uuid-info")
+    info_resp = tiering_service_pb2.InfoResponse(assets=[asset])
+    self.stub_mock.Info.return_value = info_resp
+
+    client_inst = self.client
+    await client_inst.connect()
+    assets = await client_inst.info(uuid="asset-uuid-info")
+    self.assertEqual(len(assets), 1)
+    self.assertEqual(assets[0].uuid, "asset-uuid-info")
+    self.stub_mock.Info.assert_called_once()
+    await client_inst.close()
+
+  @mock.patch("grpc.aio.insecure_channel")
   @mock.patch(
-      "orbax.checkpoint.experimental.tiering_service.client_auth.get_oauth_token"
-  )
-  @mock.patch(
-      "orbax.checkpoint.experimental.tiering_service.environment.get_gcp_zone"
-  )
-  @mock.patch(
-      "orbax.checkpoint.experimental.tiering_service.environment.get_gcp_region"
+      "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub"
   )
-  async def test_prefetch_with_uuid(
-      self,
-      mock_get_region,
-      mock_get_zone,
-      mock_get_token,
-      mock_stub_class,
-      mock_insecure_channel,
-  ):
+  async def test_release_path(self, mock_stub_class, mock_insecure_channel):
     mock_insecure_channel.return_value = self.insecure_channel_mock
     mock_stub_class.return_value = self.stub_mock
-    mock_get_token.return_value = "fake-token"
-    mock_get_zone.return_value = "us-east5-a"
-    mock_get_region.return_value = "us-east5"
 
     ready_time = timestamp_pb2.Timestamp(seconds=123456)
     backend_l0 = tiering_service_pb2.StorageBackend(level=0, prefix="/lustre")
@@ -628,7 +647,7 @@ async def test_prefetch_with_uuid(
         tier_path_uuid="tp-uuid-2",
     )
     asset = tiering_service_pb2.Asset(
-        uuid="asset-uuid-5678", tier_paths=[tp_l0]
+        uuid="asset-uuid-release", tier_paths=[tp_l0]
     )
     self.stub_mock.Prefetch.return_value = tiering_service_pb2.PrefetchResponse(
         asset=asset,
@@ -636,21 +655,15 @@ async def test_prefetch_with_uuid(
         closest_tier_path_uuid="tp-uuid-2",
     )
 
-    client_inst = client.TieringClient()
-    future = await client_inst.prefetch(uuid="asset-uuid-5678")
-    self.assertTrue(future.done())
-    self.assertEqual(await future, "/lustre/path1")
+    client_inst = self.client
+    await client_inst.connect()
+    path = await client_inst.prefetch(uuid="asset-uuid-release")
+    self.assertEqual(path, "/lustre/path1")
+    self.assertEqual(client_inst._active_asset_uuid, "asset-uuid-release")
 
-    self.stub_mock.Prefetch.assert_called_once()
-    args, kwargs = self.stub_mock.Prefetch.call_args
-    request = args[0]
-    self.assertEqual(request.uuid, "asset-uuid-5678")
-    self.assertFalse(request.HasField("path"))
-    self.assertEqual(request.zone, "us-east5-a")
-    self.assertEqual(request.region, "us-east5")
-    self.assertEqual(
-        kwargs["metadata"], [("authorization", "Bearer fake-token")]
-    )
+    await client_inst.release_path("/lustre/path1")
+    self.assertIsNone(client_inst._active_asset_uuid)
+    await client_inst.close()
 
 
 if __name__ == "__main__":
diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py
index e63fc76099..585c230ea3 100644
--- a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py
+++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py
@@ -21,9 +21,10 @@
 import datetime
 import os
 import pprint
-import sys
 import uuid
 
+from absl import app
+from absl import flags
 from absl import logging
 import fire
 import grpc
@@ -632,10 +633,10 @@ async def serve(
       await servicer.close()
 
 
-def main(argv: Sequence[str] | None = None) -> None:
+def main(argv: Sequence[str]) -> None:
   """Main entry point for CTS server."""
-  if argv is None:
-    argv = sys.argv
+  logging.use_absl_handler()
+  logging.set_verbosity(logging.INFO)
   uvloop.install()
   try:
     asyncio.get_event_loop()
@@ -647,4 +648,4 @@ def main(argv: Sequence[str] | None = None) -> None:
 
 
 if __name__ == "__main__":
-  main()
+  app.run(main, flags_parser=lambda argv: flags.FLAGS(argv, known_only=True))
diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py
index 5e083c6c0b..acfdb75665 100644
--- a/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py
+++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py
@@ -129,7 +129,7 @@ def setUp(self):
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
     # Mock metadata for OAuth token
-    self.context.invocation_metadata = mock.AsyncMock(
+    self.context.invocation_metadata = mock.Mock(
         return_value=(("authorization", "Bearer valid-mock-token"),)
     )