From 979ca0003aecf1d8c029153e12088b651cd80215 Mon Sep 17 00:00:00 2001 From: Sam Welborn Date: Sun, 12 Jul 2026 17:24:14 -0400 Subject: [PATCH] use a key to dedupe reqs we use a key here instead of transfer ID. the main idea is that the transfer ID will always be unique. the experiment ID will be unchanged per request. If a cache is up and we req for an experiment (shared mode), then it will be OK. If we have a new id or we previously shut down the cache for this particular experiment, then the key is removed from the table and we can reuse it to start up another cache --- src/fastcache_api/models.py | 6 ++- src/fastcache_api/reconcile.py | 2 + src/fastcache_api/routes/cache.py | 74 +++++++++++++++++++++++++++---- src/fastcache_api/tables.py | 11 ++++- 4 files changed, 81 insertions(+), 12 deletions(-) diff --git a/src/fastcache_api/models.py b/src/fastcache_api/models.py index 0135923..40a3e52 100644 --- a/src/fastcache_api/models.py +++ b/src/fastcache_api/models.py @@ -50,7 +50,9 @@ def to_fastcache_json(self, indent: int = 2) -> str: class CacheRequest(BaseModel): - transfer_id: str + # The cache's dedup/lookup identity. Could be per-transfer (transfer ID) + # or per-experiment (experiment name) + key: str # Human who initiated the transfer upstream (bearer token is a shared # service identity, so attribution must travel in the request body). requested_by: str @@ -62,7 +64,7 @@ class CachePublic(BaseModel): model_config = ConfigDict(from_attributes=True) id: UUID - transfer_id: str + key: str | None user: str state: CacheState exit_code: int | None diff --git a/src/fastcache_api/reconcile.py b/src/fastcache_api/reconcile.py index 02dbfca..3aa3385 100644 --- a/src/fastcache_api/reconcile.py +++ b/src/fastcache_api/reconcile.py @@ -35,6 +35,7 @@ async def sweep_dead_caches() -> int: ec = exit_code(cache.pid) cache.state = CacheState.completed if ec == 0 else CacheState.failed cache.exit_code = ec + cache.key = None # Free the key! logger.warning( "Cache %s (pid=%d) is no longer running (exit_code=%s); marking %s", cache.id, @@ -68,6 +69,7 @@ async def watch_and_record(cache_id: UUID, pid: int) -> None: return cache.state = CacheState.completed if exit_code == 0 else CacheState.failed cache.exit_code = exit_code + cache.key = None await session.commit() logger.info( "Cache %s (pid=%d) exited (code=%s); marked %s", diff --git a/src/fastcache_api/routes/cache.py b/src/fastcache_api/routes/cache.py index e3545df..1dcf2e1 100644 --- a/src/fastcache_api/routes/cache.py +++ b/src/fastcache_api/routes/cache.py @@ -1,13 +1,20 @@ from typing import Annotated from uuid import UUID, uuid4 -from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, status +from fastapi import ( + APIRouter, + BackgroundTasks, + Depends, + HTTPException, + Response, + status, +) from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from ..config import settings -from ..db import get_session +from ..db import SessionLocal, get_session from ..dependencies import TokenPayload, require_user from ..lifecycle import schedule_exit_watch from ..models import ( @@ -20,6 +27,8 @@ from ..process import ( allocate_port_pair, canonical_hostname, + exit_code, + is_alive, ports_in_use, start_cache, stop_cache, @@ -33,6 +42,29 @@ ) +async def _find_active_by_key(session: AsyncSession, key: str) -> Cache | None: + result = await session.execute( + select(Cache).where( + Cache.key == key, + Cache.state.in_([s.value for s in CacheState if not s.is_final()]), + ) + ) + candidate = result.scalar_one_or_none() + if candidate is None: + return None + if is_alive(candidate.pid, candidate.create_time): + return candidate + # The row hasn't been reconciled yet (exit watcher/sweep haven't run), + # but the process is already gone. Finalize here and free the key so + # caller creates a fresh one instead. + ec = exit_code(candidate.pid) + candidate.state = CacheState.completed if ec == 0 else CacheState.failed + candidate.exit_code = ec + candidate.key = None + await session.commit() + return None + + @router.get("/", response_model=CachesPublic) async def get_caches( session: Annotated[AsyncSession, Depends(get_session)], @@ -60,7 +92,13 @@ async def create_cache( req: CacheRequest, user: Annotated[TokenPayload, Depends(require_user)], session: Annotated[AsyncSession, Depends(get_session)], + response: Response, ): + existing = await _find_active_by_key(session, req.key) + if existing is not None: + response.status_code = status.HTTP_200_OK + return existing + # The cache always runs as a local subprocess of this api server, so the # ZMQ URIs are published under this host's canonical (FQDN) name. hostname = canonical_hostname() @@ -101,7 +139,7 @@ async def create_cache( cache = Cache( id=cache_id, - transfer_id=req.transfer_id, + key=req.key, pid=proc.pid, create_time=proc.create_time, user=req.requested_by, @@ -115,18 +153,36 @@ async def create_cache( await session.commit() except IntegrityError as exc: await session.rollback() - # Don't leave an orphaned process behind for the rejected request. + # Lost a race with a concurrent request for the same key: don't + # leave our process running, join the winner instead. await stop_cache(proc.pid, proc.create_time, timeout=0) - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=(f"Cache with transfer_id '{req.transfer_id}' already exists"), - ) from exc + winner = await _find_active_by_key(session, req.key) + if winner is None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Cache key '{req.key}' conflict", + ) from exc + response.status_code = status.HTTP_200_OK + return winner await session.refresh(cache) schedule_exit_watch(cache.id, proc.pid) + response.status_code = status.HTTP_201_CREATED return cache +async def _teardown_and_free_key( + cache_id: UUID, pid: int, create_time: float | None +) -> None: + await stop_cache(pid, create_time, settings.SHUTDOWN_GRACE_SECONDS) + # Only free the key once the process is confirmed stopped + async with SessionLocal() as session: + cache = await session.get(Cache, cache_id) + if cache is not None: + cache.key = None + await session.commit() + + @router.delete("/{cache_id}", response_model=CachePublic) async def shutdown_cache( cache_id: UUID, @@ -145,6 +201,6 @@ async def shutdown_cache( # Tear down the process tree in the background so DELETE returns promptly. background_tasks.add_task( - stop_cache, cache.pid, cache.create_time, settings.SHUTDOWN_GRACE_SECONDS + _teardown_and_free_key, cache.id, cache.pid, cache.create_time ) return cache diff --git a/src/fastcache_api/tables.py b/src/fastcache_api/tables.py index fa97719..f93b9c2 100644 --- a/src/fastcache_api/tables.py +++ b/src/fastcache_api/tables.py @@ -39,7 +39,16 @@ class Cache(DTMixin, Base): __tablename__ = "caches" id: Mapped[UUID] = mapped_column(default=uuid4, primary_key=True) - transfer_id: Mapped[str] = mapped_column(unique=True, doc="Unique transfer ID") + key: Mapped[str | None] = mapped_column( + unique=True, + default=None, + doc=( + "Dedup/lookup ID from the request (e.g. a transfer id for " + "a one-off cache, or a shared key like an experiment name for a " + "joinable one). Cleared when the cache is torn down, freeing the " + "key for reuse." + ), + ) pid: Mapped[int] = mapped_column(doc="PID of the cache process") create_time: Mapped[float | None] = mapped_column( default=None,