bootstrap_space_runtime.py

from __future__ import annotations

import os
import shutil
import sys
from pathlib import Path

DEFAULT_ENV = {
    "REMOTE_BRAIN_ONLY": "1",
    "KAGGLE_AUTO_BOOTSTRAP": "0",
    "BRAIN_AUTO_NGROK": "0",
    "BRAIN_AUTO_PUBLISH_URL_ON_STARTUP": "0",
    "BRAIN_REUSE_PUBLIC_URL_ON_RESTART": "0",
    "HF_SPACE_DOCKER": "1",
    "KAPO_COMPUTE_PROFILE": "cpu",
    "HF_ACCELERATOR": "cpu",
    "KAPO_HF_TRANSFORMERS_RUNTIME": "1",
    "KAPO_LAZY_MODEL_STARTUP": "1",
    "KAPO_LAZY_EMBED_STARTUP": "1",
    "MODEL_PROFILE_ID": "hf-debugger-qwen25-7b-instruct",
    "MODEL_REPO": "Qwen/Qwen2.5-1.5B-Instruct",
    "BRAIN_ROLES": "fallback",
    "BRAIN_LANGUAGES": "ar,en",
    "BRAIN_PLATFORM_NAME": "hf_debugger_main",
    "BRAIN_TEMPLATE": "hf-space-cpu",
    "BRAIN_PROVIDER": "huggingface",
    "FIREBASE_ENABLED": "1",
    "FIREBASE_PROJECT_ID": "citadel4travels",
    "FIREBASE_NAMESPACE": "kapo",
    "BRAIN_REUSE_PUBLIC_URL_ON_RESTART": "0",
}


def _copy_tree(source: Path, target: Path) -> None:
    if target.exists():
        shutil.rmtree(target, ignore_errors=True)
    shutil.copytree(
        source,
        target,
        ignore=shutil.ignore_patterns('__pycache__', '*.pyc', '.git', '.venv'),
    )


def _merge_overlay(overlay_root: Path, runtime_root: Path) -> None:
    if not overlay_root.exists():
        return
    for path in sorted(overlay_root.rglob('*')):
        rel = path.relative_to(overlay_root)
        dst = runtime_root / rel
        if path.is_dir():
            dst.mkdir(parents=True, exist_ok=True)
            continue
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, dst)


def main() -> None:
    source_root = Path(os.getenv('KAPO_SPACE_SOURCE_ROOT', '/app')).resolve()
    default_root = Path('/data/kapo_runtime/current') if Path('/data').exists() else Path('/tmp/kapo_runtime/current')
    runtime_root = Path(os.getenv('KAPO_RUNTIME_ROOT', str(default_root))).resolve()
    overlay_root = Path(os.getenv('KAPO_OVERLAY_ROOT', str(runtime_root.parent / 'overlay'))).resolve()
    runtime_pkg = runtime_root / 'brain_server'
    source_pkg = source_root / 'brain_server'
    for key, value in DEFAULT_ENV.items():
        os.environ.setdefault(str(key), str(value))
    space_host = str(os.getenv('SPACE_HOST', '')).strip().rstrip('/')
    if space_host:
        public_url = space_host if space_host.startswith('http://') or space_host.startswith('https://') else f'https://{space_host}'
        os.environ['BRAIN_PUBLIC_URL'] = public_url.rstrip('/')
    runtime_root.mkdir(parents=True, exist_ok=True)
    overlay_root.mkdir(parents=True, exist_ok=True)
    if not runtime_pkg.exists():
        _copy_tree(source_pkg, runtime_pkg)
    _merge_overlay(overlay_root, runtime_root)
    os.environ['KAPO_RUNTIME_ROOT'] = str(runtime_root)
    os.environ['KAPO_SYNC_ROOT'] = str(runtime_root)
    os.environ['KAPO_OVERLAY_ROOT'] = str(overlay_root)
    port = str(os.getenv('PORT', '7860') or '7860')
    os.execvp(
        sys.executable,
        [
            sys.executable,
            '-m',
            'uvicorn',
            'api.main:app',
            '--host',
            '0.0.0.0',
            '--port',
            port,
            '--app-dir',
            str(runtime_pkg),
        ],
    )


if __name__ == '__main__':
    main()