File size: 1,553 Bytes
8377188 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | #!/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()
|