Tinker Engine

class TinkerEngine:
    """Background engine for processing training requests.

    The engine handles:
    - Database operations (futures, checkpoints)
    - Request finding/scheduling
    - File I/O (download/upload checkpoints)
    - Validating requests against loaded models

    Computation and model management are delegated to the backend.
    """

    def _filter_valid_requests(
        self,
        requests: dict[str, tuple[str, BaseModel]],
    ) -> tuple[dict[str, types.ErrorResponse], dict[str, tuple[str, BaseModel]]]:
        """Filter out requests with invalid model_ids and return error results for them.

        Args:
            requests: Dict mapping request_id to (model_id, request_data) tuples

        Returns:
            Tuple of (error_results, valid_requests)
        """
        results = {}
        valid_requests = {}

        for request_id, (model_id, request_data) in requests.items():
            error = None
            if model_id and not self.backend.has_model(model_id):
                error = f"Model {model_id} not loaded"
            elif not model_id and isinstance(request_data, types.SampleInput):
                if request_data.base_model != self.config.base_model:
                    error = f"Engine is configured for '{self.config.base_model}' but request specified '{request_data.base_model}'"
                elif request_data.checkpoint_id:
                    error = "checkpoint_id must be empty for base model sampling"

            if error:
                results[request_id] = types.ErrorResponse(error=error, status="failed")
            else:
                valid_requests[request_id] = (model_id, request_data)

        return results, valid_requests

    def __init__(
        self,
        config: EngineConfig,
    ):
        """Initialize the engine with a database connection and base model."""
        self.config = config
        self.db_engine = create_engine(config.database_url, echo=False)
        enable_sqlite_wal(self.db_engine)

        if config.backend == "jax" and config.runtime_role != "combined":
            raise ValueError("Single-role runtimes require the fsdp or megatron backend")

        # Initialize the backend (handles model state, computation, and adapter management)
        use_ray = config.backend_config.get("use_ray", False)
        backend_class, backend_config_class = get_backend_classes(config.backend, use_ray=use_ray)
        backend_overrides = dict(config.backend_config)
        if config.backend in ("fsdp", "megatron"):
            backend_overrides["runtime_role"] = config.runtime_role
            if config.base_model_checkpoint_path is not None:
                backend_overrides["generator.inference_engine.served_model_name"] = config.base_model
        backend_config = backend_config_class(**backend_overrides)
        self.backend = backend_class(config.base_model_checkpoint_path or config.base_model, backend_config)

        # Backends that support async sample routing notify us when their
        # inference endpoint changes; we persist it to EngineStateDB so the
        # API process can forward sample requests directly. Backends stay
        # DB-free; only the engine owns the connection.
        if hasattr(self.backend, "set_inference_state_publisher"):
            self.backend.set_inference_state_publisher(self._write_inference_state_to_db)

        # Track last cleanup time for periodic stale session cleanup
        self._last_cleanup_time: float = time.time()

        # Continuous sampling: complete each sample future the moment its own
        # generation finishes instead of batching futures to the wave's
        # slowest member. Requires the backend to expose the awaitable
        # per-request path + main-thread engine wake (skyrl-train backends).
        # SKYRL_TINKER_CONTINUOUS_SAMPLING=0 falls back to the serial loop.
        self._sampler: Optional[_ContinuousSampler] = None
        self._continuous_sampling: bool = (
            SKYRL_TINKER_CONTINUOUS_SAMPLING
            and hasattr(self.backend, "sample_batch_async")
            and hasattr(self.backend, "prepare_for_sampling")
        )
        if self._continuous_sampling:
            logger.info(
                "Continuous sampling enabled: per-request future completion, "
                "drain-before-blocking-ops (set SKYRL_TINKER_CONTINUOUS_SAMPLING=0 for the serial loop)"
            )

        # Active torch profiling session, mirrored in memory so the optim_step
        # path never touches the DB. None when no session is running.
        self._profiling_model_id: str | None = None
        self._profiling_steps: int = 0
        self._profiling_error: str | None = None
        self._profiler_cfg = TinkerTorchProfilerConfig(**config.torch_profiler) if config.torch_profiler else None
        if self._profiler_cfg is not None:
            # Clear a session left behind by a crashed engine. Skipped when profiling
            # is off: there is nothing to reset, and the table may not exist yet.
            self._reset_profiler_row()

        logger.info(f"Initialized TinkerEngine with backend={type(self.backend).__name__}")

    @property
    def metrics(self) -> types.EngineMetrics:
        """Pass-through to backend metrics for backwards compatibility."""
        return self.backend.metrics

    def _write_inference_state_to_db(self, proxy_url: str | None) -> None:
        """Upsert the singleton EngineStateDB row.

        Wired into the backend via set_inference_state_publisher so the API
        process can resolve the engine-managed vLLM URL on the async sample
        routing path. ``proxy_url=None`` clears the row (post-teardown).
        """
        with Session(self.db_engine) as session:
            row = session.get(EngineStateDB, 1) or EngineStateDB(singleton_id=1)
            row.inference_proxy_url = proxy_url
            row.updated_at = datetime.now(timezone.utc)
            session.add(row)
            session.commit()

    @contextmanager
    def _checkpoint_status_context(self, model_id: str, checkpoint_id: str, checkpoint_type: types.CheckpointType):
        """Context manager to handle checkpoint DB status updates.

        Fetches the checkpoint entry, yields it, and updates its status to COMPLETED
        or FAILED based on whether an exception occurred.
        """
        with Session(self.db_engine) as session:
            # Fail fast if API didn't create the checkpoint row first.
            if session.get(CheckpointDB, (model_id, checkpoint_id, checkpoint_type)) is None:
                raise ValueError(
                    f"Checkpoint entry not found for model '{model_id}', checkpoint '{checkpoint_id}', type '{checkpoint_type}'"
                )

        status = CheckpointStatus.FAILED
        error_message = "checkpoint operation interrupted"
        try:
            # Run potentially slow checkpoint I/O without an open DB transaction.
            yield
            status = CheckpointStatus.COMPLETED
            error_message = None
        except Exception as e:
            logger.exception(f"Error saving checkpoint for model {model_id}, checkpoint {checkpoint_id}: {e}")
            error_message = str(e)
            raise
        finally:
            # Persist final status in a short write transaction.
            with Session(self.db_engine) as session:
                result = session.exec(
                    update(CheckpointDB)
                    .where(CheckpointDB.model_id == model_id)
                    .where(CheckpointDB.checkpoint_id == checkpoint_id)
                    .where(CheckpointDB.checkpoint_type == checkpoint_type)
                    .values(
                        status=status,
                        error_message=error_message,
                        completed_at=datetime.now(timezone.utc),
                    )
                )
                if not result.rowcount:
                    logger.warning(
                        f"Checkpoint row disappeared before status update: "
                        f"model_id={model_id}, checkpoint_id={checkpoint_id}, checkpoint_type={checkpoint_type}"
                    )
                session.commit()

    def _load_requests(self, session: Session, requests: list[tuple[Any, ...]]) -> list[tuple[Any, ...]]:
        """Append each request's request_data to its tuple, dropping any request that no longer exists.

        Chunked to stay under SQLite's limit on bound parameters per statement.
        """
        request_ids = [request_id for request_id, *_ in requests]
        payloads: dict[int, dict] = {}
        for start in range(0, len(request_ids), _MAX_IDS_PER_QUERY):
            chunk = request_ids[start : start + _MAX_IDS_PER_QUERY]
            rows = session.exec(
                select(FutureDB.request_id, FutureDB.request_data).where(FutureDB.request_id.in_(chunk))
            ).all()
            payloads.update(rows)
        return [(request_id, *args, payloads[request_id]) for request_id, *args in requests if request_id in payloads]

    def _find_destructive_barriers(self, session: Session) -> dict[str, int]:
        """Find the earliest pending destructive operation (optim_step/load_weights) per model.

        These act as scheduling barriers: model passes before them can be batched
        safely, and single requests after a blocked pass must wait.
        """
        query = (
            select(FutureDB.model_id, func.min(FutureDB.request_id).label("barrier_id"))
            .where(
                (FutureDB.request_type == types.RequestType.OPTIM_STEP)
                | (FutureDB.request_type == types.RequestType.LOAD_WEIGHTS)
            )
            .where(FutureDB.status == RequestStatus.PENDING)
            .group_by(FutureDB.model_id)
        )
        return dict(session.exec(query).all())

    def find_batchable_model_passes(
        self, session: Session, request_type: types.RequestType
    ) -> dict[str, tuple[str, types.ForwardBackwardInput]]:
        """Find all requests of the given type that come before any destructive update for their model.

        Uses look-ahead scheduling: for each model, only returns operations
        that have no optim_step or load_weights blocking them in the queue.

        Args:
            session: Database session
            request_type: The type of request to find (e.g., FORWARD or FORWARD_BACKWARD)

        Returns:
            Dict mapping request_id to (model_id, request_data) tuples
        """
        barriers = self._find_destructive_barriers(session)

        # Get all pending operations of the requested type ordered by request_id
        query = (
            select(FutureDB.request_id, FutureDB.model_id)
            .where(FutureDB.request_type == request_type)
            .where(FutureDB.status == RequestStatus.PENDING)
            .order_by(FutureDB.request_id)
        )
        ops = session.exec(query).all()

        # Filter: only include ops that come before their model's barrier
        batchable = [
            (request_id, model_id)
            for request_id, model_id in ops
            if model_id not in barriers or request_id < barriers[model_id]
        ]

        return {
            str(request_id): (model_id, types.ForwardBackwardInput.model_validate(request_data))
            for request_id, model_id, request_data in self._load_requests(session, batchable)
        }

    def find_batchable_sample(self, session: Session) -> dict[str, tuple[str, types.SampleInput]]:
        """Find all sample ops that can be safely batched together.

        Returns sample operations ensuring that each model_id has only one checkpoint_id
        to avoid loading different checkpoints for the same model in a single batch.

        If sample_max_num_sequences is configured, limits to that many requests so we don't
        produce partial batches in process_sample_batch. If num_samples > 1 for some requests,
        this may not be perfect, but it's good until we implement continuous batching.

        Args:
            session: Database session

        Returns:
            Dict mapping request_id to (model_id, request_data) tuples
        """
        sampler = getattr(self, "_sampler", None)
        inflight_ids = sampler.inflight_ids if sampler is not None else set()

        # checkpoint_id is extracted in the database so prompts stay out of this query
        sample_query = (
            select(FutureDB.request_id, FutureDB.model_id, FutureDB.request_data["checkpoint_id"].as_string())
            .where(FutureDB.request_type == types.RequestType.SAMPLE)
            .where(FutureDB.status == RequestStatus.PENDING)
            .order_by(FutureDB.request_id)
        )
        sample_ops = session.exec(sample_query).all()

        batchable = []
        model_checkpoints = {}  # Map from model_id to checkpoint_id of first request to that model
        for request_id, model_id, checkpoint_id in sample_ops:
            if str(request_id) in inflight_ids:
                continue
            # Base model requests (empty checkpoint_id) are always compatible, otherwise only
            # take only requests with one checkpoint_id for a given model_id
            if not checkpoint_id or model_checkpoints.setdefault(model_id, checkpoint_id) == checkpoint_id:
                batchable.append((request_id, model_id))

        # TODO: This leaks the abstraction by accessing backend-specific config.
        # We should find a better way to handle this going forward.
        if self.config.backend == "jax" and self.backend.config.sample_max_num_sequences > 0:
            batchable = batchable[: self.backend.config.sample_max_num_sequences]

        return {
            str(request_id): (model_id, types.SampleInput.model_validate(request_data))
            for request_id, model_id, request_data in self._load_requests(session, batchable)
        }

    def find_single_requests(self, session: Session) -> dict[str, tuple[str, types.RequestType, dict]]:
        """Find all requests that need to be processed individually (not batchable).

        Args:
            session: Database session

        Returns:
            Dict mapping request_id to (model_id, request_type, request_data) tuples
        """
        # Find the first blocked forward pass per model: a pending FORWARD/FORWARD_BACKWARD
        # that sits behind a destructive barrier and won't be batched this iteration.
        # Single requests must not jump ahead of these.
        destructive_barriers = self._find_destructive_barriers(session)
        blocked_pass_barriers: dict[str, int] = {}
        if destructive_barriers:
            pending_passes = session.exec(
                select(FutureDB.model_id, FutureDB.request_id)
                .where(
                    (FutureDB.request_type == types.RequestType.FORWARD_BACKWARD)
                    | (FutureDB.request_type == types.RequestType.FORWARD)
                )
                .where(FutureDB.status == RequestStatus.PENDING)
                .where(FutureDB.model_id.in_(destructive_barriers.keys()))
                .order_by(FutureDB.request_id)
            ).all()
            for model_id, req_id in pending_passes:
                if req_id >= destructive_barriers[model_id]:
                    blocked_pass_barriers.setdefault(model_id, req_id)

        statement = (
            select(FutureDB.request_id, FutureDB.model_id, FutureDB.request_type)
            .where(FutureDB.status == RequestStatus.PENDING)
            .where(FutureDB.request_type != types.RequestType.FORWARD_BACKWARD)
            .where(FutureDB.request_type != types.RequestType.FORWARD)
            .where(FutureDB.request_type != types.RequestType.SAMPLE)
            .where(FutureDB.request_type != types.RequestType.EXTERNAL)
            .order_by(FutureDB.request_id)
        )
        other_futures = session.exec(statement).all()

        # Filter: only include ops that come before the first blocked pass for their model
        other_futures = [
            (request_id, model_id, request_type)
            for request_id, model_id, request_type in other_futures
            if model_id not in blocked_pass_barriers or request_id < blocked_pass_barriers[model_id]
        ]

        return {
            str(request_id): (model_id, request_type, request_data)
            for request_id, model_id, request_type, request_data in self._load_requests(session, other_futures)
        }

    def process_create_model(self, model_id: str, request_data: types.CreateModelInput) -> types.CreateModelOutput:
        """Create and initialize a model."""
        # Create model in backend (allocates adapter_index, creates optimizer, and configures adapter)
        self.backend.create_model(model_id, request_data.lora_config, model_role=request_data.model_role)

        logger.info(f"Created LoRA model {model_id}")

        return types.CreateModelOutput(
            model_id=model_id,
            base_model=self.config.base_model,
            lora_config=request_data.lora_config,
        )

    def process_unload_model(self, model_id: str, request_data: types.UnloadModelInput) -> types.UnloadModelOutput:
        """Unload a model and free all resources."""
        if not self.backend.has_model(model_id):
            logger.warning(f"Ignoring unload request for model {model_id} that is not loaded.")
        else:
            self.backend.delete_model(model_id)

            # Update model status in DB
            with Session(self.db_engine) as session:
                _ = session.exec(update(ModelDB).where(ModelDB.model_id == model_id).values(status="unloaded"))
                session.commit()

            logger.info(f"Unloaded model {model_id}")

        return types.UnloadModelOutput(model_id=model_id, status="unloaded")

    def cleanup_stale_sessions(self) -> int:
        """Cleanup sessions with no recent heartbeat and unload their models.

        Returns:
            Number of models unloaded
        """
        cutoff = datetime.now(timezone.utc) - timedelta(seconds=self.config.session_timeout_sec)
        unloaded_count = 0

        with Session(self.db_engine) as session:
            # Find stale sessions (active sessions with heartbeat older than cutoff)
            stale_sessions = session.exec(
                select(SessionDB).where(
                    SessionDB.status == "active",
                    SessionDB.last_heartbeat_at < cutoff,
                )
            ).all()

            if not stale_sessions:
                return 0

            stale_session_ids = {s.session_id for s in stale_sessions}

            # Find all models for all stale sessions in one query
            models_to_process = session.exec(
                select(ModelDB).where(
                    ModelDB.session_id.in_(stale_session_ids),
                    ModelDB.status != "unloaded",
                )
            ).all()

        # Unloading can tear the shared runtime (engines included) down when
        # the last model goes; in-flight continuous samples must finish first.
        sampler = getattr(self, "_sampler", None)
        if models_to_process and sampler is not None:
            sampler.drain(reason="stale-session model unload")

        # Unload models outside DB transactions to minimize lock time.
        sessions_with_failed_unloads: set[str] = set()
        unloaded_model_ids: set[str] = set()
        for model in models_to_process:
            if self.backend.has_model(model.model_id):
                try:
                    self.backend.delete_model(model.model_id)
                    unloaded_model_ids.add(model.model_id)
                    unloaded_count += 1
                    logger.info(f"Auto-unloaded stale model {model.model_id} from session {model.session_id}")
                except Exception as e:
                    logger.error(f"Failed to auto-unload model {model.model_id}: {e}")
                    sessions_with_failed_unloads.add(model.session_id)
            else:
                # Model already missing in backend; only DB state needs cleanup.
                unloaded_model_ids.add(model.model_id)

        sessions_to_expire = [s.session_id for s in stale_sessions if s.session_id not in sessions_with_failed_unloads]

        # Apply DB status updates in one short write transaction.
        with Session(self.db_engine) as session:
            if unloaded_model_ids:
                _ = session.exec(
                    update(ModelDB).where(ModelDB.model_id.in_(unloaded_model_ids)).values(status="unloaded")
                )
            if sessions_to_expire:
                _ = session.exec(
                    update(SessionDB).where(SessionDB.session_id.in_(sessions_to_expire)).values(status="expired")
                )
            session.commit()

        for session_id in sessions_to_expire:
            logger.info(f"Expired stale session {session_id}")

        return unloaded_count

    # ------------------------------------------------------------------
    # torch.profiler session control. The API process owns the single
    # ProfilerControlDB row; the engine reconciles to it once per loop.
    # ------------------------------------------------------------------

    def _reset_profiler_row(self) -> None:
        """Clear any session left behind by a crashed engine.

        The row survives in SQLite but the workers and staging dirs do not, so a
        fresh engine must never inherit a row claiming a session is running.
        """
        with Session(self.db_engine) as session:
            row = session.get(ProfilerControlDB, 1)
            if row is None:
                row = ProfilerControlDB(singleton_id=1)
            row.desired_state = ProfilerState.STOPPED
            row.owner_model_id = None
            row.config_json = None
            row.version = 0
            row.applied_version = 0
            row.started_at = None
            row.step = 0
            row.error = None
            session.add(row)
            session.commit()

    def _write_profiler_row(self, **values) -> None:
        with Session(self.db_engine) as session:
            row = session.get(ProfilerControlDB, 1)
            if row is None:
                return
            for key, value in values.items():
                setattr(row, key, value)
            session.add(row)
            session.commit()

    def _stop_profiling(self) -> str | None:
        """Stop the live session. Returns an error string, or None on success."""
        self._profiling_model_id = None
        try:
            self.backend.stop_profile()
        except Exception as e:
            logger.warning(f"[profiler] stop failed: {e}")
            return f"stop failed: {e}"
        return None

    def reconcile_profiler(self) -> None:
        """Converge the workers to the desired state in the control row.

        Runs at the top of the request loop, which is a boundary between request
        batches, so a session never starts or stops mid-batch.
        """
        if self._profiler_cfg is None:
            return
        with Session(self.db_engine) as session:
            row = session.get(ProfilerControlDB, 1)
            if row is None:
                return
            desired, version, applied = row.desired_state, row.version, row.applied_version
            owner, config_json, started_at = row.owner_model_id, row.config_json, row.started_at

        if applied < version:
            # A new claim or release to apply.
            if desired == ProfilerState.RUNNING:
                try:
                    try:
                        worker_config = json.loads(config_json)
                    except (TypeError, ValueError) as e:
                        raise ValueError(f"unreadable profiler config in control row: {e}") from e
                    self.backend.start_profile(worker_config)
                except Exception as e:
                    logger.warning(f"[profiler] start failed: {e}")
                    # Release the slot, but still ack: the API is blocked waiting
                    # for applied_version and would otherwise hang to its timeout.
                    self._profiling_model_id = None
                    self._write_profiler_row(
                        desired_state=ProfilerState.STOPPED,
                        owner_model_id=None,
                        applied_version=version,
                        error=f"start failed: {e}",
                    )
                    return
                self._profiling_model_id = owner
                self._profiling_steps = 0
                self._profiling_error = None
                self._write_profiler_row(applied_version=version, step=0, error=None)
                logger.info(f"[profiler] session started for model {owner}")
            else:
                error = self._stop_profiling()
                self._write_profiler_row(applied_version=version, step=self._profiling_steps, error=error)
                logger.info("[profiler] session stopped")
            return

        if self._profiling_model_id is None:
            return

        # Active session: enforce the TTL, else flush progress for /profiling_status.
        if started_at is not None:
            # SQLite hands back naive datetimes even for a timezone-aware column,
            # and everything we store is UTC.
            if started_at.tzinfo is None:
                started_at = started_at.replace(tzinfo=timezone.utc)
            age = (datetime.now(timezone.utc) - started_at).total_seconds()
            if age > self._profiler_cfg.max_session_duration_sec:
                logger.warning(
                    f"[profiler] session exceeded max_session_duration_sec="
                    f"{self._profiler_cfg.max_session_duration_sec}s; finalizing"
                )
                self._stop_profiling()
                self._write_profiler_row(
                    desired_state=ProfilerState.STOPPED,
                    owner_model_id=None,
                    step=self._profiling_steps,
                    error="session terminated by max_session_duration_sec",
                )
                return
        self._write_profiler_row(step=self._profiling_steps, error=self._profiling_error)

    def process_optim_step(
        self, model_id: str, request_data: types.OptimStepInput
    ) -> types.OptimStepOutput | types.ErrorResponse:
        """Process an optim_step request and apply accumulated gradients."""
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        result = self.backend.optim_step(model_id, request_data)

        # One profiler step per optim_step, for the owning model only. Counters are
        # kept in memory and flushed by reconcile_profiler, so profiling adds no DB
        # write to this path.
        if self._profiling_model_id is not None and self._profiling_model_id == model_id:
            self._profiling_steps += 1
            try:
                self._profiling_error = self.backend.profile_step()
            except Exception as e:
                # Never fail a client's optim_step because profiling misbehaved.
                logger.warning(f"[profiler] step failed: {e}")
                self._profiling_error = f"step failed: {e}"

        return result

    def process_forward_backward(self, requests: dict[str, tuple[str, types.ForwardBackwardInput]]) -> dict:
        """Run forward and backward pass on a batch of requests."""
        prepared = prepare_model_pass_batch(requests)
        return self.backend.forward_backward(prepared)

    def process_forward(self, requests: dict[str, tuple[str, types.ForwardBackwardInput]]) -> dict:
        """Run forward-only pass on a batch of requests."""
        prepared = prepare_model_pass_batch(requests)
        return self.backend.forward(prepared)

    def process_sample(self, requests: dict[str, tuple[str, types.SampleInput]]) -> dict:
        """Generate samples for a batch of requests."""
        prepared = prepare_sample_batch(requests, self.config.checkpoints_base)
        return self.backend.sample(prepared)

    def process_load_weights(
        self, model_id: str, request_data: types.LoadWeightsInput
    ) -> types.LoadWeightsOutput | types.ErrorResponse:
        """Loads a clean, trimmed training checkpoint."""
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        checkpoint_path = (
            self.config.checkpoints_base / request_data.source_model_id / f"{request_data.checkpoint_id}.tar.gz"
        )

        self.backend.load_checkpoint(checkpoint_path, model_id, load_optimizer=request_data.load_optimizer)

        return types.LoadWeightsOutput(type="load_weights")

    def process_save_weights(
        self, model_id: str, request_data: types.SaveWeightsInput
    ) -> types.SaveWeightsOutput | types.ErrorResponse:
        """
        Saves a clean training checkpoint by converting the trimmed NNX graph
        to a pure dictionary before serialization, following official Flax docs.
        """
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        checkpoint_id = request_data.path
        output_path = self.config.checkpoints_base / model_id / f"{checkpoint_id}.tar.gz"

        with self._checkpoint_status_context(model_id, checkpoint_id, types.CheckpointType.TRAINING):
            self.backend.save_checkpoint(output_path, model_id)
            logger.info(f"Saved trimmed training checkpoint for model {model_id} to {output_path}")

        return types.SaveWeightsOutput(
            path=f"tinker://{model_id}/weights/{checkpoint_id}",
            type="save_weights",
        )

    def process_save_weights_for_sampler(
        self, model_id: str, request_data: types.SaveWeightsForSamplerInput
    ) -> types.SaveWeightsForSamplerOutput | types.ErrorResponse:
        """Process a save_weights_for_sampler request and save model weights."""
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        # Make sure the user cannot store checkpoints in places like ../../<important file>
        checkpoint_id = Path(request_data.path).name
        output_path = self.config.checkpoints_base / model_id / "sampler_weights" / f"{checkpoint_id}.tar.gz"

        # When the caller provides a sampling_session_seq_id the save is
        # transient — weights only need to reach the inference engines, not
        # disk.  Backends can skip the expensive write in that case.
        persist = request_data.sampling_session_seq_id is None

        with self._checkpoint_status_context(model_id, checkpoint_id, types.CheckpointType.SAMPLER):
            self.backend.save_sampler_checkpoint(output_path, model_id, persist=persist)
            logger.info(f"Saved sampler checkpoint for model {model_id} to {output_path}")

        # Return path=None when using sampling_session_seq_id and seq_id (SDK expects this)
        if request_data.sampling_session_seq_id is not None and request_data.seq_id is not None:
            output_path_str = None
        else:
            output_path_str = f"tinker://{model_id}/{checkpoint_id}"

        return types.SaveWeightsForSamplerOutput(
            path=output_path_str,
            type="save_weights_for_sampler",
            sampling_session_id=request_data.sampling_session_id,
        )

    def _complete_futures(self, results: dict[str, BaseModel]):
        """Helper method to complete multiple futures in the database.

        Args:
            results: Dict mapping request_id to result (Pydantic BaseModel)
        """
        completed_at = datetime.now(timezone.utc)
        params = [
            {
                "request_id": int(request_id),
                # `result_data` holds JSON text, so serialize straight to it:
                # sample results may carry top-k logprobs for every prompt token,
                # and the dict `model_dump()` builds would only exist for
                # `json.dumps` to walk again -- several times the cost of
                # `model_dump_json()`.
                "result_data": result.model_dump_json(),
                "status": RequestStatus.FAILED if isinstance(result, types.ErrorResponse) else RequestStatus.COMPLETED,
                "completed_at": completed_at,
            }
            for request_id, result in results.items()
        ]

        with Session(self.db_engine) as session:
            session.execute(update(FutureDB), params)
            session.commit()

    def process_single_request(self, request_type: types.RequestType, model_id: str, request_data: dict) -> BaseModel:
        match request_type:
            case types.RequestType.CREATE_MODEL:
                return self.process_create_model(model_id, types.CreateModelInput.model_validate(request_data))
            case types.RequestType.OPTIM_STEP:
                return self.process_optim_step(model_id, types.OptimStepInput.model_validate(request_data))
            case types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER:
                return self.process_save_weights_for_sampler(
                    model_id, types.SaveWeightsForSamplerInput.model_validate(request_data)
                )
            case types.RequestType.SAVE_WEIGHTS:
                return self.process_save_weights(model_id, types.SaveWeightsInput.model_validate(request_data))
            case types.RequestType.LOAD_WEIGHTS:
                return self.process_load_weights(model_id, types.LoadWeightsInput.model_validate(request_data))
            case types.RequestType.UNLOAD_MODEL:
                return self.process_unload_model(model_id, types.UnloadModelInput.model_validate(request_data))
            case _:
                raise ValueError(f"Unknown request type: {request_type}")

    def process_single_requests(self, requests: dict[str, tuple[str, types.RequestType, dict]]):
        """Process a collection of single (non-batchable) requests.

        Args:
            requests: Dict mapping request_id to (model_id, request_type, request_data) tuples
        """
        if not requests:
            return
        results = {}
        for request_id, (model_id, request_type, request_data) in requests.items():
            with log_timing(f"process_single_request({request_type.value})"):
                try:
                    result = self.process_single_request(request_type, model_id, request_data)
                except Exception as e:
                    logger.exception(f"Error processing request {request_id}: {e}")
                    result = types.ErrorResponse(error=str(e), status="failed")
            results[request_id] = result
        self._complete_futures(results)

    def process_batch_requests(
        self,
        requests: dict[str, tuple[str, BaseModel]],
        processor: Callable[[dict[str, tuple[str, BaseModel]]], dict[str, BaseModel]],
        name: str,
        per_model: bool = False,
    ):
        """Process a batch of requests with error handling and future completion.

        Args:
            requests: Dict mapping request_id to (model_id, request_data) tuples
            processor: Function that processes requests and returns results dict
            name: Name for logging
            per_model: Process one model's requests at a time, completing each
                model's futures as soon as its sub-batch finishes (GPU execution
                is serialized per model anyway; this only changes completion
                granularity, not batching within a model).
        """
        if not requests:
            return
        error_results, valid_requests = self._filter_valid_requests(requests)
        if error_results:
            self._complete_futures(error_results)
        if per_model:
            grouped: dict[str, dict] = defaultdict(dict)
            for request_id, item in valid_requests.items():
                grouped[item[0]][request_id] = item
            groups = list(grouped.values())
        else:
            groups = [valid_requests] if valid_requests else []
        for group in groups:
            with log_timing(f"process_batch_requests({name}, n={len(group)})"):
                try:
                    results = processor(group)
                except Exception as e:
                    logger.exception(f"Error processing batch: {e}")
                    results = {request_id: types.ErrorResponse(error=str(e), status="failed") for request_id in group}
            self._complete_futures(results)

    def _admit_samples_continuous(self, sample_requests: dict[str, tuple[str, types.SampleInput]]) -> None:
        """Hand pending sample requests to the background sampler.

        Invalid model_ids are failed up front (same semantics as the serial
        path's _filter_valid_requests — e.g. stale requests from a previous
        server must not trigger an engine build). Engines are brought up /
        woken on this (main) thread so the sampler's coroutines never touch
        engine lifecycle. If that fails (e.g. engine build error), the
        candidate futures are failed like a serial batch failure.
        """
        error_results, valid_requests = self._filter_valid_requests(sample_requests)
        if error_results:
            self._complete_futures(error_results)
        if not valid_requests:
            return
        if self._sampler is None:
            self._sampler = _ContinuousSampler(self)
        try:
            self.backend.prepare_for_sampling()
        except Exception as e:  # noqa: BLE001 - match serial batch failure semantics
            logger.exception(f"prepare_for_sampling failed; failing {len(valid_requests)} sample request(s): {e}")
            self._complete_futures({rid: types.ErrorResponse(error=str(e), status="failed") for rid in valid_requests})
            return
        for request_id, (model_id, request_data) in valid_requests.items():
            self._sampler.submit(request_id, model_id, request_data)
        logger.debug(f"Admitted {len(valid_requests)} sample request(s); {self._sampler.inflight_count()} in flight")

    def process_pending_requests_once(self) -> None:
        """One scheduling iteration; see ``process_pending_requests``.

        With continuous sampling, sample requests are admitted to the
        background sampler (per-request future completion) whenever no
        blocking work is pending. Anything that may sleep the engines, swap
        adapters, mutate the model registry, or tear the runtime down —
        model passes, single requests, session cleanup — first drains the
        in-flight samples, preserving the serial loop's safety invariant.
        Samples queued ahead of the earliest blocking request (lower request
        id) are admitted first so the drain covers them and the blocking
        request runs after them, as in the serial loop. Later samples are
        deferred (left in the DB) and admitted on the next iteration instead
        of being processed as a convoy batch.
        """
        # Converge torch profiling to the control row before picking up work,
        # so a session never starts or stops in the middle of a batch.
        self.reconcile_profiler()

        # Query for pending requests and extract data within session context
        with Session(self.db_engine) as session:
            # Use look-ahead scheduling to find batchable forward_backward and forward model passes
            forward_backward_requests = self.find_batchable_model_passes(session, types.RequestType.FORWARD_BACKWARD)
            forward_requests = self.find_batchable_model_passes(session, types.RequestType.FORWARD)
            # Find pending sample requests that can be batched
            sample_requests = self.find_batchable_sample(session)
            # Get other pending requests (non forward_backward and non sampling)
            other_requests = self.find_single_requests(session)

        cleanup_enabled = self.config.session_cleanup_interval_sec >= 0 and self.config.session_timeout_sec >= 0
        cleanup_due = (
            cleanup_enabled and time.time() - self._last_cleanup_time > self.config.session_cleanup_interval_sec
        )

        if self._continuous_sampling:
            blocking_work = bool(forward_backward_requests or forward_requests or other_requests)
            if blocking_work:
                first_blocking_id = min(
                    int(rid) for rid in (*forward_backward_requests, *forward_requests, *other_requests)
                )
                earlier_samples = {rid: req for rid, req in sample_requests.items() if int(rid) < first_blocking_id}
                if earlier_samples:
                    self._admit_samples_continuous(earlier_samples)
                if self._sampler is not None:
                    blockers = [
                        name
                        for name, present in (
                            ("forward_backward", forward_backward_requests),
                            ("forward", forward_requests),
                            ("single requests", other_requests),
                        )
                        if present
                    ]
                    self._sampler.drain(reason=", ".join(blockers))
            elif sample_requests:
                self._admit_samples_continuous(sample_requests)
            # Samples never go through the serial batch path in this mode:
            # deferred requests are re-fetched (still PENDING) next iteration.
            # Session cleanup is not treated as blocking here; it drains
            # internally only when it actually unloads models (the common
            # no-op sweep must not stall sampling every interval).
            sample_requests = {}

        # Process batches outside of session context
        self.process_batch_requests(
            forward_backward_requests, self.process_forward_backward, "forward_backward", per_model=True
        )
        self.process_batch_requests(forward_requests, self.process_forward, "forward", per_model=True)
        self.process_batch_requests(sample_requests, self.process_sample, "sample")

        # Process other request types individually (in the future we can also batch independent optim_steps)
        self.process_single_requests(other_requests)

        # Periodically cleanup stale sessions (disabled if either config is negative)
        if cleanup_due:
            _ = self.cleanup_stale_sessions()
            self._last_cleanup_time = time.time()

    def process_pending_requests(self):
        """Main loop to process pending requests."""
        while True:
            self.process_pending_requests_once()
            # Poll every 100ms
            time.sleep(0.1)

    def run(self):
        """Entry point to start the engine."""
        logger.info("Starting background engine...")
        self.process_pending_requests()

    def find_batchable_model_passes(
        self, session: Session, request_type: types.RequestType
    ) -> dict[str, tuple[str, types.ForwardBackwardInput]]:
        """Find all requests of the given type that come before any destructive update for their model.

        Uses look-ahead scheduling: for each model, only returns operations
        that have no optim_step or load_weights blocking them in the queue.

        Args:
            session: Database session
            request_type: The type of request to find (e.g., FORWARD or FORWARD_BACKWARD)

        Returns:
            Dict mapping request_id to (model_id, request_data) tuples
        """
        barriers = self._find_destructive_barriers(session)

        # Get all pending operations of the requested type ordered by request_id
        query = (
            select(FutureDB.request_id, FutureDB.model_id)
            .where(FutureDB.request_type == request_type)
            .where(FutureDB.status == RequestStatus.PENDING)
            .order_by(FutureDB.request_id)
        )
        ops = session.exec(query).all()

        # Filter: only include ops that come before their model's barrier
        batchable = [
            (request_id, model_id)
            for request_id, model_id in ops
            if model_id not in barriers or request_id < barriers[model_id]
        ]

        return {
            str(request_id): (model_id, types.ForwardBackwardInput.model_validate(request_data))
            for request_id, model_id, request_data in self._load_requests(session, batchable)
        }

    def find_batchable_sample(self, session: Session) -> dict[str, tuple[str, types.SampleInput]]:
        """Find all sample ops that can be safely batched together.

        Returns sample operations ensuring that each model_id has only one checkpoint_id
        to avoid loading different checkpoints for the same model in a single batch.

        If sample_max_num_sequences is configured, limits to that many requests so we don't
        produce partial batches in process_sample_batch. If num_samples > 1 for some requests,
        this may not be perfect, but it's good until we implement continuous batching.

        Args:
            session: Database session

        Returns:
            Dict mapping request_id to (model_id, request_data) tuples
        """
        sampler = getattr(self, "_sampler", None)
        inflight_ids = sampler.inflight_ids if sampler is not None else set()

        # checkpoint_id is extracted in the database so prompts stay out of this query
        sample_query = (
            select(FutureDB.request_id, FutureDB.model_id, FutureDB.request_data["checkpoint_id"].as_string())
            .where(FutureDB.request_type == types.RequestType.SAMPLE)
            .where(FutureDB.status == RequestStatus.PENDING)
            .order_by(FutureDB.request_id)
        )
        sample_ops = session.exec(sample_query).all()

        batchable = []
        model_checkpoints = {}  # Map from model_id to checkpoint_id of first request to that model
        for request_id, model_id, checkpoint_id in sample_ops:
            if str(request_id) in inflight_ids:
                continue
            # Base model requests (empty checkpoint_id) are always compatible, otherwise only
            # take only requests with one checkpoint_id for a given model_id
            if not checkpoint_id or model_checkpoints.setdefault(model_id, checkpoint_id) == checkpoint_id:
                batchable.append((request_id, model_id))

        # TODO: This leaks the abstraction by accessing backend-specific config.
        # We should find a better way to handle this going forward.
        if self.config.backend == "jax" and self.backend.config.sample_max_num_sequences > 0:
            batchable = batchable[: self.backend.config.sample_max_num_sequences]

        return {
            str(request_id): (model_id, types.SampleInput.model_validate(request_data))
            for request_id, model_id, request_data in self._load_requests(session, batchable)
        }

    def find_single_requests(self, session: Session) -> dict[str, tuple[str, types.RequestType, dict]]:
        """Find all requests that need to be processed individually (not batchable).

        Args:
            session: Database session

        Returns:
            Dict mapping request_id to (model_id, request_type, request_data) tuples
        """
        # Find the first blocked forward pass per model: a pending FORWARD/FORWARD_BACKWARD
        # that sits behind a destructive barrier and won't be batched this iteration.
        # Single requests must not jump ahead of these.
        destructive_barriers = self._find_destructive_barriers(session)
        blocked_pass_barriers: dict[str, int] = {}
        if destructive_barriers:
            pending_passes = session.exec(
                select(FutureDB.model_id, FutureDB.request_id)
                .where(
                    (FutureDB.request_type == types.RequestType.FORWARD_BACKWARD)
                    | (FutureDB.request_type == types.RequestType.FORWARD)
                )
                .where(FutureDB.status == RequestStatus.PENDING)
                .where(FutureDB.model_id.in_(destructive_barriers.keys()))
                .order_by(FutureDB.request_id)
            ).all()
            for model_id, req_id in pending_passes:
                if req_id >= destructive_barriers[model_id]:
                    blocked_pass_barriers.setdefault(model_id, req_id)

        statement = (
            select(FutureDB.request_id, FutureDB.model_id, FutureDB.request_type)
            .where(FutureDB.status == RequestStatus.PENDING)
            .where(FutureDB.request_type != types.RequestType.FORWARD_BACKWARD)
            .where(FutureDB.request_type != types.RequestType.FORWARD)
            .where(FutureDB.request_type != types.RequestType.SAMPLE)
            .where(FutureDB.request_type != types.RequestType.EXTERNAL)
            .order_by(FutureDB.request_id)
        )
        other_futures = session.exec(statement).all()

        # Filter: only include ops that come before the first blocked pass for their model
        other_futures = [
            (request_id, model_id, request_type)
            for request_id, model_id, request_type in other_futures
            if model_id not in blocked_pass_barriers or request_id < blocked_pass_barriers[model_id]
        ]

        return {
            str(request_id): (model_id, request_type, request_data)
            for request_id, model_id, request_type, request_data in self._load_requests(session, other_futures)
        }

    def process_create_model(self, model_id: str, request_data: types.CreateModelInput) -> types.CreateModelOutput:
        """Create and initialize a model."""
        # Create model in backend (allocates adapter_index, creates optimizer, and configures adapter)
        self.backend.create_model(model_id, request_data.lora_config, model_role=request_data.model_role)

        logger.info(f"Created LoRA model {model_id}")

        return types.CreateModelOutput(
            model_id=model_id,
            base_model=self.config.base_model,
            lora_config=request_data.lora_config,
        )

    def process_unload_model(self, model_id: str, request_data: types.UnloadModelInput) -> types.UnloadModelOutput:
        """Unload a model and free all resources."""
        if not self.backend.has_model(model_id):
            logger.warning(f"Ignoring unload request for model {model_id} that is not loaded.")
        else:
            self.backend.delete_model(model_id)

            # Update model status in DB
            with Session(self.db_engine) as session:
                _ = session.exec(update(ModelDB).where(ModelDB.model_id == model_id).values(status="unloaded"))
                session.commit()

            logger.info(f"Unloaded model {model_id}")

        return types.UnloadModelOutput(model_id=model_id, status="unloaded")

    def cleanup_stale_sessions(self) -> int:
        """Cleanup sessions with no recent heartbeat and unload their models.

        Returns:
            Number of models unloaded
        """
        cutoff = datetime.now(timezone.utc) - timedelta(seconds=self.config.session_timeout_sec)
        unloaded_count = 0

        with Session(self.db_engine) as session:
            # Find stale sessions (active sessions with heartbeat older than cutoff)
            stale_sessions = session.exec(
                select(SessionDB).where(
                    SessionDB.status == "active",
                    SessionDB.last_heartbeat_at < cutoff,
                )
            ).all()

            if not stale_sessions:
                return 0

            stale_session_ids = {s.session_id for s in stale_sessions}

            # Find all models for all stale sessions in one query
            models_to_process = session.exec(
                select(ModelDB).where(
                    ModelDB.session_id.in_(stale_session_ids),
                    ModelDB.status != "unloaded",
                )
            ).all()

        # Unloading can tear the shared runtime (engines included) down when
        # the last model goes; in-flight continuous samples must finish first.
        sampler = getattr(self, "_sampler", None)
        if models_to_process and sampler is not None:
            sampler.drain(reason="stale-session model unload")

        # Unload models outside DB transactions to minimize lock time.
        sessions_with_failed_unloads: set[str] = set()
        unloaded_model_ids: set[str] = set()
        for model in models_to_process:
            if self.backend.has_model(model.model_id):
                try:
                    self.backend.delete_model(model.model_id)
                    unloaded_model_ids.add(model.model_id)
                    unloaded_count += 1
                    logger.info(f"Auto-unloaded stale model {model.model_id} from session {model.session_id}")
                except Exception as e:
                    logger.error(f"Failed to auto-unload model {model.model_id}: {e}")
                    sessions_with_failed_unloads.add(model.session_id)
            else:
                # Model already missing in backend; only DB state needs cleanup.
                unloaded_model_ids.add(model.model_id)

        sessions_to_expire = [s.session_id for s in stale_sessions if s.session_id not in sessions_with_failed_unloads]

        # Apply DB status updates in one short write transaction.
        with Session(self.db_engine) as session:
            if unloaded_model_ids:
                _ = session.exec(
                    update(ModelDB).where(ModelDB.model_id.in_(unloaded_model_ids)).values(status="unloaded")
                )
            if sessions_to_expire:
                _ = session.exec(
                    update(SessionDB).where(SessionDB.session_id.in_(sessions_to_expire)).values(status="expired")
                )
            session.commit()

        for session_id in sessions_to_expire:
            logger.info(f"Expired stale session {session_id}")

        return unloaded_count

    def reconcile_profiler(self) -> None:
        """Converge the workers to the desired state in the control row.

        Runs at the top of the request loop, which is a boundary between request
        batches, so a session never starts or stops mid-batch.
        """
        if self._profiler_cfg is None:
            return
        with Session(self.db_engine) as session:
            row = session.get(ProfilerControlDB, 1)
            if row is None:
                return
            desired, version, applied = row.desired_state, row.version, row.applied_version
            owner, config_json, started_at = row.owner_model_id, row.config_json, row.started_at

        if applied < version:
            # A new claim or release to apply.
            if desired == ProfilerState.RUNNING:
                try:
                    try:
                        worker_config = json.loads(config_json)
                    except (TypeError, ValueError) as e:
                        raise ValueError(f"unreadable profiler config in control row: {e}") from e
                    self.backend.start_profile(worker_config)
                except Exception as e:
                    logger.warning(f"[profiler] start failed: {e}")
                    # Release the slot, but still ack: the API is blocked waiting
                    # for applied_version and would otherwise hang to its timeout.
                    self._profiling_model_id = None
                    self._write_profiler_row(
                        desired_state=ProfilerState.STOPPED,
                        owner_model_id=None,
                        applied_version=version,
                        error=f"start failed: {e}",
                    )
                    return
                self._profiling_model_id = owner
                self._profiling_steps = 0
                self._profiling_error = None
                self._write_profiler_row(applied_version=version, step=0, error=None)
                logger.info(f"[profiler] session started for model {owner}")
            else:
                error = self._stop_profiling()
                self._write_profiler_row(applied_version=version, step=self._profiling_steps, error=error)
                logger.info("[profiler] session stopped")
            return

        if self._profiling_model_id is None:
            return

        # Active session: enforce the TTL, else flush progress for /profiling_status.
        if started_at is not None:
            # SQLite hands back naive datetimes even for a timezone-aware column,
            # and everything we store is UTC.
            if started_at.tzinfo is None:
                started_at = started_at.replace(tzinfo=timezone.utc)
            age = (datetime.now(timezone.utc) - started_at).total_seconds()
            if age > self._profiler_cfg.max_session_duration_sec:
                logger.warning(
                    f"[profiler] session exceeded max_session_duration_sec="
                    f"{self._profiler_cfg.max_session_duration_sec}s; finalizing"
                )
                self._stop_profiling()
                self._write_profiler_row(
                    desired_state=ProfilerState.STOPPED,
                    owner_model_id=None,
                    step=self._profiling_steps,
                    error="session terminated by max_session_duration_sec",
                )
                return
        self._write_profiler_row(step=self._profiling_steps, error=self._profiling_error)

    def process_optim_step(
        self, model_id: str, request_data: types.OptimStepInput
    ) -> types.OptimStepOutput | types.ErrorResponse:
        """Process an optim_step request and apply accumulated gradients."""
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        result = self.backend.optim_step(model_id, request_data)

        # One profiler step per optim_step, for the owning model only. Counters are
        # kept in memory and flushed by reconcile_profiler, so profiling adds no DB
        # write to this path.
        if self._profiling_model_id is not None and self._profiling_model_id == model_id:
            self._profiling_steps += 1
            try:
                self._profiling_error = self.backend.profile_step()
            except Exception as e:
                # Never fail a client's optim_step because profiling misbehaved.
                logger.warning(f"[profiler] step failed: {e}")
                self._profiling_error = f"step failed: {e}"

        return result

    def process_forward_backward(self, requests: dict[str, tuple[str, types.ForwardBackwardInput]]) -> dict:
        """Run forward and backward pass on a batch of requests."""
        prepared = prepare_model_pass_batch(requests)
        return self.backend.forward_backward(prepared)

    def process_forward(self, requests: dict[str, tuple[str, types.ForwardBackwardInput]]) -> dict:
        """Run forward-only pass on a batch of requests."""
        prepared = prepare_model_pass_batch(requests)
        return self.backend.forward(prepared)

    def process_sample(self, requests: dict[str, tuple[str, types.SampleInput]]) -> dict:
        """Generate samples for a batch of requests."""
        prepared = prepare_sample_batch(requests, self.config.checkpoints_base)
        return self.backend.sample(prepared)

    def process_load_weights(
        self, model_id: str, request_data: types.LoadWeightsInput
    ) -> types.LoadWeightsOutput | types.ErrorResponse:
        """Loads a clean, trimmed training checkpoint."""
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        checkpoint_path = (
            self.config.checkpoints_base / request_data.source_model_id / f"{request_data.checkpoint_id}.tar.gz"
        )

        self.backend.load_checkpoint(checkpoint_path, model_id, load_optimizer=request_data.load_optimizer)

        return types.LoadWeightsOutput(type="load_weights")

    def process_save_weights(
        self, model_id: str, request_data: types.SaveWeightsInput
    ) -> types.SaveWeightsOutput | types.ErrorResponse:
        """
        Saves a clean training checkpoint by converting the trimmed NNX graph
        to a pure dictionary before serialization, following official Flax docs.
        """
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        checkpoint_id = request_data.path
        output_path = self.config.checkpoints_base / model_id / f"{checkpoint_id}.tar.gz"

        with self._checkpoint_status_context(model_id, checkpoint_id, types.CheckpointType.TRAINING):
            self.backend.save_checkpoint(output_path, model_id)
            logger.info(f"Saved trimmed training checkpoint for model {model_id} to {output_path}")

        return types.SaveWeightsOutput(
            path=f"tinker://{model_id}/weights/{checkpoint_id}",
            type="save_weights",
        )

    def process_save_weights_for_sampler(
        self, model_id: str, request_data: types.SaveWeightsForSamplerInput
    ) -> types.SaveWeightsForSamplerOutput | types.ErrorResponse:
        """Process a save_weights_for_sampler request and save model weights."""
        if not self.backend.has_model(model_id):
            return _model_not_found_error(model_id)

        # Make sure the user cannot store checkpoints in places like ../../<important file>
        checkpoint_id = Path(request_data.path).name
        output_path = self.config.checkpoints_base / model_id / "sampler_weights" / f"{checkpoint_id}.tar.gz"

        # When the caller provides a sampling_session_seq_id the save is
        # transient — weights only need to reach the inference engines, not
        # disk.  Backends can skip the expensive write in that case.
        persist = request_data.sampling_session_seq_id is None

        with self._checkpoint_status_context(model_id, checkpoint_id, types.CheckpointType.SAMPLER):
            self.backend.save_sampler_checkpoint(output_path, model_id, persist=persist)
            logger.info(f"Saved sampler checkpoint for model {model_id} to {output_path}")

        # Return path=None when using sampling_session_seq_id and seq_id (SDK expects this)
        if request_data.sampling_session_seq_id is not None and request_data.seq_id is not None:
            output_path_str = None
        else:
            output_path_str = f"tinker://{model_id}/{checkpoint_id}"

        return types.SaveWeightsForSamplerOutput(
            path=output_path_str,
            type="save_weights_for_sampler",
            sampling_session_id=request_data.sampling_session_id,
        )

    def process_single_request(self, request_type: types.RequestType, model_id: str, request_data: dict) -> BaseModel:
        match request_type:
            case types.RequestType.CREATE_MODEL:
                return self.process_create_model(model_id, types.CreateModelInput.model_validate(request_data))
            case types.RequestType.OPTIM_STEP:
                return self.process_optim_step(model_id, types.OptimStepInput.model_validate(request_data))
            case types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER:
                return self.process_save_weights_for_sampler(
                    model_id, types.SaveWeightsForSamplerInput.model_validate(request_data)
                )
            case types.RequestType.SAVE_WEIGHTS:
                return self.process_save_weights(model_id, types.SaveWeightsInput.model_validate(request_data))
            case types.RequestType.LOAD_WEIGHTS:
                return self.process_load_weights(model_id, types.LoadWeightsInput.model_validate(request_data))
            case types.RequestType.UNLOAD_MODEL:
                return self.process_unload_model(model_id, types.UnloadModelInput.model_validate(request_data))
            case _:
                raise ValueError(f"Unknown request type: {request_type}")

    def process_single_requests(self, requests: dict[str, tuple[str, types.RequestType, dict]]):
        """Process a collection of single (non-batchable) requests.

        Args:
            requests: Dict mapping request_id to (model_id, request_type, request_data) tuples
        """
        if not requests:
            return
        results = {}
        for request_id, (model_id, request_type, request_data) in requests.items():
            with log_timing(f"process_single_request({request_type.value})"):
                try:
                    result = self.process_single_request(request_type, model_id, request_data)
                except Exception as e:
                    logger.exception(f"Error processing request {request_id}: {e}")
                    result = types.ErrorResponse(error=str(e), status="failed")
            results[request_id] = result
        self._complete_futures(results)

    def process_batch_requests(
        self,
        requests: dict[str, tuple[str, BaseModel]],
        processor: Callable[[dict[str, tuple[str, BaseModel]]], dict[str, BaseModel]],
        name: str,
        per_model: bool = False,
    ):
        """Process a batch of requests with error handling and future completion.

        Args:
            requests: Dict mapping request_id to (model_id, request_data) tuples
            processor: Function that processes requests and returns results dict
            name: Name for logging
            per_model: Process one model's requests at a time, completing each
                model's futures as soon as its sub-batch finishes (GPU execution
                is serialized per model anyway; this only changes completion
                granularity, not batching within a model).
        """
        if not requests:
            return
        error_results, valid_requests = self._filter_valid_requests(requests)
        if error_results:
            self._complete_futures(error_results)
        if per_model:
            grouped: dict[str, dict] = defaultdict(dict)
            for request_id, item in valid_requests.items():
                grouped[item[0]][request_id] = item
            groups = list(grouped.values())
        else:
            groups = [valid_requests] if valid_requests else []
        for group in groups:
            with log_timing(f"process_batch_requests({name}, n={len(group)})"):
                try:
                    results = processor(group)
                except Exception as e:
                    logger.exception(f"Error processing batch: {e}")
                    results = {request_id: types.ErrorResponse(error=str(e), status="failed") for request_id in group}
            self._complete_futures(results)

    def process_pending_requests_once(self) -> None:
        """One scheduling iteration; see ``process_pending_requests``.

        With continuous sampling, sample requests are admitted to the
        background sampler (per-request future completion) whenever no
        blocking work is pending. Anything that may sleep the engines, swap
        adapters, mutate the model registry, or tear the runtime down —
        model passes, single requests, session cleanup — first drains the
        in-flight samples, preserving the serial loop's safety invariant.
        Samples queued ahead of the earliest blocking request (lower request
        id) are admitted first so the drain covers them and the blocking
        request runs after them, as in the serial loop. Later samples are
        deferred (left in the DB) and admitted on the next iteration instead
        of being processed as a convoy batch.
        """
        # Converge torch profiling to the control row before picking up work,
        # so a session never starts or stops in the middle of a batch.
        self.reconcile_profiler()

        # Query for pending requests and extract data within session context
        with Session(self.db_engine) as session:
            # Use look-ahead scheduling to find batchable forward_backward and forward model passes
            forward_backward_requests = self.find_batchable_model_passes(session, types.RequestType.FORWARD_BACKWARD)
            forward_requests = self.find_batchable_model_passes(session, types.RequestType.FORWARD)
            # Find pending sample requests that can be batched
            sample_requests = self.find_batchable_sample(session)
            # Get other pending requests (non forward_backward and non sampling)
            other_requests = self.find_single_requests(session)

        cleanup_enabled = self.config.session_cleanup_interval_sec >= 0 and self.config.session_timeout_sec >= 0
        cleanup_due = (
            cleanup_enabled and time.time() - self._last_cleanup_time > self.config.session_cleanup_interval_sec
        )

        if self._continuous_sampling:
            blocking_work = bool(forward_backward_requests or forward_requests or other_requests)
            if blocking_work:
                first_blocking_id = min(
                    int(rid) for rid in (*forward_backward_requests, *forward_requests, *other_requests)
                )
                earlier_samples = {rid: req for rid, req in sample_requests.items() if int(rid) < first_blocking_id}
                if earlier_samples:
                    self._admit_samples_continuous(earlier_samples)
                if self._sampler is not None:
                    blockers = [
                        name
                        for name, present in (
                            ("forward_backward", forward_backward_requests),
                            ("forward", forward_requests),
                            ("single requests", other_requests),
                        )
                        if present
                    ]
                    self._sampler.drain(reason=", ".join(blockers))
            elif sample_requests:
                self._admit_samples_continuous(sample_requests)
            # Samples never go through the serial batch path in this mode:
            # deferred requests are re-fetched (still PENDING) next iteration.
            # Session cleanup is not treated as blocking here; it drains
            # internally only when it actually unloads models (the common
            # no-op sweep must not stall sampling every interval).
            sample_requests = {}

        # Process batches outside of session context
        self.process_batch_requests(
            forward_backward_requests, self.process_forward_backward, "forward_backward", per_model=True
        )
        self.process_batch_requests(forward_requests, self.process_forward, "forward", per_model=True)
        self.process_batch_requests(sample_requests, self.process_sample, "sample")

        # Process other request types individually (in the future we can also batch independent optim_steps)
        self.process_single_requests(other_requests)

        # Periodically cleanup stale sessions (disabled if either config is negative)
        if cleanup_due:
            _ = self.cleanup_stale_sessions()
            self._last_cleanup_time = time.time()

    def process_pending_requests(self):
        """Main loop to process pending requests."""
        while True:
            self.process_pending_requests_once()
            # Poll every 100ms
            time.sleep(0.1)

    def run(self):
        """Entry point to start the engine."""
        logger.info("Starting background engine...")
        self.process_pending_requests()