cpiuk commited on
Commit
cee36d1
·
verified ·
1 Parent(s): 202aaad

Upload runpod/handler.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. runpod/handler.py +14 -3
runpod/handler.py CHANGED
@@ -440,9 +440,20 @@ def main():
440
  state.status = "failed"
441
  sys.exit(1)
442
 
443
- # Phase 2: Run training
444
- logger.info("Phase 2: Training")
445
- run_training(state)
 
 
 
 
 
 
 
 
 
 
 
446
 
447
  # Brief grace for final /status poll
448
  grace = 30 if state.status == "completed" else 120
 
440
  state.status = "failed"
441
  sys.exit(1)
442
 
443
+ # Phase 2: Run training in a BACKGROUND thread so the control API stays responsive.
444
+ # The main thread just waits. This ensures /status, /logs, /stop remain accessible
445
+ # throughout model loading, training, and export.
446
+ import threading
447
+
448
+ logger.info("Phase 2: Training (background thread)")
449
+ training_thread = threading.Thread(target=run_training, args=(state,), name="training")
450
+ training_thread.start()
451
+
452
+ # Main thread: keep alive while training runs, poll for completion
453
+ while training_thread.is_alive():
454
+ training_thread.join(timeout=10)
455
+
456
+ logger.info(f"Training thread finished. Status: {state.status}")
457
 
458
  # Brief grace for final /status poll
459
  grace = 30 if state.status == "completed" else 120