Upload runner.py with huggingface_hub

runner.py

@@ -6,8 +6,9 @@ import sys
 import threading
 from http.server import HTTPServer, SimpleHTTPRequestHandler
 from huggingface_hub import HfApi
+import torch
 
-# Start a dummy HTTP server on 7860 so Hugging Face marks the space as 'Healthy' and doesn't kill it.
+# Start a dummy HTTP server on 7860 so Hugging Face marks the space as 'Healthy'.
 def start_health_server():
     class HealthCheckHandler(SimpleHTTPRequestHandler):
         def do_GET(self):
@@ -28,12 +29,14 @@ def main():
     
     log("===== STARTING TRAINING RUNNER =====")
     
-    # Verify GPU exists
-    try:
-        subprocess.run(["nvidia-smi"], check=True)
-        log("GPU detected. Proceeding with training setup...")
-    except:
-        log("No GPU detected. Idling to prevent loops.")
+    # RELIABLE GPU DETECTION using PyTorch
+    is_gpu = torch.cuda.is_available()
+    
+    if is_gpu:
+        gpu_name = torch.cuda.get_device_name(0)
+        log(f"GPU DETECTED: {gpu_name}. Proceeding with training setup...")
+    else:
+        log("No GPU detected by PyTorch! Idling to prevent loops on CPU.")
         while True: time.sleep(3600)
         
     if os.path.exists("training_done.txt"):
@@ -70,8 +73,11 @@ def main():
         f.write("done")
         
     log("DOWNGRADING HARDWARE TO CPU TO STOP BILLING...")
-    api.request_space_hardware(repo_id=repo_id, hardware="cpu-basic")
-    log("Hardware downgraded successfully.")
+    try:
+        api.request_space_hardware(repo_id=repo_id, hardware="cpu-basic")
+        log("Hardware downgraded successfully.")
+    except Exception as e:
+        log(f"Failed to downgrade hardware: {e}")
     
     while True:
         time.sleep(3600)