File size: 6,957 Bytes
38c9982
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93e8468
38c9982
 
 
 
 
 
93e8468
38c9982
 
 
 
 
 
 
 
 
 
 
 
 
93e8468
 
 
 
 
 
 
 
 
 
38c9982
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93e8468
 
 
38c9982
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
from __future__ import annotations

import argparse
import os
import shutil
import sys
import tempfile
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.executive_assistant.deployment import (
    DEFAULT_SPACE_TITLE,
    HFSpaceDeployConfig,
    parse_hf_usernames,
    stage_space_bundle,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Create or update a Hugging Face Space from this repository in one command."
    )
    parser.add_argument(
        "--repo-id",
        default=os.environ.get("HF_SPACE_REPO", "").strip(),
        help="Target Space repo in owner/name form. Defaults to HF_SPACE_REPO.",
    )
    parser.add_argument(
        "--token",
        default=os.environ.get("HF_TOKEN", "").strip(),
        help="Hugging Face token. Defaults to HF_TOKEN.",
    )
    parser.add_argument(
        "--title",
        default=os.environ.get("HF_SPACE_TITLE", DEFAULT_SPACE_TITLE),
        help="Space title used in the generated HF README.",
    )
    parser.add_argument(
        "--team-name",
        default=os.environ.get("HF_SPACE_TEAM_NAME", "Team Epsilon"),
        help="Team name shown in the generated HF README.",
    )
    parser.add_argument(
        "--hf-usernames",
        default=os.environ.get(
            "HF_SPACE_TEAM_USERNAMES",
            "flickinshots,ShreyaKhatik,itsayushdey",
        ),
        help="Comma-separated HF usernames for the HF README placeholders.",
    )
    parser.add_argument(
        "--checkpoint-name",
        default=os.environ.get("HF_SPACE_CHECKPOINT_NAME", "q_policy_notebook.json"),
        help="Checkpoint filename staged into artifacts/checkpoints/ for RL replay.",
    )
    parser.add_argument(
        "--openrouter-api-key",
        default=os.environ.get("OPENROUTER_API_KEY", "").strip(),
        help="Optional secret to set on the Space during deployment.",
    )
    parser.add_argument(
        "--api-base-url",
        default=os.environ.get("API_BASE_URL", "https://openrouter.ai/api/v1").strip(),
        help="OpenAI-compatible API base URL for inference.py. Defaults to OpenRouter.",
    )
    parser.add_argument(
        "--model-name",
        default=os.environ.get("MODEL_NAME", "google/gemma-4-31b-it").strip(),
        help="OpenRouter model id for inference.py. Defaults to Gemma 4.",
    )
    parser.add_argument(
        "--private",
        action="store_true",
        default=os.environ.get("HF_SPACE_PRIVATE", "").strip().lower() == "true",
        help="Create or keep the Space private.",
    )
    parser.add_argument(
        "--skip-checkpoint",
        action="store_true",
        help="Skip bundling the RL checkpoint.",
    )
    parser.add_argument(
        "--keep-stage-dir",
        default="",
        help="Optional local folder where the prepared Space bundle should be copied after upload.",
    )
    return parser


def require_huggingface_hub():
    try:
        from huggingface_hub import HfApi  # type: ignore
    except ImportError as exc:
        raise SystemExit(
            "huggingface_hub is required for deployment. Install the training environment "
            "or run `python -m pip install huggingface_hub` first."
        ) from exc
    return HfApi


def maybe_set_space_secret(api, repo_id: str, key: str, value: str) -> str:
    if not value.strip():
        return f"Skipped secret {key} because no value was provided."
    add_secret = getattr(api, "add_space_secret", None)
    if add_secret is None:
        return f"Upload succeeded, but this huggingface_hub version cannot set {key} automatically."
    add_secret(repo_id=repo_id, key=key, value=value)
    return f"Set Space secret {key}."


def maybe_set_space_variable(api, repo_id: str, key: str, value: str) -> str:
    add_variable = getattr(api, "add_space_variable", None)
    if add_variable is None:
        return f"Upload succeeded, but this huggingface_hub version cannot set variable {key} automatically."
    add_variable(repo_id=repo_id, key=key, value=value)
    return f"Set Space variable {key}={value}."


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if not args.repo_id:
        parser.error("A Space repo id is required. Pass --repo-id or set HF_SPACE_REPO.")
    if "/" not in args.repo_id:
        parser.error("Space repo id must be in owner/name form.")
    if not args.token:
        parser.error("A Hugging Face token is required. Pass --token or set HF_TOKEN.")

    config = HFSpaceDeployConfig(
        repo_id=args.repo_id,
        title=args.title,
        team_name=args.team_name,
        hf_usernames=parse_hf_usernames(args.hf_usernames),
        checkpoint_name=args.checkpoint_name,
        private=args.private,
        include_checkpoint=not args.skip_checkpoint,
    )

    HfApi = require_huggingface_hub()
    api = HfApi(token=args.token)

    with tempfile.TemporaryDirectory(prefix="hf-space-stage-") as tmp_dir:
        stage_dir = Path(tmp_dir)
        checkpoint_path = stage_space_bundle(config, stage_dir)

        api.create_repo(
            repo_id=config.repo_id,
            repo_type="space",
            space_sdk="docker",
            private=config.private,
            exist_ok=True,
        )
        api.upload_folder(
            folder_path=str(stage_dir),
            repo_id=config.repo_id,
            repo_type="space",
            commit_message="Deploy Project Epsilon Space bundle",
            delete_patterns=["*", "**/*"],
        )

        messages = [
            f"Uploaded Space bundle to {config.space_url}",
            f"App URL: {config.app_url}",
        ]
        if checkpoint_path is not None:
            messages.append(f"Bundled RL checkpoint: {checkpoint_path.relative_to(stage_dir)}")
        messages.append(maybe_set_space_secret(api, config.repo_id, "OPENROUTER_API_KEY", args.openrouter_api_key))
        messages.append(maybe_set_space_secret(api, config.repo_id, "OPENAI_API_KEY", args.openrouter_api_key))
        messages.append(maybe_set_space_variable(api, config.repo_id, "API_BASE_URL", args.api_base_url))
        messages.append(maybe_set_space_variable(api, config.repo_id, "MODEL_NAME", args.model_name))
        messages.append(maybe_set_space_variable(api, config.repo_id, "OPENROUTER_APP_NAME", config.title))
        messages.append(maybe_set_space_variable(api, config.repo_id, "OPENROUTER_SITE_URL", config.app_url))

        if args.keep_stage_dir:
            target_dir = Path(args.keep_stage_dir).resolve()
            if target_dir.exists():
                shutil.rmtree(target_dir)
            shutil.copytree(stage_dir, target_dir)
            messages.append(f"Saved staged bundle to {target_dir}")

    for message in messages:
        print(message)
    return 0


if __name__ == "__main__":
    sys.exit(main())