hf_sync.py

import os
import subprocess
import sys
import shutil

# Configuration
REPO_ID = os.environ.get("DATASET_REPO_ID")
HF_TOKEN = os.environ.get("HF_TOKEN")

def run_command(command):
    print(f"Running: {' '.join(command)}")
    # Ensure HF_TOKEN is in environment for the command
    env = os.environ.copy()
    if HF_TOKEN:
        env["HF_TOKEN"] = HF_TOKEN
    
    result = subprocess.run(command, capture_output=True, text=True, env=env)
    if result.returncode != 0:
        print(f"Error: {result.stderr}")
    else:
        print(f"Output: {result.stdout}")
    return result.returncode == 0

def download():
    if not REPO_ID:
        print("DATASET_REPO_ID not set, skipping download.")
        return
    
    print(f"Downloading data from {REPO_ID}...")
    # hf download REPO_ID --repo-type dataset --local-dir data_repo
    # Using --local-dir-use-symlinks False to avoid issues in some environments
    success = run_command(["hf", "download", REPO_ID, "--repo-type", "dataset", "--local-dir", "data_repo", "--local-dir-use-symlinks", "False"])
    if success:
        print("Download successful.")
    else:
        print("Download failed or repository is empty.")

def upload():
    if not REPO_ID:
        print("DATASET_REPO_ID not set, skipping upload.")
        return
    if not HF_TOKEN:
        print("HF_TOKEN not set, skipping upload.")
        return

    print(f"Uploading data to {REPO_ID}...")
    # hf upload REPO_ID data_repo / --repo-type dataset
    # We upload the contents of data_repo to the root of the dataset
    success = run_command(["hf", "upload", REPO_ID, "data_repo", ".", "--repo-type", "dataset"])
    if success:
        print("Upload successful.")
    else:
        print("Upload failed.")

def init_local():
    """Ensure data_repo has the necessary structure if download failed or it's new."""
    os.makedirs("data_repo/output", exist_ok=True)
    os.makedirs("data_repo/processed", exist_ok=True)
    os.makedirs("data_repo/uploads", exist_ok=True)
    # database.db will be created by the app if it doesn't exist, 
    # but we should ensure it's in data_repo.
    # We'll handle this in entrypoint.sh by symlinking.

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python hf_sync.py [download|upload|init]")
        sys.exit(1)
    
    action = sys.argv[1]
    if action == "download":
        download()
    elif action == "upload":
        upload()
    elif action == "init":
        init_local()
    else:
        print(f"Unknown action: {action}")