src/infer_runtime/settings.py

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path

from infer_runtime.checkpoints import resolve_checkpoint_layout


@dataclass
class InferSettings:
    config_path: str
    ckpt_path: str
    rewrite_model: str
    openai_api_key: str | None
    openai_base_url: str | None
    default_seed: int


def _materialize_config_if_needed(config_path: Path, ckpt_root: Path) -> Path:
    """Make a real config file inside ckpt_root when the source config is a symlink.

    Some Hugging Face cache layouts expose snapshot files as symlinks into a shared
    blobs directory. The upstream infer_config.py uses Path(__file__).resolve().parent
    to locate the checkpoint root, so loading the symlink directly makes it resolve to
    the blobs directory instead of the snapshot root. Materializing a real file inside
    the checkpoint root preserves the intended relative layout.
    """
    if not config_path.exists():
        raise FileNotFoundError(f"Configuration file not found: {config_path}")

    if not config_path.is_symlink():
        return config_path

    materialized = ckpt_root / '_space_runtime_infer_config.py'
    print(f"[Model] Materializing symlinked config for Spaces: {config_path} -> {materialized}")
    source_text = config_path.read_text(encoding='utf-8')
    current_text = materialized.read_text(encoding='utf-8') if materialized.exists() else None
    if current_text != source_text:
        tmp_path = materialized.with_suffix('.tmp')
        tmp_path.write_text(source_text, encoding='utf-8')
        tmp_path.replace(materialized)
    return materialized


def load_settings(
    *,
    ckpt_root: str,
    config_path: str | None = None,
    rewrite_model: str | None = None,
    default_seed: int = 42,
) -> InferSettings:
    layout = resolve_checkpoint_layout(ckpt_root)
    default_config = layout.root / 'infer_config.py'
    if config_path is None and not default_config.exists():
        raise FileNotFoundError(
            f"Missing inference config: {default_config}. Pass --config explicitly to choose a config file."
        )

    chosen_config = Path(config_path).expanduser() if config_path is not None else default_config
    chosen_config = _materialize_config_if_needed(chosen_config, layout.root)

    return InferSettings(
        config_path=str(chosen_config),
        ckpt_path=str(layout.transformer_ckpt),
        rewrite_model=rewrite_model or 'gpt-5',
        openai_api_key=os.environ.get('OPENAI_API_KEY'),
        openai_base_url=os.environ.get('OPENAI_BASE_URL'),
        default_seed=default_seed,
    )