diff --git a/checkpoint/orbax/__init__.py b/checkpoint/orbax/__init__.py index af16c599e..82a91a199 100644 --- a/checkpoint/orbax/__init__.py +++ b/checkpoint/orbax/__init__.py @@ -19,6 +19,7 @@ import contextlib import functools + 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 af16c599e..82a91a199 100644 --- a/checkpoint/orbax/checkpoint/__init__.py +++ b/checkpoint/orbax/checkpoint/__init__.py @@ -19,6 +19,7 @@ import contextlib import functools + 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 435de5e35..086ef0212 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 += ( @@ -957,6 +962,7 @@ def _configure_checkpointer_common( options: CheckpointManagerOptions, use_async: bool, ) -> Checkpointer: + kwargs = {} if use_async: return async_checkpointer.AsyncCheckpointer( handler, @@ -965,6 +971,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( @@ -973,6 +980,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( @@ -1516,6 +1524,17 @@ 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 + save_directory = epath.Path(lustre_path) + logging.info( '[process=%s] Saving checkpoint at step %d', process_index, step ) @@ -1737,8 +1756,24 @@ 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) + if self._tiering_client is not None: + logical_restore_dir = self._get_read_step_directory(step, directory) + self._run_async( + self._tiering_client.release_path(str(logical_restore_dir)) + ) step_stats.checkpointer_duration_secs = ( time.time() - step_stats.checkpointer_start_time ) @@ -2153,6 +2188,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( @@ -2179,6 +2220,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 faf792561..82036289f 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py @@ -22,6 +22,7 @@ from collections.abc import Collection, Sequence import dataclasses import datetime +import uuid from absl import logging from orbax.checkpoint.experimental.tiering_service import db_schema @@ -33,13 +34,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. @@ -105,6 +114,10 @@ def _get_location_kwargs(sb: db_schema.StorageBackend): expires_at_pb = timestamp_pb2.Timestamp() expires_at_pb.FromDatetime(tier_path.expires_at) + state_val = tier_path.state or db_schema.TierPathState.UNSPECIFIED + proto_state_name = f"TIER_PATH_STATE_{state_val.name}" + state_pb = tiering_service_pb2.TierPathState.Value(proto_state_name) + return tiering_service_pb2.TierPath( id=tier_path.id, path=tier_path.path, @@ -112,6 +125,7 @@ def _get_location_kwargs(sb: db_schema.StorageBackend): ready_at=ready_at_pb, expires_at=expires_at_pb, tier_path_uuid=tier_path.tier_path_uuid, + state=state_pb, ) @@ -271,6 +285,9 @@ async def create_or_fetch_asset( request: tiering_service_pb2.ReserveRequest, backend: db_schema.StorageBackend, config: tiering_service_pb2.ServerConfig, + *, + tier_path_uuid: str, + storage_path: str, ) -> db_schema.Asset: """Creates a new asset or fetches an existing one on unique constraint conflict. @@ -284,6 +301,8 @@ async def create_or_fetch_asset( request: The ReserveRequest containing path, user, and tags. backend: The StorageBackend to associate the asset with. config: The ServerConfig to get the keep-alive interval. + tier_path_uuid: The pre-generated UUID for the new TierPath. + storage_path: The pre-generated physical storage path. Returns: The created or fetched Asset object. @@ -302,10 +321,11 @@ async def create_or_fetch_asset( ) ), ) - storage_path = storage_backend_lib.get_storage_path(backend, request.path) + backend = await session.merge(backend) tier_path = db_schema.TierPath( storage_backend=backend, path=storage_path, + tier_path_uuid=tier_path_uuid, ) db_asset.tier_paths.append(tier_path) @@ -371,7 +391,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. @@ -383,6 +404,14 @@ async def finalize_asset( Raises: ValueError: If the asset is not in ACTIVE_WRITE state. """ + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == db_asset.asset_uuid) + .options(sqlalchemy.orm.joinedload(db_schema.Asset.tier_paths)) + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + if db_asset.state != db_schema.AssetState.ASSET_STATE_ACTIVE_WRITE: raise ValueError( f"Asset {db_asset.asset_uuid} is in state {db_asset.state.name}, but" @@ -396,8 +425,42 @@ async def finalize_asset( for tier_path in db_asset.tier_paths: tier_path.ready_at = now + tier_path.state = db_schema.TierPathState.READY # 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_tp_uuid = uuid.uuid4().hex + gcs_path = storage_backend_lib.get_storage_path( + gcs_backend, db_asset.path, gcs_tp_uuid + ) + new_gcs_tp = db_schema.TierPath( + storage_backend=gcs_backend, + path=gcs_path, + tier_path_uuid=gcs_tp_uuid, + 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 @@ -453,6 +516,7 @@ async def create_prefetch_job( *, backend: db_schema.StorageBackend, storage_path: str, + tier_path_uuid: str, client_keep_alive_interval: datetime.timedelta, ) -> CreatePrefetchJobResult: """Queues a prefetch job for the given asset to the target backend. @@ -468,6 +532,7 @@ async def create_prefetch_job( db_asset: The asset to prefetch. backend: The target storage backend (level 0). storage_path: The storage path to use for the new TierPath. + tier_path_uuid: The pre-generated unique identifier for the TierPath. client_keep_alive_interval: The interval to set for the initial expires_at of the TierPath. @@ -479,6 +544,14 @@ async def create_prefetch_job( DeletionPendingError: If the asset is already marked for deletion. """ + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == db_asset.asset_uuid) + .options(sqlalchemy.orm.joinedload(db_schema.Asset.tier_paths)) + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + # Check if there is already a preceding delete job if await is_delete_pending(session, asset_uuid=db_asset.asset_uuid): raise DeletionPendingError( @@ -494,9 +567,11 @@ async def create_prefetch_job( db_asset.asset_uuid, backend.id, ) + backend = await session.merge(backend) new_tp = db_schema.TierPath( storage_backend=backend, path=storage_path, + tier_path_uuid=tier_path_uuid, expires_at=calculate_expires_at(client_keep_alive_interval), ) db_asset.tier_paths.append(new_tp) @@ -550,6 +625,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 +633,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 +724,93 @@ async def queue_delete_asset_job( session.add(db_job) # pyrefly: ignore[missing-attribute] await session.commit() + + +async def begin_delete_tier_path( + session: AsyncSession, + tier_path: db_schema.TierPath, +) -> db_schema.TierPath: + """Transitions a tier path state to DELETE_IN_PROCESS and clears read access.""" + tier_path.state = db_schema.TierPathState.DELETE_IN_PROCESS + tier_path.ready_at = None + tier_path.expires_at = None + session.add(tier_path) # pyrefly: ignore[missing-attribute] + return tier_path + + +async def begin_delete_asset( + session: AsyncSession, + db_asset: db_schema.Asset, +) -> db_schema.Asset: + """Transitions asset state to DELETED and all its tier paths to DELETE_IN_PROCESS.""" + db_asset.state = db_schema.AssetState.ASSET_STATE_DELETED + db_asset.deleted_at = datetime.datetime.now(datetime.timezone.utc) + session.add(db_asset) # pyrefly: ignore[missing-attribute] + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == db_asset.asset_uuid) + .options(sqlalchemy.orm.joinedload(db_schema.Asset.tier_paths)) + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + for tp in db_asset.tier_paths: + tp.state = db_schema.TierPathState.DELETE_IN_PROCESS + tp.ready_at = None + tp.expires_at = None + session.add(tp) # pyrefly: ignore[missing-attribute] + return db_asset + + +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. + """ + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == db_asset.asset_uuid) + .options(sqlalchemy.orm.joinedload(db_schema.Asset.tier_paths)) + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + + now = datetime.datetime.now(datetime.timezone.utc) + db_asset.state = db_schema.AssetState.ASSET_STATE_DELETED + db_asset.deleted_at = now + session.add(db_asset) # pyrefly: ignore[missing-attribute] + + for tp in db_asset.tier_paths: + tp.state = db_schema.TierPathState.DELETED + tp.ready_at = None + tp.expires_at = None + session.add(tp) # pyrefly: ignore[missing-attribute] + + 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 + session.add(tier_path) # pyrefly: ignore[missing-attribute] + 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 b624c85be..6165f2bb8 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): @@ -119,6 +123,7 @@ def test_proto_from_db_asset_tier_path_and_backend(self): ready_at=dt_ready, expires_at=dt_expires, storage_backend=db_backend, + state=db_schema.TierPathState.READY, ) db_asset = db_schema.Asset( asset_uuid="test-uuid", @@ -135,6 +140,10 @@ def test_proto_from_db_asset_tier_path_and_backend(self): self.assertEqual(tp_proto.path, "/mnt/lustre/test/path") self.assertEqual(tp_proto.ready_at, expected_ts_ready) self.assertEqual(tp_proto.expires_at, expected_ts_expires) + self.assertEqual( + tp_proto.state, + tiering_service_pb2.TierPathState.TIER_PATH_STATE_READY, + ) sb_proto = tp_proto.storage_backend self.assertEqual(sb_proto.id, 1) @@ -253,16 +262,28 @@ async def test_create_or_fetch_asset_and_queries(self): client_keep_alive_interval_seconds=600 ) + tp_uuid = "test-tp-uuid" + storage_path = "test/storage/path" # Create asset. asset = await assets.create_or_fetch_asset( - session, request, backend, config + session, + request, + backend, + config, + tier_path_uuid=tp_uuid, + storage_path=storage_path, ) self.assertEqual(asset.path, "test/path") self.assertLen(asset.tier_paths, 1) # Try creating it again (triggers unique conflict fetch fallback). asset2 = await assets.create_or_fetch_asset( - session, request, backend, config + session, + request, + backend, + config, + tier_path_uuid=tp_uuid, + storage_path=storage_path, ) self.assertEqual(asset2.asset_uuid, asset.asset_uuid) @@ -311,8 +332,17 @@ async def test_mutations_keep_alive_finalize(self): config = tiering_service_pb2.ServerConfig( client_keep_alive_interval_seconds=600 ) + tp_uuid = "test-tp-uuid" + storage_path = storage_backend_lib.get_storage_path( + backend, request.path, tp_uuid + ) asset = await assets.create_or_fetch_asset( - session, request, backend, config + session, + request, + backend, + config, + tier_path_uuid=tp_uuid, + storage_path=storage_path, ) # Keep alive. @@ -337,6 +367,9 @@ async def test_mutations_keep_alive_finalize(self): self.assertEqual(finalized.state, db_schema.AssetState.ASSET_STATE_STORED) self.assertLen(finalized.tier_paths, 1) self.assertEqual(finalized.tier_paths[0].ready_at, finalized.finalized_at) + self.assertEqual( + finalized.tier_paths[0].state, db_schema.TierPathState.READY + ) # Verify finalize persistence. async with self.session_maker() as session3: @@ -348,6 +381,10 @@ async def test_mutations_keep_alive_finalize(self): self._assert_date_time_equal( fetched3[0].finalized_at, finalized.finalized_at ) + self.assertLen(fetched3[0].tier_paths, 1) + self.assertEqual( + fetched3[0].tier_paths[0].state, db_schema.TierPathState.READY + ) async def test_queries_filtering(self): async with self.session_maker() as session: @@ -370,8 +407,17 @@ async def test_queries_filtering(self): user="user-a", zone="us-central1-a", ) + tp_uuid_a = "test-tp-uuid-a" + storage_path_a = storage_backend_lib.get_storage_path( + backend, request_a.path, tp_uuid_a + ) asset_a = await assets.create_or_fetch_asset( - session, request_a, backend, config + session, + request_a, + backend, + config, + tier_path_uuid=tp_uuid_a, + storage_path=storage_path_a, ) # Create Asset B. @@ -380,8 +426,17 @@ async def test_queries_filtering(self): user="user-b", zone="us-central1-a", ) + tp_uuid_b = "test-tp-uuid-b" + storage_path_b = storage_backend_lib.get_storage_path( + backend, request_b.path, tp_uuid_b + ) asset_b = await assets.create_or_fetch_asset( - session, request_b, backend, config + session, + request_b, + backend, + config, + tier_path_uuid=tp_uuid_b, + storage_path=storage_path_b, ) # Verify fetch_asset_by_path only returns the matched asset. @@ -444,6 +499,10 @@ async def _set_a_finalized_asset( user="test-user", zone="us-central1-a", ) + tp_uuid = "test-tp-uuid-b1" + storage_path = storage_backend_lib.get_storage_path( + b1, request.path, tp_uuid + ) reserved_asset = await assets.create_or_fetch_asset( session, request, @@ -451,6 +510,8 @@ async def _set_a_finalized_asset( tiering_service_pb2.ServerConfig( client_keep_alive_interval_seconds=600 ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, ) finalized_asset = await assets.finalize_asset(session, reserved_asset) return finalized_asset, b1, b2 @@ -458,12 +519,16 @@ async def _set_a_finalized_asset( async def test_create_prefetch_job_returns_created_and_updated_asset(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) self.assertTrue(result.created) @@ -472,12 +537,16 @@ async def test_create_prefetch_job_returns_created_and_updated_asset(self): async def test_create_prefetch_job_updates_tier_paths(self): async with self.session_maker() as session: asset, b1, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -486,7 +555,9 @@ async def test_create_prefetch_job_updates_tier_paths(self): self.assertCountEqual( paths, [ - storage_backend_lib.get_storage_path(b1, asset.path), + storage_backend_lib.get_storage_path( + b1, asset.path, "test-tp-uuid-b1" + ), storage_path, ], ) @@ -494,12 +565,16 @@ async def test_create_prefetch_job_updates_tier_paths(self): async def test_create_prefetch_job_db_tier_path_not_ready(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) @@ -515,12 +590,16 @@ async def test_create_prefetch_job_db_tier_path_not_ready(self): async def test_create_prefetch_job_db_queues_copy_job(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) @@ -584,6 +663,7 @@ async def test_create_prefetch_job_concurrent_fails_gracefully( asset1, backend=sb2, storage_path=attempted_path, + tier_path_uuid="test-prefetch-attempt-uuid", client_keep_alive_interval=datetime.timedelta(seconds=600), ) @@ -601,12 +681,16 @@ async def test_create_prefetch_job_concurrent_fails_gracefully( async def test_create_prefetch_job_sets_expires_at_and_uuid(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -617,18 +701,22 @@ async def test_create_prefetch_job_sets_expires_at_and_uuid(self): for tp in updated_asset.tier_paths if tp.storage_backend_id == b2.id ) - self.assertIsNotNone(tp_b.tier_path_uuid) + self.assertEqual(tp_b.tier_path_uuid, tp_uuid) self.assertIsNotNone(tp_b.expires_at) async def test_prefetch_keep_alive_success(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -665,12 +753,16 @@ async def test_prefetch_keep_alive_not_found(self): async def test_prefetch_keep_alive_no_op_when_permanent(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -699,12 +791,16 @@ async def test_prefetch_keep_alive_no_op_when_permanent(self): async def test_prefetch_keep_alive_only_extends(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -732,12 +828,16 @@ async def test_prefetch_keep_alive_only_extends(self): async def test_prefetch_keep_alive_fails_if_deleting(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -762,12 +862,16 @@ async def test_prefetch_keep_alive_fails_if_deleting(self): async def test_prefetch_keep_alive_fails_if_instance_deleting(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -799,7 +903,10 @@ async def test_prefetch_keep_alive_fails_if_instance_deleting(self): async def test_create_prefetch_job_fails_if_deleting(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) await assets.queue_delete_asset_job(session, asset) @@ -811,6 +918,7 @@ async def test_create_prefetch_job_fails_if_deleting(self): asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) @@ -842,12 +950,16 @@ async def test_is_delete_pending_false(self): async def test_is_tier_path_delete_pending_true(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -878,12 +990,16 @@ async def test_is_tier_path_delete_pending_true(self): async def test_is_tier_path_delete_pending_false(self): async with self.session_maker() as session: asset, _, b2 = await self._set_a_finalized_asset(session) - storage_path = storage_backend_lib.get_storage_path(b2, asset.path) + tp_uuid = "test-prefetch-tp-uuid-b2" + storage_path = storage_backend_lib.get_storage_path( + b2, asset.path, tp_uuid + ) result = await assets.create_prefetch_job( session, asset, backend=b2, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta(seconds=600), ) updated_asset = result.asset @@ -959,6 +1075,162 @@ async def test_queue_delete_asset_job_already_deleted(self): jobs = res.scalars().all() self.assertEmpty(jobs) + async def test_begin_delete_asset_lazy_load_success(self): + async with self.session_maker() as session: + asset, _, _ = await self._set_a_finalized_asset(session) + asset_uuid = asset.asset_uuid + + # Now load asset in a new session without eagerly loading tier_paths. + async with self.session_maker() as session: + stmt = select(db_schema.Asset).where( + db_schema.Asset.asset_uuid == asset_uuid + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + self.assertIsNotNone(db_asset) + + # Should succeed after our fix! + await assets.begin_delete_asset(session, db_asset) + await session.commit() + + async with self.session_maker() as session: + db_assets = await assets.fetch_asset_by_uuid(session, asset_uuid) + db_asset = db_assets[0] + self.assertEqual(db_asset.state, db_schema.AssetState.ASSET_STATE_DELETED) + self.assertNotEmpty(db_asset.tier_paths) + for tp in db_asset.tier_paths: + self.assertEqual(tp.state, db_schema.TierPathState.DELETE_IN_PROCESS) + + async def test_finalize_asset_lazy_load_success(self): + async with self.session_maker() as session: + b1 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + session.add(b1) # pyrefly: ignore[missing-attribute] + await session.commit() + + request = tiering_service_pb2.ReserveRequest( + path="test/path/unfinalized_asset", + user="test-user", + zone="us-central1-a", + ) + tp_uuid = "test-tp-uuid-unfinalized" + storage_path = storage_backend_lib.get_storage_path( + b1, request.path, tp_uuid + ) + reserved_asset = await assets.create_or_fetch_asset( + session, + request, + b1, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, + ) + asset_uuid = reserved_asset.asset_uuid + + async with self.session_maker() as session: + stmt = select(db_schema.Asset).where( + db_schema.Asset.asset_uuid == asset_uuid + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + self.assertIsNotNone(db_asset) + + # Should succeed now! + await assets.finalize_asset(session, db_asset) + await session.commit() + + async with self.session_maker() as session: + db_assets = await assets.fetch_asset_by_uuid(session, asset_uuid) + db_asset = db_assets[0] + self.assertEqual(db_asset.state, db_schema.AssetState.ASSET_STATE_STORED) + self.assertNotEmpty(db_asset.tier_paths) + for tp in db_asset.tier_paths: + self.assertIsNotNone(tp.ready_at) + self.assertEqual(tp.state, db_schema.TierPathState.READY) + + async def test_create_prefetch_job_lazy_load_success(self): + async with self.session_maker() as session: + asset, _, b2 = await self._set_a_finalized_asset(session) + asset_uuid = asset.asset_uuid + b2_id = b2.id + + # Now load asset in a new session without eagerly loading tier_paths. + async with self.session_maker() as session: + stmt = select(db_schema.Asset).where( + db_schema.Asset.asset_uuid == asset_uuid + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + self.assertIsNotNone(db_asset) + + res_b2 = await session.execute( + select(db_schema.StorageBackend).where( + db_schema.StorageBackend.id == b2_id + ) + ) + db_backend = res_b2.scalars().first() + + tp_uuid = "test-prefetch-tp-uuid-lazy" + storage_path = storage_backend_lib.get_storage_path( + db_backend, db_asset.path, tp_uuid + ) + + # Should succeed now! + result = await assets.create_prefetch_job( + session, + db_asset, + backend=db_backend, + storage_path=storage_path, + tier_path_uuid=tp_uuid, + client_keep_alive_interval=datetime.timedelta(seconds=600), + ) + self.assertTrue(result.created) + await session.commit() + + async with self.session_maker() as session: + db_assets = await assets.fetch_asset_by_uuid(session, asset_uuid) + db_asset = db_assets[0] + # Verify the new tier path is in the DB. + self.assertLen( + db_asset.tier_paths, 2 + ) # L0 reserved (1), L0 prefetched (2) + tp_prefetched = next( + tp for tp in db_asset.tier_paths if tp.tier_path_uuid == tp_uuid + ) + self.assertEqual(tp_prefetched.state, db_schema.TierPathState.PENDING) + + async def test_complete_delete_asset_lazy_load_success(self): + async with self.session_maker() as session: + asset, _, _ = await self._set_a_finalized_asset(session) + asset_uuid = asset.asset_uuid + + # Now load asset in a new session without eagerly loading tier_paths. + async with self.session_maker() as session: + stmt = select(db_schema.Asset).where( + db_schema.Asset.asset_uuid == asset_uuid + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + self.assertIsNotNone(db_asset) + + # Should succeed now! + await assets.complete_delete_asset(session, db_asset) + await session.commit() + + async with self.session_maker() as session: + db_assets = await assets.fetch_asset_by_uuid(session, asset_uuid) + db_asset = db_assets[0] + self.assertEqual(db_asset.state, db_schema.AssetState.ASSET_STATE_DELETED) + self.assertNotEmpty(db_asset.tier_paths) + for tp in db_asset.tier_paths: + self.assertEqual(tp.state, db_schema.TierPathState.DELETED) + if __name__ == "__main__": absltest.main() diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py index 9d0380254..12d96a477 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 index 6570a95f0..55d1fac4c 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client.py @@ -48,7 +48,7 @@ def __init__( self.interval = interval self.tier_path_uuid = tier_path_uuid self.loop = asyncio.get_running_loop() - self.next_run = asyncio.get_running_loop().time() + interval + self.next_run = self.loop.time() + interval class TieringClient: @@ -65,16 +65,23 @@ def __init__( """ self._server_address = server_address self._secure = secure - self._channel = None - self._stub = None + 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_task: asyncio.Task[None] | None = None - self._keep_alive_event: asyncio.Event = asyncio.Event() - self._prefetch_futures: dict[str, asyncio.Future[str]] = {} + self._keep_alive_manager_tasks: dict[ + asyncio.AbstractEventLoop, asyncio.Task[None] + ] = {} + self._keep_alive_events: dict[asyncio.AbstractEventLoop, asyncio.Event] = {} + self._prefetch_futures: dict[ + asyncio.AbstractEventLoop, dict[str, asyncio.Future[str]] + ] = {} + self._path_to_uuid: dict[str, str] = {} async def __aenter__(self) -> "TieringClient": await self.connect() @@ -83,55 +90,98 @@ async def __aenter__(self) -> "TieringClient": async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: await self.close() - async def connect(self) -> None: - """Establishes an async gRPC channel with the server.""" - if self._channel is not None: - return - - if self._secure: - is_local = ( - "localhost" in self._server_address - or "127.0.0.1" in self._server_address - ) - if is_local: - try: - # Secure channel setup. Fall back to SSL if local creds not supported. - creds = grpc.local_channel_credentials() - except AttributeError: + 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: - creds = grpc.ssl_channel_credentials() - self._channel = grpc.aio.secure_channel(self._server_address, creds) - else: - self._channel = grpc.aio.insecure_channel(self._server_address) + channel = grpc.aio.insecure_channel(self._server_address) - self._stub = tiering_service_pb2_grpc.TieringServiceStub(self._channel) + 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.""" - # Release pending prefetches and cancel futures - for asset_uuid, fut in list(self._prefetch_futures.items()): - if not fut.done(): - self._stop_prefetch_keep_alive(asset_uuid) - fut.cancel() - self._prefetch_futures.clear() - - # Cancel manager task and clear job list - if self._keep_alive_manager_task is not None: - self._keep_alive_manager_task.cancel() + 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 self._keep_alive_manager_task + await task except asyncio.CancelledError: pass - self._keep_alive_manager_task = None - - self._keep_alives.clear() - self._keep_alive_event.clear() + del self._keep_alive_manager_tasks[current_loop] + + # Cancel manager tasks for other loops (or all if no current_loop) + for loop_val, task in list(self._keep_alive_manager_tasks.items()): + task.cancel() + del self._keep_alive_manager_tasks[loop_val] + + # Clean up jobs belonging to current loop (or all if current_loop is None) + 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] + else: + self._keep_alives.clear() + self._keep_alive_events.clear() + + # Release pending prefetches and cancel futures belonging to current loop + if current_loop and current_loop in self._prefetch_futures: + for asset_uuid, fut in list(self._prefetch_futures[current_loop].items()): + if not fut.done(): + self._stop_prefetch_keep_alive(asset_uuid) + fut.cancel() + del self._prefetch_futures[current_loop] + + # Clean up remaining loops' futures (cancel without awaiting) + for loop_val, fut_dict in list(self._prefetch_futures.items()): + for asset_uuid, fut in list(fut_dict.items()): + if not fut.done(): + fut.cancel() + del self._prefetch_futures[loop_val] + + 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: + asyncio.run(channel.close()) + except Exception: # pylint: disable=broad-except + pass + del self._channels[loop_val] - if self._channel is not None: - await self._channel.close() - self._channel = None - self._stub = None + 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.""" @@ -154,11 +204,14 @@ async def _get_auth_metadata(self) -> list[tuple[str, str]]: 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 ( - self._keep_alive_manager_task is None - or self._keep_alive_manager_task.done() + loop not in self._keep_alive_manager_tasks + or self._keep_alive_manager_tasks[loop].done() ): - self._keep_alive_manager_task = asyncio.create_task( + self._keep_alive_manager_tasks[loop] = asyncio.create_task( self._keep_alive_manager_loop() ) @@ -171,13 +224,13 @@ def _start_write_keep_alive(self, asset_uuid: str, interval: int) -> None: ) self._keep_alives[(asset_uuid, JobType.WRITE)] = job self._ensure_manager_running() - self._keep_alive_event.set() + 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_event.set() + self._keep_alive_events[job.loop].set() def _start_prefetch_keep_alive( self, asset_uuid: str, tier_path_uuid: str, interval: int @@ -191,52 +244,73 @@ def _start_prefetch_keep_alive( ) self._keep_alives[(asset_uuid, JobType.PREFETCH)] = job self._ensure_manager_running() - self._keep_alive_event.set() + 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) + loop = None if job: - self._keep_alive_event.set() - fut = self._prefetch_futures.pop(asset_uuid, None) - if fut and not fut.done(): - fut.cancel() + self._keep_alive_events[job.loop].set() + loop = job.loop + else: + try: + loop = asyncio.get_running_loop() + except RuntimeError: + pass + + if loop and loop in self._prefetch_futures: + fut = self._prefetch_futures[loop].pop(asset_uuid, None) + if fut and not fut.done(): + fut.cancel() - def _get_earliest_job(self) -> tuple[_KeepAliveJob | None, float | None]: + 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 earliest_time is None or job.next_run < earliest_time: - earliest_time = job.next_run - earliest_job = job + 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, timeout: float) -> bool: + 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.""" + event_task = loop.create_task(self._keep_alive_events[loop].wait()) + sleep_task = loop.create_task(asyncio.sleep(timeout)) try: - await asyncio.wait_for(self._keep_alive_event.wait(), timeout=timeout) - return True - except asyncio.TimeoutError: - return False + done, _ = await asyncio.wait( + [event_task, sleep_task], + return_when=asyncio.FIRST_COMPLETED + ) + return event_task in done + finally: + event_task.cancel() + sleep_task.cancel() 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_event.clear() - earliest_job, earliest_time = self._get_earliest_job() + 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_event.wait() + await self._keep_alive_events[loop].wait() continue - now = asyncio.get_running_loop().time() + 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(sleep_duration): + if await self._wait_for_next_job(loop, sleep_duration): continue await self._run_keep_alive_job(earliest_job) @@ -311,13 +385,14 @@ async def _run_prefetch_keep_alive_job( 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) + loop = asyncio.get_running_loop() + if loop in self._prefetch_futures: + fut = self._prefetch_futures[loop].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 + ) except grpc.aio.AioRpcError as e: logging.warning( @@ -325,12 +400,18 @@ async def _run_prefetch_keep_alive_job( job.asset_uuid, e.details(), ) - if e.code() == grpc.StatusCode.NOT_FOUND: - fut = self._prefetch_futures.get(job.asset_uuid) - if fut and not fut.done(): - fut.set_exception( - RuntimeError(f"Prefetch failed: asset {job.asset_uuid} not found") - ) + if e.code() in ( + grpc.StatusCode.NOT_FOUND, + grpc.StatusCode.FAILED_PRECONDITION, + grpc.StatusCode.ABORTED, + ): + loop = asyncio.get_running_loop() + if loop in self._prefetch_futures: + fut = self._prefetch_futures[loop].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) @@ -340,19 +421,16 @@ async def _run_prefetch_keep_alive_job( job.asset_uuid, e, ) - fut = self._prefetch_futures.get(job.asset_uuid) - if fut and not fut.done(): - fut.set_exception(e) + loop = asyncio.get_running_loop() + if loop in self._prefetch_futures: + fut = self._prefetch_futures[loop].get(job.asset_uuid) + if fut and not fut.done(): + fut.set_exception(e) self._keep_alives.pop((job.asset_uuid, JobType.PREFETCH), None) async def _run_keep_alive_job(self, job: _KeepAliveJob) -> None: """Executes a single keep-alive heartbeat request.""" - stub = self._stub - if stub is None: - logging.error("Stub is not initialized in keep-alive manager.") - job.next_run = asyncio.get_running_loop().time() + job.interval - return - + stub = self._get_or_create_stub() now = asyncio.get_running_loop().time() if job.job_type == JobType.WRITE: @@ -379,12 +457,7 @@ async def reserve( Raises: RuntimeError: If gRPC call fails or no Tier 0 path is returned. """ - if not self._stub: - await self.connect() - - stub = self._stub - if stub is None: - raise RuntimeError("Stub is not initialized after connect.") + stub = self._get_or_create_stub() if user is None: user = environment.get_current_user() @@ -444,12 +517,7 @@ async def finalize(self, uuid: str) -> None: Raises: RuntimeError: If gRPC call fails. """ - if not self._stub: - await self.connect() - - stub = self._stub - if stub is None: - raise RuntimeError("Stub is not initialized after connect.") + stub = self._get_or_create_stub() request = tiering_service_pb2.FinalizeRequest(uuid=uuid) metadata = await self._get_auth_metadata() @@ -486,15 +554,11 @@ async def prefetch( 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] + loop = asyncio.get_running_loop() + if loop in self._prefetch_futures and uuid in self._prefetch_futures[loop]: + return self._prefetch_futures[loop][uuid] - if not self._stub: - await self.connect() - - stub = self._stub - if stub is None: - raise RuntimeError("Stub is not initialized after connect.") + stub = self._get_or_create_stub() zone, region = await self._get_gcp_zone_and_region() @@ -521,11 +585,16 @@ async def prefetch( asset_uuid = asset.uuid interval = response.keep_alive_interval_seconds - if asset_uuid in self._prefetch_futures: - return self._prefetch_futures[asset_uuid] + if ( + loop in self._prefetch_futures + and asset_uuid in self._prefetch_futures[loop] + ): + return self._prefetch_futures[loop][asset_uuid] - future = asyncio.get_running_loop().create_future() - self._prefetch_futures[asset_uuid] = future + if loop not in self._prefetch_futures: + self._prefetch_futures[loop] = {} + future = loop.create_future() + self._prefetch_futures[loop][asset_uuid] = future closest_tp = None if response.closest_tier_path_uuid: @@ -540,7 +609,8 @@ async def prefetch( ) if closest_tp is None: - self._prefetch_futures.pop(asset_uuid, None) + if loop in self._prefetch_futures: + self._prefetch_futures[loop].pop(asset_uuid, None) raise RuntimeError( "Prefetch response did not contain closest TierPath matching " f"{response.closest_tier_path_uuid} for asset {asset_uuid}" @@ -550,9 +620,11 @@ async def prefetch( asset_uuid, closest_tp.tier_path_uuid, interval ) + self._path_to_uuid[asset.path] = asset_uuid + self._path_to_uuid[closest_tp.path] = asset_uuid + if closest_tp.HasField("ready_at"): future.set_result(closest_tp.path) - return future async def release(self, uuid: str) -> None: @@ -563,6 +635,19 @@ async def release(self, uuid: str) -> None: """ self._stop_prefetch_keep_alive(uuid) + async def release_path(self, path: str) -> None: + """Client-side release of prefetch keep-alive loop by path. + + Args: + path: Logical path or physical Lustre path of the asset. + """ + uuid = self._path_to_uuid.pop(path, None) + if uuid: + keys_to_remove = [k for k, v in self._path_to_uuid.items() if v == uuid] + for k in keys_to_remove: + self._path_to_uuid.pop(k, None) + await self.release(uuid) + async def delete( self, path: str | None = None, @@ -583,12 +668,7 @@ async def delete( if path is not None and uuid is not None: raise ValueError("Only one of path or uuid can be specified.") - if not self._stub: - await self.connect() - - stub = self._stub - if stub is None: - raise RuntimeError("Stub is not initialized after connect.") + stub = self._get_or_create_stub() if uuid is not None: request = tiering_service_pb2.DeleteRequest(uuid=uuid) @@ -626,12 +706,7 @@ async def info( if path is not None and uuid is not None: raise ValueError("Only one of path or uuid can be specified.") - if not self._stub: - await self.connect() - - stub = self._stub - if stub is None: - raise RuntimeError("Stub is not initialized after connect.") + stub = self._get_or_create_stub() if uuid is not None: request = tiering_service_pb2.InfoRequest(uuid=uuid) diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py index 4e72cfc83..4e2fa66e9 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/client_test.py @@ -131,11 +131,12 @@ async def test_connect_and_close( client_inst = client.TieringClient() await client_inst.connect() - self.assertEqual(client_inst._stub, self.stub_mock) + loop = asyncio.get_running_loop() + self.assertEqual(client_inst._stubs[loop], self.stub_mock) await client_inst.close() - self.assertIsNone(client_inst._channel) - self.assertIsNone(client_inst._stub) + 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") @@ -393,7 +394,9 @@ async def mock_sleep_fn(delay): self.assertEqual(uuid, "asset-1") self.assertIn(("asset-1", client.JobType.WRITE), client_inst._keep_alives) - manager_task = client_inst._keep_alive_manager_task + manager_task = client_inst._keep_alive_manager_tasks.get( + asyncio.get_running_loop() + ) assert manager_task is not None self.assertFalse(manager_task.done()) @@ -581,8 +584,7 @@ async def mock_wait_for_fn(fut, timeout=None): asyncio.CancelledError(), ] - await asyncio.sleep(0) - await asyncio.sleep(0) + await original_sleep(0.1) self.assertTrue(future.done()) self.assertEqual(await future, "/lustre/path1") diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/db_lib.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/db_lib.py index a29d53ad5..adf48d2ee 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/db_lib.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/db_lib.py @@ -56,7 +56,9 @@ def set_sqlite_pragma(dbapi_connection, connection_record): cursor.close() -def get_async_engine(config: tiering_service_pb2.ServerConfig) -> AsyncEngine: +def get_async_engine( + config: tiering_service_pb2.ServerConfig, **kwargs +) -> AsyncEngine: """Returns an AsyncEngine configured from ServerConfig.""" input_url = config.db_connection_str @@ -67,7 +69,7 @@ def get_async_engine(config: tiering_service_pb2.ServerConfig) -> AsyncEngine: url = input_url.replace("sqlite://", "sqlite+aiosqlite://", 1) else: url = input_url - return create_async_engine(url) + return create_async_engine(url, **kwargs) @contextlib.asynccontextmanager diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/db_schema.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/db_schema.py index 2f9474987..eedfc904b 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/db_schema.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/db_schema.py @@ -72,6 +72,7 @@ class TierPathState(enum.IntEnum): READY = 3 FAILED = 4 DELETED = 5 + DELETE_IN_PROCESS = 6 class Asset(Base): 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 1d63c5b78..179bc9865 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client.py @@ -18,8 +18,13 @@ import asyncio import dataclasses import datetime +import email.parser 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 @@ -115,6 +120,11 @@ def __init__( self._credentials = None self._async_client = None + @property + def zone(self) -> str | None: + """Returns the GCP zone/location of the client.""" + return self.location + @property def async_client(self) -> httpx.AsyncClient: if self._async_client is None: @@ -178,6 +188,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 +336,153 @@ 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 using GCS JSON Batch API.""" + bucket, prefix = _parse_gcs_path(path) + token, _ = await self._get_token_and_project() + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + } + + items = await self._list_objects(bucket, prefix, headers) + if not items: + return + + await self._delete_objects_batch(bucket, items, token) + + async def _list_objects( + self, bucket: str, prefix: str, headers: dict[str, str] + ) -> list[dict[str, Any]]: + """Lists GCS objects under the given prefix.""" + list_url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o" + params = {"prefix": prefix, "projection": "noAcl"} + all_items = [] + page_token = None + + while True: + if page_token: + params["pageToken"] = page_token + response = await self.async_client.get( + list_url, params=params, headers=headers + ) + 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", []) + normalized_prefix = prefix.rstrip("/") + filtered_items = [ + item + for item in items + if item["name"] == prefix + or item["name"].startswith(normalized_prefix + "/") + ] + all_items.extend(filtered_items) + page_token = data.get("nextPageToken") + if not page_token: + break + + return all_items + + async def _delete_objects_batch( + self, bucket: str, items: list[dict[str, Any]], token: str + ) -> None: + """Chunks and executes GCS batch delete requests.""" + chunk_size = 100 + chunks = [ + items[i : i + chunk_size] for i in range(0, len(items), chunk_size) + ] + boundary = "===============CTS_BATCH_DELETE_BOUNDARY====" + + for chunk in chunks: + body = self._serialize_batch_request(bucket, chunk, boundary) + response_content, content_type = await self._send_batch_request( + body, token, boundary + ) + self._deserialize_batch_response(response_content, content_type) + + def _serialize_batch_request( + self, bucket: str, chunk: list[dict[str, Any]], boundary: str + ) -> bytes: + """Serializes a chunk of delete requests into multipart/mixed body bytes.""" + body_parts = [] + for item in chunk: + name = item["name"] + encoded_name = urllib.parse.quote(name, safe="") + body_parts.append( + f"--{boundary}\r\n" + "Content-Type: application/http\r\n" + "\r\n" + f"DELETE /storage/v1/b/{bucket}/o/{encoded_name} HTTP/1.1\r\n" + "\r\n" + ) + body_parts.append(f"--{boundary}--\r\n") + return "".join(body_parts).encode("utf-8") + + async def _send_batch_request( + self, body: bytes, token: str, boundary: str + ) -> tuple[bytes, str]: + """Sends the serialized batch POST request to GCS.""" + batch_url = "https://storage.googleapis.com/batch/storage/v1" + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": f"multipart/mixed; boundary={boundary}", + } + resp = await self.async_client.post( + batch_url, headers=headers, content=body + ) + if resp.status_code != 200: + raise RuntimeError( + f"GCS Batch request failed: {resp.status_code} - {resp.text}" + ) + return resp.content, resp.headers.get("Content-Type", "") + + def _deserialize_batch_response( + self, content: bytes, content_type: str + ) -> None: + """Parses GCS multipart response and validates nested HTTP status codes.""" + msg_bytes = ( + f"Content-Type: {content_type}\r\n\r\n".encode("utf-8") + content + ) + msg = email.parser.BytesParser().parsebytes(msg_bytes) + if not msg.is_multipart(): + raise RuntimeError("GCS batch response is not multipart/mixed") + + payload = msg.get_payload() + if not isinstance(payload, list): + raise RuntimeError( + "Expected GCS batch response payload to be a list of parts" + ) + + for part in payload: + if isinstance(part, str): + continue + part_payload = part.get_payload(decode=True) + if not isinstance(part_payload, bytes): + continue + part_text = part_payload.decode("utf-8") + status_line = part_text.split("\r\n")[0] + parts = status_line.split(" ") + if len(parts) >= 2: + try: + status_code = int(parts[1]) + if status_code not in (200, 204, 404): + raise RuntimeError( + "Batch object deletion failed with nested status:" + f" {status_line}" + ) + except ValueError as exc: + raise RuntimeError( + f"Failed to parse nested status code from line: {status_line}" + ) from exc + else: + raise RuntimeError( + f"Invalid status line in nested response: {status_line}" + ) + class GcpLustreBaseClient(GCPStorageClient): """Base client interface to interact with GCP Managed Lustre API via REST.""" @@ -338,10 +499,10 @@ def __init__( super().__init__( project=project, + location=zone, instance=instance, service_account=service_account, ) - self.zone = zone async def poll_operation( self, @@ -398,8 +559,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 +592,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 +781,49 @@ 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.islink(path): + os.remove(path) + elif os.path.isdir(path): + shutil.rmtree(path) + elif os.path.exists(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/gcp_storage_client_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client_test.py index 27e2ad964..365be088c 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client_test.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/gcp_storage_client_test.py @@ -226,6 +226,111 @@ async def test_lustre_poll_operation_in_progress(self): ) self.assertEqual(result.detail_info, {"percent_complete": 42}) + async def test_gcs_delete_path_batch_success(self): + client = gcp_storage_client.GcsToGcsClient(project="test-project") + + # 1. Mock list response (150 objects, triggers 2 chunks: 100 and 50) + mock_list_resp = mock.MagicMock(spec=httpx.Response) + mock_list_resp.status_code = 200 + mock_list_resp.json.return_value = { + "items": [{"name": f"test/dir/file_{i}.txt"} for i in range(150)] + } + self.mock_client.get.return_value = mock_list_resp + + # 2. Mock batch response content + boundary = "dummy_response_boundary" + + # Construct a multipart response with 100 parts for the first chunk, and + # 50 parts for the second chunk + def make_mock_batch_resp(count: int) -> bytes: + parts = [] + for _ in range(count): + parts.append( + f"--{boundary}\r\n" + "Content-Type: application/http\r\n" + "\r\n" + "HTTP/1.1 204 No Content\r\n" + ) + parts.append(f"--{boundary}--\r\n") + return "".join(parts).encode("utf-8") + + mock_batch_resp_1 = mock.MagicMock(spec=httpx.Response) + mock_batch_resp_1.status_code = 200 + mock_batch_resp_1.headers = { + "Content-Type": f"multipart/mixed; boundary={boundary}" + } + mock_batch_resp_1.content = make_mock_batch_resp(100) + + mock_batch_resp_2 = mock.MagicMock(spec=httpx.Response) + mock_batch_resp_2.status_code = 200 + mock_batch_resp_2.headers = { + "Content-Type": f"multipart/mixed; boundary={boundary}" + } + mock_batch_resp_2.content = make_mock_batch_resp(50) + + self.mock_client.post.side_effect = [mock_batch_resp_1, mock_batch_resp_2] + + # Run delete + await client.delete_path("gs://test-bucket/test/dir") + + # Assertions + self.mock_client.get.assert_called_once() + self.assertEqual(self.mock_client.post.call_count, 2) + + # Check request payload structure of the first post + first_call_args = self.mock_client.post.call_args_list[0] + called_content = first_call_args.kwargs["content"].decode("utf-8") + self.assertIn( + "DELETE /storage/v1/b/test-bucket/o/test%2Fdir%2Ffile_0.txt", + called_content, + ) + self.assertIn( + "DELETE /storage/v1/b/test-bucket/o/test%2Fdir%2Ffile_99.txt", + called_content, + ) + self.assertNotIn( + "DELETE /storage/v1/b/test-bucket/o/test%2Fdir%2Ffile_100.txt", + called_content, + ) + + async def test_gcs_delete_path_batch_nested_failure(self): + client = gcp_storage_client.GcsToGcsClient(project="test-project") + + # Mock list response + mock_list_resp = mock.MagicMock(spec=httpx.Response) + mock_list_resp.status_code = 200 + mock_list_resp.json.return_value = { + "items": [{"name": "test/dir/file_1.txt"}] + } + self.mock_client.get.return_value = mock_list_resp + + # Mock batch response with a nested 403 Forbidden status + boundary = "dummy_response_boundary" + response_content = ( + f"--{boundary}\r\n" + "Content-Type: application/http\r\n" + "\r\n" + "HTTP/1.1 403 Forbidden\r\n" + "\r\n" + "Permission denied\r\n" + f"--{boundary}--\r\n" + ).encode("utf-8") + + mock_batch_resp = mock.MagicMock(spec=httpx.Response) + mock_batch_resp.status_code = 200 + mock_batch_resp.headers = { + "Content-Type": f"multipart/mixed; boundary={boundary}" + } + mock_batch_resp.content = response_content + self.mock_client.post.return_value = mock_batch_resp + + with self.assertRaisesRegex( + RuntimeError, + "Batch object deletion failed with nested status: HTTP/1.1 403" + " Forbidden", + ): + await client.delete_path("gs://test-bucket/test/dir") + @mock.patch("google.auth.impersonated_credentials.Credentials") async def test_service_account_token_impersonation( self, mock_impersonated_creds_class diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py index 1091bc8be..e594c2e6c 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py @@ -69,6 +69,8 @@ def __init__( self._shutdown_event = asyncio.Event() self._backends = None self._backends_to_try = [] + self._delete_queue = asyncio.Queue() + self._delete_worker_task = None async def start(self): """Starts the background worker loops.""" @@ -91,11 +93,15 @@ async def start(self): self._tasks.append(asyncio.create_task(self._run_acquisition_loop())) self._tasks.append(asyncio.create_task(self._run_polling_loop())) + self._delete_worker_task = asyncio.create_task(self._run_delete_worker()) + self._tasks.append(self._delete_worker_task) async def stop(self): """Stops the background worker loops gracefully.""" logging.info("Stopping TieringServiceWorker...") self._shutdown_event.set() + if self._delete_worker_task: + self._delete_worker_task.cancel() if self._tasks: # Wait for tasks to finish await asyncio.gather(*self._tasks, return_exceptions=True) @@ -155,6 +161,23 @@ async def _poll_single_job( """Polls a single active copy job and updates its status.""" logging.info("Polling job %d", job.id) + if job.request_type in ( + db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_INSTANCE, + db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS, + ): + async with self._session_maker() as session: # pyrefly: ignore[not-callable] + async with session.begin(): + merged_job = await session.get( + db_schema.AssetJob, job.id, with_for_update=True + ) + if merged_job and ( + merged_job.status == db_schema.JobStatus.JOB_STATUS_PROCESSING + and merged_job.worker_host == self._hostname + and merged_job.worker_pid == self._pid + ): + await self._extend_lease(merged_job, now) + return + status_dict = job.transfer_status if not status_dict or "request_id" not in status_dict: logging.warning("Job %d in PROCESSING but has no request_id", job.id) @@ -311,10 +334,17 @@ 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, + ): + # Synchronously transition paths to DELETE_IN_PROCESS to block reads + # instantly before it waits in the delete queue. + await self._transition_to_deleting_in_db(job) + await self._delete_queue.put(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 +458,17 @@ 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 == prefix: + return "/" + elif path.startswith(prefix.rstrip("/") + "/"): + return path[len(prefix.rstrip("/")) :] + 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 +496,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,7 +513,7 @@ 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 _get_target_tier_path( @@ -542,29 +585,30 @@ 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] - - logging.info( - "Completed job %d, target TierPath %s marked ready", - job.id, - target_tp.path if target_tp else "None", - ) + target_tp = None + if job.request_type == db_schema.RequestType.REQUEST_TYPE_COPY: + # 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] + logging.info( + "Completed job %d, target TierPath %s marked ready", + job.id, + target_tp.path if target_tp else "None", + ) async def _extend_lease( self, @@ -574,6 +618,175 @@ async def _extend_lease( """Extends the lease of the job (heartbeat).""" job.expiration_at = now + self._lease_duration + async def _transition_to_deleting_in_db(self, job: db_schema.AssetJob): + """Transitions targeted TierPaths to DELETE_IN_PROCESS state immediately.""" + async with self._session_maker() as session: # pyrefly: ignore[not-callable] + async with session.begin(): + if job.target_tier_path_id: + stmt = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == job.target_tier_path_id) + .with_for_update() + ) + res = await session.execute(stmt) + tp = res.scalars().first() + if tp: + await assets.begin_delete_tier_path(session, tp) + else: + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == job.asset_uuid) + .options(joinedload(db_schema.Asset.tier_paths)) + .with_for_update() + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + if db_asset: + await assets.begin_delete_asset(session, db_asset) + + async def _run_delete_worker(self): + """Processes enqueued delete jobs sequentially.""" + while not self._shutdown_event.is_set(): + try: + job = await self._delete_queue.get() + try: + await self._execute_delete_job(job) + finally: + self._delete_queue.task_done() + except asyncio.CancelledError: + break + except Exception: # pylint: disable=broad-except + logging.exception("Error in background delete worker loop.") + await asyncio.sleep(5) + + async def _execute_delete_job(self, job: db_schema.AssetJob): + """Executes path deletion in the background and updates DB status.""" + logging.info("Starting background delete execution for job %d", job.id) + now = datetime.datetime.now(datetime.timezone.utc) + + # 1. Fetch paths to delete (already in DELETE_IN_PROCESS state) + target_tp_id = job.target_tier_path_id + async with self._session_maker() as session: # pyrefly: ignore[not-callable] + if target_tp_id: + stmt = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == target_tp_id) + .options(joinedload(db_schema.TierPath.storage_backend)) + ) + result = await session.execute(stmt) + tier_paths = [result.scalars().first()] + else: + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == job.asset_uuid) + .options( + joinedload(db_schema.Asset.tier_paths).joinedload( + db_schema.TierPath.storage_backend + ) + ) + ) + result = await session.execute(stmt) + db_asset = result.scalars().first() + tier_paths = db_asset.tier_paths if db_asset else [] + + if not tier_paths or any(tp is None for tp in tier_paths): + logging.warning("No tier paths found for delete job %d", job.id) + await self._fail_job_by_id(job.id, "No paths to delete") + return + + # 2. Perform deletions asynchronously + try: + for tp in tier_paths: + client = gcp_storage_client.determine_delete_client( + tp, + project=self._config.gcp_project or None, + service_account=self._config.service_account or None, + ) + try: + logging.info( + "Deleting path %s using client %s", + tp.path, + type(client).__name__, + ) + await client.delete_path(tp.path) + finally: + await client.close() + + # 3. Transition status in database to completed + async with self._session_maker() as session: # pyrefly: ignore[not-callable] + async with session.begin(): + merged_job = await session.get( + db_schema.AssetJob, job.id, with_for_update=True + ) + if merged_job and ( + merged_job.status == db_schema.JobStatus.JOB_STATUS_PROCESSING + and merged_job.worker_host == self._hostname + and merged_job.worker_pid == self._pid + ): + if ( + merged_job.request_type + == db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS + ): + stmt_asset = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == job.asset_uuid) + .options(joinedload(db_schema.Asset.tier_paths)) + .with_for_update() + ) + res_asset = await session.execute(stmt_asset) + db_asset = res_asset.scalars().first() + logging.info( + "DELETE_FROM_ALL_TIERS: db_asset found: %s", + str(db_asset is not None), + ) + if db_asset: + logging.info( + "DELETE_FROM_ALL_TIERS: db_asset %s state before" + " complete: %s", + db_asset.asset_uuid, + db_asset.state, + ) + await assets.complete_delete_asset(session, db_asset) + logging.info( + "DELETE_FROM_ALL_TIERS: db_asset %s state after" + " complete: %s", + db_asset.asset_uuid, + db_asset.state, + ) + else: + for tp in tier_paths: + stmt_tp = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == tp.id) + .with_for_update() + ) + res_tp = await session.execute(stmt_tp) + db_tp = res_tp.scalars().first() + if db_tp: + logging.info( + "DELETE_FROM_INSTANCE: db_tp %d state before" + " complete: %s", + db_tp.id, + db_tp.state, + ) + await assets.complete_delete_tier_path(session, db_tp) + logging.info( + "DELETE_FROM_INSTANCE: db_tp %d state after complete: %s", + db_tp.id, + db_tp.state, + ) + + await self._complete_job(session, merged_job, now) + logging.info( + "Successfully completed delete job %d in background", job.id + ) + + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Failed to execute background delete for job %d", job.id + ) + await self._fail_job_by_id(job.id, str(e)) + async def run_tiering_service_worker_loop( session_maker: sessionmaker | None, diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker_test.py index 075b40759..e52df3ede 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker_test.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker_test.py @@ -22,6 +22,7 @@ from absl.testing import absltest import aiosqlite # pylint: disable=unused-import import greenlet # pylint: disable=unused-import +from orbax.checkpoint.experimental.tiering_service import assets from orbax.checkpoint.experimental.tiering_service import db_lib from orbax.checkpoint.experimental.tiering_service import db_schema from orbax.checkpoint.experimental.tiering_service import gcp_storage_client @@ -30,6 +31,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool class DummyGcpParallelstoreClient(gcp_storage_client.GCPStorageClient): @@ -85,11 +87,35 @@ async def poll_operation( }, ) + async def delete_path(self, path: str) -> None: + """Mock delete_path.""" + logging.info("Dummy deleted path: %s", path) + class TieringServiceWorkerTest( absltest.TestCase, unittest.IsolatedAsyncioTestCase ): + async def _wait_for_condition( + self, check_fn, max_attempts=50, sleep_time=0.1 + ) -> None: + """Helper to poll a condition with database lock retries.""" + for i in range(max_attempts): + try: + if await check_fn(): + return + except Exception as e: # pylint: disable=broad-exception-caught + if "database table is locked" in str(e): + logging.info( + "Database locked during poll, retrying... (attempt %d)", i + ) + else: + raise + await asyncio.sleep(sleep_time) + # Final check without exception catching + if not await check_fn(): + self.fail("Condition not met after polling timeout") + async def asyncSetUp(self): await super().asyncSetUp() storage_backends_config = [ @@ -122,7 +148,7 @@ async def asyncSetUp(self): f"sqlite+aiosqlite:///file:{db_name}?mode=memory&cache=shared&uri=true" ) - self.engine = db_lib.get_async_engine(self.config) + self.engine = db_lib.get_async_engine(self.config, poolclass=NullPool) # Open and keep a connection alive to prevent SQLite from deleting the # in-memory shared database when engines are disposed/recreated. self._keep_alive_conn = await self.engine.connect() @@ -137,6 +163,13 @@ async def asyncSetUp(self): ) self.determine_client_mock.return_value = self.gcp_client + self.determine_delete_client_mock = self.enter_context( + mock.patch.object( + gcp_storage_client, "determine_delete_client", autospec=True + ) + ) + self.determine_delete_client_mock.return_value = self.gcp_client + def mock_get_client(status_dict, project=None, service_account=None): del project, service_account # Unused client_type = status_dict.get("client_type") @@ -185,6 +218,7 @@ async def _create_asset_and_job( target_backend_id, source_backend_id, job_status=db_schema.JobStatus.JOB_STATUS_QUEUED, + request_type=db_schema.RequestType.REQUEST_TYPE_COPY, ): """Helper to create an asset, its tier paths, and a prefetch job.""" del self @@ -214,11 +248,16 @@ async def _create_asset_and_job( session.add(target_tp) await session.flush() + if request_type == db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS: + target_tp_id = None + else: + target_tp_id = target_tp.id + job = db_schema.AssetJob( asset_uuid=asset_uuid, - request_type=db_schema.RequestType.REQUEST_TYPE_COPY, + request_type=request_type, status=job_status, - target_tier_path_id=target_tp.id, + target_tier_path_id=target_tp_id, ) session.add(job) await session.commit() @@ -243,14 +282,13 @@ async def test_job_acquisition_success(self): # Start worker to process the job await self.worker.start() - # Wait for job to complete with timeout - for _ in range(10): - await asyncio.sleep(1.0) + async def _check(): async with self.session_maker() as session: result = await session.execute(select(db_schema.AssetJob)) job = result.scalars().first() - if job.status == db_schema.JobStatus.JOB_STATUS_COMPLETED: - break + return job.status == db_schema.JobStatus.JOB_STATUS_COMPLETED + + await self._wait_for_condition(_check, max_attempts=10, sleep_time=1.0) await self.worker.stop() @@ -308,17 +346,16 @@ async def test_concurrency_limit_respected(self): ), ): await self.worker.start() - # Wait for at least one job to transition to PROCESSING - for _ in range(50): + async def _check(): async with self.session_maker() as session: result = await session.execute(select(db_schema.AssetJob)) jobs_list = result.scalars().all() - if any( + return any( j.status == db_schema.JobStatus.JOB_STATUS_PROCESSING for j in jobs_list - ): - break - await asyncio.sleep(0.1) + ) + + await self._wait_for_condition(_check, max_attempts=50, sleep_time=0.1) await self.worker.stop() async with self.session_maker() as session: @@ -369,17 +406,16 @@ async def test_different_backends_concurrency(self): ), ): await self.worker.start() - # Wait for both jobs to transition to PROCESSING - for _ in range(50): + async def _check(): async with self.session_maker() as session: result = await session.execute(select(db_schema.AssetJob)) jobs_list = result.scalars().all() - if all( + return all( j.status == db_schema.JobStatus.JOB_STATUS_PROCESSING for j in jobs_list - ): - break - await asyncio.sleep(0.1) + ) + + await self._wait_for_condition(_check, max_attempts=50, sleep_time=0.1) await self.worker.stop() async with self.session_maker() as session: @@ -417,14 +453,13 @@ async def test_job_trigger_failure(self): side_effect=RuntimeError("Mocked trigger failure"), ): await self.worker.start() - # Wait for job to fail - for _ in range(5): - await asyncio.sleep(0.5) + async def _check(): async with self.session_maker() as session: result = await session.execute(select(db_schema.AssetJob)) job = result.scalars().first() - if job.status == db_schema.JobStatus.JOB_STATUS_FAILED: - break + return job.status == db_schema.JobStatus.JOB_STATUS_FAILED + + await self._wait_for_condition(_check, max_attempts=5, sleep_time=0.5) await self.worker.stop() async with self.session_maker() as session: @@ -473,14 +508,13 @@ async def test_job_failure_clean_up_target_tier_path(self): ), ): await self.worker.start() - # Wait for job to fail - for _ in range(10): - await asyncio.sleep(1) + async def _check(): async with self.session_maker() as session: result = await session.execute(select(db_schema.AssetJob)) job = result.scalars().first() - if job.status == db_schema.JobStatus.JOB_STATUS_FAILED: - break + return job.status == db_schema.JobStatus.JOB_STATUS_FAILED + + await self._wait_for_condition(_check, max_attempts=10, sleep_time=1) await self.worker.stop() async with self.session_maker() as session: @@ -548,14 +582,13 @@ async def test_crash_recovery_on_lease_expiration(self): # Start worker await self.worker.start() - # Wait for job to complete with timeout - for _ in range(10): - await asyncio.sleep(1) + async def _check(): async with self.session_maker() as session: result = await session.execute(select(db_schema.AssetJob)) recovered_job = result.scalars().first() - if recovered_job.status == db_schema.JobStatus.JOB_STATUS_COMPLETED: - break + return recovered_job.status == db_schema.JobStatus.JOB_STATUS_COMPLETED + + await self._wait_for_condition(_check, max_attempts=10, sleep_time=1) await self.worker.stop() @@ -705,6 +738,43 @@ async def test_update_job_status_after_poll_extends_lease(self): expected_expiration.replace(tzinfo=None), ) + async def test_poll_delete_job_extends_lease(self): + async with self.session_maker() as session: + result = await session.execute(select(db_schema.StorageBackend)) + backends = result.scalars().all() + lustre_a = next(b for b in backends if b.zone == "us-central1-a") + gcs = next(b for b in backends if b.region == "us-central1") + + job, _ = await self._create_asset_and_job( + session, + asset_uuid="asset-delete-lease", + path="path/delete/lease", + target_backend_id=lustre_a.id, + source_backend_id=gcs.id, + job_status=db_schema.JobStatus.JOB_STATUS_PROCESSING, + request_type=db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS, + ) + job.worker_host = self.worker._hostname + job.worker_pid = self.worker._pid + initial_expiration = datetime.datetime.now( + datetime.timezone.utc + ) - datetime.timedelta(seconds=10) + job.expiration_at = initial_expiration + await session.commit() + + now = datetime.datetime.now(datetime.timezone.utc) + # Trigger polling manually + await self.worker._poll_single_job(job, now) + + # Verify expiration_at is now updated to now + lease_duration + async with self.session_maker() as session: + db_job = await session.get(db_schema.AssetJob, job.id) + expected_expiration = now + self.worker._lease_duration + self.assertEqual( + db_job.expiration_at.replace(tzinfo=None), + expected_expiration.replace(tzinfo=None), + ) + async def test_custom_transfer_status_details_preserved(self): async with self.session_maker() as session: result = await session.execute(select(db_schema.StorageBackend)) @@ -735,20 +805,208 @@ async def test_custom_transfer_status_details_preserved(self): await self.worker.start() custom_metric_val = None - for _ in range(20): - await asyncio.sleep(0.5) + async def _check(): + nonlocal custom_metric_val async with self.session_maker() as session: db_job = await session.get(db_schema.AssetJob, job.id) - if ( - db_job.transfer_status - and db_job.transfer_status.get("custom_metric") + if db_job.transfer_status and db_job.transfer_status.get( + "custom_metric" ): custom_metric_val = db_job.transfer_status.get("custom_metric") - break + return True + return False + + await self._wait_for_condition(_check, max_attempts=20, sleep_time=0.5) await self.worker.stop() self.assertEqual(custom_metric_val, "custom_value") + async def test_delete_from_instance_job_lifecycle(self): + async with self.session_maker() as session: + result = await session.execute(select(db_schema.StorageBackend)) + backends = result.scalars().all() + gcs = next(b for b in backends if b.region == "us-central1") + + # Create Asset and ready TierPath + asset = db_schema.Asset( + asset_uuid="asset-del-inst-1", + path="path/del/inst/1", + user="test-user", + state=db_schema.AssetState.ASSET_STATE_STORED, + ) + session.add(asset) + tp = db_schema.TierPath( + asset_uuid="asset-del-inst-1", + storage_backend_id=gcs.id, + path="gs://my-bucket/path/del/inst/1", + ready_at=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + state=db_schema.TierPathState.READY, + ) + session.add(tp) + await session.commit() + + # Queue DELETE_FROM_INSTANCE job + db_job = db_schema.AssetJob( + asset_uuid=asset.asset_uuid, + target_tier_path_id=tp.id, + request_type=db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_INSTANCE, + status=db_schema.JobStatus.JOB_STATUS_QUEUED, + ) + session.add(db_job) + await session.commit() + + # Start worker - it should pick up the job, transition state to + # DELETE_IN_PROCESS, and complete the deletion. + await self.worker.start() + + # Verify that the tier path has transitioned to DELETE_IN_PROCESS + # almost immediately (or upon completion). + async def _check(): + async with self.session_maker() as session: + stmt = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == tp.id) + .execution_options(populate_existing=True) + ) + res = await session.execute(stmt) + db_tp = res.scalars().first() + return db_tp and db_tp.state == db_schema.TierPathState.DELETED + + await self._wait_for_condition(_check, max_attempts=20, sleep_time=0.2) + + await self.worker.stop() + + # Final DB assertion checks + async with self.session_maker() as session: + async with session.begin(): + stmt_tp = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == tp.id) + .execution_options(populate_existing=True) + ) + res_tp = await session.execute(stmt_tp) + db_tp = res_tp.scalars().first() + self.assertEqual(db_tp.state, db_schema.TierPathState.DELETED) + self.assertIsNone(db_tp.ready_at) + + stmt_job = ( + select(db_schema.AssetJob) + .where(db_schema.AssetJob.id == db_job.id) + .execution_options(populate_existing=True) + ) + res_job = await session.execute(stmt_job) + db_job_res = res_job.scalars().first() + self.assertEqual( + db_job_res.status, db_schema.JobStatus.JOB_STATUS_COMPLETED + ) + + async def test_delete_from_all_tiers_job_lifecycle(self): + async with self.session_maker() as session: + result = await session.execute(select(db_schema.StorageBackend)) + backends = result.scalars().all() + gcs = next(b for b in backends if b.region == "us-central1") + lustre = next(b for b in backends if b.zone == "us-central1-a") + + # Create Asset and 2 ready TierPaths (Lustre + GCS) + asset = db_schema.Asset( + asset_uuid="asset-del-all-1", + path="path/del/all/1", + user="test-user", + state=db_schema.AssetState.ASSET_STATE_STORED, + ) + session.add(asset) + tp_gcs = db_schema.TierPath( + asset_uuid="asset-del-all-1", + storage_backend_id=gcs.id, + path="gs://my-bucket/path/del/all/1", + ready_at=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + state=db_schema.TierPathState.READY, + ) + tp_lustre = db_schema.TierPath( + asset_uuid="asset-del-all-1", + storage_backend_id=lustre.id, + path="/mnt/lustre/path/del/all/1", + ready_at=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + state=db_schema.TierPathState.READY, + ) + session.add(tp_gcs) + session.add(tp_lustre) + await session.commit() + + # Queue DELETE_FROM_ALL_TIERS job + await assets.queue_delete_asset_job(session, asset) + + # Find queued job + stmt = select(db_schema.AssetJob).where( + db_schema.AssetJob.asset_uuid == asset.asset_uuid, + db_schema.AssetJob.request_type + == db_schema.RequestType.REQUEST_TYPE_DELETE_FROM_ALL_TIERS, + ) + res_job = await session.execute(stmt) + db_job = res_job.scalars().first() + + # Start worker + await self.worker.start() + + # Verify both paths deleted + async def _check(): + async with self.session_maker() as session: + stmt = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == tp_gcs.id) + .execution_options(populate_existing=True) + ) + res = await session.execute(stmt) + db_tp = res.scalars().first() + return db_tp and db_tp.state == db_schema.TierPathState.DELETED + + await self._wait_for_condition(_check, max_attempts=20, sleep_time=0.2) + + await self.worker.stop() + + # Assertions + async with self.session_maker() as session: + async with session.begin(): + stmt_tp_g = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == tp_gcs.id) + .execution_options(populate_existing=True) + ) + res_tp_g = await session.execute(stmt_tp_g) + db_tp_g = res_tp_g.scalars().first() + self.assertEqual(db_tp_g.state, db_schema.TierPathState.DELETED) + + stmt_tp_l = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == tp_lustre.id) + .execution_options(populate_existing=True) + ) + res_tp_l = await session.execute(stmt_tp_l) + db_tp_l = res_tp_l.scalars().first() + self.assertEqual(db_tp_l.state, db_schema.TierPathState.DELETED) + + stmt_asset = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == asset.asset_uuid) + .execution_options(populate_existing=True) + ) + res_asset = await session.execute(stmt_asset) + db_asset = res_asset.scalars().first() + self.assertEqual( + db_asset.state, db_schema.AssetState.ASSET_STATE_DELETED + ) + + stmt_job = ( + select(db_schema.AssetJob) + .where(db_schema.AssetJob.id == db_job.id) + .execution_options(populate_existing=True) + ) + res_job = await session.execute(stmt_job) + db_job_res = res_job.scalars().first() + self.assertEqual( + db_job_res.status, db_schema.JobStatus.JOB_STATUS_COMPLETED + ) + if __name__ == "__main__": absltest.main() 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 4948946ec..9adc1265b 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/proto/tiering_service.proto @@ -18,6 +18,16 @@ enum AssetState { ASSET_STATE_INCOMPLETE = 4; } +enum TierPathState { + TIER_PATH_STATE_UNSPECIFIED = 0; + TIER_PATH_STATE_PENDING = 1; + TIER_PATH_STATE_IN_PROGRESS = 2; + TIER_PATH_STATE_READY = 3; + TIER_PATH_STATE_FAILED = 4; + TIER_PATH_STATE_DELETED = 5; + TIER_PATH_STATE_DELETE_IN_PROCESS = 6; +} + enum BackendType { BACKEND_TYPE_UNSPECIFIED = 0; BACKEND_TYPE_LUSTRE = 1; @@ -59,6 +69,9 @@ message TierPath { // A unique identifier for this tier path. string tier_path_uuid = 6; + + // The state of the tier path. + TierPathState state = 7; } message Asset { diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py index 25b722803..230d89887 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py @@ -19,9 +19,11 @@ from concurrent import futures import contextlib import datetime +import logging as std_logging import os import pprint import sys +import uuid from absl import logging import fire @@ -133,7 +135,10 @@ async def Reserve( return tiering_service_pb2.ReserveResponse() # Calculate resolved path and check GCS permission. - storage_path = storage_backend.get_storage_path(backend, request.path) + tp_uuid = str(uuid.uuid4()) + storage_path = storage_backend.get_storage_path( + backend, request.path, tp_uuid + ) token = await auth.get_oauth_token(context) if not await auth.has_write_permission( token, backend=backend, path=storage_path @@ -151,7 +156,12 @@ async def Reserve( async with self._session_scope() as session: try: db_asset = await assets.create_or_fetch_asset( - session, request, backend, self._config + session, + request, + backend, + self._config, + tier_path_uuid=tp_uuid, + storage_path=storage_path, ) except ValueError as e: logging.exception("Failed to reserve asset for path: %s", request.path) @@ -324,6 +334,12 @@ async def Prefetch( return tiering_service_pb2.PrefetchResponse() for tp in db_asset.tier_paths: + if tp.state not in ( + db_schema.TierPathState.PENDING, + db_schema.TierPathState.IN_PROGRESS, + db_schema.TierPathState.READY, + ): + continue if tp.storage_backend_id == closest_backend.id: # TODO: b/503445463 - Extend the expiration of the existing TierPath # if needed. @@ -341,8 +357,9 @@ async def Prefetch( ) # No existing TierPath, we need to prefetch + tp_uuid = str(uuid.uuid4()) storage_path = storage_backend.get_storage_path( - closest_backend, db_asset.path + closest_backend, db_asset.path, tp_uuid ) token = await auth.get_oauth_token(context) @@ -361,6 +378,7 @@ async def Prefetch( db_asset, backend=closest_backend, storage_path=storage_path, + tier_path_uuid=tp_uuid, client_keep_alive_interval=datetime.timedelta( seconds=self._config.client_keep_alive_interval_seconds ), @@ -395,6 +413,12 @@ async def Prefetch( closest_tier_path_uuid = "" for tp in db_asset.tier_paths: + if tp.state not in ( + db_schema.TierPathState.PENDING, + db_schema.TierPathState.IN_PROGRESS, + db_schema.TierPathState.READY, + ): + continue if tp.storage_backend_id == closest_backend.id: closest_tier_path_uuid = tp.tier_path_uuid break @@ -433,6 +457,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", @@ -605,6 +634,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 f6b64f2cc..5e083c6c0 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( @@ -485,14 +489,29 @@ async def test_prefetch_success_rpc_response(self): self.assertEqual(prefetch_res.asset.uuid, asset_uuid) paths = [tp.path for tp in prefetch_res.asset.tier_paths] - self.assertCountEqual( - paths, - ["/mnt/lustre-a/test/path", "/mnt/lustre-b/test/path"], + self.assertLen(paths, 3) + lustre_a_paths = [ + p for p in paths if p.startswith("/mnt/lustre-a/test/path/") + ] + lustre_b_paths = [ + p for p in paths if p.startswith("/mnt/lustre-b/test/path/") + ] + gcs_paths = [ + p for p in paths if p.startswith("gs://my-bucket/test/path/") + ] + self.assertLen(lustre_a_paths, 1) + self.assertLen(lustre_b_paths, 1) + self.assertLen(gcs_paths, 1) + + # 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.startswith("/mnt/lustre-b/test/path/") ) - # In the setup, index 1 corresponds to zone B's Lustre backend target self.assertEqual( prefetch_res.closest_tier_path_uuid, - prefetch_res.asset.tier_paths[1].tier_path_uuid, + expected_tp.tier_path_uuid, ) self.assertTrue(prefetch_res.closest_tier_path_uuid) @@ -505,25 +524,33 @@ 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, + res_backend = await session.execute( + select(db_schema.StorageBackend).filter_by(zone="us-central1-b") ) - 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 + backend_b = res_backend.scalars().first() + self.assertIsNotNone(backend_b) stmt_tp = select(db_schema.TierPath).filter_by( - asset_uuid=asset_uuid, path="/mnt/lustre-b/test/path" + asset_uuid=asset_uuid, storage_backend_id=backend_b.id ) 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() @@ -551,7 +578,64 @@ 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_while_delete_in_process(self): + servicer, asset_uuid = await self._setup_servicer_and_asset() + + # 1. Prefetch from B (creates first TierPath and COPY job). + res1 = await servicer.Prefetch( + tiering_service_pb2.PrefetchRequest( + uuid=asset_uuid, zone="us-central1-b" + ), + self.context, + ) + first_tp_uuid = res1.closest_tier_path_uuid + + # Verify first TierPath has the UUID in its path! + async with servicer._session_scope() as session: + tp1 = await session.execute( + select(db_schema.TierPath).filter_by(tier_path_uuid=first_tp_uuid) + ) + tp1_obj = tp1.scalars().first() + self.assertIsNotNone(tp1_obj) + self.assertTrue(tp1_obj.path.endswith(first_tp_uuid)) + first_path = tp1_obj.path + + # 2. Transition first TierPath to DELETE_IN_PROCESS to simulate eviction. + tp1_obj.state = db_schema.TierPathState.DELETE_IN_PROCESS + await session.commit() + + # 3. Call Prefetch again from B. + # Since the first one is DELETE_IN_PROCESS, it should be ignored, + # and a NEW TierPath (and COPY job) should be created. + res2 = await servicer.Prefetch( + tiering_service_pb2.PrefetchRequest( + uuid=asset_uuid, zone="us-central1-b" + ), + self.context, + ) + second_tp_uuid = res2.closest_tier_path_uuid + self.assertNotEqual(first_tp_uuid, second_tp_uuid) + + # Verify second TierPath has a different UUID and path! + async with servicer._session_scope() as session: + tp2 = await session.execute( + select(db_schema.TierPath).filter_by(tier_path_uuid=second_tp_uuid) + ) + tp2_obj = tp2.scalars().first() + self.assertIsNotNone(tp2_obj) + self.assertTrue(tp2_obj.path.endswith(second_tp_uuid)) + self.assertNotEqual(first_path, tp2_obj.path) + + # Verify two COPY jobs now exist + result = await session.execute( + select(db_schema.AssetJob).filter_by( + asset_uuid=asset_uuid, + request_type=db_schema.RequestType.REQUEST_TYPE_COPY, + ) + ) + self.assertLen(result.scalars().all(), 3) async def test_prefetch_already_ready(self): # If we prefetch to the same zone where it was reserved and finalized, @@ -569,18 +653,30 @@ 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).where( + db_schema.TierPath.asset_uuid == asset_uuid, + db_schema.TierPath.path.like("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() @@ -677,7 +773,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 c3fd2d562..ab88aa56e 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py @@ -97,20 +97,30 @@ def locate_closest_backend( def get_storage_path( backend: db_schema.StorageBackend, relative_path: str, + tier_path_uuid: str | None = None, ) -> str: """Builds the absolute storage path for the given backend and relative path. Combines the backend's prefix with the relative path, ensuring proper - formatting. + formatting. If a tier_path_uuid is provided, append it to ensure uniqueness. Args: backend: The StorageBackend target. relative_path: The relative path of the asset. + tier_path_uuid: Optional unique identifier for the tier path. Returns: The absolute storage path. """ - return f"{backend.prefix.rstrip('/')}/{relative_path.lstrip('/')}" + if relative_path.startswith(backend.prefix): + base_path = relative_path + else: + base_path = f"{backend.prefix.rstrip('/')}/{relative_path.lstrip('/')}" + + if tier_path_uuid is not None: + if not base_path.endswith(f"/{tier_path_uuid}"): + return f"{base_path.rstrip('/')}/{tier_path_uuid}" + return base_path def get_backend_name(backend: db_schema.StorageBackend) -> str: diff --git a/checkpoint/orbax/checkpoint/options.py b/checkpoint/orbax/checkpoint/options.py index a3f636319..a62d9c208 100644 --- a/checkpoint/orbax/checkpoint/options.py +++ b/checkpoint/orbax/checkpoint/options.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Configuration options for APIs like CheckpointManager and Checkpointer.""" +"""Orbax Checkpoint Options.""" + +from __future__ import annotations import dataclasses +import typing from typing import Callable, Optional, Set from orbax.checkpoint._src.multihost import multihost