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"""Push an OpenEnv environment to Hugging Face Spaces."""8 9from __future__ import annotations10 11import shutil12import sys13import tempfile14from fnmatch import fnmatch15from pathlib import Path16from typing import Annotated17 18import typer19import yaml20from huggingface_hub import HfApi, login, whoami21 22from .._cli_utils import console, validate_env_structure23 24app = typer.Typer(help="Push an OpenEnv environment to Hugging Face Spaces")25 26 27DEFAULT_PUSH_IGNORE_PATTERNS = [".*", "__pycache__", "*.pyc"]28 29 30def _path_matches_pattern(relative_path: Path, pattern: str) -> bool:31    """Return True if a relative path matches an exclude pattern."""32    normalized_pattern = pattern.strip()33    if normalized_pattern.startswith("!"):34        return False35 36    while normalized_pattern.startswith("./"):37        normalized_pattern = normalized_pattern[2:]38 39    if normalized_pattern.startswith("/"):40        normalized_pattern = normalized_pattern[1:]41 42    if not normalized_pattern:43        return False44 45    posix_path = relative_path.as_posix()46    pattern_candidates = [normalized_pattern]47    if normalized_pattern.startswith("**/"):48        # Gitignore-style "**/" can also match directly at the root.49        pattern_candidates.append(normalized_pattern[3:])50 51    # Support directory patterns such as "artifacts/" and "**/outputs/".52    if normalized_pattern.endswith("/"):53        dir_pattern_candidates: list[str] = []54        for candidate in pattern_candidates:55            base = candidate.rstrip("/")56            if not base:57                continue58            dir_pattern_candidates.extend([base, f"{base}/*"])59 60        return any(61            fnmatch(posix_path, candidate) for candidate in dir_pattern_candidates62        )63 64    # Match both full relative path and basename for convenience.65    return any(66        fnmatch(posix_path, candidate) for candidate in pattern_candidates67    ) or any(fnmatch(relative_path.name, candidate) for candidate in pattern_candidates)68 69 70def _should_exclude_path(relative_path: Path, ignore_patterns: list[str]) -> bool:71    """Return True when the path should be excluded from staging/upload."""72    return any(73        _path_matches_pattern(relative_path, pattern) for pattern in ignore_patterns74    )75 76 77def _read_ignore_file(ignore_path: Path) -> tuple[list[str], int]:78    """Read ignore patterns from a file and return (patterns, ignored_negations)."""79    patterns: list[str] = []80    ignored_negations = 081 82    for line in ignore_path.read_text().splitlines():83        stripped = line.strip()84        if not stripped or stripped.startswith("#"):85            continue86        if stripped.startswith("!"):87            ignored_negations += 188            continue89        patterns.append(stripped)90 91    return patterns, ignored_negations92 93 94def _load_ignore_patterns(env_dir: Path, exclude_file: str | None) -> list[str]:95    """Load ignore patterns from defaults and an optional ignore file."""96    patterns = list(DEFAULT_PUSH_IGNORE_PATTERNS)97    ignored_negations = 098 99    def _merge_ignore_file(ignore_path: Path, *, source_label: str) -> None:100        nonlocal ignored_negations101        file_patterns, skipped_negations = _read_ignore_file(ignore_path)102        patterns.extend(file_patterns)103        ignored_negations += skipped_negations104        console.print(105            f"[bold green]✓[/bold green] Loaded {len(file_patterns)} ignore patterns from {source_label}: {ignore_path}"106        )107 108    # Optional source: explicit exclude file from CLI.109    if exclude_file:110        ignore_path = Path(exclude_file)111        if not ignore_path.is_absolute():112            ignore_path = env_dir / ignore_path113        ignore_path = ignore_path.resolve()114 115        if not ignore_path.exists() or not ignore_path.is_file():116            raise typer.BadParameter(117                f"Exclude file not found or not a file: {ignore_path}"118            )119 120        _merge_ignore_file(ignore_path, source_label="--exclude")121 122    # Keep stable order while removing duplicates.123    patterns = list(dict.fromkeys(patterns))124 125    if ignored_negations > 0:126        console.print(127            f"[bold yellow]⚠[/bold yellow] Skipped {ignored_negations} negated ignore patterns ('!') because negation is not supported for push excludes"128        )129 130    return patterns131 132 133def _copytree_ignore_factory(env_dir: Path, ignore_patterns: list[str]):134    """Build a shutil.copytree ignore callback from path-based patterns."""135 136    def _ignore(path: str, names: list[str]) -> set[str]:137        current_dir = Path(path)138        ignored: set[str] = set()139 140        for name in names:141            candidate = current_dir / name142            try:143                relative_path = candidate.relative_to(env_dir)144            except ValueError:145                # candidate is not under env_dir (e.g. symlink or146                # copytree root differs from env_dir); skip filtering.147                continue148            if _should_exclude_path(relative_path, ignore_patterns):149                ignored.add(name)150 151        return ignored152 153    return _ignore154 155 156def _validate_openenv_directory(directory: Path) -> tuple[str, dict]:157    """158    Validate that the directory is an OpenEnv environment.159 160    Returns:161        Tuple of (env_name, manifest_data)162    """163    # Use the comprehensive validation function164    try:165        warnings = validate_env_structure(directory)166        for warning in warnings:167            console.print(f"[bold yellow]⚠[/bold yellow] {warning}")168    except FileNotFoundError as e:169        raise typer.BadParameter(f"Invalid OpenEnv environment structure: {e}") from e170 171    # Load and validate manifest172    manifest_path = directory / "openenv.yaml"173    try:174        with open(manifest_path, "r") as f:175            manifest = yaml.safe_load(f)176    except Exception as e:177        raise typer.BadParameter(f"Failed to parse openenv.yaml: {e}") from e178 179    if not isinstance(manifest, dict):180        raise typer.BadParameter("openenv.yaml must be a YAML dictionary")181 182    env_name = manifest.get("name")183    if not env_name:184        raise typer.BadParameter("openenv.yaml must contain a 'name' field")185 186    return env_name, manifest187 188 189def _ensure_hf_authenticated() -> str:190    """191    Ensure user is authenticated with Hugging Face.192 193    Returns:194        Username of authenticated user195    """196    try:197        # Try to get current user198        user_info = whoami()199        # Handle both dict and object return types200        if isinstance(user_info, dict):201            username = (202                user_info.get("name")203                or user_info.get("fullname")204                or user_info.get("username")205            )206        else:207            # If it's an object, try to get name attribute208            username = (209                getattr(user_info, "name", None)210                or getattr(user_info, "fullname", None)211                or getattr(user_info, "username", None)212            )213 214        if not username:215            raise ValueError("Could not extract username from whoami response")216 217        console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")218        return username219    except Exception:220        # Not authenticated, prompt for login221        console.print(222            "[bold yellow]Not authenticated with Hugging Face. Please login...[/bold yellow]"223        )224 225        try:226            login()227            # Verify login worked228            user_info = whoami()229            # Handle both dict and object return types230            if isinstance(user_info, dict):231                username = (232                    user_info.get("name")233                    or user_info.get("fullname")234                    or user_info.get("username")235                )236            else:237                username = (238                    getattr(user_info, "name", None)239                    or getattr(user_info, "fullname", None)240                    or getattr(user_info, "username", None)241                )242 243            if not username:244                raise ValueError("Could not extract username from whoami response")245 246            console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")247            return username248        except Exception as e:249            raise typer.BadParameter(250                f"Hugging Face authentication failed: {e}. Please run login manually."251            ) from e252 253 254def _prepare_staging_directory(255    env_dir: Path,256    env_name: str,257    staging_dir: Path,258    ignore_patterns: list[str],259    base_image: str | None = None,260    enable_interface: bool = True,261) -> None:262    """263    Prepare files for deployment.264 265    This includes:266    - Copying necessary files267    - Modifying Dockerfile to optionally enable web interface and update base image268    - Ensuring README has proper HF frontmatter (if interface enabled)269    """270    # Create staging directory structure271    staging_dir.mkdir(parents=True, exist_ok=True)272 273    # Copy all files from env directory274    copy_ignore = _copytree_ignore_factory(env_dir, ignore_patterns)275    for item in env_dir.iterdir():276        relative_path = item.relative_to(env_dir)277        if _should_exclude_path(relative_path, ignore_patterns):278            continue279 280        dest = staging_dir / item.name281        if item.is_dir():282            shutil.copytree(item, dest, dirs_exist_ok=True, ignore=copy_ignore)283        else:284            shutil.copy2(item, dest)285 286    # Dockerfile must be at repo root for Hugging Face. Prefer root if present287    # (it was copied there); otherwise move server/Dockerfile to root.288    dockerfile_server_path = staging_dir / "server" / "Dockerfile"289    dockerfile_root_path = staging_dir / "Dockerfile"290    dockerfile_path: Path | None = None291 292    if dockerfile_root_path.exists():293        dockerfile_path = dockerfile_root_path294    elif dockerfile_server_path.exists():295        dockerfile_server_path.rename(dockerfile_root_path)296        console.print(297            "[bold cyan]Moved Dockerfile to repository root for deployment[/bold cyan]"298        )299        dockerfile_path = dockerfile_root_path300 301    # Modify Dockerfile to optionally enable web interface and update base image302    if dockerfile_path and dockerfile_path.exists():303        dockerfile_content = dockerfile_path.read_text()304        lines = dockerfile_content.split("\n")305        new_lines = []306        cmd_found = False307        base_image_updated = False308        web_interface_env_exists = "ENABLE_WEB_INTERFACE" in dockerfile_content309        last_instruction = None310 311        for line in lines:312            stripped = line.strip()313            token = stripped.split(maxsplit=1)[0] if stripped else ""314            current_instruction = token.upper()315 316            is_healthcheck_continuation = last_instruction == "HEALTHCHECK"317 318            # Update base image if specified319            if base_image and stripped.startswith("FROM") and not base_image_updated:320                new_lines.append(f"FROM {base_image}")321                base_image_updated = True322                last_instruction = "FROM"323                continue324 325            if (326                stripped.startswith("CMD")327                and not cmd_found328                and not web_interface_env_exists329                and enable_interface330                and not is_healthcheck_continuation331            ):332                new_lines.append("ENV ENABLE_WEB_INTERFACE=true")333                cmd_found = True334 335            new_lines.append(line)336 337            if current_instruction:338                last_instruction = current_instruction339 340        if not cmd_found and not web_interface_env_exists and enable_interface:341            new_lines.append("ENV ENABLE_WEB_INTERFACE=true")342 343        if base_image and not base_image_updated:344            new_lines.insert(0, f"FROM {base_image}")345 346        dockerfile_path.write_text("\n".join(new_lines))347 348        changes = []349        if base_image and base_image_updated:350            changes.append("updated base image")351        if enable_interface and not web_interface_env_exists:352            changes.append("enabled web interface")353        if changes:354            console.print(355                f"[bold green]✓[/bold green] Updated Dockerfile: {', '.join(changes)}"356            )357    else:358        console.print(359            "[bold yellow]⚠[/bold yellow] No Dockerfile at server/ or repo root"360        )361 362    # Ensure README has proper HF frontmatter (only if interface enabled)363    if enable_interface:364        readme_path = staging_dir / "README.md"365        if readme_path.exists():366            readme_content = readme_path.read_text()367            if "base_path: /web" not in readme_content:368                # Check if frontmatter exists369                if readme_content.startswith("---"):370                    # Add base_path to existing frontmatter371                    lines = readme_content.split("\n")372                    new_lines = []373                    _in_frontmatter = True374                    for i, line in enumerate(lines):375                        new_lines.append(line)376                        if line.strip() == "---" and i > 0:377                            # End of frontmatter, add base_path before this line378                            if "base_path:" not in "\n".join(new_lines):379                                new_lines.insert(-1, "base_path: /web")380                            _in_frontmatter = False381                    readme_path.write_text("\n".join(new_lines))382                else:383                    # No frontmatter, add it384                    frontmatter = f"""---385title: {env_name.replace("_", " ").title()} Environment Server386emoji: 🔊387colorFrom: '#00C9FF'388colorTo: '#1B2845'389sdk: docker390pinned: false391app_port: 8000392base_path: /web393tags:394  - openenv395---396 397"""398                    readme_path.write_text(frontmatter + readme_content)399                console.print(400                    "[bold green]✓[/bold green] Updated README with HF Space frontmatter"401                )402        else:403            console.print("[bold yellow]⚠[/bold yellow] No README.md found")404 405 406def _create_hf_space(407    repo_id: str,408    api: HfApi,409    private: bool = False,410) -> None:411    """Create a Hugging Face Space if it doesn't exist."""412    console.print(f"[bold cyan]Creating/verifying space: {repo_id}[/bold cyan]")413 414    try:415        api.create_repo(416            repo_id=repo_id,417            repo_type="space",418            space_sdk="docker",419            private=private,420            exist_ok=True,421        )422        console.print(f"[bold green]✓[/bold green] Space {repo_id} is ready")423    except Exception as e:424        # Space might already exist, which is okay with exist_ok=True425        # But if there's another error, log it426        console.print(f"[bold yellow]⚠[/bold yellow] Space creation: {e}")427 428 429def _upload_to_hf_space(430    repo_id: str,431    staging_dir: Path,432    api: HfApi,433    ignore_patterns: list[str],434    private: bool = False,435    create_pr: bool = False,436    commit_message: str | None = None,437) -> None:438    """Upload files to Hugging Face Space."""439    if create_pr:440        console.print(441            f"[bold cyan]Uploading files to {repo_id} (will open a Pull Request)...[/bold cyan]"442        )443    else:444        console.print(f"[bold cyan]Uploading files to {repo_id}...[/bold cyan]")445 446    upload_kwargs: dict = {447        "folder_path": str(staging_dir),448        "repo_id": repo_id,449        "repo_type": "space",450        "create_pr": create_pr,451        "ignore_patterns": ignore_patterns,452    }453    if commit_message:454        upload_kwargs["commit_message"] = commit_message455 456    try:457        result = api.upload_folder(**upload_kwargs)458        console.print("[bold green]✓[/bold green] Upload completed successfully")459        if create_pr and result is not None and hasattr(result, "pr_url"):460            console.print(f"[bold]Pull request:[/bold] {result.pr_url}")461        console.print(462            f"[bold]Space URL:[/bold] https://huggingface.co/spaces/{repo_id}"463        )464    except Exception as e:465        console.print(f"[bold red]✗[/bold red] Upload failed: {e}")466        raise typer.Exit(1) from e467 468 469@app.command()470def push(471    directory: Annotated[472        str | None,473        typer.Argument(474            help="Directory containing the OpenEnv environment (default: current directory)"475        ),476    ] = None,477    repo_id: Annotated[478        str | None,479        typer.Option(480            "--repo-id",481            "-r",482            help="Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)",483        ),484    ] = None,485    base_image: Annotated[486        str | None,487        typer.Option(488            "--base-image",489            "-b",490            help="Base Docker image to use (overrides Dockerfile FROM)",491        ),492    ] = None,493    interface: Annotated[494        bool,495        typer.Option(496            "--interface",497            help="Enable web interface (default: True if no registry specified)",498        ),499    ] = None,500    no_interface: Annotated[501        bool,502        typer.Option(503            "--no-interface",504            help="Disable web interface",505        ),506    ] = False,507    registry: Annotated[508        str | None,509        typer.Option(510            "--registry",511            help="Custom registry URL (e.g., docker.io/username). Disables web interface by default.",512        ),513    ] = None,514    private: Annotated[515        bool,516        typer.Option(517            "--private",518            help="Deploy the space as private",519        ),520    ] = False,521    create_pr: Annotated[522        bool,523        typer.Option(524            "--create-pr",525            help="Create a Pull Request instead of pushing to the default branch",526        ),527    ] = False,528    exclude: Annotated[529        str | None,530        typer.Option(531            "--exclude",532            help="Optional additional ignore file with newline-separated glob patterns to exclude from Hugging Face uploads",533        ),534    ] = None,535) -> None:536    """537    Push an OpenEnv environment to Hugging Face Spaces or a custom Docker registry.538 539    This command:540    1. Validates that the directory is an OpenEnv environment (openenv.yaml present)541    2. Builds and pushes to Hugging Face Spaces or custom Docker registry542    3. Optionally enables web interface for deployment543 544    The web interface is enabled by default when pushing to HuggingFace Spaces,545    but disabled by default when pushing to a custom Docker registry.546 547    Examples:548        # Push to HuggingFace Spaces from current directory (web interface enabled)549        $ cd my_env550        $ openenv push551 552        # Push to HuggingFace repo and open a Pull Request553        $ openenv push my-org/my-env --create-pr554        $ openenv push --repo-id my-org/my-env --create-pr555 556        # Push to HuggingFace without web interface557        $ openenv push --no-interface558 559        # Push to Docker Hub560        $ openenv push --registry docker.io/myuser561 562        # Push to GitHub Container Registry563        $ openenv push --registry ghcr.io/myorg564 565        # Push to custom registry with web interface566        $ openenv push --registry myregistry.io/path1/path2 --interface567 568        # Push to specific HuggingFace repo569        $ openenv push --repo-id my-org/my-env570 571        # Push privately with custom base image572        $ openenv push --private --base-image ghcr.io/meta-pytorch/openenv-base:latest573    """574    # Handle interface flag logic575    if no_interface and interface:576        console.print(577            "[bold red]Error:[/bold red] Cannot specify both --interface and --no-interface",578            file=sys.stderr,579        )580        raise typer.Exit(1)581 582    # Determine if web interface should be enabled583    if no_interface:584        enable_interface = False585    elif interface is not None:586        enable_interface = interface587    elif registry is not None:588        # Custom registry: disable interface by default589        enable_interface = False590    else:591        # HuggingFace: enable interface by default592        enable_interface = True593 594    # Determine directory595    if directory:596        env_dir = Path(directory).resolve()597    else:598        env_dir = Path.cwd().resolve()599 600    if not env_dir.exists() or not env_dir.is_dir():601        raise typer.BadParameter(f"Directory does not exist: {env_dir}")602 603    # Check for openenv.yaml to confirm this is an environment directory604    openenv_yaml = env_dir / "openenv.yaml"605    if not openenv_yaml.exists():606        console.print(607            f"[bold red]Error:[/bold red] Not an OpenEnv environment directory (missing openenv.yaml): {env_dir}",608        )609        console.print(610            "[yellow]Hint:[/yellow] Run this command from the environment root directory",611        )612        raise typer.Exit(1)613 614    # Validate OpenEnv environment615    console.print(616        f"[bold cyan]Validating OpenEnv environment in {env_dir}...[/bold cyan]"617    )618    env_name, manifest = _validate_openenv_directory(env_dir)619    console.print(f"[bold green]✓[/bold green] Found OpenEnv environment: {env_name}")620 621    # Handle custom registry push622    if registry:623        console.print("[bold cyan]Preparing to push to custom registry...[/bold cyan]")624        if enable_interface:625            console.print("[bold cyan]Web interface will be enabled[/bold cyan]")626 627        # Import build functions628        from .build import _build_docker_image, _push_docker_image629 630        # Prepare build args for custom registry deployment631        build_args = {}632        if enable_interface:633            build_args["ENABLE_WEB_INTERFACE"] = "true"634 635        # Build Docker image from the environment directory636        tag = f"{registry}/{env_name}"637        console.print(f"[bold cyan]Building Docker image: {tag}[/bold cyan]")638 639        success = _build_docker_image(640            env_path=env_dir,641            tag=tag,642            build_args=build_args if build_args else None,643        )644 645        if not success:646            console.print("[bold red]✗ Docker build failed[/bold red]")647            raise typer.Exit(1)648 649        console.print("[bold green]✓ Docker build successful[/bold green]")650 651        # Push to registry652        console.print(f"[bold cyan]Pushing to registry: {registry}[/bold cyan]")653 654        success = _push_docker_image(655            tag, registry=None656        )  # Tag already includes registry657 658        if not success:659            console.print("[bold red]✗ Docker push failed[/bold red]")660            raise typer.Exit(1)661 662        console.print("\n[bold green]✓ Deployment complete![/bold green]")663        console.print(f"[bold]Image:[/bold] {tag}")664        return665 666    ignore_patterns = _load_ignore_patterns(env_dir, exclude)667 668    # Ensure authentication for HuggingFace669    username = _ensure_hf_authenticated()670 671    # Determine repo_id672    if not repo_id:673        repo_id = f"{username}/{env_name}"674 675    # Validate repo_id format676    if "/" not in repo_id or repo_id.count("/") != 1:677        raise typer.BadParameter(678            f"Invalid repo-id format: {repo_id}. Expected format: 'username/repo-name'"679        )680 681    # Initialize Hugging Face API682    api = HfApi()683 684    # Prepare staging directory685    deployment_type = (686        "with web interface" if enable_interface else "without web interface"687    )688    console.print(689        f"[bold cyan]Preparing files for Hugging Face deployment ({deployment_type})...[/bold cyan]"690    )691    with tempfile.TemporaryDirectory() as tmpdir:692        staging_dir = Path(tmpdir) / "staging"693        _prepare_staging_directory(694            env_dir,695            env_name,696            staging_dir,697            ignore_patterns=ignore_patterns,698            base_image=base_image,699            enable_interface=enable_interface,700        )701 702        # Create/verify space (no-op if exists; needed when pushing to own new repo)703        if not create_pr:704            _create_hf_space(repo_id, api, private=private)705        # When create_pr we rely on upload_folder to create branch and PR706 707        # Upload files708        _upload_to_hf_space(709            repo_id,710            staging_dir,711            api,712            private=private,713            create_pr=create_pr,714            ignore_patterns=ignore_patterns,715        )716 717        console.print("\n[bold green]✓ Deployment complete![/bold green]")718        console.print(f"Visit your space at: https://huggingface.co/spaces/{repo_id}")719