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/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 )