From f6a76c63e05087900ce63625713f1b23cbc86f99 Mon Sep 17 00:00:00 2001 From: Daniel Ng Date: Wed, 24 Jun 2026 19:00:33 -0700 Subject: [PATCH] Internal PiperOrigin-RevId: 937685763 --- checkpoint/orbax/__init__.py | 5 +- checkpoint/orbax/checkpoint/__init__.py | 5 +- .../orbax/checkpoint/checkpoint_manager.py | 53 +- .../experimental/tiering_service/assets.py | 93 ++- .../tiering_service/assets_test.py | 6 +- .../experimental/tiering_service/auth.py | 8 +- .../experimental/tiering_service/client.py | 667 ++++++++++++++++++ .../tiering_service/client_auth.py | 51 ++ .../tiering_service/client_test.py | 579 +++++++++++++++ .../tiering_service/cts_client_cli.py | 241 +++++++ .../tiering_service/cts_e2e_test_runner.py | 272 +++++++ .../tiering_service/cts_integration_run.py | 114 +++ .../tiering_service/environment.py | 71 ++ .../tiering_service/environment_test.py | 115 +++ .../tiering_service/gcp_storage_client.py | 99 ++- .../tiering_service/job_worker.py | 194 ++++- .../proto/tiering_service.proto | 2 + .../experimental/tiering_service/server.py | 25 + .../tiering_service/server_test.py | 74 +- .../tiering_service/storage_backend.py | 2 + checkpoint/orbax/checkpoint/options.py | 1 + 21 files changed, 2622 insertions(+), 55 deletions(-) create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/client.py create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/client_auth.py create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/cts_client_cli.py create mode 100755 checkpoint/orbax/checkpoint/experimental/tiering_service/cts_e2e_test_runner.py create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/cts_integration_run.py create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/environment.py create mode 100644 checkpoint/orbax/checkpoint/experimental/tiering_service/environment_test.py diff --git a/checkpoint/orbax/__init__.py b/checkpoint/orbax/__init__.py index 4ae0da9732..243b7aed7e 100644 --- a/checkpoint/orbax/__init__.py +++ b/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 diff --git a/checkpoint/orbax/checkpoint/__init__.py b/checkpoint/orbax/checkpoint/__init__.py index 4ae0da9732..243b7aed7e 100644 --- a/checkpoint/orbax/checkpoint/__init__.py +++ b/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 diff --git a/checkpoint/orbax/checkpoint/checkpoint_manager.py b/checkpoint/orbax/checkpoint/checkpoint_manager.py index db47bb7f36..c0ea2227d4 100644 --- a/checkpoint/orbax/checkpoint/checkpoint_manager.py +++ b/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() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py index faf792561d..4137572217 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py +++ b/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,6 +588,7 @@ 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) @@ -557,6 +596,9 @@ async def prefetch_keep_alive( 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 diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/assets_test.py index b624c85be9..36451be042 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets_test.py +++ b/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): diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py index 9d0380254f..12d96a4777 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py +++ b/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") diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py new file mode 100644 index 0000000000..a562b581d2 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py @@ -0,0 +1,667 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint Tiering Service (CTS) client library implementation.""" + +import asyncio +from collections.abc import Sequence +import enum +from typing import Any +from absl import logging +import grpc +from orbax.checkpoint.experimental.tiering_service import client_auth +from orbax.checkpoint.experimental.tiering_service import environment +from orbax.checkpoint.experimental.tiering_service.proto import tiering_service_pb2 +from orbax.checkpoint.experimental.tiering_service.proto import tiering_service_pb2_grpc + + +class JobType(enum.Enum): + """Job types managed by the centralized keep-alive manager.""" + + 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 = self.loop.time() + interval + + +class TieringClient: + """Client library to communicate with the Checkpoint Tiering Service (CTS).""" + + def __init__( + self, server_address: str = "localhost:50051", secure: bool = False + ): + """Initializes the TieringClient. + + Args: + server_address: Address of the gRPC server. + secure: If True, establishes a secure gRPC channel. + """ + self._server_address = server_address + self._secure = secure + self._channels: dict[asyncio.AbstractEventLoop, grpc.aio.Channel] = {} + self._stubs: dict[ + asyncio.AbstractEventLoop, tiering_service_pb2_grpc.TieringServiceStub + ] = {} + 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_tasks: dict[ + asyncio.AbstractEventLoop, asyncio.Task[None] + ] = {} + self._keep_alive_events: dict[asyncio.AbstractEventLoop, asyncio.Event] = {} + self._prefetch_futures: dict[str, asyncio.Future[str]] = {} + + async def __aenter__(self) -> "TieringClient": + await self.connect() + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.close() + + def _get_or_create_stub(self) -> tiering_service_pb2_grpc.TieringServiceStub: + """Gets or creates the gRPC stub for the current event loop.""" + loop = asyncio.get_running_loop() + if loop not in self._stubs: + if self._secure: + is_local = ( + "localhost" in self._server_address + or "127.0.0.1" in self._server_address + ) + if is_local: + try: + # Secure channel setup. Fall back to SSL if local creds not + # supported. + creds = grpc.local_channel_credentials() + except AttributeError: + creds = grpc.ssl_channel_credentials() + else: + creds = grpc.ssl_channel_credentials() + channel = grpc.aio.secure_channel(self._server_address, creds) + else: + channel = grpc.aio.insecure_channel(self._server_address) + + self._channels[loop] = channel + self._stubs[loop] = tiering_service_pb2_grpc.TieringServiceStub(channel) + + return self._stubs[loop] + + async def connect(self) -> None: + """Establishes an async gRPC channel with the server.""" + self._get_or_create_stub() + + async def close(self) -> None: + """Closes the gRPC channel.""" + try: + current_loop = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + # Cancel manager task for current loop + if current_loop and current_loop in self._keep_alive_manager_tasks: + task = self._keep_alive_manager_tasks[current_loop] + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + del self._keep_alive_manager_tasks[current_loop] + + # Clean up jobs belonging to current loop + if current_loop: + for key, job in list(self._keep_alives.items()): + if job.loop == current_loop: + del self._keep_alives[key] + if current_loop in self._keep_alive_events: + del self._keep_alive_events[current_loop] + + # Release pending prefetches and cancel futures belonging to current loop + for asset_uuid, fut in list(self._prefetch_futures.items()): + if not fut.done(): + if current_loop and fut.get_loop() == current_loop: + self._stop_prefetch_keep_alive(asset_uuid) + fut.cancel() + + if current_loop: + for uuid, fut in list(self._prefetch_futures.items()): + if fut.get_loop() == current_loop: + del self._prefetch_futures[uuid] + + for loop_val, channel in list(self._channels.items()): + if current_loop and loop_val == current_loop: + await channel.close() + del self._channels[loop_val] + elif not current_loop: + try: + await channel.close() + except Exception: # pylint: disable=broad-except + pass + del self._channels[loop_val] + + if not self._channels: + self._stubs.clear() + + async def _get_gcp_zone_and_region(self) -> tuple[str | None, str | None]: + """Retrieves and caches GCP zone and region.""" + lock = self._env_lock + if lock is None: + lock = asyncio.Lock() + self._env_lock = lock + async with lock: + if not self._env_queried: + self._zone = await environment.get_gcp_zone() + self._region = await environment.get_gcp_region() + self._env_queried = True + return self._zone, self._region + + async def _get_auth_metadata(self) -> list[tuple[str, str]]: + """Retrieves GCP OAuth token and formats it as gRPC metadata.""" + token = await client_auth.get_oauth_token() + if token: + return [("authorization", f"Bearer {token}")] + return [] + + def _ensure_manager_running(self) -> None: + loop = asyncio.get_running_loop() + if loop not in self._keep_alive_events: + self._keep_alive_events[loop] = asyncio.Event() + if ( + loop not in self._keep_alive_manager_tasks + or self._keep_alive_manager_tasks[loop].done() + ): + self._keep_alive_manager_tasks[loop] = 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_events[job.loop].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_events[job.loop].set() + + def _start_prefetch_keep_alive( + self, asset_uuid: str, tier_path_uuid: str, interval: int + ) -> 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_events[job.loop].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_events[job.loop].set() + fut = self._prefetch_futures.pop(asset_uuid, None) + if fut and not fut.done(): + fut.cancel() + + def _get_earliest_job( + self, loop: asyncio.AbstractEventLoop + ) -> 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 job.loop == loop: + 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, loop: asyncio.AbstractEventLoop, timeout: float + ) -> bool: + """Waits for next job or early wakeup. Returns True if woken up early.""" + + try: + await asyncio.wait_for( + self._keep_alive_events[loop].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.") + loop = asyncio.get_running_loop() + while True: + try: + self._keep_alive_events[loop].clear() + earliest_job, earliest_time = self._get_earliest_job(loop) + if earliest_job is None or earliest_time is None: + # Wait indefinitely for a new job. + await self._keep_alive_events[loop].wait() + continue + + now = 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(loop, sleep_duration): + continue + + await self._run_keep_alive_job(earliest_job) + + 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) + logging.info( + "Prefetch completed and resolved for asset %s", job.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() in ( + grpc.StatusCode.NOT_FOUND, + grpc.StatusCode.FAILED_PRECONDITION, + grpc.StatusCode.ABORTED, + ): + fut = self._prefetch_futures.get(job.asset_uuid) + if fut and not fut.done(): + fut.set_exception( + RuntimeError(f"Prefetch failed: {e.details()}") + ) + 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._get_or_create_stub() + 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( + 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. + """ + stub = self._get_or_create_stub() + + 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, + ) + if zone is not None: + request.zone = zone + if region is not None: + request.region = region + + metadata = await self._get_auth_metadata() + try: + response = await stub.Reserve(request, metadata=metadata) + except grpc.aio.AioRpcError as e: + raise RuntimeError( + f"Reserve RPC failed: {e.details()} ({e.code()})" + ) from e + + 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) + + 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}" + ) + + 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. + """ + 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) + 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. + + 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] + + stub = self._get_or_create_stub() + + zone, region = await self._get_gcp_zone_and_region() + + request = tiering_service_pb2.PrefetchRequest() + if uuid is not None: + request.uuid = uuid + else: + request.path = path + + if zone is not None: + request.zone = zone + if region is not None: + request.region = region + + metadata = await self._get_auth_metadata() + try: + response = await stub.Prefetch(request, metadata=metadata) + except grpc.aio.AioRpcError as e: + raise RuntimeError( + f"Prefetch RPC failed: {e.details()} ({e.code()})" + ) from e + + 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: + raise RuntimeError( + "Prefetch succeeded but returned no closest_tier_path_uuid for asset" + f" {asset.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}" + ) + + self._start_prefetch_keep_alive( + asset_uuid, closest_tp.tier_path_uuid, interval + ) + + if closest_tp.HasField("ready_at"): + future.set_result(closest_tp.path) + return future + + async def release(self, uuid: str) -> None: + """Client-side release of prefetch keep-alive loop. + + Args: + uuid: Asset UUID to release. + """ + self._stop_prefetch_keep_alive(uuid) + + async def 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. + """ + 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.") + + stub = self._get_or_create_stub() + + if uuid is not None: + request = tiering_service_pb2.DeleteRequest(uuid=uuid) + else: + request = tiering_service_pb2.DeleteRequest(path=path) + + metadata = await self._get_auth_metadata() + try: + await stub.Delete(request, metadata=metadata) + except grpc.aio.AioRpcError as e: + raise RuntimeError( + f"Delete 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. + + Args: + path: Logical path of the asset. + uuid: Asset UUID. + + Returns: + A list of matching Asset configurations. + + 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.") + + stub = self._get_or_create_stub() + + if uuid is not None: + request = tiering_service_pb2.InfoRequest(uuid=uuid) + else: + request = tiering_service_pb2.InfoRequest(path=path) + + metadata = await self._get_auth_metadata() + try: + response = await stub.Info(request, metadata=metadata) + return list(response.assets) + except grpc.aio.AioRpcError as e: + raise RuntimeError(f"Info RPC failed: {e.details()} ({e.code()})") from e diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/client_auth.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_auth.py new file mode 100644 index 0000000000..80196ac264 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_auth.py @@ -0,0 +1,51 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint Tiering Service (CTS) client authentication utilities.""" + +import asyncio +from absl import logging +import google.auth +import google.auth.transport.requests + +# Note: asyncio.Lock is not thread-safe. This lock only serializes concurrent +# async tasks running on the same event loop. +_LOCK = None +_CREDENTIALS = None + + +async def get_oauth_token() -> str | None: + """Fetches the Google ADC OAuth 2.0 access token asynchronously with caching.""" + global _LOCK, _CREDENTIALS + if _LOCK is None: + _LOCK = asyncio.Lock() + + async with _LOCK: + try: + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + if _CREDENTIALS is None: + # Discover default credentials. + _CREDENTIALS, _ = await asyncio.to_thread( + google.auth.default, scopes=scopes + ) + + # If credentials are not valid (e.g. expired), refresh asynchronously. + if not _CREDENTIALS.valid: + request = google.auth.transport.requests.Request() + await asyncio.to_thread(_CREDENTIALS.refresh, request) + + return _CREDENTIALS.token + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning("Failed to retrieve GCP OAuth token: %s", e) + return None diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py new file mode 100644 index 0000000000..1641d49a4f --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py @@ -0,0 +1,579 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the CTS Client library and utility modules.""" + +import asyncio +import unittest +from unittest import mock + +import grpc +from orbax.checkpoint.experimental.tiering_service import client +from orbax.checkpoint.experimental.tiering_service import client_auth +from orbax.checkpoint.experimental.tiering_service import environment +from orbax.checkpoint.experimental.tiering_service.proto import tiering_service_pb2 + +from google.protobuf import timestamp_pb2 + + +class EnvironmentTest(unittest.IsolatedAsyncioTestCase): + + @mock.patch("os.environ", {}) + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_zone_metadata_server(self, mock_get): + mock_response = mock.MagicMock() + mock_response.status_code = 200 + mock_response.text = "projects/123456/zones/us-east5-a" + mock_get.return_value = mock_response + + zone = await environment.get_gcp_zone() + self.assertEqual(zone, "us-east5-a") + + @mock.patch("os.environ", {"GCP_ZONE": "us-west1-b"}) + async def test_get_gcp_zone_env_override(self): + zone = await environment.get_gcp_zone() + self.assertEqual(zone, "us-west1-b") + + @mock.patch("os.environ", {}) + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_zone_timeout(self, mock_get): + mock_get.side_effect = Exception("Connection timeout") + zone = await environment.get_gcp_zone() + self.assertIsNone(zone) + + @mock.patch("os.environ", {}) + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_region_metadata_server(self, mock_get): + mock_response = mock.MagicMock() + mock_response.status_code = 200 + mock_response.text = "projects/123456/zones/us-east5-a" + mock_get.return_value = mock_response + + region = await environment.get_gcp_region() + self.assertEqual(region, "us-east5") + + @mock.patch("os.environ", {"GCP_REGION": "us-west1"}) + async def test_get_gcp_region_env_override(self): + region = await environment.get_gcp_region() + self.assertEqual(region, "us-west1") + + +class ClientAuthTest(unittest.IsolatedAsyncioTestCase): + + def setUp(self): + super().setUp() + client_auth._CREDENTIALS = None + + def tearDown(self): + client_auth._CREDENTIALS = None + super().tearDown() + + @mock.patch("google.auth.default") + @mock.patch("google.auth.transport.requests.Request") + async def test_get_oauth_token_success(self, _, mock_default): + mock_creds = mock.MagicMock() + mock_creds.valid = False + mock_creds.token = "fake-access-token" + + def mock_refresh(_): + mock_creds.valid = True + + 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") + mock_creds.refresh.assert_not_called() + + @mock.patch("google.auth.default") + async def test_get_oauth_token_failure(self, mock_default): + mock_default.side_effect = Exception("No ADC credentials") + token = await client_auth.get_oauth_token() + self.assertIsNone(token) + + +class TieringClientTest(unittest.IsolatedAsyncioTestCase): + + def setUp(self): + super().setUp() + 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 + + @mock.patch("grpc.aio.insecure_channel") + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub" + ) + async def test_connect_and_close( + self, mock_stub_class, mock_insecure_channel + ): + mock_insecure_channel.return_value = self.insecure_channel_mock + mock_stub_class.return_value = self.stub_mock + + client_inst = client.TieringClient() + await client_inst.connect() + loop = asyncio.get_running_loop() + self.assertEqual(client_inst._stubs[loop], self.stub_mock) + + await client_inst.close() + self.assertNotIn(loop, client_inst._channels) + self.assertNotIn(loop, client_inst._stubs) + self.channel_close_mock.assert_called_once() + + @mock.patch("grpc.aio.secure_channel") + @mock.patch("grpc.local_channel_credentials") + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub" + ) + async def test_connect_secure_local( + self, mock_stub_class, mock_local_creds, mock_secure_channel + ): + mock_local_creds.return_value = "local-creds" + 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 + ) + await client_inst.connect() + mock_local_creds.assert_called_once() + mock_secure_channel.assert_called_once_with( + "localhost:50051", "local-creds" + ) + + @mock.patch("grpc.aio.secure_channel") + @mock.patch("grpc.ssl_channel_credentials") + @mock.patch("grpc.local_channel_credentials") + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub" + ) + async def test_connect_secure_remote( + self, + mock_stub_class, + mock_local_creds, + mock_ssl_creds, + mock_secure_channel, + ): + mock_ssl_creds.return_value = "ssl-creds" + 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 + ) + 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") + + @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" + ) + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.environment.get_current_user" + ) + async def test_reserve_success( + self, + mock_get_user, + mock_get_region, + mock_get_zone, + mock_get_token, + 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" + mock_get_user.return_value = "test-user" + + 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-uuid-1234", tier_paths=[tp_l0] + ) + reserve_resp = tiering_service_pb2.ReserveResponse( + asset=asset, + keep_alive_interval_seconds=60, + tier_path_uuid="tp-uuid-1", + ) + self.stub_mock.Reserve.return_value = reserve_resp + + client_inst = client.TieringClient() + uuid, path = await client_inst.reserve(path="logical/path", tags=["my-tag"]) + + self.assertEqual(uuid, "asset-uuid-1234") + self.assertEqual(path, "/lustre/path1") + self.stub_mock.Reserve.assert_called_once() + args, kwargs = self.stub_mock.Reserve.call_args + request = args[0] + self.assertEqual(request.user, "test-user") + self.assertEqual(request.zone, "us-east5-a") + self.assertEqual(request.region, "us-east5") + self.assertEqual(list(request.tags), ["my-tag"]) + self.assertEqual( + kwargs["metadata"], [("authorization", "Bearer fake-token")] + ) + + @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" + ) + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.environment.get_current_user" + ) + async def test_reserve_caching_behavior( + self, + mock_get_user, + mock_get_region, + mock_get_zone, + mock_get_token, + 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" + mock_get_user.return_value = "test-user" + + 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-uuid-1234", tier_paths=[tp_l0] + ) + reserve_resp = tiering_service_pb2.ReserveResponse( + asset=asset, + keep_alive_interval_seconds=60, + tier_path_uuid="tp-uuid-1", + ) + self.stub_mock.Reserve.return_value = reserve_resp + + client_inst = client.TieringClient() + # Call reserve twice + await client_inst.reserve(path="logical/path") + 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() + + @mock.patch("grpc.aio.insecure_channel") + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.proto.tiering_service_pb2_grpc.TieringServiceStub" + ) + async def test_reserve_rpc_failure( + self, mock_stub_class, mock_insecure_channel + ): + mock_insecure_channel.return_value = self.insecure_channel_mock + mock_stub_class.return_value = self.stub_mock + + rpc_error = grpc.aio.AioRpcError( + code=grpc.StatusCode.INTERNAL, + initial_metadata=grpc.aio.Metadata(), + trailing_metadata=grpc.aio.Metadata(), + details="database error", + ) + self.stub_mock.Reserve.side_effect = rpc_error + + client_inst = client.TieringClient() + with self.assertRaises(RuntimeError) as context: + await client_inst.reserve(path="logical/path") + self.assertIn("Reserve RPC failed", str(context.exception)) + + @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" + ) + @mock.patch( + "orbax.checkpoint.experimental.tiering_service.environment.get_current_user" + ) + async def test_reserve_starts_and_stops_keep_alive( + self, + mock_get_user, + mock_get_region, + mock_get_zone, + mock_get_token, + 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" + mock_get_user.return_value = "test-user" + + original_sleep = asyncio.sleep + with mock.patch( + "orbax.checkpoint.experimental.tiering_service.client.asyncio.sleep" + ) as mock_sleep: + + async def mock_sleep_fn(delay): + if delay == 0: + await original_sleep(0) + else: + return None + + mock_sleep.side_effect = mock_sleep_fn + + 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=10, + tier_path_uuid="tp-uuid-1", + ) + + client_inst = client.TieringClient() + 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_tasks.get( + asyncio.get_running_loop() + ) + assert manager_task is not None + self.assertFalse(manager_task.done()) + + 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 + + self.assertNotIn( + ("asset-1", client.JobType.WRITE), client_inst._keep_alives + ) + + @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_resolves_immediately_if_ready( + self, + mock_get_region, + mock_get_zone, + mock_get_token, + 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") + 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", + ) + + client_inst = client.TieringClient() + future = await client_inst.prefetch(path="logical/path") + self.assertTrue(future.done()) + self.assertEqual(await future, "/lustre/path1") + self.assertIn( + ("asset-2", client.JobType.PREFETCH), client_inst._keep_alives + ) + await client_inst.release("asset-2") + self.assertNotIn( + ("asset-2", 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" + ) + @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, + ): + 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 + + mock_sleep.side_effect = mock_sleep_fn + + async def mock_wait_for_fn(fut, timeout=None): + if timeout == 0: + return await fut + fut.close() + raise asyncio.TimeoutError() + + mock_wait_for.side_effect = mock_wait_for_fn + + 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() + future = await client_inst.prefetch(path="logical/path") + + self.assertFalse(future.done()) + self.assertIn( + ("asset-3", client.JobType.PREFETCH), client_inst._keep_alives + ) + + resp_not_ready = tiering_service_pb2.PrefetchKeepAliveResponse( + keep_alive_interval_seconds=10, asset=asset_not_ready + ) + + 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 + ) + + self.stub_mock.PrefetchKeepAlive.side_effect = [ + resp_not_ready, + resp_ready, + asyncio.CancelledError(), + ] + + await asyncio.sleep(0) + await asyncio.sleep(0) + + self.assertTrue(future.done()) + self.assertEqual(await future, "/lustre/path1") + + await client_inst.release("asset-3") + self.assertNotIn( + ("asset-3", client.JobType.PREFETCH), client_inst._keep_alives + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_client_cli.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_client_cli.py new file mode 100644 index 0000000000..1aeb2af5a7 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_client_cli.py @@ -0,0 +1,241 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint Tiering Service (CTS) Client CLI Tool.""" + +import asyncio +import logging +import os +import shutil +import sys +from typing import Sequence + +import fire +from orbax.checkpoint.experimental.tiering_service import client +import uvloop + + +class CtsClientCli: + """CLI tool for Checkpoint Tiering Service (CTS) Client.""" + + def __init__(self, server_address: str = "localhost:50051"): + self._server_address = server_address + + def reserve( + self, path: str, user: str | None = None, auto_finalize: bool = False + ) -> None: + """Reserves an asset path on Tier 0 storage. + + Args: + path: Unique checkpoint logical path. + user: Optional owner user. + auto_finalize: If True, finalizes the asset immediately after reserving. + """ + + async def _run(): + c = client.TieringClient(self._server_address) + try: + uuid, t0_path = await c.reserve(path, user=user) + print("Reserve succeeded:") + print(f" Asset UUID: {uuid}") + print(f" Tier 0 Path: {t0_path}") + + if auto_finalize: + await c.finalize(uuid) + print(f"Auto-finalized asset UUID: {uuid}") + else: + print("Keep-alive task is active. Maintaining reservation lease...") + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, + input, + "\nWrite dummy data to Tier 0 Path. Press Enter to finalize and" + " exit...\n", + ) + await c.finalize(uuid) + print(f"Finalized asset UUID: {uuid}") + finally: + await c.close() + + asyncio.run(_run()) + + def finalize(self, uuid: str) -> None: + """Finalizes the asset, marking it stored and immutable. + + Args: + uuid: Asset UUID to finalize. + """ + + async def _run(): + c = client.TieringClient(self._server_address) + try: + await c.finalize(uuid) + print(f"Finalize succeeded for asset UUID: {uuid}") + finally: + await c.close() + + asyncio.run(_run()) + + def prefetch(self, path_or_uuid: str) -> None: + """Prefetches the asset to Tier 0 storage and waits for completion. + + Args: + path_or_uuid: Logical path or asset UUID. + """ + + async def _run(): + c = client.TieringClient(self._server_address) + try: + print(f"Initiating prefetch for {path_or_uuid}...") + future = await c.prefetch(path_or_uuid) + print("Waiting for prefetch to resolve to Tier 0 path...") + t0_path = await future + print(f"Prefetch resolved successfully! Tier 0 path: {t0_path}") + + # Keep keep-alive running until user releases + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, + input, + "\nPrefetch lease is active. Press Enter to release and exit...\n", + ) + + is_uuid = len(path_or_uuid) == 36 and "-" in path_or_uuid + if is_uuid: + uuid = path_or_uuid + else: + assets_list = await c.info(path_or_uuid) + if assets_list: + uuid = assets_list[0].uuid + else: + print("Could not retrieve asset UUID to release.") + return + await c.release(uuid) + print(f"Released prefetch for asset: {uuid}") + finally: + await c.close() + + asyncio.run(_run()) + + def release(self, uuid: str) -> None: + """Releases client-side prefetch keep-alive loop. + + Args: + uuid: Asset UUID to release. + """ + + async def _run(): + c = client.TieringClient(self._server_address) + try: + await c.release(uuid) + print(f"Released prefetch keep-alive for asset UUID: {uuid}") + finally: + await c.close() + + asyncio.run(_run()) + + def delete(self, path_or_uuid: str) -> None: + """Queues a delete job for the asset (deletes from all tiers). + + Args: + path_or_uuid: Logical path or asset UUID to delete. + """ + + async def _run(): + c = client.TieringClient(self._server_address) + try: + await c.delete(path_or_uuid) + print(f"Delete job queued successfully for {path_or_uuid}") + finally: + await c.close() + + asyncio.run(_run()) + + def info(self, path_or_uuid: str) -> None: + """Retrieves and prints metadata info for an asset. + + Args: + path_or_uuid: Logical path or asset UUID. + """ + + async def _run(): + c = client.TieringClient(self._server_address) + try: + assets_list = await c.info(path_or_uuid) + if not assets_list: + print("No matching assets found.") + return + for i, asset in enumerate(assets_list): + print(f"Asset #{i+1}:") + print(f" UUID: {asset.uuid}") + print(f" Path: {asset.path}") + print(f" User: {asset.user}") + print(f" State: {asset.state}") + if asset.HasField("created_at"): + print(f" Created At: {asset.created_at.ToDatetime()}") + if asset.HasField("deleted_at"): + print(f" Deleted At: {asset.deleted_at.ToDatetime()}") + print(" Tier Paths:") + for tp in asset.tier_paths: + print(f" - Path: {tp.path}") + print(f" TierPath UUID: {tp.tier_path_uuid}") + print( + f" Backend: Level {tp.storage_backend.level}" + f" ({tp.storage_backend.backend_type})" + ) + if tp.HasField("ready_at"): + print( + f" Status: READY (Ready At: {tp.ready_at.ToDatetime()})" + ) + else: + print(" Status: PENDING / NOT READY") + if tp.HasField("expires_at"): + print(f" Expires At: {tp.expires_at.ToDatetime()}") + finally: + await c.close() + + asyncio.run(_run()) + + def evict(self, path: str) -> None: + """Simulates cache eviction by manually deleting a local Lustre path. + + Args: + path: Local file or directory path to delete. + """ + if not os.path.exists(path): + print(f"Error: Path {path} does not exist.") + return + if os.path.isdir(path): + shutil.rmtree(path) + print(f"Evicted directory: {path}") + else: + os.remove(path) + print(f"Evicted file: {path}") + + +def main(argv: Sequence[str] | None = None) -> None: + if argv is None: + argv = sys.argv + uvloop.install() + try: + asyncio.get_event_loop() + except RuntimeError: + loop = uvloop.new_event_loop() + asyncio.set_event_loop(loop) + logging.basicConfig(level=logging.WARNING) + fire.Fire(CtsClientCli, command=argv[1:]) + + +if __name__ == "__main__": + main() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_e2e_test_runner.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_e2e_test_runner.py new file mode 100755 index 0000000000..b90c8549bb --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_e2e_test_runner.py @@ -0,0 +1,272 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/usr/bin/env python3 +"""Automated CTS E2E Integration Test Runner.""" + +import argparse +import os +import subprocess +import time + +POD_NAME = "cts-test-pod" +NAMESPACE = "default" +WORKSPACE_DIR = "/google/src/cloud/dnlng/implement-cts-client-library/" +PACKAGE_PATH = "orbax/checkpoint/experimental/tiering_service" + + +def run_cmd(cmd, check=True, capture_output=True, text=True): + """Helper to run shell commands on host.""" + print(f"[HOST EXEC] Running: {cmd}") + try: + result = subprocess.run( + cmd, + shell=True, + check=check, + stdout=subprocess.PIPE if capture_output else None, + stderr=subprocess.PIPE if capture_output else None, + text=text, + ) + if capture_output and result.stdout: + print(f"Stdout:\n{result.stdout.strip()}") + if capture_output and result.stderr: + print(f"Stderr:\n{result.stderr.strip()}") + return result + except subprocess.CalledProcessError as e: + if capture_output and e.stdout: + print(f"FAILED Command Stdout:\n{e.stdout.strip()}") + if capture_output and e.stderr: + print(f"FAILED Command Stderr:\n{e.stderr.strip()}") + raise e + + +def run_pod_cmd(cmd, check=True, background=False): + """Helper to run commands inside GKE test pod.""" + kube_cmd = f"kubectl exec {POD_NAME} -n {NAMESPACE} -- {cmd}" + if background: + print(f"[POD BG EXEC] Running: {kube_cmd}") + return subprocess.Popen(kube_cmd, shell=True) + else: + print(f"[POD EXEC] Running: {kube_cmd}") + return run_cmd(kube_cmd, check=check) + + +def bootstrap(): + """Bootstrap the pod environment.""" + print("=== Bootstrapping GKE pod ===") + # Create directory structure in pod + run_pod_cmd("mkdir -p /app") + + # Install pip packages + run_pod_cmd( + "pip install httpx grpcio uvloop fire sqlalchemy aiosqlite pyyaml" + " pytimeparse greenlet protobuf grpcio-tools absl-py jax jaxlib" + " google-auth google-cloud-storage orbax-checkpoint" + ) + + # Create site-packages experimental directory + run_pod_cmd( + "mkdir -p" + " /usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental" + ) + # Remove any old version to prevent nested copy issues + run_pod_cmd( + "rm -rf" + " /usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/tiering_service" + ) + + # Copy source files recursively into site-packages + local_src = os.path.join(WORKSPACE_DIR, PACKAGE_PATH) + run_cmd( + f"kubectl cp {local_src}" + f" {POD_NAME}:/usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/tiering_service" + ) + + # Copy modified checkpoint_manager.py to overwrite standard one in + # site-packages + local_cm = os.path.join( + WORKSPACE_DIR, "orbax/checkpoint/checkpoint_manager.py" + ) + run_cmd( + f"kubectl cp {local_cm}" + f" {POD_NAME}:/usr/local/lib/python3.11/site-packages/orbax/checkpoint/checkpoint_manager.py" + ) + + # Touch __init__.py files + run_pod_cmd( + "touch" + " /usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/__init__.py" + ) + + # Strip datapol from proto + run_pod_cmd( + "sed -i '/datapol/d'" + " /usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto" + ) + + # Compile proto inside site-packages + run_pod_cmd( + "python3 -m grpc_tools.protoc -I/usr/local/lib/python3.11/site-packages " + "-I/usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/tiering_service/proto " + "--python_out=/usr/local/lib/python3.11/site-packages " + "--grpc_python_out=/usr/local/lib/python3.11/site-packages " + "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto" + ) + + # Write Server Config yaml file locally on the host + config_content = """db_connection_str: "sqlite+aiosqlite:///app/cts.db" +gcp_project: "orbax-checkpoint" +client_keep_alive_interval: "30s" +max_active_jobs_per_backend: 5 +storage_backends: + - level: 0 + backend_type: "LUSTRE" + prefix: "/lustre" + zone: "us-east1-b" + - level: 1 + backend_type: "GCS" + prefix: "gs://orbax-benchmarks-us-east5/dnlng/tiering_service/e2e_test/" + zone: "us-east1-b" +""" + local_config_path = "/tmp/cts_server_config.yaml" + with open(local_config_path, "w", encoding="utf-8") as f: + f.write(config_content) + + # Copy config file to the pod + run_cmd(f"kubectl cp {local_config_path} {POD_NAME}:/app/server_config.yaml") + + # Clean up local config file + if os.path.exists(local_config_path): + os.remove(local_config_path) + + print("Bootstrap completed successfully!") + + +def run_tests(): + """Run E2E integration test sequence.""" + print("=== Running E2E Integration Tests ===") + # Clean up old database files and checkpoints + run_pod_cmd("rm -f /app/cts.db") + run_pod_cmd("rm -rf /lustre/checkpoints") + + # Copy integration scripts from site-packages to app root for easy execution + run_pod_cmd( + "cp /usr/local/lib/python3.11/site-packages/orbax/checkpoint/experimental/tiering_service/cts_integration_run.py" + " /app/cts_integration_run.py" + ) + + # Modify checkpoints folder path to match GKE mount point + run_pod_cmd( + "sed -i 's|/tmp/lustre-mount/checkpoints|/lustre/checkpoints|g'" + " /app/cts_integration_run.py" + ) + + # Initialize SQLite database schema + run_pod_cmd( + "python3 -m orbax.checkpoint.experimental.tiering_service.db_cli" + " initialize_db /app/server_config.yaml" + ) + + # Start Server and Worker process in background + server_proc = run_pod_cmd( + "python3 -u -m orbax.checkpoint.experimental.tiering_service.server serve" + " /app/server_config.yaml --start-tiering-service-worker=True", + background=True, + ) + time.sleep(3) # Wait for server to start up + + try: + # 1. Execute JAX Save + print("Executing JAX save...") + run_pod_cmd("python3 -u /app/cts_integration_run.py --mode=save --step=1") + + # 2. Wait for GCS transfer job to complete + print("Waiting for GCS copy job to complete...") + for _ in range(120): # Wait up to 10 minutes + # Query DB to check if the copy job status is COMPLETED (value 2) + res = run_pod_cmd( + "python3 -c \"import sqlite3; conn = sqlite3.connect('/app/cts.db');" + " print(conn.execute('SELECT status FROM asset_jobs LIMIT" + " 1;').fetchone()[0])\"" + ) + status = res.stdout.strip() + print(f"Current GCS copy job status: {status}") + if status == "JOB_STATUS_COMPLETED" or status == "2": + print("GCS copy job completed successfully!") + break + time.sleep(5) + else: + raise TimeoutError("Timed out waiting for GCS copy job to complete.") + + # 3. Simulate Eviction from local Lustre and update DB to simulate + # GCS-only state + print("Simulating eviction from Lustre...") + # Clean checkpoints directory + run_pod_cmd("rm -rf /lustre/checkpoints/1") + + # Copy and seed the DB for prefetch test + local_seed = "/usr/local/google/home/dnlng/.gemini/jetski/brain/de63b75e-9e4f-4555-bac6-8ef451aedc60/scratch/seed_db.py" + run_cmd(f"kubectl cp {local_seed} {POD_NAME}:/app/seed_db.py") + + # Update seed_db paths using sed + run_pod_cmd("sed -i 's|/tmp/cts_db.db|/app/cts.db|g' /app/seed_db.py") + run_pod_cmd( + "sed -i 's|/tmp/lustre-mount/checkpoints|/lustre/checkpoints|g'" + " /app/seed_db.py" + ) + + # Run database seeder + run_pod_cmd("python3 -u /app/seed_db.py") + + # 4. Execute JAX Restore (triggers prefetch and polls LRO to completion) + print("Executing JAX restore...") + res = run_pod_cmd( + "python3 -u /app/cts_integration_run.py --mode=restore --step=1" + ) + if "INTEGRATION_TEST_SUCCESS" in res.stdout: + print("\n====================================") + print("E2E INTEGRATION TEST SUCCESSFUL!") + print("====================================\n") + else: + raise RuntimeError(f"Restore verification failed: {res.stdout}") + + finally: + # Cleanup background server process + print("Cleaning up server process...") + # Kill the server python process in the pod + run_pod_cmd("pkill -f 'tiering_service.server'", check=False) + server_proc.terminate() + + +def main(): + parser = argparse.ArgumentParser(description="CTS E2E Test Runner") + parser.add_argument( + "--action", + choices=["bootstrap", "run", "all"], + default="all", + help="Action to perform: bootstrap pod, run integration tests, or all.", + ) + args = parser.parse_args() + + if args.action == "bootstrap": + bootstrap() + elif args.action == "run": + run_tests() + elif args.action == "all": + bootstrap() + run_tests() + + +if __name__ == "__main__": + main() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_integration_run.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_integration_run.py new file mode 100644 index 0000000000..22ec8fc6ae --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/cts_integration_run.py @@ -0,0 +1,114 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration test script for CTS Client with CheckpointManager.""" + +import asyncio +from absl import app +from absl import flags +from absl import logging +import jax +import jax.numpy as jnp +import orbax.checkpoint as ocp +from orbax.checkpoint.experimental.tiering_service import client as cts_client + +FLAGS = flags.FLAGS +flags.DEFINE_string( + "mode", None, "Execution mode: save or restore.", required=True +) +flags.DEFINE_integer("step", 1, "Step number to save or restore.") + + +def main(argv): + del argv # Unused + + # Initialize client + client = cts_client.TieringClient( + server_address="localhost:50051", secure=False + ) + + # CheckpointManager options with client + options = ocp.CheckpointManagerOptions( + create=True, + tiering_client=client, + ) + + mngr = ocp.CheckpointManager( + "/tmp/lustre-mount/checkpoints", + options=options, + ) + + print( + "DEBUG CheckpointManager class location:", + ocp.CheckpointManager.__module__, + ocp.CheckpointManager.__init__.__code__.co_filename, + ) + print( + "DEBUG options tiering_client:", + getattr(options, "tiering_client", "MISSING"), + ) + print("DEBUG options dir:", [x for x in dir(options) if "tier" in x]) + print( + "DEBUG mngr tiering_client:", getattr(mngr, "_tiering_client", "MISSING") + ) + + try: + if FLAGS.mode == "save": + logging.info("Starting JAX save...") + x = jnp.arange(10, dtype=jnp.int32) + y = jnp.ones((2, 5), dtype=jnp.float32) + + # Save step using ArraySave (as standard single arrays are saved this way) + save_args = ocp.args.Composite( + x=ocp.args.ArraySave(x), + y=ocp.args.ArraySave(y), + ) + mngr.save(FLAGS.step, args=save_args) + logging.info("Saved step %d successfully.", FLAGS.step) + + elif FLAGS.mode == "restore": + logging.info("Starting JAX restore...") + # Define structure to restore + abstract_x = jax.ShapeDtypeStruct((10,), jnp.int32) + abstract_y = jax.ShapeDtypeStruct((2, 5), jnp.float32) + + restore_args = ocp.args.Composite( + x=ocp.args.ArrayRestore(abstract_x), + y=ocp.args.ArrayRestore(abstract_y), + ) + + restored = mngr.restore(FLAGS.step, args=restore_args) + + # Assert values + if not jnp.array_equal(restored.x, jnp.arange(10, dtype=jnp.int32)): + raise ValueError( + f"Restored X mismatch: expected {jnp.arange(10)}, got {restored.x}" + ) + if not jnp.array_equal(restored.y, jnp.ones((2, 5), dtype=jnp.float32)): + raise ValueError( + f"Restored Y mismatch: expected ones, got {restored.y}" + ) + print("INTEGRATION_TEST_SUCCESS") + logging.info( + "Restored step %d successfully and verified content.", FLAGS.step + ) + + finally: + mngr.close() + # Close client channel connection + asyncio.run(client.close()) + + +if __name__ == "__main__": + app.run(main) diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/environment.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/environment.py new file mode 100644 index 0000000000..6e332e29ba --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/environment.py @@ -0,0 +1,71 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint Tiering Service (CTS) client environment discovery utilities.""" + +import getpass +import os +from absl import logging +import httpx + +_METADATA_URL = "http://metadata.google.internal/computeMetadata/v1" +_HEADERS = {"Metadata-Flavor": "Google"} + + +def get_current_user() -> str: + """Returns the username of the current user.""" + try: + return os.getlogin() + except OSError: + return getpass.getuser() + + +async def get_gcp_zone() -> str | None: + """Gets the current zone from environment variable or GCP metadata server. + + Returns: + The zone name (e.g., 'us-east5-a'), or None if not found/unavailable. + """ + zone = os.environ.get("GCP_ZONE") + if zone: + return zone + try: + async with httpx.AsyncClient(timeout=2.0) as client: + response = await client.get( + f"{_METADATA_URL}/instance/zone", headers=_HEADERS + ) + if response.status_code == 200: + # The metadata server returns 'projects//zones/' + zone_path = response.text.strip() + return zone_path.split("/")[-1] + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning("Failed to query GCP metadata server for zone: %s", e) + return None + + +async def get_gcp_region() -> str | None: + """Gets the current region from environment variable or derived from zone. + + Returns: + The region name (e.g., 'us-east5'), or None if not found. + """ + region = os.environ.get("GCP_REGION") + if region: + return region + zone = await get_gcp_zone() + if zone: + parts = zone.split("-") + if len(parts) >= 2: + return "-".join(parts[:-1]) + return None diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/environment_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/environment_test.py new file mode 100644 index 0000000000..79fc4e6ea5 --- /dev/null +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/environment_test.py @@ -0,0 +1,115 @@ +# Copyright 2026 The Orbax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for environment discovery utilities.""" + +import os +import unittest +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +import httpx +from orbax.checkpoint.experimental.tiering_service import environment + + +class EnvironmentTest(parameterized.TestCase, unittest.IsolatedAsyncioTestCase): + + def setUp(self): + super().setUp() + # Clear env variables to ensure tests are hermetic + self.original_environ = dict(os.environ) + if "GCP_ZONE" in os.environ: + del os.environ["GCP_ZONE"] + if "GCP_REGION" in os.environ: + del os.environ["GCP_REGION"] + + def tearDown(self): + os.environ.clear() + os.environ.update(self.original_environ) + super().tearDown() + + @mock.patch("os.getlogin") + def test_get_current_user_os_login(self, mock_login): + mock_login.return_value = "env-user" + self.assertEqual(environment.get_current_user(), "env-user") + + @mock.patch("os.getlogin") + @mock.patch("getpass.getuser") + def test_get_current_user_fallback(self, mock_getuser, mock_login): + mock_login.side_effect = OSError("No tty") + mock_getuser.return_value = "fallback-user" + self.assertEqual(environment.get_current_user(), "fallback-user") + + async def test_get_gcp_zone_from_env(self): + os.environ["GCP_ZONE"] = "us-east5-a" + zone = await environment.get_gcp_zone() + self.assertEqual(zone, "us-east5-a") + + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_zone_from_metadata(self, mock_get): + mock_response = mock.MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.text = "projects/12345/zones/us-central1-c" + + async def mock_get_fn(*unused_args, **unused_kwargs): + return mock_response + + mock_get.side_effect = mock_get_fn + + zone = await environment.get_gcp_zone() + self.assertEqual(zone, "us-central1-c") + + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_zone_metadata_server_unavailable(self, mock_get): + async def mock_get_fn(*unused_args, **unused_kwargs): + raise httpx.ConnectError("Network unreachable") + + mock_get.side_effect = mock_get_fn + + zone = await environment.get_gcp_zone() + self.assertIsNone(zone) + + async def test_get_gcp_region_from_env(self): + os.environ["GCP_REGION"] = "us-east5" + region = await environment.get_gcp_region() + self.assertEqual(region, "us-east5") + + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_region_derived_from_metadata(self, mock_get): + mock_response = mock.MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.text = "projects/12345/zones/us-central1-c" + + async def mock_get_fn(*unused_args, **unused_kwargs): + return mock_response + + mock_get.side_effect = mock_get_fn + + region = await environment.get_gcp_region() + self.assertEqual(region, "us-central1") + + @mock.patch("httpx.AsyncClient.get") + async def test_get_gcp_region_metadata_server_unavailable(self, mock_get): + async def mock_get_fn(*unused_args, **unused_kwargs): + raise httpx.ConnectError("Network unreachable") + + mock_get.side_effect = mock_get_fn + + region = await environment.get_gcp_region() + self.assertIsNone(region) + + +if __name__ == "__main__": + absltest.main() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client.py index 1d63c5b788..5631016590 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client.py @@ -19,7 +19,11 @@ import dataclasses import datetime import enum +import os +import shutil from typing import Any, Callable +import urllib.parse + import google.auth from google.auth import exceptions as auth_exceptions @@ -178,6 +182,10 @@ async def poll_operation( """Polls operation status and returns a Result object.""" pass + async def delete_path(self, path: str) -> None: + """Deletes a path from the storage backend.""" + raise NotImplementedError() + def _parse_gcs_path(gcs_path: str) -> tuple[str, str]: """Parses a GCS path like gs://bucket/prefix/file into (bucket, prefix).""" @@ -322,6 +330,43 @@ async def poll_operation( detail_info={"error": error_msg}, ) + async def delete_path(self, path: str) -> None: + """Deletes a GCS object or prefix (directory) recursively.""" + bucket, prefix = _parse_gcs_path(path) + token, _ = await self._get_token_and_project() + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + + # List all objects under the prefix. + list_url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o" + params = {"prefix": prefix, "projection": "noAcl"} + + response = await self.async_client.get( + list_url, params=params, headers=headers + ) + if response.status_code == 404: + return + if response.status_code != 200: + raise RuntimeError( + f"Failed to list GCS objects for deletion: {response.status_code} -" + f" {response.text}" + ) + + data = response.json() + items = data.get("items", []) + for item in items: + name = item["name"] + encoded_name = urllib.parse.quote(name, safe="") + del_url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_name}" + del_resp = await self.async_client.delete(del_url, headers=headers) + if del_resp.status_code not in (200, 204, 404): + raise RuntimeError( + f"Failed to delete GCS object {name}: {del_resp.status_code} -" + f" {del_resp.text}" + ) + class GcpLustreBaseClient(GCPStorageClient): """Base client interface to interact with GCP Managed Lustre API via REST.""" @@ -398,8 +443,9 @@ async def trigger_copy( "Authorization": f"Bearer {token}", "Content-Type": "application/json", } + gcs_uri = source_path if source_path.endswith("/") else source_path + "/" payload = { - "gcsPath": {"uri": source_path}, + "gcsPath": {"uri": gcs_uri}, "lustrePath": {"path": destination_path}, "requestId": request_id, } @@ -430,9 +476,14 @@ async def trigger_copy( "Authorization": f"Bearer {token}", "Content-Type": "application/json", } + gcs_uri = ( + destination_path + if destination_path.endswith("/") + else destination_path + "/" + ) payload = { "lustrePath": {"path": source_path}, - "gcsPath": {"uri": destination_path}, + "gcsPath": {"uri": gcs_uri}, "requestId": request_id, } response = await self.async_client.post(url, json=payload, headers=headers) @@ -614,3 +665,47 @@ def get_client_from_status( ) else: raise ValueError(f"Unknown or missing client type: {client_type}") + + +class LustreClient(GCPStorageClient): + """Client implementation for Lustre filesystem operations.""" + + async def trigger_copy( + self, + request_id: str, + source_path: str, + destination_path: str, + ) -> str: + raise NotImplementedError() + + async def poll_operation( + self, + operation_name: str, + ) -> Result: + raise NotImplementedError() + + async def delete_path(self, path: str) -> None: + """Deletes a Lustre file or directory tree recursively.""" + + def _delete(): + if os.path.isdir(path): + shutil.rmtree(path) + elif os.path.exists(path) or os.path.islink(path): + os.remove(path) + + await asyncio.to_thread(_delete) + + +def determine_delete_client( + tier_path: db_schema.TierPath, + project: str | None = None, + service_account: str | None = None, +) -> GCPStorageClient: + """Determines and returns the GCP Storage client for deleting a TierPath.""" + backend_type = tier_path.storage_backend.backend_type + if backend_type == db_schema.BackendType.BACKEND_TYPE_GCS: + return GcsToGcsClient(project=project, service_account=service_account) + elif backend_type == db_schema.BackendType.BACKEND_TYPE_LUSTRE: + return LustreClient(project=project, service_account=service_account) + else: + raise ValueError(f"Unsupported backend type for delete: {backend_type}") diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py index 1091bc8be3..bbd1deaa59 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py @@ -311,10 +311,15 @@ async def _process_job(self, job: db_schema.AssetJob): if job.request_type == db_schema.RequestType.REQUEST_TYPE_COPY: await self._process_copy_job(job) + elif job.request_type in ( + db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_INSTANCE, + db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS, + ): + await self._process_delete_job(job) else: - # TODO: b/503445463 - Support DELETE jobs. logging.warning("Unsupported job request type: %s", job.request_type) - # Mark as failed for now if not COPY + + # Mark as failed for now if not COPY or DELETE async with self._session_maker() as session: # pyrefly: ignore[not-callable] async with session.begin(): merged_job = await session.get( @@ -428,6 +433,18 @@ async def _save_triggered_transfer_status( operation_name, ) + def _get_transfer_path(self, tp: db_schema.TierPath) -> str: + backend = tp.storage_backend + if backend.backend_type == db_schema.BackendType.BACKEND_TYPE_LUSTRE: + prefix = backend.prefix + path = tp.path + if path.startswith(prefix): + rel_path = path[len(prefix) :] + if not rel_path.startswith("/"): + rel_path = "/" + rel_path + return rel_path + return tp.path + async def _process_copy_job(self, job: db_schema.AssetJob): """Processes a COPY job (eg. GCS -> Lustre or Lustre -> GCS).""" target_tp = job.target_tier_path @@ -455,9 +472,11 @@ async def _process_copy_job(self, job: db_schema.AssetJob): error_msg = None operation_name = None + source_path = self._get_transfer_path(source_tp) + target_path = self._get_transfer_path(target_tp) try: operation_name = await determined_client.trigger_copy( - job.request_id, source_tp.path, target_tp.path + job.request_id, source_path, target_path ) except Exception as e: # pylint: disable=broad-except logging.exception("Failed to trigger transfer for job %d", job.id) @@ -470,9 +489,110 @@ async def _process_copy_job(self, job: db_schema.AssetJob): error_msg=error_msg, operation_name=operation_name, client_type=determined_client.__class__.__name__, - zone=getattr(determined_client, "location", None), + zone=getattr(determined_client, "zone", None), ) + async def _process_delete_job(self, job: db_schema.AssetJob): + """Processes a DELETE job by deleting the files from GCS/Lustre and updating DB.""" + tps_to_delete_info = [] + async with self._session_maker() as session: + async with session.begin(): + stmt = ( + select(db_schema.Asset) + .options( + joinedload(db_schema.Asset.tier_paths).joinedload( + db_schema.TierPath.storage_backend + ) + ) + .where(db_schema.Asset.asset_uuid == job.asset_uuid) + ) + result = await session.execute(stmt) + db_asset = result.scalars().first() + + if not db_asset: + asset_found = False + else: + asset_found = True + if ( + job.request_type + == db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_INSTANCE + ): + tps = [ + tp + for tp in db_asset.tier_paths + if tp.id == job.target_tier_path_id + ] + else: + tps = [ + tp + for tp in db_asset.tier_paths + if tp.state + in ( + db_schema.TierPathState.PENDING, + db_schema.TierPathState.IN_PROGRESS, + db_schema.TierPathState.READY, + ) + ] + for tp in tps: + tps_to_delete_info.append( + (tp.path, tp.storage_backend.backend_type, tp.id) + ) + + if not asset_found: + await self._fail_job_by_id(job.id, "Asset not found") + return + + errors = [] + for path, backend_type, _ in tps_to_delete_info: + try: + dummy_sb = db_schema.StorageBackend(backend_type=backend_type) + dummy_tp = db_schema.TierPath(path=path, storage_backend=dummy_sb) + + client = gcp_storage_client.determine_delete_client( + dummy_tp, + project=self._config.gcp_project or None, + service_account=self._config.service_account or None, + ) + try: + await client.delete_path(path) + finally: + await client.close() + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Failed to delete path %s for asset %s", path, job.asset_uuid + ) + errors.append(f"Failed to delete path {path}: {e}") + + async with self._session_maker() as session: + async with session.begin(): + merged_job = await session.get( + db_schema.AssetJob, job.id, with_for_update=True + ) + if not merged_job: + logging.warning("Job %d was deleted", job.id) + return + if ( + merged_job.status != db_schema.JobStatus.JOB_STATUS_PROCESSING + or merged_job.worker_host != self._hostname + or merged_job.worker_pid != self._pid + ): + logging.warning( + "Job %d is no longer owned by this worker (%s:%d).", + job.id, + self._hostname, + self._pid, + ) + return + + now = datetime.datetime.now(datetime.timezone.utc) + if errors: + error_msg = "; ".join(errors) + await self._fail_job( + session, merged_job, error_msg, {"status": "FAILED"}, now + ) + else: + await self._complete_job(session, merged_job, now) + async def _get_target_tier_path( self, session: AsyncSession, @@ -542,29 +662,51 @@ async def _complete_job( job.expiration_at = None session.add(job) # pyrefly: ignore[missing-attribute] - # Mark target TierPath as ready - target_tp = await self._get_target_tier_path( - session, job, load_backend=True - ) - if target_tp: - target_tp.state = db_schema.TierPathState.READY - target_tp.ready_at = now - # Calculate expiration for TierPath if it is Level 0 (Lustre) - # "the checkpoint could be removed from this location after it expires." - # GCS (Level 1) paths usually don't expire. - if ( - target_tp.storage_backend.backend_type - == db_schema.BackendType.BACKEND_TYPE_LUSTRE - ): - ttl = datetime.timedelta(seconds=60 * 60) - target_tp.expires_at = assets.calculate_expires_at(ttl) - session.add(target_tp) # pyrefly: ignore[missing-attribute] + if job.request_type == db_schema.RequestType.REQUEST_TYPE_COPY: + target_tp = await self._get_target_tier_path( + session, job, load_backend=True + ) + if target_tp: + target_tp.state = db_schema.TierPathState.READY + target_tp.ready_at = now + if ( + target_tp.storage_backend.backend_type + == db_schema.BackendType.BACKEND_TYPE_LUSTRE + ): + ttl = datetime.timedelta(seconds=60 * 60) + target_tp.expires_at = assets.calculate_expires_at(ttl) + session.add(target_tp) # pyrefly: ignore[missing-attribute] + logging.info( + "Completed job %d, target TierPath %s marked ready", + job.id, + target_tp.path if target_tp else "None", + ) - logging.info( - "Completed job %d, target TierPath %s marked ready", - job.id, - target_tp.path if target_tp else "None", - ) + elif ( + job.request_type + == db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_INSTANCE + ): + target_tp = await self._get_target_tier_path(session, job) + if target_tp: + await assets.complete_delete_tier_path(session, target_tp) + logging.info( + "Completed job %d, target TierPath %s marked deleted", + job.id, + target_tp.path if target_tp else "None", + ) + + elif ( + job.request_type + == db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS + ): + db_assets = await assets.fetch_asset_by_uuid(session, job.asset_uuid) + if db_assets: + await assets.complete_delete_asset(session, db_assets[0]) + logging.info( + "Completed job %d, asset %s marked deleted", + job.id, + job.asset_uuid, + ) async def _extend_lease( self, diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto b/checkpoint/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto index 862bc0e59b..4948946ecb 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto @@ -85,6 +85,7 @@ message ReserveRequest { message ReserveResponse { Asset asset = 1; int64 keep_alive_interval_seconds = 2; + string tier_path_uuid = 3; } message ReserveKeepAliveRequest { @@ -116,6 +117,7 @@ message PrefetchRequest { message PrefetchResponse { Asset asset = 1; int64 keep_alive_interval_seconds = 2; + string closest_tier_path_uuid = 3; } message PrefetchKeepAliveRequest { diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py index c63235f685..c7b0807c49 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py @@ -19,6 +19,7 @@ from concurrent import futures import contextlib import datetime +import logging as std_logging import os import pprint import sys @@ -160,9 +161,15 @@ async def Reserve( ) return tiering_service_pb2.ReserveResponse() + tier_path_uuid = "" + for tp in db_asset.tier_paths: + if tp.storage_backend_id == backend.id: + tier_path_uuid = tp.tier_path_uuid + break return tiering_service_pb2.ReserveResponse( asset=assets.proto_from_db_asset(db_asset), keep_alive_interval_seconds=self._config.client_keep_alive_interval_seconds, + tier_path_uuid=tier_path_uuid, ) async def ReserveKeepAlive( @@ -331,6 +338,7 @@ async def Prefetch( keep_alive_interval_seconds=( self._config.client_keep_alive_interval_seconds ), + closest_tier_path_uuid=tp.tier_path_uuid, ) # No existing TierPath, we need to prefetch @@ -386,9 +394,16 @@ async def Prefetch( await context.abort(grpc.StatusCode.NOT_FOUND, "Asset not found") return tiering_service_pb2.PrefetchResponse() + closest_tier_path_uuid = "" + for tp in db_asset.tier_paths: + if tp.storage_backend_id == closest_backend.id: + closest_tier_path_uuid = tp.tier_path_uuid + break + return tiering_service_pb2.PrefetchResponse( asset=assets.proto_from_db_asset(db_asset), keep_alive_interval_seconds=self._config.client_keep_alive_interval_seconds, + closest_tier_path_uuid=closest_tier_path_uuid, ) async def PrefetchKeepAlive( @@ -419,6 +434,11 @@ async def PrefetchKeepAlive( logging.warning(error_msg) await context.abort(grpc.StatusCode.FAILED_PRECONDITION, error_msg) return tiering_service_pb2.PrefetchKeepAliveResponse() + except assets.PrefetchFailedError as e: + error_msg = f"PrefetchKeepAlive: Prefetch failed for TierPath: {e}" + logging.warning(error_msg) + await context.abort(grpc.StatusCode.ABORTED, error_msg) + return tiering_service_pb2.PrefetchKeepAliveResponse() except Exception: # pylint: disable=broad-except logging.exception( "Failed to keep alive prefetch for TierPath: %s", @@ -591,6 +611,11 @@ def main(argv: Sequence[str] | None = None) -> None: """Main entry point for CTS server.""" if argv is None: argv = sys.argv + std_logging.basicConfig( + format="%(asctime)s %(levelname)s %(message)s", + level=std_logging.INFO, + force=True, + ) uvloop.install() try: asyncio.get_event_loop() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py index 0c0ea9a4d7..486aa815e0 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/server_test.py @@ -29,7 +29,11 @@ from orbax.checkpoint.experimental.tiering_service.proto import tiering_service_pb2 from sqlalchemy.future import select -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 async def _setup_prefetch_req( @@ -210,6 +214,10 @@ async def test_reserve_success(self): ) self.assertLen(response.asset.tier_paths, 1) self.assertStartsWith(response.asset.tier_paths[0].path, "/mnt/lustre") + self.assertEqual( + response.tier_path_uuid, response.asset.tier_paths[0].tier_path_uuid + ) + self.assertTrue(response.tier_path_uuid) async def test_reserve_twice(self): request1 = tiering_service_pb2.ReserveRequest( @@ -483,8 +491,23 @@ async def test_prefetch_success_rpc_response(self): paths = [tp.path for tp in prefetch_res.asset.tier_paths] self.assertCountEqual( paths, - ["/mnt/lustre-a/test/path", "/mnt/lustre-b/test/path"], + [ + "/mnt/lustre-a/test/path", + "/mnt/lustre-b/test/path", + "gs://my-bucket/test/path", + ], + ) + # In the setup, find the tier path corresponding to zone B's Lustre backend + # target. + expected_tp = next( + tp for tp in prefetch_res.asset.tier_paths + if tp.path == "/mnt/lustre-b/test/path" + ) + self.assertEqual( + prefetch_res.closest_tier_path_uuid, + expected_tp.tier_path_uuid, ) + self.assertTrue(prefetch_res.closest_tier_path_uuid) async def test_prefetch_success_db_job_creation(self): servicer, asset_uuid = await self._setup_servicer_and_asset() @@ -495,25 +518,27 @@ async def test_prefetch_success_db_job_creation(self): await servicer.Prefetch(prefetch_req, self.context) async with servicer._session_scope() as session: - stmt = select(db_schema.AssetJob).filter_by( - asset_uuid=asset_uuid, - request_type=db_schema.RequestType.REQUEST_TYPE_COPY, - ) - result = await session.execute(stmt) - jobs = result.scalars().all() - self.assertLen(jobs, 1) - self.assertEqual(jobs[0].status, db_schema.JobStatus.JOB_STATUS_QUEUED) - target_tp_id = jobs[0].target_tier_path_id - stmt_tp = select(db_schema.TierPath).filter_by( asset_uuid=asset_uuid, path="/mnt/lustre-b/test/path" ) result_tp = await session.execute(stmt_tp) tp_b = result_tp.scalars().first() self.assertIsNotNone(tp_b) - self.assertEqual(target_tp_id, tp_b.id) self.assertIsNone(tp_b.ready_at) + stmt = select(db_schema.AssetJob).filter_by( + asset_uuid=asset_uuid, + request_type=db_schema.RequestType.REQUEST_TYPE_COPY, + ) + result = await session.execute(stmt) + jobs = result.scalars().all() + self.assertLen(jobs, 2) + prefetch_job = next(j for j in jobs if j.target_tier_path_id == tp_b.id) + self.assertEqual( + prefetch_job.status, db_schema.JobStatus.JOB_STATUS_QUEUED + ) + self.assertEqual(prefetch_job.target_tier_path_id, tp_b.id) + async def test_prefetch_idempotent(self): servicer, asset_uuid = await self._setup_servicer_and_asset() @@ -541,7 +566,7 @@ async def test_prefetch_idempotent(self): ) result = await session.execute(stmt) jobs = result.scalars().all() - self.assertLen(jobs, 1) + self.assertLen(jobs, 2) async def test_prefetch_already_ready(self): # If we prefetch to the same zone where it was reserved and finalized, @@ -559,18 +584,29 @@ async def test_prefetch_already_ready(self): # Verify response self.assertEqual(response.asset.uuid, asset_uuid) - self.assertLen(response.asset.tier_paths, 1) - self.assertIsNotNone(response.asset.tier_paths[0].ready_at.ToDatetime()) + self.assertLen(response.asset.tier_paths, 2) + lustre_a_tp = next( + tp for tp in response.asset.tier_paths if "lustre" in tp.path + ) + self.assertIsNotNone(lustre_a_tp.ready_at.ToDatetime()) - # Verify NO job was created + # Verify NO prefetch job was created (only GCS copy job exists) async with self.servicer._session_scope() as session: + stmt_tp = select(db_schema.TierPath).filter_by( + asset_uuid=asset_uuid, path="gs://my-bucket/test/path" + ) + result_tp = await session.execute(stmt_tp) + gcs_tp = result_tp.scalars().first() + self.assertIsNotNone(gcs_tp) + stmt = select(db_schema.AssetJob).filter_by( asset_uuid=asset_uuid, request_type=db_schema.RequestType.REQUEST_TYPE_COPY, ) result = await session.execute(stmt) jobs = result.scalars().all() - self.assertEmpty(jobs) + self.assertLen(jobs, 1) + self.assertEqual(jobs[0].target_tier_path_id, gcs_tp.id) async def test_prefetch_permission_denied(self): servicer, asset_uuid = await self._setup_servicer_and_asset() @@ -667,7 +703,7 @@ async def test_prefetch_keep_alive_grpc_success(self): ), self.context, ) - self.assertLen(prefetch_res.asset.tier_paths, 2) + self.assertLen(prefetch_res.asset.tier_paths, 3) tp_b = next( tp for tp in prefetch_res.asset.tier_paths if "lustre-b" in tp.path ) diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py index c3fd2d5624..2c7676b48c 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py @@ -110,6 +110,8 @@ def get_storage_path( Returns: The absolute storage path. """ + if relative_path.startswith(backend.prefix): + return relative_path return f"{backend.prefix.rstrip('/')}/{relative_path.lstrip('/')}" diff --git a/checkpoint/orbax/checkpoint/options.py b/checkpoint/orbax/checkpoint/options.py index a3f636319a..80750598f8 100644 --- a/checkpoint/orbax/checkpoint/options.py +++ b/checkpoint/orbax/checkpoint/options.py @@ -15,6 +15,7 @@ """Configuration options for APIs like CheckpointManager and Checkpointer.""" import dataclasses +import typing from typing import Callable, Optional, Set from orbax.checkpoint._src.multihost import multihost