@@ -64,6 +64,7 @@ class _Job:
6464 checkpoint : str | None = None
6565 final_loss : float | None = None
6666 method : str | None = None
67+ worker : str | None = None # trained on this machine; inference routes there
6768 created : float = 0.0
6869 logs : list [str ] = field (default_factory = list ) # captured console lines
6970 live : str = "" # the in-progress (post-\r) line, e.g. the progress bar
@@ -74,7 +75,7 @@ class _Job:
7475 def record (self ) -> dict :
7576 """The serializable job record persisted to disk (survives restarts)."""
7677 return {"job_id" : self .job_id , "base_model" : self .base_model ,
77- "name" : self .name ,
78+ "name" : self .name , "worker" : self . worker ,
7879 "status" : self .status , "method" : self .method ,
7980 "error" : self .error , "checkpoint" : self .checkpoint ,
8081 "final_loss" : self .final_loss , "created" : self .created ,
@@ -84,7 +85,7 @@ def record(self) -> dict:
8485 @classmethod
8586 def from_record (cls , d : dict ) -> "_Job" :
8687 job = cls (job_id = d ["job_id" ], base_model = d .get ("base_model" , "?" ),
87- name = d .get ("name" , "" ),
88+ name = d .get ("name" , "" ), worker = d . get ( "worker" ),
8889 status = d .get ("status" , "succeeded" ), method = d .get ("method" ),
8990 error = d .get ("error" ), checkpoint = d .get ("checkpoint" ),
9091 final_loss = d .get ("final_loss" ), created = d .get ("created" , 0.0 ))
@@ -371,12 +372,18 @@ class _Worker:
371372 backend : str = "?"
372373 device : str = "?"
373374 gpus : int = 0
375+ gpu_name : str = ""
376+ vram_gb : float = 0.0
377+ ram_gb : float = 0.0
378+ cores : int = 0
374379 last_seen : float = 0.0
375380 inbox : "queue.Queue[str]" = field (default_factory = queue .Queue )
376381
377382 def info (self ) -> dict :
378383 return {"name" : self .name , "backend" : self .backend , "device" : self .device ,
379- "gpus" : self .gpus , "last_seen" : int (self .last_seen ),
384+ "gpus" : self .gpus , "gpu_name" : self .gpu_name ,
385+ "vram_gb" : self .vram_gb , "ram_gb" : self .ram_gb ,
386+ "cores" : self .cores , "last_seen" : int (self .last_seen ),
380387 "online" : (time .time () - self .last_seen ) < 90 ,
381388 "queued" : self .inbox .qsize ()}
382389
@@ -409,6 +416,7 @@ def __init__(self, *, backend: str, accelerator: str, device: str,
409416 self .jobs : dict [str , _Job ] = {}
410417 self .workers : dict [str , _Worker ] = {} # remote executors, keyed by name
411418 self .worker_socks : dict [str , object ] = {} # name → live WSConn
419+ self ._infer_waiters : dict [str , "queue.Queue[dict]" ] = {} # req id → reply
412420 self .datasets = DatasetStore (work_root / "datasets" )
413421 self .queue : "queue.Queue[str]" = queue .Queue ()
414422 self ._lock = threading .Lock () # job-store mutations
@@ -768,6 +776,7 @@ def submit(self, payload: dict) -> str:
768776 job = _Job (job_id = job_id , base_model = payload ["base_model" ],
769777 name = (payload .get ("name" ) or "" ).strip (),
770778 method = (payload .get ("config" ) or {}).get ("method" ),
779+ worker = (payload .get ("worker" ) or None ),
771780 created = int (time .time ()))
772781 job ._payload = payload
773782 with self ._lock :
@@ -792,6 +801,10 @@ def register_worker(self, body: dict) -> _Worker:
792801 w .backend = body .get ("backend" , w .backend )
793802 w .device = body .get ("device" , w .device )
794803 w .gpus = int (body .get ("gpus" ) or 0 )
804+ w .gpu_name = str (body .get ("gpu_name" ) or "" )
805+ w .vram_gb = float (body .get ("vram_gb" ) or 0 )
806+ w .ram_gb = float (body .get ("ram_gb" ) or 0 )
807+ w .cores = int (body .get ("cores" ) or 0 )
795808 w .last_seen = time .time ()
796809 return w
797810
@@ -875,6 +888,9 @@ def sender() -> None:
875888 if self .ingest_events (job , msg )["cancel" ]:
876889 conn .send_json ({"type" : "cancel" ,
877890 "job_id" : job .job_id })
891+ elif msg .get ("type" ) == "infer_result" :
892+ if (q := self ._infer_waiters .pop (msg .get ("id" , "" ), None )):
893+ q .put (msg )
878894 except (ConnectionError , OSError , json .JSONDecodeError ):
879895 pass # a dropped worker is normal — the studio just shows offline
880896 finally :
@@ -885,6 +901,38 @@ def sender() -> None:
885901 conn .close ()
886902 print (f"[hub] worker '{ name } , flush = True )
887903
904+ def worker_infer (self , job : _Job , req : dict , * , timeout : float = 240.0 ) -> str :
905+ """Run generate/chat on the machine that trained `job`, over its socket.
906+
907+ An mlx-trained adapter can't load into this box's torch stack (and vice
908+ versa) — the weights live where they trained, so inference goes there.
909+ """
910+ conn = self .worker_socks .get (job .worker )
911+ if conn is None :
912+ raise RuntimeError (
913+ f"this shadow was trained on machine '{ job .worker }
914+ f"offline — start `shadowlm worker` there to chat with it" )
915+ rid = uuid .uuid4 ().hex
916+ q : "queue.Queue[dict]" = queue .Queue ()
917+ self ._infer_waiters [rid ] = q
918+ try :
919+ conn .send_json ({** req , "id" : rid , "job_id" : job .job_id ,
920+ "base_model" : job .base_model })
921+ reply = q .get (timeout = timeout )
922+ except queue .Empty :
923+ raise RuntimeError (
924+ f"machine '{ job .worker } { timeout :.0f}
925+ ) from None
926+ except OSError as e :
927+ raise RuntimeError (
928+ f"machine '{ job .worker } { e }
929+ ) from None
930+ finally :
931+ self ._infer_waiters .pop (rid , None )
932+ if reply .get ("error" ):
933+ raise RuntimeError (f"machine '{ job .worker } { reply ['error' ]} )
934+ return reply .get ("text" , "" )
935+
888936 def push_cancel (self , job : _Job ) -> None :
889937 """Tell the executing worker to stop, right now, over its socket."""
890938 name = (getattr (job , "_payload" , None ) or {}).get ("worker" )
@@ -1280,6 +1328,16 @@ def do_POST(self): # noqa: N802
12801328 self ._send (200 , {"ok" : True })
12811329 elif parts == ["v1" , "generate" ]:
12821330 b = self ._body ()
1331+ # a worker-trained shadow answers on its own machine — the
1332+ # adapter format matches that backend, not necessarily ours
1333+ wjob = server .jobs .get (b .get ("adapter" ) or "" )
1334+ if wjob is not None and wjob .worker :
1335+ self ._send (200 , {"text" : server .worker_infer (wjob , {
1336+ "type" : "generate" , "prompt" : b ["prompt" ],
1337+ "max_new_tokens" : b .get ("max_new_tokens" , 256 ),
1338+ "temperature" : b .get ("temperature" , 0.7 ),
1339+ "top_p" : b .get ("top_p" , 0.95 )})})
1340+ return
12831341 with server ._model_lock :
12841342 m = server ._infer_model (b ["model" ], b .get ("adapter" ), b .get ("checkpoint" ))
12851343 text = m .generate (
@@ -1290,6 +1348,14 @@ def do_POST(self): # noqa: N802
12901348 self ._send (200 , {"text" : text })
12911349 elif parts == ["v1" , "chat" ]:
12921350 b = self ._body ()
1351+ wjob = server .jobs .get (b .get ("adapter" ) or "" )
1352+ if wjob is not None and wjob .worker :
1353+ self ._send (200 , {"text" : server .worker_infer (wjob , {
1354+ "type" : "chat" , "messages" : b ["messages" ],
1355+ "max_new_tokens" : b .get ("max_new_tokens" , 512 ),
1356+ "temperature" : b .get ("temperature" , 0.7 ),
1357+ "top_p" : b .get ("top_p" , 0.95 )})})
1358+ return
12931359 with server ._model_lock :
12941360 m = server ._infer_model (b ["model" ], b .get ("adapter" ), b .get ("checkpoint" ))
12951361 reply = m .chat (
0 commit comments