Internal

checkpoint/orbax/__init__.py

@@ -19,7 +19,10 @@
 import contextlib
 import functools
 
-from orbax.checkpoint.experimental import v1
+try:
+  from orbax.checkpoint.experimental import v1
+except (ImportError, ModuleNotFoundError):
+  pass
 from orbax.checkpoint import arrays
 from orbax.checkpoint import aggregate_handlers
 from orbax.checkpoint import args

checkpoint/orbax/checkpoint/__init__.py

@@ -19,7 +19,10 @@
 import contextlib
 import functools
 
-from orbax.checkpoint.experimental import v1
+try:
+  from orbax.checkpoint.experimental import v1
+except (ImportError, ModuleNotFoundError):
+  pass
 from orbax.checkpoint import arrays
 from orbax.checkpoint import aggregate_handlers
 from orbax.checkpoint import args

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import concurrent
 import dataclasses
 import datetime
@@ -61,10 +62,10 @@
 from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.path import temporary_paths
 from orbax.checkpoint._src.path import utils as path_utils
+from orbax.checkpoint.experimental.tiering_service import client as cts_client
 from typing_extensions import Self  # for Python version < 3.11
 
 
-
 PyTree = Any
 CheckpointDirs = Tuple[str, str]
 SaveParams = Mapping[str, Any]
@@ -96,6 +97,7 @@
 MultiprocessingOptions = options_lib.MultiprocessingOptions
 FileOptions = options_lib.FileOptions
 
+
 DEFAULT_ITEM_NAME = 'default'
 METRIC_ITEM_NAME = 'metrics'
 METADATA_ITEM_NAME = 'metadata'
@@ -407,10 +409,11 @@ class CheckpointManagerOptions:
       None
   )
   prevent_write_metrics: bool = False
-  # TODO(b/428061876) Remove this option.
   enable_should_save_is_saving_in_progress_check: bool = True
+  # TODO(b/428061876) Remove this option.
   enable_per_process_directory_creation: bool = False
   lightweight_initialize: bool = False
+  tiering_client: Optional[cts_client.TieringClient] = None
 
   def __post_init__(self):
     step_name_format_single_host_load_and_broadcast = (
@@ -719,6 +722,8 @@ def __init__(
 
     self._options = options or CheckpointManagerOptions()
     self._multiprocessing_options = self._options.multiprocessing_options
+    self._tiering_client = self._options.tiering_client
+    self._tiering_uuids = {}  # maps step (int) -> uuid (str)
 
     if self._options.enable_per_process_directory_creation:
       future.AwaitableSignalsContract.awaitable_signals_contract_prefix += (
@@ -953,6 +958,7 @@ def _configure_checkpointer_common(
       options: CheckpointManagerOptions,
       use_async: bool,
   ) -> Checkpointer:
+    kwargs = {}
     if use_async:
       return async_checkpointer.AsyncCheckpointer(
           handler,
@@ -961,6 +967,7 @@ def _configure_checkpointer_common(
           file_options=options.file_options,
           checkpoint_metadata_store=self._non_blocking_metadata_store,
           temporary_path_class=options.temporary_path_class,
+          **kwargs,
       )
     else:
       return Checkpointer(
@@ -969,6 +976,7 @@ def _configure_checkpointer_common(
           file_options=options.file_options,
           checkpoint_metadata_store=self._blocking_metadata_store,
           temporary_path_class=options.temporary_path_class,
+          **kwargs,
       )
 
   def _configure_checkpointer_legacy_init(
@@ -1511,6 +1519,16 @@ def save(
     args = args_lib.Composite(**args_dict)
 
     save_directory = self._get_write_step_directory(step, self.directory)
+    if self._tiering_client is not None:
+      logging.info('[CTS] Reserving path: %s', save_directory)
+      uuid, lustre_path = self._run_async(
+          self._tiering_client.reserve(str(save_directory))
+      )
+      logging.info(
+          '[CTS] Reserved path. UUID: %s, Lustre Path: %s', uuid, lustre_path
+      )
+      self._tiering_uuids[step] = uuid
+
     logging.info(
         '[process=%s] Saving checkpoint at step %d', process_index, step
     )
@@ -1719,6 +1737,17 @@ def restore(
         args = typing.cast(args_lib.Composite, args)
 
     restore_directory = self._get_read_step_directory(step, directory)
+    if self._tiering_client is not None:
+      logging.info('[CTS] Prefetching path: %s', restore_directory)
+
+      async def _do_prefetch():
+        fut = await self._tiering_client.prefetch(str(restore_directory))
+        return await fut
+
+      lustre_path = self._run_async(_do_prefetch())
+      logging.info('[CTS] Prefetch complete. Lustre Path: %s', lustre_path)
+      restore_directory = epath.Path(lustre_path)
+
     step_stats.checkpointer_start_time = time.time()
     restored = self._checkpointer.restore(restore_directory, args=args)
     step_stats.checkpointer_duration_secs = (
@@ -2135,6 +2164,12 @@ def _finalize(self, step: int, steps_to_remove: List[int]):
       # If an error is encountered while waiting for commit futures to complete,
       # we will not proceed past this point.
       self._finalize_checkpoint(step)
+      if self._tiering_client is not None and step in self._tiering_uuids:
+        uuid = self._tiering_uuids[step]
+        logging.info('[CTS] Finalizing step %d, UUID: %s', step, uuid)
+        self._run_async(self._tiering_client.finalize(uuid))
+        del self._tiering_uuids[step]
+
       remove_steps_start_time = time.time()
       self._checkpoint_deleter.delete_steps(steps_to_remove)
       jax.monitoring.record_event_duration_secs(
@@ -2161,6 +2196,20 @@ def _finalize(self, step: int, steps_to_remove: List[int]):
       # This time is tracked for metric purposes only.
       self._last_save_time = time.time()
 
+  def _run_async(self, coro):
+
+    try:
+      loop = asyncio.get_event_loop()
+    except RuntimeError:
+      loop = asyncio.new_event_loop()
+      asyncio.set_event_loop(loop)
+    if loop.is_running():
+      with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        fut = executor.submit(asyncio.run, coro)
+        return fut.result()
+    else:
+      return loop.run_until_complete(coro)
+
   def close(self):
     """Waits for outstanding operations to finish and closes internal objects."""
     self.wait_until_finished()

checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py

@@ -33,13 +33,21 @@
 from sqlalchemy.future import select
 import sqlalchemy.orm
 
-from google.protobuf import timestamp_pb2
+try:
+  from google.protobuf import timestamp_pb2  # pylint: disable=g-import-not-at-top
+except ImportError:
+  # pytype: disable=import-error
+  from google.protobuf import timestamp_pb2  # pylint: disable=g-import-not-at-top
 
 
 class DeletionPendingError(ValueError):
   """Raised when an operation is attempted on an asset/TierPath marked for deletion."""
 
 
+class PrefetchFailedError(ValueError):
+  """Raised when a prefetch operation has failed."""
+
+
 @dataclasses.dataclass
 class CreatePrefetchJobResult:
   """Result of creating a prefetch job.
@@ -371,7 +379,8 @@ async def finalize_asset(
   """Finalizes asset status, transitions state to STORED inside a transaction.
 
   Updates the asset state, sets the finalized timestamp, and marks the
-  associated tier path as ready.
+  associated tier path as ready. Also queues a GCS copy job if a level 1
+  backend exists.
 
   Args:
     session: The database session.
@@ -398,6 +407,35 @@ async def finalize_asset(
     tier_path.ready_at = now
     # TODO: b/503445463 - Set expires_at when policy is supported.
 
+  # Look up level 1 (GCS) storage backend to queue the copy job
+  stmt = select(db_schema.StorageBackend).where(
+      db_schema.StorageBackend.level == 1
+  )
+  res = await session.execute(stmt)
+  gcs_backend = res.scalars().first()
+
+  if gcs_backend is not None:
+    gcs_path = storage_backend_lib.get_storage_path(gcs_backend, db_asset.path)
+    new_gcs_tp = db_schema.TierPath(
+        storage_backend=gcs_backend,
+        path=gcs_path,
+        state=db_schema.TierPathState.PENDING,
+    )
+    db_asset.tier_paths.append(new_gcs_tp)
+
+    db_job = db_schema.AssetJob(
+        asset_uuid=db_asset.asset_uuid,
+        request_type=db_schema.RequestType.REQUEST_TYPE_COPY,
+        status=db_schema.JobStatus.JOB_STATUS_QUEUED,
+        target_tier_path=new_gcs_tp,
+    )
+    session.add(db_job)
+    logging.info(
+        "Finalize: Queued GCS copy job for asset %s, target path: %s",
+        db_asset.asset_uuid,
+        gcs_path,
+    )
+
   await session.commit()
   await session.refresh(db_asset, attribute_names=["updated_at"])
   return db_asset
@@ -550,13 +588,17 @@ async def prefetch_keep_alive(
     DeletionPendingError: If the asset associated with the TierPath is marked
       for deletion, or if the specific TierPath instance is marked for
       deletion.
+    PrefetchFailedError: If the prefetch operation on the TierPath failed.
   """
   stmt = select(db_schema.TierPath).filter_by(tier_path_uuid=tier_path_uuid)
   result = await session.execute(stmt)
   tp = result.scalars().first()
   if tp is None:
     return None
 
+  if tp.state == db_schema.TierPathState.FAILED:
+    raise PrefetchFailedError(f"Prefetch failed for TierPath {tier_path_uuid}.")
+
   if await is_delete_pending(session, asset_uuid=tp.asset_uuid):
     raise DeletionPendingError(f"Asset {tp.asset_uuid} is marked for deletion.")
 
@@ -645,3 +687,50 @@ async def queue_delete_asset_job(
   session.add(db_job)  # pyrefly: ignore[missing-attribute]
 
   await session.commit()
+
+
+async def complete_delete_asset(
+    session: AsyncSession,
+    db_asset: db_schema.Asset,
+) -> db_schema.Asset:
+  """Transitions asset state to DELETED and marks all tier paths as deleted.
+
+  Args:
+    session: The database session.
+    db_asset: The Asset model instance to delete.
+
+  Returns:
+    The updated Asset object.
+  """
+  now = datetime.datetime.now(datetime.timezone.utc)
+  db_asset.state = db_schema.AssetState.ASSET_STATE_DELETED
+  db_asset.deleted_at = now
+
+  for tp in db_asset.tier_paths:
+    tp.state = db_schema.TierPathState.DELETED
+    tp.ready_at = None
+    tp.expires_at = None
+
+  await session.commit()
+  await session.refresh(db_asset, attribute_names=["updated_at"])
+  return db_asset
+
+
+async def complete_delete_tier_path(
+    session: AsyncSession,
+    tier_path: db_schema.TierPath,
+) -> db_schema.TierPath:
+  """Transitions a tier path state to DELETED.
+
+  Args:
+    session: The database session.
+    tier_path: The TierPath model instance to update.
+
+  Returns:
+    The updated TierPath object.
+  """
+  tier_path.state = db_schema.TierPathState.DELETED
+  tier_path.ready_at = None
+  tier_path.expires_at = None
+  await session.commit()
+  return tier_path

checkpoint/orbax/checkpoint/experimental/tiering_service/assets_test.py

@@ -28,7 +28,11 @@
 from sqlalchemy.future import select
 from sqlalchemy.orm import sessionmaker
 
-from google.protobuf import timestamp_pb2
+try:
+  from google.protobuf import timestamp_pb2  # pylint: disable=g-import-not-at-top
+except ImportError:
+  # pytype: disable=import-error
+  from google.protobuf import timestamp_pb2  # pylint: disable=g-import-not-at-top
 
 
 class AssetsProtoTest(absltest.TestCase):

checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py

@@ -19,6 +19,8 @@
 """
 
 from collections.abc import Collection
+import inspect
+
 from absl import logging
 import grpc
 from orbax.checkpoint.experimental.tiering_service import db_schema
@@ -39,7 +41,11 @@ 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()
+  if inspect.isawaitable(raw_metadata):
+    metadata = dict(await raw_metadata)  # pyrefly: ignore[not-async]
+  else:
+    metadata = dict(raw_metadata)
   # Standard header for OAuth tokens in gRPC is 'authorization'.
   auth_header = metadata.get("authorization")