diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/assets.py index faf792561..70dc17f1d 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 @@ -40,6 +41,10 @@ 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 +110,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 +121,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 +281,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 +297,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 +317,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) @@ -383,6 +399,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,6 +420,7 @@ 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. await session.commit() @@ -403,6 +428,141 @@ async def finalize_asset( return db_asset +def find_matching_backend( + backends: Sequence[db_schema.StorageBackend], + source_backend: db_schema.StorageBackend, +) -> db_schema.StorageBackend | None: + """Finds a storage backend in `backends` that matches `source_backend`. + + The matching logic prioritizes: + 1. Exact zone match (if both source and target have zone). + 2. Region match (if they share the same region, or region extracted from + zone). + 3. Multi-regions match (if target has multi_regions containing source's + region). + + Args: + backends: A list of candidate StorageBackend objects. + source_backend: The source StorageBackend object. + + Returns: + The matching StorageBackend object, or None if no match is found. + """ + + def _get_region_from_zone(zone: str) -> str: + return zone.rsplit("-", 1)[0] + + # Get source region + source_region = None + if source_backend.zone: + source_region = _get_region_from_zone(source_backend.zone) + elif source_backend.region: + source_region = source_backend.region + + # 1. Match by zone first (if both have zone) + if source_backend.zone: + for b in backends: + if b.zone and b.zone == source_backend.zone: + return b + + # 2. Match by region + if source_region: + # Try exact region match first + for b in backends: + if b.region and b.region == source_region: + return b + if b.zone and _get_region_from_zone(b.zone) == source_region: + return b + + # Try multi-regions match + for b in backends: + if b.multi_regions and source_region in b.multi_regions: + return b + + # 3. Handle source having multi_regions (if L0 is multi-regional) + if source_backend.multi_regions: + # Check if target has region that is in source's multi_regions + for b in backends: + if b.region and b.region in source_backend.multi_regions: + return b + if ( + b.zone + and _get_region_from_zone(b.zone) in source_backend.multi_regions + ): + return b + if b.multi_regions: + # Check overlap + overlap = set(source_backend.multi_regions) & set(b.multi_regions) + if overlap: + return b + + return None + + +async def trigger_l0_to_l1_copy( + session: AsyncSession, + db_asset: db_schema.Asset, +) -> None: + """Triggers copying the asset from Level 0 to Level 1 storage. + + We assume always preserve policy, so trigger L0-L1 immediately. + Generically calls them L0 and L1 as they are not always Lustre and GCS. + + Args: + session: The database session. + db_asset: The Asset model instance to copy. + """ + # Find the Level 0 backend from the asset's existing tier paths + l0_backend = None + for tp in db_asset.tier_paths: + if tp.storage_backend.level == 0: + l0_backend = tp.storage_backend + break + + if l0_backend is None: + raise ValueError(f"Asset {db_asset.asset_uuid} has no Level 0 backend.") + + # Look up Level 1 storage backend to queue the copy job + stmt = select(db_schema.StorageBackend).where( + db_schema.StorageBackend.level == 1 + ) + res = await session.execute(stmt) + level_1_backends = res.scalars().all() + + l1_backend = find_matching_backend(level_1_backends, l0_backend) + if level_1_backends and l1_backend is None: + raise ValueError( + f"No matching Level 1 backend found for Level 0 backend: {l0_backend}" + ) + + if l1_backend is not None: + l1_tp_uuid = str(uuid.uuid4()) + l1_path = storage_backend_lib.get_storage_path( + l1_backend, db_asset.path, l1_tp_uuid + ) + new_l1_tp = db_schema.TierPath( + storage_backend=l1_backend, + path=l1_path, + tier_path_uuid=l1_tp_uuid, + state=db_schema.TierPathState.PENDING, + ) + db_asset.tier_paths.append(new_l1_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_l1_tp, + ) + session.add(db_job) + logging.info( + "Finalize: Queued copy job to Level 1 for asset %s, target path: %s", + db_asset.asset_uuid, + l1_path, + ) + await session.commit() + + async def is_delete_pending(session: AsyncSession, *, asset_uuid: str) -> bool: """Checks if there is a pending delete job for the asset.""" stmt = ( @@ -453,6 +613,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 +629,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 +641,17 @@ 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() + + if db_asset is None: + return CreatePrefetchJobResult(asset=None, created=False) + # 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 +667,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 +725,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 +733,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 +824,97 @@ 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 | None: + """Transitions asset state to DELETED and all its tier paths to DELETE_IN_PROCESS.""" + 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 is None: + return None + + db_asset.state = db_schema.AssetState.ASSET_STATE_DELETED + db_asset.deleted_at = datetime.datetime.now(datetime.timezone.utc) + + for tp in db_asset.tier_paths: + tp.state = db_schema.TierPathState.DELETE_IN_PROCESS + tp.ready_at = None + tp.expires_at = None + return db_asset + + +async def complete_delete_asset( + session: AsyncSession, + db_asset: db_schema.Asset, +) -> db_schema.Asset | None: + """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, or None if it was concurrently deleted. + """ + 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 is None: + return None + + now = datetime.datetime.now(datetime.timezone.utc) + db_asset.state = db_schema.AssetState.ASSET_STATE_DELETED + db_asset.deleted_at = now + + for tp in db_asset.tier_paths: + tp.state = db_schema.TierPathState.DELETED + tp.ready_at = None + tp.expires_at = None + + 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..865ca73d0 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/assets_test.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/assets_test.py @@ -26,6 +26,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.future import select +from sqlalchemy.orm import joinedload from sqlalchemy.orm import sessionmaker from google.protobuf import timestamp_pb2 @@ -119,6 +120,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 +137,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 +259,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 +329,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 +364,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 +378,131 @@ 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_trigger_l0_to_l1_copy_success(self): + async with self.session_maker() as session: + # Setup L0 backend (Lustre) + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre", + zone="us-central1-a", + ) + # Setup L1 backend (GCS) + b1 = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://my-bucket", + region="us-central1", + ) + session.add_all([b0, b1]) + await session.commit() + + request = tiering_service_pb2.ReserveRequest( + path="test/path/copy_test", + user="test-user", + zone="us-central1-a", + ) + tp_uuid = "test-tp-uuid-b0" + storage_path = storage_backend_lib.get_storage_path( + b0, request.path, tp_uuid + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, + ) + # Finalize the asset first + finalized = await assets.finalize_asset(session, asset) + + # Now trigger copy + await assets.trigger_l0_to_l1_copy(session, finalized) + + # Verify L1 tier path and copy job were created + async with self.session_maker() as session2: + fetched = await assets.fetch_asset_by_uuid(session2, asset.asset_uuid) + self.assertLen(fetched, 1) + db_asset = fetched[0] + # Should have 2 tier paths now: b0 (Lustre) and b1 (GCS) + self.assertLen(db_asset.tier_paths, 2) + + l1_tp = next( + tp for tp in db_asset.tier_paths if tp.storage_backend_id == b1.id + ) + self.assertIsNotNone(l1_tp) + self.assertEqual(l1_tp.state, db_schema.TierPathState.PENDING) + self.assertTrue( + l1_tp.path.startswith("gs://my-bucket/test/path/copy_test/") + ) + + # Verify copy job + stmt = select(db_schema.AssetJob).filter_by( + asset_uuid=asset.asset_uuid, + request_type=db_schema.RequestType.REQUEST_TYPE_COPY, + status=db_schema.JobStatus.JOB_STATUS_QUEUED, + target_tier_path_id=l1_tp.id, + ) + result = await session2.execute(stmt) + job = result.scalars().first() + self.assertIsNotNone(job) + + async def test_trigger_l0_to_l1_copy_multiple_l1_backends_raises_error(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre", + zone="us-central1-a", + ) + b1_a = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-a", + region="us-west1", + ) + b1_b = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-b", + region="us-east1", + ) + session.add_all([b0, b1_a, b1_b]) + await session.commit() + + request = tiering_service_pb2.ReserveRequest( + path="test/path/copy_test_fail", + user="test-user", + zone="us-central1-a", + ) + tp_uuid = "test-tp-uuid-b0-fail" + storage_path = storage_backend_lib.get_storage_path( + b0, request.path, tp_uuid + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, + ) + finalized = await assets.finalize_asset(session, asset) + + with self.assertRaisesRegex( + ValueError, "No matching Level 1 backend found" + ): + await assets.trigger_l0_to_l1_copy(session, finalized) async def test_queries_filtering(self): async with self.session_maker() as session: @@ -370,8 +525,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 +544,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 +617,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 +628,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 +637,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 +655,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 +673,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 +683,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 +708,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 +781,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 +799,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 +819,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 +871,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 +909,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 +946,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 +980,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 +1021,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 +1036,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 +1068,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 +1108,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 +1193,788 @@ 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) + + 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) + + 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 + ) + + 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) + + 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) + + async def test_finalize_asset_success_with_multiple_assets(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + session.add(b0) + await session.commit() + + # Create Asset 1 + request1 = tiering_service_pb2.ReserveRequest( + path="test/path/asset1", + user="test-user", + zone="us-central1-a", + ) + asset1 = await assets.create_or_fetch_asset( + session, + request1, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-1", + storage_path="/mnt/lustre-a/asset1", + ) + + # Create Asset 2 + request2 = tiering_service_pb2.ReserveRequest( + path="test/path/asset2", + user="test-user", + zone="us-central1-a", + ) + asset2 = await assets.create_or_fetch_asset( + session, + request2, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-2", + storage_path="/mnt/lustre-a/asset2", + ) + + # Finalize Asset 2 + finalized = await assets.finalize_asset(session, asset2) + await session.commit() + + # Verify we finalized Asset 2 and NOT Asset 1 + self.assertEqual(finalized.asset_uuid, asset2.asset_uuid) + + async with self.session_maker() as session2: + fetched1 = await assets.fetch_asset_by_uuid(session2, asset1.asset_uuid) + self.assertEqual( + fetched1[0].state, + db_schema.AssetState.ASSET_STATE_ACTIVE_WRITE, + ) + + fetched2 = await assets.fetch_asset_by_uuid(session2, asset2.asset_uuid) + self.assertEqual( + fetched2[0].state, db_schema.AssetState.ASSET_STATE_STORED + ) + + async def test_create_prefetch_job_success_with_multiple_assets(self): + """Verifies prefetch job creation isolates TierPaths to the correct asset.""" + async with self.session_maker() as session: + # Setup L0 backend (Lustre) and L1 backend (GCS) + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + b1 = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://my-bucket", + region="us-central1", + ) + session.add_all([b0, b1]) + await session.commit() + + # Create and finalize Asset 1 + request1 = tiering_service_pb2.ReserveRequest( + path="test/path/asset1", + user="test-user", + zone="us-central1-a", + ) + asset1 = await assets.create_or_fetch_asset( + session, + request1, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-1", + storage_path="/mnt/lustre-a/asset1", + ) + await assets.finalize_asset(session, asset1) + + # Create and finalize Asset 2 (the one we will prefetch) + request2 = tiering_service_pb2.ReserveRequest( + path="test/path/asset2", + user="test-user", + zone="us-central1-a", + ) + asset2 = await assets.create_or_fetch_asset( + session, + request2, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-2", + storage_path="/mnt/lustre-a/asset2", + ) + await assets.finalize_asset(session, asset2) + # Trigger L0-L1 copy + await assets.trigger_l0_to_l1_copy(session, asset2) + await session.commit() + + # Simulate copy completed and L0 eviction + async with self.session_maker() as session: + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == asset2.asset_uuid) + .options(joinedload(db_schema.Asset.tier_paths)) + ) + res = await session.execute(stmt) + db_asset = res.scalars().first() + for tp in db_asset.tier_paths: + if tp.storage_backend_id == b0.id: + tp.state = db_schema.TierPathState.DELETED + elif tp.storage_backend_id == b1.id: + tp.state = db_schema.TierPathState.READY + tp.ready_at = datetime.datetime.now(datetime.timezone.utc) + await session.commit() + + # Now in a new session, prefetch Asset 2 + async with self.session_maker() as session2: + # Fetch L0 backend again for prefetch target + res_b0 = await session2.execute( + select(db_schema.StorageBackend).where( + db_schema.StorageBackend.id == b0.id + ) + ) + db_b0 = res_b0.scalars().first() + + # Fetch Asset 2 again to pass to prefetch + db_assets2 = await assets.fetch_asset_by_uuid(session2, asset2.asset_uuid) + db_asset2 = db_assets2[0] + + tp_uuid = "prefetch-tp-uuid-2" + storage_path = storage_backend_lib.get_storage_path( + db_b0, db_asset2.path, tp_uuid + ) + + # Trigger prefetch on Asset 2 + result = await assets.create_prefetch_job( + session2, + db_asset2, + backend=db_b0, + storage_path=storage_path, + tier_path_uuid=tp_uuid, + client_keep_alive_interval=datetime.timedelta(seconds=600), + ) + self.assertTrue(result.created) + asset_res = result.asset + assert asset_res is not None + self.assertEqual(asset_res.asset_uuid, asset2.asset_uuid) + await session2.commit() + + # Verify that Asset 2 got the new tier path, but Asset 1 did not + async with self.session_maker() as session3: + fetched1 = await assets.fetch_asset_by_uuid(session3, asset1.asset_uuid) + self.assertNotIn( + tp_uuid, [tp.tier_path_uuid for tp in fetched1[0].tier_paths] + ) + + fetched2 = await assets.fetch_asset_by_uuid(session3, asset2.asset_uuid) + # Asset 2 should have the prefetch path + self.assertIn( + tp_uuid, [tp.tier_path_uuid for tp in fetched2[0].tier_paths] + ) + + async def test_begin_delete_tier_path(self): + async with self.session_maker() as session: + # Setup asset with a tier path + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + session.add(b0) + await session.commit() + request = tiering_service_pb2.ReserveRequest( + path="test/path/delete_tp", + user="test-user", + zone="us-central1-a", + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-del", + storage_path="/mnt/lustre-a/delete_tp", + ) + tp = asset.tier_paths[0] + # Set it to ready first + tp.state = db_schema.TierPathState.READY + tp.ready_at = datetime.datetime.now(datetime.timezone.utc) + await session.commit() + + # Now begin delete + updated_tp = await assets.begin_delete_tier_path(session, tp) + await session.commit() + + self.assertEqual( + updated_tp.state, db_schema.TierPathState.DELETE_IN_PROCESS + ) + self.assertIsNone(updated_tp.ready_at) + self.assertIsNone(updated_tp.expires_at) + + # Verify persistence + async with self.session_maker() as session2: + stmt = select(db_schema.TierPath).where(db_schema.TierPath.id == tp.id) + res = await session2.execute(stmt) + db_tp = res.scalars().first() + self.assertEqual(db_tp.state, db_schema.TierPathState.DELETE_IN_PROCESS) + + async def test_complete_delete_tier_path(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + session.add(b0) + await session.commit() + request = tiering_service_pb2.ReserveRequest( + path="test/path/complete_del_tp", + user="test-user", + zone="us-central1-a", + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-comp-del", + storage_path="/mnt/lustre-a/complete_del_tp", + ) + tp = asset.tier_paths[0] + tp.state = db_schema.TierPathState.DELETE_IN_PROCESS + await session.commit() + + # Now complete delete + updated_tp = await assets.complete_delete_tier_path(session, tp) + await session.commit() + + self.assertEqual(updated_tp.state, db_schema.TierPathState.DELETED) + self.assertIsNone(updated_tp.ready_at) + self.assertIsNone(updated_tp.expires_at) + + # Verify persistence + async with self.session_maker() as session2: + stmt = select(db_schema.TierPath).where(db_schema.TierPath.id == tp.id) + res = await session2.execute(stmt) + db_tp = res.scalars().first() + self.assertEqual(db_tp.state, db_schema.TierPathState.DELETED) + + async def test_begin_delete_asset_success_with_multiple_assets(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + session.add(b0) + await session.commit() + + # Create Asset 1 + request1 = tiering_service_pb2.ReserveRequest( + path="test/path/asset1", + user="test-user", + zone="us-central1-a", + ) + asset1 = await assets.create_or_fetch_asset( + session, + request1, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-1", + storage_path="/mnt/lustre-a/asset1", + ) + await assets.finalize_asset(session, asset1) + + # Create Asset 2 (the one we will delete) + request2 = tiering_service_pb2.ReserveRequest( + path="test/path/asset2", + user="test-user", + zone="us-central1-a", + ) + asset2 = await assets.create_or_fetch_asset( + session, + request2, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-2", + storage_path="/mnt/lustre-a/asset2", + ) + await assets.finalize_asset(session, asset2) + await session.commit() + + # Now call begin_delete_asset on Asset 2 + async with self.session_maker() as session2: + db_assets2 = await assets.fetch_asset_by_uuid(session2, asset2.asset_uuid) + db_asset2 = db_assets2[0] + + deleted = await assets.begin_delete_asset(session2, db_asset2) + await session2.commit() + + self.assertIsNotNone(deleted) + self.assertEqual(deleted.asset_uuid, asset2.asset_uuid) + self.assertEqual(deleted.state, db_schema.AssetState.ASSET_STATE_DELETED) + + # Verify that Asset 1 is UNTOUCHED (still STORED) + async with self.session_maker() as session3: + fetched1 = await assets.fetch_asset_by_uuid(session3, asset1.asset_uuid) + self.assertEqual( + fetched1[0].state, db_schema.AssetState.ASSET_STATE_STORED + ) + # Also verify Asset 1's tier paths are still READY + for tp in fetched1[0].tier_paths: + self.assertEqual(tp.state, db_schema.TierPathState.READY) + + # Verify Asset 2 is DELETED and its tier paths are DELETE_IN_PROCESS + fetched2 = await assets.fetch_asset_by_uuid(session3, asset2.asset_uuid) + self.assertEqual( + fetched2[0].state, db_schema.AssetState.ASSET_STATE_DELETED + ) + for tp in fetched2[0].tier_paths: + self.assertEqual(tp.state, db_schema.TierPathState.DELETE_IN_PROCESS) + + async def test_complete_delete_asset_success_with_multiple_assets(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre-a", + zone="us-central1-a", + ) + session.add(b0) + await session.commit() + + # Create Asset 1 + request1 = tiering_service_pb2.ReserveRequest( + path="test/path/asset1", + user="test-user", + zone="us-central1-a", + ) + asset1 = await assets.create_or_fetch_asset( + session, + request1, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-1", + storage_path="/mnt/lustre-a/asset1", + ) + await assets.finalize_asset(session, asset1) + + # Create Asset 2 (the one we will delete) + request2 = tiering_service_pb2.ReserveRequest( + path="test/path/asset2", + user="test-user", + zone="us-central1-a", + ) + asset2 = await assets.create_or_fetch_asset( + session, + request2, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid="tp-uuid-2", + storage_path="/mnt/lustre-a/asset2", + ) + await assets.finalize_asset(session, asset2) + await session.commit() + + # Transition Asset 2 to DELETE_IN_PROCESS via begin_delete_asset + async with self.session_maker() as session2: + db_assets2 = await assets.fetch_asset_by_uuid(session2, asset2.asset_uuid) + db_asset2 = db_assets2[0] + await assets.begin_delete_asset(session2, db_asset2) + await session2.commit() + + # Now call complete_delete_asset on Asset 2 + async with self.session_maker() as session3: + db_assets2 = await assets.fetch_asset_by_uuid(session3, asset2.asset_uuid) + db_asset2 = db_assets2[0] + + completed = await assets.complete_delete_asset(session3, db_asset2) + await session3.commit() + + self.assertIsNotNone(completed) + self.assertEqual(completed.asset_uuid, asset2.asset_uuid) + + # Verify that Asset 1 is UNTOUCHED (still STORED and tier paths READY) + async with self.session_maker() as session4: + fetched1 = await assets.fetch_asset_by_uuid(session4, asset1.asset_uuid) + self.assertEqual( + fetched1[0].state, db_schema.AssetState.ASSET_STATE_STORED + ) + for tp in fetched1[0].tier_paths: + self.assertEqual(tp.state, db_schema.TierPathState.READY) + + # Verify Asset 2 is DELETED and its tier paths are DELETED + fetched2 = await assets.fetch_asset_by_uuid(session4, asset2.asset_uuid) + self.assertEqual( + fetched2[0].state, db_schema.AssetState.ASSET_STATE_DELETED + ) + for tp in fetched2[0].tier_paths: + self.assertEqual(tp.state, db_schema.TierPathState.DELETED) + + async def test_trigger_l0_to_l1_copy_zone_match(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre", + zone="us-central1-a", + ) + b1_a = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-a", + zone="us-central1-a", + ) + b1_b = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-b", + zone="us-east1-a", + ) + session.add_all([b0, b1_a, b1_b]) + await session.commit() + + request = tiering_service_pb2.ReserveRequest( + path="test/path/copy_zone_match", + user="test-user", + zone="us-central1-a", + ) + tp_uuid = "tp-uuid-b0" + storage_path = storage_backend_lib.get_storage_path( + b0, request.path, tp_uuid + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, + ) + finalized = await assets.finalize_asset(session, asset) + + await assets.trigger_l0_to_l1_copy(session, finalized) + await session.commit() + + # Verify it selected b1_a (same zone) and NOT b1_b + async with self.session_maker() as session2: + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == finalized.asset_uuid) + .options(joinedload(db_schema.Asset.tier_paths)) + ) + res = await session2.execute(stmt) + db_asset = res.scalars().first() + + # Should have 2 tier paths: L0 (b0) and L1 (b1_a) + self.assertLen(db_asset.tier_paths, 2) + backends = [tp.storage_backend_id for tp in db_asset.tier_paths] + self.assertIn(b0.id, backends) + self.assertIn(b1_a.id, backends) + self.assertNotIn(b1_b.id, backends) + + async def test_trigger_l0_to_l1_copy_region_match(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre", + zone="us-central1-a", + ) + b1_a = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-a", + region="us-central1", + ) + b1_b = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-b", + region="us-east1", + ) + session.add_all([b0, b1_a, b1_b]) + await session.commit() + + request = tiering_service_pb2.ReserveRequest( + path="test/path/copy_region_match", + user="test-user", + zone="us-central1-a", + ) + tp_uuid = "tp-uuid-b0-reg" + storage_path = storage_backend_lib.get_storage_path( + b0, request.path, tp_uuid + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, + ) + finalized = await assets.finalize_asset(session, asset) + + await assets.trigger_l0_to_l1_copy(session, finalized) + await session.commit() + + # Verify it selected b1_a (region matches L0 zone's region) + async with self.session_maker() as session2: + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == finalized.asset_uuid) + .options(joinedload(db_schema.Asset.tier_paths)) + ) + res = await session2.execute(stmt) + db_asset = res.scalars().first() + + self.assertLen(db_asset.tier_paths, 2) + backends = [tp.storage_backend_id for tp in db_asset.tier_paths] + self.assertIn(b0.id, backends) + self.assertIn(b1_a.id, backends) + self.assertNotIn(b1_b.id, backends) + + async def test_trigger_l0_to_l1_copy_multi_region_match(self): + async with self.session_maker() as session: + b0 = db_schema.StorageBackend( + level=0, + backend_type=db_schema.BackendType.BACKEND_TYPE_LUSTRE, + prefix="/mnt/lustre", + zone="us-central1-a", + ) + b1_a = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-a", + multi_regions=["us-central1", "us-east1"], + ) + b1_b = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://bucket-b", + region="us-west1", + ) + session.add_all([b0, b1_a, b1_b]) + await session.commit() + + request = tiering_service_pb2.ReserveRequest( + path="test/path/copy_multi_region_match", + user="test-user", + zone="us-central1-a", + ) + tp_uuid = "tp-uuid-b0-mreg" + storage_path = storage_backend_lib.get_storage_path( + b0, request.path, tp_uuid + ) + asset = await assets.create_or_fetch_asset( + session, + request, + b0, + tiering_service_pb2.ServerConfig( + client_keep_alive_interval_seconds=600 + ), + tier_path_uuid=tp_uuid, + storage_path=storage_path, + ) + finalized = await assets.finalize_asset(session, asset) + + await assets.trigger_l0_to_l1_copy(session, finalized) + await session.commit() + + # Verify it selected b1_a (L0 region is in b1_a's multi-regions) + async with self.session_maker() as session2: + stmt = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == finalized.asset_uuid) + .options(joinedload(db_schema.Asset.tier_paths)) + ) + res = await session2.execute(stmt) + db_asset = res.scalars().first() + + self.assertLen(db_asset.tier_paths, 2) + backends = [tp.storage_backend_id for tp in db_asset.tier_paths] + self.assertIn(b0.id, backends) + self.assertIn(b1_a.id, backends) + self.assertNotIn(b1_b.id, backends) + if __name__ == "__main__": absltest.main() 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..f94a58cd2 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,12 @@ 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 @@ -95,8 +99,8 @@ def data(self) -> bytes: raise auth_exceptions.TransportError(f"Request error: {e}") -class GCPStorageClient(abc.ABC): - """Client interface to interact with GCP storage backend (e.g. +class GCPStorageClient: + """Base helper class to interact with GCP storage backend (e.g. Lustre, GCS). """ @@ -115,6 +119,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: @@ -160,6 +169,10 @@ async def _get_token_and_project(self) -> tuple[str, str]: return self._credentials.token, self.project # pyrefly: ignore[missing-attribute] + +class StorageTransferClient(GCPStorageClient, abc.ABC): + """Interface for transferring data between storage backends.""" + @abc.abstractmethod async def trigger_copy( self, @@ -179,6 +192,19 @@ async def poll_operation( pass +class StorageBackendClient(abc.ABC): + """Interface for single-backend storage operations.""" + + @abc.abstractmethod + async def delete_path(self, path: str) -> None: + """Deletes a path from the storage backend.""" + pass + + async def close(self) -> None: + """Closes the client and releases resources.""" + pass + + def _parse_gcs_path(gcs_path: str) -> tuple[str, str]: """Parses a GCS path like gs://bucket/prefix/file into (bucket, prefix).""" path_no_scheme = gcs_path.replace("gs://", "") @@ -188,7 +214,7 @@ def _parse_gcs_path(gcs_path: str) -> tuple[str, str]: return bucket, prefix -class GcsToGcsClient(GCPStorageClient): +class GcsToGcsTransferClient(StorageTransferClient): """Client implementation for GCS-to-GCS operations using Storage Transfer Service.""" def __init__( @@ -323,7 +349,173 @@ async def poll_operation( ) -class GcpLustreBaseClient(GCPStorageClient): +class GcsClient(GCPStorageClient, StorageBackendClient): + """Client implementation for GCS backend operations (like deletion).""" + + _BATCH_DELETE_BOUNDARY = "cts_delete" + + 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. + + This design reduces network overhead and latency by packing up to 100 + deletions into a single HTTP request (the GCS batch API limit). It + efficiently handles bulk deletions without flooding the API with individual + calls. + + Args: + bucket: The GCS bucket name. + items: A list of object dictionaries to delete. + token: The OAuth2 access token for authorization. + """ + chunk_size = 100 + chunks = [ + items[i : i + chunk_size] for i in range(0, len(items), chunk_size) + ] + + for chunk in chunks: + body = self._serialize_batch_request(bucket, chunk) + response_content, content_type = await self._send_batch_request( + body, token + ) + self._deserialize_batch_response(response_content, content_type) + + def _serialize_batch_request( + self, bucket: str, chunk: list[dict[str, Any]] + ) -> bytes: + """Serializes a chunk of delete requests into multipart/mixed body bytes.""" + boundary = self._BATCH_DELETE_BOUNDARY + 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 + ) -> 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={self._BATCH_DELETE_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(StorageTransferClient): """Base client interface to interact with GCP Managed Lustre API via REST.""" def __init__( @@ -338,10 +530,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 +590,11 @@ async def trigger_copy( "Authorization": f"Bearer {token}", "Content-Type": "application/json", } + # Assets are directories. The Lustre API requires a trailing slash for + # directory GCS URIs. + 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 +625,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) @@ -482,7 +682,7 @@ def deserialize_transfer_status( tuple[db_schema.BackendType, db_schema.BackendType], Callable[ [db_schema.TierPath, db_schema.TierPath, str | None, str | None], - GCPStorageClient, + StorageTransferClient, ], ] = {} @@ -492,9 +692,11 @@ def register_client( target_type: db_schema.BackendType, ): """Decorator to register a client builder for a specific backend pair.""" + def decorator(builder_func): _CLIENT_BUILDERS[(source_type, target_type)] = builder_func return builder_func + return decorator @@ -507,10 +709,12 @@ def _build_gcs_to_gcs( target_tp: db_schema.TierPath, project: str | None, service_account: str | None, -) -> GcsToGcsClient: - """Builds GcsToGcsClient from TierPaths.""" +) -> GcsToGcsTransferClient: + """Builds GcsToGcsTransferClient from TierPaths.""" del source_tp, target_tp # Unused - return GcsToGcsClient(project=project, service_account=service_account) + return GcsToGcsTransferClient( + project=project, service_account=service_account + ) @register_client( @@ -566,7 +770,7 @@ def determine_client( target_tp: db_schema.TierPath, project: str | None = None, service_account: str | None = None, -) -> GCPStorageClient: +) -> StorageTransferClient: """Determines and returns the GCP Storage client based on backends.""" source_type = source_tp.storage_backend.backend_type target_type = target_tp.storage_backend.backend_type @@ -582,14 +786,16 @@ def get_client_from_status( status_dict: dict[str, Any], project: str | None = None, service_account: str | None = None, -) -> GCPStorageClient: +) -> StorageTransferClient: """Resolves and instantiates the client based on status metadata.""" client_type = status_dict.get("client_type") if not client_type: raise ValueError("Unknown or missing client type in status") - if client_type == "GcsToGcsClient": - return GcsToGcsClient(project=project, service_account=service_account) + if client_type == "GcsToGcsClient" or client_type == "GcsToGcsTransferClient": + return GcsToGcsTransferClient( + project=project, service_account=service_account + ) elif client_type == "LustreToGcsClient": zone = status_dict.get("zone") if not zone: @@ -614,3 +820,45 @@ def get_client_from_status( ) else: raise ValueError(f"Unknown or missing client type: {client_type}") + + +class LustreClient(StorageBackendClient): + """Client implementation for Lustre filesystem operations.""" + + def __init__(self, *args, **kwargs): + del args, kwargs + super().__init__() + + async def delete_path(self, path: str) -> None: + """Deletes a Lustre file or directory tree recursively.""" + + def _delete(): + try: + if os.path.islink(path): + os.remove(path) + elif os.path.isdir(path): + shutil.rmtree(path) + elif os.path.exists(path): + os.remove(path) + except FileNotFoundError: + # Catch FileNotFoundError to handle potential race conditions where the + # path is removed by another process/thread between checking and + # deletion. + pass + + await asyncio.to_thread(_delete) + + +def determine_delete_client( + tier_path: db_schema.TierPath, + project: str | None = None, + service_account: str | None = None, +) -> StorageBackendClient: + """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 GcsClient(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..17e9e9db6 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 @@ -14,6 +14,8 @@ """Unit tests for GCP Storage Clients.""" +import os +import tempfile import unittest from unittest import mock import httpx @@ -46,7 +48,7 @@ def tearDown(self): super().tearDown() async def test_gcs_to_gcs_trigger_copy_success(self): - client = gcp_storage_client.GcsToGcsClient(project="test-project") + client = gcp_storage_client.GcsToGcsTransferClient(project="test-project") # 1. Mock the first POST call to create transfer job mock_post_resp_1 = mock.MagicMock(spec=httpx.Response) @@ -70,7 +72,7 @@ async def test_gcs_to_gcs_trigger_copy_success(self): self.assertEqual(self.mock_client.post.call_count, 2) async def test_gcs_to_gcs_trigger_copy_sts_fail(self): - client = gcp_storage_client.GcsToGcsClient(project="test-project") + client = gcp_storage_client.GcsToGcsTransferClient(project="test-project") mock_post_resp = mock.MagicMock(spec=httpx.Response) mock_post_resp.status_code = 400 @@ -86,7 +88,7 @@ async def test_gcs_to_gcs_trigger_copy_sts_fail(self): self.assertIn("Failed to create Storage Transfer Job", str(ctx.exception)) async def test_gcs_to_gcs_poll_operation_in_progress(self): - client = gcp_storage_client.GcsToGcsClient(project="test-project") + client = gcp_storage_client.GcsToGcsTransferClient(project="test-project") mock_get_resp = mock.MagicMock(spec=httpx.Response) mock_get_resp.status_code = 200 @@ -109,7 +111,7 @@ async def test_gcs_to_gcs_poll_operation_in_progress(self): self.assertEqual(result.detail_info["total_bytes"], 1000) async def test_gcs_to_gcs_poll_operation_success(self): - client = gcp_storage_client.GcsToGcsClient(project="test-project") + client = gcp_storage_client.GcsToGcsTransferClient(project="test-project") mock_get_resp = mock.MagicMock(spec=httpx.Response) mock_get_resp.status_code = 200 @@ -123,7 +125,7 @@ async def test_gcs_to_gcs_poll_operation_success(self): self.assertEqual(result.status, gcp_storage_client.OperationStatus.SUCCESS) async def test_gcs_to_gcs_poll_operation_failed(self): - client = gcp_storage_client.GcsToGcsClient(project="test-project") + client = gcp_storage_client.GcsToGcsTransferClient(project="test-project") mock_get_resp = mock.MagicMock(spec=httpx.Response) mock_get_resp.status_code = 200 @@ -226,6 +228,113 @@ 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): + """Tests batch chunking limits and multipart HTTP request serialization.""" + client = gcp_storage_client.GcsClient(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): + """Tests parsing multipart HTTP responses to catch nested API error codes.""" + client = gcp_storage_client.GcsClient(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 @@ -235,7 +344,7 @@ async def test_service_account_token_impersonation( mock_imp_creds.token = "impersonated_token" mock_impersonated_creds_class.return_value = mock_imp_creds - client = gcp_storage_client.GcsToGcsClient( + client = gcp_storage_client.GcsToGcsTransferClient( project="test-project", service_account="sa@test.iam.gserviceaccount.com", ) @@ -251,6 +360,26 @@ async def test_service_account_token_impersonation( target_scopes=["https://www.googleapis.com/auth/cloud-platform"], ) + async def test_lustre_delete_path_file(self): + client = gcp_storage_client.LustreClient() + with tempfile.NamedTemporaryFile(delete=False) as f: + temp_path = f.name + + self.assertTrue(os.path.exists(temp_path)) + await client.delete_path(temp_path) + self.assertFalse(os.path.exists(temp_path)) + + async def test_lustre_delete_path_dir(self): + client = gcp_storage_client.LustreClient() + temp_dir = tempfile.mkdtemp() + temp_file = os.path.join(temp_dir, "file.txt") + with open(temp_file, "w") as f: + f.write("test") + + self.assertTrue(os.path.exists(temp_file)) + await client.delete_path(temp_dir) + self.assertFalse(os.path.exists(temp_dir)) + class GCPStorageClientHelpersTest(unittest.TestCase): @@ -268,7 +397,7 @@ def test_determine_client_gcs_to_gcs(self): client = gcp_storage_client.determine_client( source_tp, target_tp, project="my-proj", service_account="my-sa" ) - self.assertIsInstance(client, gcp_storage_client.GcsToGcsClient) + self.assertIsInstance(client, gcp_storage_client.GcsToGcsTransferClient) self.assertEqual(client.project, "my-proj") self.assertEqual(client.service_account, "my-sa") @@ -354,12 +483,44 @@ def test_determine_client_gcs_to_lustre_missing_zone_raises(self): with self.assertRaisesRegex(ValueError, "Lustre zone is missing"): gcp_storage_client.determine_client(source_tp, target_tp) + def test_determine_delete_client_gcs(self): + tier_path = mock.MagicMock() + tier_path.storage_backend.backend_type = ( + gcp_storage_client.db_schema.BackendType.BACKEND_TYPE_GCS + ) + client = gcp_storage_client.determine_delete_client( + tier_path, project="my-proj", service_account="my-sa" + ) + self.assertIsInstance(client, gcp_storage_client.GcsClient) + self.assertEqual(client.project, "my-proj") + self.assertEqual(client.service_account, "my-sa") + + def test_determine_delete_client_lustre(self): + tier_path = mock.MagicMock() + tier_path.storage_backend.backend_type = ( + gcp_storage_client.db_schema.BackendType.BACKEND_TYPE_LUSTRE + ) + client = gcp_storage_client.determine_delete_client( + tier_path, project="my-proj", service_account="my-sa" + ) + self.assertIsInstance(client, gcp_storage_client.LustreClient) + + def test_determine_delete_client_invalid_raises(self): + tier_path = mock.MagicMock() + tier_path.storage_backend.backend_type = ( + gcp_storage_client.db_schema.BackendType.BACKEND_TYPE_UNSPECIFIED + ) + with self.assertRaisesRegex( + ValueError, "Unsupported backend type for delete" + ): + gcp_storage_client.determine_delete_client(tier_path) + def test_get_client_from_status_gcs_to_gcs(self): status = {"client_type": "GcsToGcsClient"} client = gcp_storage_client.get_client_from_status( status, project="my-proj" ) - self.assertIsInstance(client, gcp_storage_client.GcsToGcsClient) + self.assertIsInstance(client, gcp_storage_client.GcsToGcsTransferClient) self.assertEqual(client.project, "my-proj") def test_get_client_from_status_gcs_to_lustre(self): diff --git a/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/job_worker.py index 1091bc8be..0262d91e3 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,209 @@ 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 _get_paths_to_delete( + self, job: db_schema.AssetJob + ) -> list[db_schema.TierPath]: + """Fetches tier paths associated with the delete job.""" + async with self._session_maker() as session: # pyrefly: ignore[not-callable] + if job.target_tier_path_id: + stmt = ( + select(db_schema.TierPath) + .where(db_schema.TierPath.id == job.target_tier_path_id) + .options(joinedload(db_schema.TierPath.storage_backend)) + ) + result = await session.execute(stmt) + tp = result.scalars().first() + return [tp] if tp else [] + 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() + return db_asset.tier_paths if db_asset else [] + + async def _perform_physical_deletions( + self, + tier_paths: list[db_schema.TierPath], + ) -> None: + """Invokes storage clients to delete paths physically.""" + 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() + + async def _complete_delete_all_tiers( + self, + session: AsyncSession, + asset_uuid: str, + ) -> None: + """Completes deletion of all tiers for an asset.""" + stmt_asset = ( + select(db_schema.Asset) + .where(db_schema.Asset.asset_uuid == 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, + ) + + async def _complete_delete_instances( + self, + session: AsyncSession, + tier_paths: list[db_schema.TierPath], + ) -> None: + """Completes deletion of specific tier paths.""" + 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, + ) + + async def _complete_delete_job_db_transition( + self, + job: db_schema.AssetJob, + tier_paths: list[db_schema.TierPath], + now: datetime.datetime, + ) -> None: + """Updates database states for completed deletion job.""" + 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 + ): + await self._complete_delete_all_tiers(session, job.asset_uuid) + else: + await self._complete_delete_instances(session, tier_paths) + + await self._complete_job(session, merged_job, now) + + 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) + + try: + # 1. Fetch paths to delete (already in DELETE_IN_PROCESS state) + tier_paths = await self._get_paths_to_delete(job) + 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 + await self._perform_physical_deletions(tier_paths) + + # 3. Transition status in database to completed + await self._complete_delete_job_db_transition(job, tier_paths, 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..aa6f38e09 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,11 +163,19 @@ 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") if client_type not in { "GcsToGcsClient", + "GcsToGcsTransferClient", "LustreToGcsClient", "GcsToLustreClient", "DummyGcpParallelstoreClient", @@ -185,6 +219,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 +249,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 +283,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 +347,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 +407,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 +454,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 +509,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 +583,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 +739,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 +806,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..e63fc7609 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/server.py @@ -22,6 +22,7 @@ import os import pprint import sys +import uuid from absl import logging import fire @@ -133,7 +134,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 +155,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) @@ -261,6 +270,9 @@ async def Finalize( if db_asset is None: # This is unlikely to happen since we just finalized the asset. raise ValueError("Asset not found after finalize") + + # default policy to preserve it from L0 to L1 + await assets.trigger_l0_to_l1_copy(session, db_asset) except ValueError as e: logging.exception("Finalize failed for UUID: %s", request.uuid) await context.abort( @@ -324,6 +336,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 +359,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 +380,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 +415,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 +459,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", 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..9b54d9818 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend.py @@ -97,20 +97,31 @@ 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('/')}" + prefix = backend.prefix.rstrip("/") + if relative_path == prefix or relative_path.startswith(prefix + "/"): + base_path = relative_path + else: + base_path = f"{prefix}/{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/experimental/tiering_service/storage_backend_test.py b/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend_test.py index 8237327d4..06ad47ab9 100644 --- a/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend_test.py +++ b/checkpoint/orbax/checkpoint/experimental/tiering_service/storage_backend_test.py @@ -98,5 +98,42 @@ async def test_locate_closest_backend(self): self.assertIsNone(match_fallback) +class StorageBackendPathTest(absltest.TestCase): + + def test_get_storage_path_safe_prefix(self): + backend = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://my-bucket", + ) + self.assertEqual( + storage_backend.get_storage_path(backend, "gs://my-bucket/file"), + "gs://my-bucket/file", + ) + self.assertEqual( + storage_backend.get_storage_path(backend, "gs://my-bucket"), + "gs://my-bucket", + ) + self.assertEqual( + storage_backend.get_storage_path(backend, "gs://my-bucket-2/file"), + "gs://my-bucket/gs://my-bucket-2/file", + ) + self.assertEqual( + storage_backend.get_storage_path(backend, "file"), + "gs://my-bucket/file", + ) + + def test_get_storage_path_with_trailing_slash(self): + backend = db_schema.StorageBackend( + level=1, + backend_type=db_schema.BackendType.BACKEND_TYPE_GCS, + prefix="gs://my-bucket/", + ) + self.assertEqual( + storage_backend.get_storage_path(backend, "gs://my-bucket/file"), + "gs://my-bucket/file", + ) + + if __name__ == "__main__": absltest.main()