88import argparse
99import json
1010import os
11- import queue
1211import signal
1312import sys
1413import threading
2120import backend_pb2_grpc
2221
2322_ONE_DAY_IN_SECONDS = 60 * 60 * 24
24- MAX_WORKERS = int (os .environ .get ('PYTHON_GRPC_MAX_WORKERS' , '1 ' ))
23+ MAX_WORKERS = int (os .environ .get ('PYTHON_GRPC_MAX_WORKERS' , '2 ' ))
2524
2625
2726class ProgressCallback :
@@ -38,14 +37,61 @@ def get_callback(self):
3837 parent = self
3938
4039 class _Callback (TrainerCallback ):
40+ def __init__ (self ):
41+ super ().__init__ ()
42+ self ._train_start_time = None
43+ self ._last_step_event_time = 0
44+
45+ def on_train_begin (self , args , state , control , ** kwargs ):
46+ self ._train_start_time = time .time ()
47+ # Emit the total step count as soon as the Trainer has computed it.
48+ # Pre-training events (loading_model, etc.) don't know this value.
49+ total_steps = state .max_steps if state .max_steps > 0 else 0
50+ update = backend_pb2 .FineTuneProgressUpdate (
51+ job_id = parent .job_id ,
52+ current_step = 0 ,
53+ total_steps = total_steps ,
54+ total_epochs = float (parent .total_epochs ),
55+ status = "training" ,
56+ message = f"Training started — { total_steps } ,
57+ )
58+ parent .progress_queue .push (update )
59+
60+ def on_step_end (self , args , state , control , ** kwargs ):
61+ # Lightweight step-count update so the UI's progress bar stays
62+ # in sync with the tqdm bar even between on_log events.
63+ # Throttle to at most once per second to avoid flooding the queue.
64+ now = time .time ()
65+ if now - self ._last_step_event_time < 1.0 :
66+ return
67+ self ._last_step_event_time = now
68+ total_steps = state .max_steps if state .max_steps > 0 else 0
69+ progress = (state .global_step / total_steps * 100 ) if total_steps > 0 else 0
70+ eta = 0.0
71+ if state .global_step > 0 and total_steps > 0 and self ._train_start_time is not None :
72+ elapsed = time .time () - self ._train_start_time
73+ remaining_steps = total_steps - state .global_step
74+ eta = remaining_steps * (elapsed / state .global_step )
75+
76+ update = backend_pb2 .FineTuneProgressUpdate (
77+ job_id = parent .job_id ,
78+ current_step = state .global_step ,
79+ total_steps = total_steps ,
80+ total_epochs = float (parent .total_epochs ),
81+ progress_percent = float (progress ),
82+ eta_seconds = float (eta ),
83+ status = "training" ,
84+ )
85+ parent .progress_queue .push (update )
86+
4187 def on_log (self , args , state , control , logs = None , ** kwargs ):
4288 if logs is None :
4389 return
4490 total_steps = state .max_steps if state .max_steps > 0 else 0
4591 progress = (state .global_step / total_steps * 100 ) if total_steps > 0 else 0
4692 eta = 0.0
47- if state .global_step > 0 and total_steps > 0 :
48- elapsed = time .time () - state . logging_steps # approximate
93+ if state .global_step > 0 and total_steps > 0 and self . _train_start_time is not None :
94+ elapsed = time .time () - self . _train_start_time
4995 remaining_steps = total_steps - state .global_step
5096 if state .global_step > 1 :
5197 eta = remaining_steps * (elapsed / state .global_step )
@@ -70,39 +116,91 @@ def on_log(self, args, state, control, logs=None, **kwargs):
70116 status = "training" ,
71117 extra_metrics = extra_metrics ,
72118 )
73- parent .progress_queue .put (update )
119+ parent .progress_queue .push (update )
74120
75121 def on_save (self , args , state , control , ** kwargs ):
122+ total_steps = state .max_steps if state .max_steps > 0 else 0
76123 checkpoint_path = os .path .join (args .output_dir , f"checkpoint-{ state .global_step } )
77124 update = backend_pb2 .FineTuneProgressUpdate (
78125 job_id = parent .job_id ,
79126 current_step = state .global_step ,
127+ total_steps = total_steps ,
80128 status = "saving" ,
81129 message = f"Checkpoint saved at step { state .global_step } ,
82130 checkpoint_path = checkpoint_path ,
83131 )
84- parent .progress_queue .put (update )
132+ parent .progress_queue .push (update )
85133
86134 def on_train_end (self , args , state , control , ** kwargs ):
87135 update = backend_pb2 .FineTuneProgressUpdate (
88136 job_id = parent .job_id ,
89137 current_step = state .global_step ,
90- total_steps = state .max_steps ,
138+ total_steps = state .max_steps if state . max_steps > 0 else 0 ,
91139 progress_percent = 100.0 ,
92140 status = "completed" ,
93141 message = "Training completed" ,
94142 )
95- parent .progress_queue .put (update )
143+ parent .progress_queue .push (update )
96144
97145 return _Callback ()
98146
99147
148+ class EventLog :
149+ """Append-only event log that supports multiple concurrent readers without losing events.
150+
151+ Unlike queue.Queue (where get() is destructive), readers iterate the log
152+ from any position. Reconnecting clients can resume from where they left off,
153+ or start from 0 to replay everything.
154+ """
155+
156+ def __init__ (self ):
157+ self ._events = []
158+ self ._lock = threading .Lock ()
159+ self ._condition = threading .Condition (threading .Lock ())
160+ self ._done = False # Set when the sentinel (None) is pushed
161+
162+ def push (self , update ):
163+ """Append an event. If update is None, mark the log as done."""
164+ with self ._lock :
165+ self ._events .append (update )
166+ if update is None :
167+ self ._done = True
168+ with self ._condition :
169+ self ._condition .notify_all ()
170+
171+ def iter_from (self , start = 0 , timeout = 1.0 , is_alive = None ):
172+ """Yield events from *start* onward, blocking when caught up.
173+
174+ *is_alive* is an optional callable checked on each timeout.
175+ If it returns False the iterator ends (e.g. pass ``context.is_active``).
176+ Ends iteration when the sentinel (None) is reached.
177+ """
178+ idx = start
179+ while True :
180+ with self ._lock :
181+ if idx < len (self ._events ):
182+ event = self ._events [idx ]
183+ idx += 1
184+ if event is None :
185+ return # sentinel — stream complete
186+ yield event
187+ continue
188+ done = self ._done
189+ if done :
190+ return
191+ if is_alive is not None and not is_alive ():
192+ return
193+ # Wait for new events
194+ with self ._condition :
195+ self ._condition .wait (timeout = timeout )
196+
197+
100198class ActiveJob :
101199 """Represents an active fine-tuning job."""
102200
103201 def __init__ (self , job_id ):
104202 self .job_id = job_id
105- self .progress_queue = queue . Queue ()
203+ self .progress_queue = EventLog ()
106204 self .trainer = None
107205 self .thread = None
108206 self .model = None
@@ -179,9 +277,9 @@ def _run_training(self, request, job):
179277 status = "failed" ,
180278 message = msg ,
181279 )
182- job .progress_queue .put (update )
280+ job .progress_queue .push (update )
183281 # Send sentinel
184- job .progress_queue .put (None )
282+ job .progress_queue .push (None )
185283
186284 def _do_training (self , request , job ):
187285 import torch
@@ -193,7 +291,7 @@ def _do_training(self, request, job):
193291 training_type = request .training_type or "lora"
194292
195293 # Send loading status
196- job .progress_queue .put (backend_pb2 .FineTuneProgressUpdate (
294+ job .progress_queue .push (backend_pb2 .FineTuneProgressUpdate (
197295 job_id = job .job_id , status = "loading_model" , message = f"Loading model { request .model } ,
198296 ))
199297
@@ -241,7 +339,7 @@ def _do_training(self, request, job):
241339 model = get_peft_model (model , peft_config )
242340
243341 # Load dataset
244- job .progress_queue .put (backend_pb2 .FineTuneProgressUpdate (
342+ job .progress_queue .push (backend_pb2 .FineTuneProgressUpdate (
245343 job_id = job .job_id , status = "loading_dataset" , message = "Loading dataset" ,
246344 ))
247345
@@ -534,10 +632,8 @@ def _do_training(self, request, job):
534632
535633 job .trainer = trainer
536634
537- # Start training
538- job .progress_queue .put (backend_pb2 .FineTuneProgressUpdate (
539- job_id = job .job_id , status = "training" , message = "Training started" ,
540- ))
635+ # Note: the "training started" event with total_steps is emitted by
636+ # on_train_begin in ProgressCallback, once the Trainer has computed max_steps.
541637
542638 resume_ckpt = request .resume_from_checkpoint if request .resume_from_checkpoint else None
543639 trainer .train (resume_from_checkpoint = resume_ckpt )
@@ -549,7 +645,7 @@ def _do_training(self, request, job):
549645
550646 job .completed = True
551647 # Sentinel to signal stream end
552- job .progress_queue .put (None )
648+ job .progress_queue .push (None )
553649
554650 def FineTuneProgress (self , request , context ):
555651 if self .active_job is None or self .active_job .job_id != request .job_id :
@@ -558,20 +654,14 @@ def FineTuneProgress(self, request, context):
558654 return
559655
560656 job = self .active_job
561- while True :
562- try :
563- update = job .progress_queue .get (timeout = 1.0 )
564- if update is None :
565- break
566- yield update
567- if update .status in ("completed" , "failed" , "stopped" ):
568- break
569- except queue .Empty :
570- if job .completed or job .stopped :
571- break
572- if not context .is_active ():
573- break
574- continue
657+ # Each stream gets its own iterator over the append-only event log.
658+ # Reconnecting clients replay from event 0, so no events are lost.
659+ for update in job .progress_queue .iter_from (
660+ start = 0 , timeout = 1.0 , is_alive = context .is_active ,
661+ ):
662+ yield update
663+ if update .status in ("completed" , "failed" , "stopped" ):
664+ break
575665
576666 def StopFineTune (self , request , context ):
577667 # No-op: stopping is handled by killing the backend process from Go.
0 commit comments