openenv/atari_env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""Fork (duplicate) a Hugging Face Space using the Hub API."""8 9from __future__ import annotations10 11from typing import Annotated12 13import typer14from huggingface_hub import HfApi, login, whoami15 16from .._cli_utils import console17 18app = typer.Typer(19    help="Fork (duplicate) an OpenEnv environment on Hugging Face to your account"20)21 22 23def _parse_key_value(s: str) -> tuple[str, str]:24    """Parse KEY=VALUE string. Raises BadParameter if no '='."""25    if "=" not in s:26        raise typer.BadParameter(27            f"Expected KEY=VALUE format, got: {s!r}. "28            "Use --set-env KEY=VALUE or --set-secret KEY=VALUE"29        )30    key, _, value = s.partition("=")31    key = key.strip()32    if not key:33        raise typer.BadParameter(f"Empty key in: {s!r}")34    return key, value.strip()35 36 37def _ensure_hf_authenticated() -> str:38    """Ensure user is authenticated with Hugging Face. Returns username."""39    try:40        user_info = whoami()41        if isinstance(user_info, dict):42            username = (43                user_info.get("name")44                or user_info.get("fullname")45                or user_info.get("username")46            )47        else:48            username = (49                getattr(user_info, "name", None)50                or getattr(user_info, "fullname", None)51                or getattr(user_info, "username", None)52            )53        if not username:54            raise ValueError("Could not extract username from whoami response")55        console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")56        return username57    except Exception:58        console.print(59            "[bold yellow]Not authenticated with Hugging Face. Please login...[/bold yellow]"60        )61        try:62            login()63            user_info = whoami()64            if isinstance(user_info, dict):65                username = (66                    user_info.get("name")67                    or user_info.get("fullname")68                    or user_info.get("username")69                )70            else:71                username = (72                    getattr(user_info, "name", None)73                    or getattr(user_info, "fullname", None)74                    or getattr(user_info, "username", None)75                )76            if not username:77                raise ValueError("Could not extract username from whoami response")78            console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")79            return username80        except Exception as e:81            raise typer.BadParameter(82                f"Hugging Face authentication failed: {e}. Please run login manually."83            ) from e84 85 86@app.command()87def fork(88    source_space: Annotated[89        str,90        typer.Argument(91            help="Source Space ID in format 'owner/space-name' (e.g. org/my-openenv-space)"92        ),93    ],94    repo_id: Annotated[95        str | None,96        typer.Option(97            "--repo-id",98            "-r",99            help="Target repo ID for the fork (default: created under your account with same name)",100        ),101    ] = None,102    private: Annotated[103        bool,104        typer.Option("--private", help="Create the forked Space as private"),105    ] = False,106    set_env: Annotated[107        list[str],108        typer.Option(109            "--set-env",110            "-e",111            help="Set Space variable (public). Can be repeated. Format: KEY=VALUE",112        ),113    ] = [],114    set_secret: Annotated[115        list[str],116        typer.Option(117            "--set-secret",118            "--secret",119            "-s",120            help="Set Space secret. Can be repeated. Format: KEY=VALUE",121        ),122    ] = [],123    hardware: Annotated[124        str | None,125        typer.Option(126            "--hardware",127            "-H",128            help="Request hardware (e.g. t4-medium, cpu-basic). See Hub docs for options.",129        ),130    ] = None,131) -> None:132    """133    Fork (duplicate) a Hugging Face Space to your account using the Hub API.134 135    Uses the Hugging Face duplicate_space API. You can set environment variables136    and secrets, and request hardware/storage/sleep time at creation time.137 138    Examples:139        $ openenv fork owner/source-space140        $ openenv fork owner/source-space --private141        $ openenv fork owner/source-space --repo-id myuser/my-fork142        $ openenv fork owner/source-space --set-env MODEL_ID=user/model --set-secret HF_TOKEN=hf_xxx143        $ openenv fork owner/source-space --hardware t4-medium144    """145    if "/" not in source_space or source_space.count("/") != 1:146        raise typer.BadParameter(147            f"Invalid source Space ID: {source_space!r}. Expected format: 'owner/space-name'"148        )149 150    _ensure_hf_authenticated()151    api = HfApi()152 153    # Build kwargs for duplicate_space (only pass what we have)154    dup_kwargs: dict = {155        "from_id": source_space,156        "private": private,157    }158    if set_env:159        dup_kwargs["variables"] = [160            {"key": k, "value": v} for k, v in (_parse_key_value(x) for x in set_env)161        ]162    if set_secret:163        dup_kwargs["secrets"] = [164            {"key": k, "value": v} for k, v in (_parse_key_value(x) for x in set_secret)165        ]166    # HF API requires hardware when duplicating; default to free cpu-basic167    dup_kwargs["hardware"] = hardware if hardware is not None else "cpu-basic"168    if repo_id is not None:169        if "/" not in repo_id or repo_id.count("/") != 1:170            raise typer.BadParameter(171                f"Invalid --repo-id: {repo_id!r}. Expected format: 'username/repo-name'"172            )173        dup_kwargs["to_id"] = repo_id174 175    console.print(f"[bold cyan]Forking Space {source_space}...[/bold cyan]")176    try:177        result = api.duplicate_space(**dup_kwargs)178    except Exception as e:179        console.print(f"[bold red]✗[/bold red] Fork failed: {e}")180        raise typer.Exit(1) from e181 182    # result is RepoUrl (str-like) or similar; get repo_id for display183    if hasattr(result, "repo_id"):184        new_repo_id = result.repo_id185    elif isinstance(result, str):186        # URL like https://huggingface.co/spaces/owner/name -> owner/name187        if "/spaces/" in result:188            new_repo_id = result.split("/spaces/")[-1].rstrip("/")189        else:190            new_repo_id = result191    else:192        new_repo_id = getattr(result, "repo_id", str(result))193 194    console.print("[bold green]✓[/bold green] Space forked successfully")195    console.print(196        f"[bold]Space URL:[/bold] https://huggingface.co/spaces/{new_repo_id}"197    )198