import os
from huggingface_hub import login
import subprocess
import shutil

def main():
    # 書き込み可能なキャッシュディレクトリを指定
    os.environ["HF_HOME"] = "/tmp/.cache/huggingface"
    os.makedirs("/tmp/.cache/huggingface", exist_ok=True)  # ディレクトリ作成
    
    # Hugging Face トークンの取得
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("Hugging Face token is not set. Please provide it as a Secret in your environment.")
    
    print(f"Using Hugging Face token: {hf_token[:5]}...")  # トークンの先頭5文字を表示
    
    # Hugging Faceへのログイン処理
    try:
        login(token=hf_token)
        print("Login successful.")
    except Exception as e:
        print(f"Failed to log in to Hugging Face: {e}")
        raise

    # vLLMサーバーを起動
    try:
        print("Starting vLLM server...")

        # vllm serve コマンドを設定
        command = [
            "vllm", "serve",
            "--model", "mistralai/Pixtral-Large-Instruct-2411",  # モデル名が正しく指定されているか確認
            "--tensor-parallel-size", "1",
            "--device", "cpu"  # GPUを使用せず、必ずCPUモードで実行
        ]
        
        print(f"Running command: {' '.join(command)}")  # 実行するコマンドを表示
        
        # コマンド実行
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Failed to start vLLM server: {e}")
        raise

if __name__ == "__main__":
    main()