Update submit_job.py

submit_job.py

@@ -1,42 +1,48 @@
 """
 Submit the Data-Centric AI training job to HF infrastructure.
 
-Cost: T4 small = $0.40/hr × ~5 hrs ≈ $2-3 from your $30 credits.
+⚠️  RECOMMENDED: Use the Colab notebook instead — it's more reliable.
+    https://colab.research.google.com/github/CelestialWorthyOfHeavenAndEarth/data-centric-env/blob/main/train_colab.ipynb
 
-Usage:
-    set HF_TOKEN=hf_yourtoken   (Windows)
+HF Jobs is provided as an alternative for automated / unattended runs.
+
+Usage (Windows):
+    set HF_TOKEN=hf_yourtoken
     python submit_job.py
+
+Usage (Linux/Mac):
+    HF_TOKEN=hf_yourtoken python submit_job.py
 """
 import os, sys
 from huggingface_hub import HfApi
 
 TOKEN = os.environ.get("HF_TOKEN") or input("Enter your HF token (hf_...): ").strip()
 
-# We use a stable PyTorch image (2.5.1) to avoid Unsloth/TorchAO incompatibility in 2.6
-SPACE_IMAGE  = "pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel"
-ENV_URL      = "https://aswini-kumar-data-centric-env.hf.space"
-REPO_URL     = "https://huggingface.co/spaces/Aswini-Kumar/data-centric-env"
+ENV_URL  = "https://aswini-kumar-data-centric-env.hf.space"
+REPO_URL = "https://huggingface.co/spaces/Aswini-Kumar/data-centric-env"
+
+# Use official Unsloth Docker image — has torch 2.4.1 + compatible torchao pre-installed
+# See: https://hub.docker.com/r/unsloth/unsloth/tags
+DOCKER_IMAGE = "unsloth/unsloth:latest-torch241"
 
 api = HfApi(token=TOKEN)
 
 print("Submitting HF training job...")
-print(f"  Docker image: {SPACE_IMAGE}")
+print(f"  Docker image: {DOCKER_IMAGE}")
 print(f"  ENV_URL     : {ENV_URL}")
-print(f"  Hardware    : a10g-large (Fast GPU)")
+print(f"  Hardware    : a10g-large")
 
-# Command installs git, clones your repository, and runs training
+# Clone repo + run training (torchao + unsloth are pre-installed in the image)
 bash_cmd = f"""
-apt-get update && apt-get install -y git && \\
+apt-get update -qq && apt-get install -y -qq git && \\
 git clone {REPO_URL} /app && cd /app && \\
-pip install -q torchao==0.6.1 && \\
-pip install -q 'unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git' \\
-trl>=0.15.0 datasets>=2.0.0 transformers>=4.40.0 accelerate>=0.30.0 matplotlib && \\
+pip install -q openenv-core[core]>=0.2.1 scikit-learn>=1.3.0 pandas>=2.0.0 numpy matplotlib && \\
 pip install -e . && \\
 python hf_job_train.py
 """
 
 job = api.run_job(
-    image=SPACE_IMAGE,
+    image=DOCKER_IMAGE,
     command=["bash", "-c", bash_cmd],
     env={
         "ENV_URL":  ENV_URL,
@@ -45,12 +51,9 @@ job = api.run_job(
     flavor="a10g-large",
 )
 
-print(f"\nJob submitted successfully!")
+print(f"\nJob submitted!")
 print(f"   Job ID  : {job.id}")
 print(f"   Status  : {job.status}")
-print(f"   Monitor : https://huggingface.co/settings/jobs")
-print(f"\nWhen DONE (~5 hrs), run:")
-print(f"  git pull hf main")
-print(f"  git add plots/ logs/")
-print(f'  git commit -m "Add training results"')
-print(f"  git push origin main && git push hf main")
+print(f"   Monitor : https://huggingface.co/jobs/Aswini-Kumar/{job.id}")
+print(f"\n⚡ Alternatively, use Colab for a more reliable run:")
+print(f"   https://colab.research.google.com/github/CelestialWorthyOfHeavenAndEarth/data-centric-env/blob/main/train_colab.ipynb")