train_uv.py

#!/usr/bin/env python3
"""
Elasticsearch Expert Model Training Script (UV Wrapper)

This script uses `uv` to manage dependencies on Hugging Face Jobs,
providing much faster and more reliable environment setup.
"""

import subprocess
import sys
import os

def run_command(cmd):
    print(f"Executing: {cmd}")
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    for line in process.stdout:
        print(line, end="")
    process.wait()
    return process.returncode

def main():
    print("=" * 50)
    print("Elasticsearch Training Job - Environment Setup (UV)")
    print("=" * 50)

    # 1. Install uv
    print("\nInstalling uv...")
    if run_command("curl -LsSf https://astral.sh/uv/install.sh | sh") != 0:
        print("Failed to install uv")
        sys.exit(1)
    
    # Add uv to path
    os.environ["PATH"] = f"{os.path.expanduser('~/.local/bin')}:{os.environ['PATH']}"

    # 2. Run training using uv
    print("\nLaunching training script with uv...")
    # We use 'uv run' which handles the virtualenv and dependencies automatically
    # based on the requirements_train.txt or inline metadata.
    # Here we'll pass the requirements file.
    cmd = "uv run --with-requirements requirements_train.txt python train.py"
    
    exit_code = run_command(cmd)
    
    if exit_code == 0:
        print("\nTraining completed successfully!")
    else:
        print(f"\nTraining failed with exit code {exit_code}")
    
    sys.exit(exit_code)

if __name__ == "__main__":
    main()