Skip to content

Commit

Permalink
[CLI] Add flagscale pull command (FlagOpen#364)
Browse files Browse the repository at this point in the history
# Pull Instruction for FlagScale

This PR introduces the `flagscale pull` command, which allows users to
pull Docker images and download checkpoints (ckpt) from a specified Git
repository.

## Installation

First, please install FlagScale:

```bash
cd FlagScale && pip install .
```

## Usage

The command for pulling images and checkpoints, along with setting the
storage path for checkpoints, is as follows:

```bash
flagscale pull [--image <image>] [--ckpt <ckpt>] [--ckpt-path <ckpt_save_path>]
```

### Parameters

- `--image`: Optional. Docker image address that you want to pull.
- `--ckpt`: Optional. The Git repository address of the checkpoint you
wish to download.
- `--ckpt-path`: Optional. If provided, the checkpoint will be stored in
the specified path. If not set, it will be downloaded to a folder named
`model_download` in the current directory.
  • Loading branch information
phoenixdong authored Feb 23, 2025
1 parent 7cf2c6b commit fe6c78c
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 1 deletion.
2 changes: 1 addition & 1 deletion docker/Dockerfile.cuda
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ RUN apt-get update && \
autotools-dev \
build-essential \
ca-certificates ccache cmake curl emacs \
gcc git g++ \
gcc git git-lfs g++ \
htop \
iftop iotop \
libcairo2-dev libfontconfig-dev libibverbs1 libibverbs-dev libnuma-dev libx11-dev lsb-release \
Expand Down
64 changes: 64 additions & 0 deletions flagscale/cli.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
import subprocess
import sys

import click
Expand Down Expand Up @@ -85,5 +86,68 @@ def serve(model_name, yaml_path=None):
run_main()


@flagscale.command()
# Input the name of the Docker image (required)
@click.option(
"--image",
"image_name",
required=True,
type=str,
help="The name of the Docker image",
)
# Input the address of the Git repository (required)
@click.option(
"--ckpt",
"ckpt_name",
required=True,
type=str,
help="The address of the ckpt's git repository",
)
# Input the address of the local directory (optional)
@click.option(
"--ckpt-path",
"ckpt_path",
type=click.Path(),
required=False,
help="The path to save ckpt",
)
def pull(image_name, ckpt_name, ckpt_path):
# If ckpt_path is not provided, use the default download directory
if ckpt_path is None:
ckpt_path = os.path.join(os.getcwd(), "model_download")

# Check and create the directory
if not os.path.exists(ckpt_path):
os.makedirs(ckpt_path)
print(f"Directory {ckpt_path} created.")

# Pull the Docker image
try:
print(f"Pulling Docker image: {image_name}...")
subprocess.run(["docker", "pull", image_name], check=True)
print(f"Successfully pulled Docker image: {image_name}")
except subprocess.CalledProcessError:
print(f"Failed to pull Docker image: {image_name}")
return

# Clone the Git repository
try:
print(f"Cloning Git repository: {ckpt_name} into {ckpt_path}...")
subprocess.run(["git", "clone", ckpt_name, ckpt_path], check=True)
print(f"Successfully cloned Git repository: {ckpt_name}")
except subprocess.CalledProcessError:
print(f"Failed to clone Git repository: {ckpt_name}")
return

# Pull large files using Git LFS
print("Pulling Git LFS files...")
try:
subprocess.run(["git", "lfs", "pull"], cwd=ckpt_path, check=True)
print("Successfully pulled Git LFS files")
except subprocess.CalledProcessError:
print("Failed to pull Git LFS files")
return


if __name__ == "__main__":
flagscale()

0 comments on commit fe6c78c

Please sign in to comment.