From fe6c78c6a26cbaeb5d1a137e70b65fa8746017f8 Mon Sep 17 00:00:00 2001
From: phoenixDong <42799804+phoenixdong@users.noreply.github.com>
Date: Sun, 23 Feb 2025 17:42:01 +0800
Subject: [PATCH] [CLI] Add flagscale pull command (#364)

# 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.
---
 docker/Dockerfile.cuda |  2 +-
 flagscale/cli.py       | 64 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda
index 1256cd1d5..ad9adb46b 100644
--- a/docker/Dockerfile.cuda
+++ b/docker/Dockerfile.cuda
@@ -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 \
diff --git a/flagscale/cli.py b/flagscale/cli.py
index caabbe3a0..5171a3c8b 100644
--- a/flagscale/cli.py
+++ b/flagscale/cli.py
@@ -1,4 +1,5 @@
 import os
+import subprocess
 import sys
 
 import click
@@ -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()