1use std::collections::HashMap;
34use std::ffi::OsString;
35use std::process::Stdio;
36use std::sync::atomic::{AtomicU64, Ordering};
37use std::sync::Arc;
38
39use std::time::Duration;
40
41use car_inference::stream::StreamEvent;
42use car_inference::tasks::generate::GenerateRequest;
43use car_inference::{
44 InferenceConfig, InferenceEngine, InferenceError, InferenceResult, InferenceTerminationAck,
45 LocalGenerationOffload, LocalLoadPreflight, LocalOffloadResult, LocalOffloadStream,
46 LocalWorkerAdmission, LocalWorkerResidency,
47};
48use serde::{Deserialize, Serialize};
49use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines};
50use tokio::process::{Child, ChildStdin, ChildStdout, Command};
51use tokio::sync::{mpsc, oneshot, Mutex};
52use tokio::time::Instant;
53
54type WorkerCancelSender = mpsc::UnboundedSender<oneshot::Sender<bool>>;
55
56struct WorkerControlRegistration {
57 inference_id: String,
58 controls: Arc<std::sync::Mutex<HashMap<String, WorkerCancelSender>>>,
59}
60
61impl WorkerControlRegistration {
62 fn install(
63 controls: Arc<std::sync::Mutex<HashMap<String, WorkerCancelSender>>>,
64 inference_id: String,
65 ) -> (Self, mpsc::UnboundedReceiver<oneshot::Sender<bool>>) {
66 let (sender, receiver) = mpsc::unbounded_channel();
67 controls
68 .lock()
69 .unwrap_or_else(std::sync::PoisonError::into_inner)
70 .insert(inference_id.clone(), sender);
71 (
72 Self {
73 inference_id,
74 controls,
75 },
76 receiver,
77 )
78 }
79}
80
81impl Drop for WorkerControlRegistration {
82 fn drop(&mut self) {
83 self.controls
84 .lock()
85 .unwrap_or_else(std::sync::PoisonError::into_inner)
86 .remove(&self.inference_id);
87 }
88}
89
90pub const WORKER_ENV: &str = "CAR_INFERENCE_WORKER";
95
96#[derive(Serialize, Deserialize)]
98enum WorkerRequest {
99 Generate {
101 request: Box<GenerateRequest>,
102 admission: LocalWorkerAdmission,
103 },
104 Stream {
106 request: Box<GenerateRequest>,
107 admission: LocalWorkerAdmission,
108 },
109}
110
111#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
112struct WorkerResidencyAck {
113 model_id: String,
114 measured_weights_bytes: u64,
115 retention: car_inference::backend_cache::BackendRetention,
116}
117
118#[derive(Serialize, Deserialize)]
120enum WorkerResponse {
121 Result {
123 result: Box<InferenceResult>,
124 residency: WorkerResidencyAck,
125 },
126 StreamStarted {
127 residency: WorkerResidencyAck,
128 },
129 Event(Box<StreamEvent>),
131 StreamEnd,
133 Error(String),
137 LocalResourceBlocked {
138 preflight: LocalLoadPreflight,
139 recovery: String,
140 },
141}
142
143enum Exchange<T> {
146 Ok(T),
147 Reported(InferenceError),
149 Dead(String),
151}
152
153struct WorkerProc {
155 child: Child,
156 stdin: ChildStdin,
157 stdout: Lines<BufReader<ChildStdout>>,
158 policy_generation: u64,
159 state_root: Option<std::path::PathBuf>,
160}
161
162#[derive(Clone)]
163struct WorkerResident {
164 allocation_id: String,
165 coordinator: Arc<car_inference::resource_policy::LocalAdmissionCoordinator>,
166 last_used: Option<Instant>,
169}
170
171type WorkerResidentMap = HashMap<(std::path::PathBuf, String), WorkerResident>;
172
173#[derive(Clone)]
174struct WorkerResidentOwner {
175 root: std::path::PathBuf,
176 logical_model_id: String,
177 allocation_id: String,
178 coordinator: Arc<car_inference::resource_policy::LocalAdmissionCoordinator>,
179}
180
181struct WorkerProcessGuard {
182 worker: Option<WorkerProc>,
183 residents: Vec<WorkerResidentOwner>,
184 candidate: Option<WorkerResidentOwner>,
185 resident_models: Arc<std::sync::Mutex<WorkerResidentMap>>,
186 teardown_started: bool,
187}
188
189impl WorkerProcessGuard {
190 fn new(
191 worker: WorkerProc,
192 resident_models: Arc<std::sync::Mutex<WorkerResidentMap>>,
193 candidate: Option<(std::path::PathBuf, String, String)>,
194 ) -> Self {
195 let mut residents = resident_models
196 .lock()
197 .unwrap_or_else(std::sync::PoisonError::into_inner)
198 .iter()
199 .map(|((root, logical_model_id), resident)| WorkerResidentOwner {
200 root: root.clone(),
201 logical_model_id: logical_model_id.clone(),
202 allocation_id: resident.allocation_id.clone(),
203 coordinator: resident.coordinator.clone(),
204 })
205 .collect::<Vec<_>>();
206 let mut candidate_owner = None;
207 if let Some((root, model_id, allocation_id)) =
208 candidate.filter(|(_, model_id, _)| !model_id.is_empty())
209 {
210 let root = car_inference::resource_policy::normalized_state_root_key(&root);
211 if !residents
212 .iter()
213 .any(|resident| resident.root == root && resident.logical_model_id == model_id)
214 {
215 if let Some(coordinator) =
216 car_inference::resource_policy::local_admission_for_scope(&root)
217 {
218 let owner = WorkerResidentOwner {
219 root,
220 allocation_id,
221 logical_model_id: model_id,
222 coordinator,
223 };
224 residents.push(owner.clone());
225 candidate_owner = Some(owner);
226 }
227 }
228 }
229 Self {
230 worker: Some(worker),
231 residents,
232 candidate: candidate_owner,
233 resident_models,
234 teardown_started: false,
235 }
236 }
237
238 fn worker_mut(&mut self) -> &mut WorkerProc {
239 self.worker.as_mut().expect("worker process owned")
240 }
241
242 fn charge_candidate(&self, measured_weights_bytes: u64) {
243 if let Some(candidate) = &self.candidate {
244 candidate
245 .coordinator
246 .mark_teardown_pending_allocation_with_charge(
247 &candidate.logical_model_id,
248 &candidate.allocation_id,
249 measured_weights_bytes,
250 );
251 }
252 }
253
254 fn clear_candidate(&mut self) {
255 let Some(candidate) = self.candidate.take() else {
256 return;
257 };
258 candidate
259 .coordinator
260 .finish_teardown_allocation(&candidate.logical_model_id, &candidate.allocation_id);
261 self.residents.retain(|resident| {
262 resident.root != candidate.root
263 || resident.logical_model_id != candidate.logical_model_id
264 || resident.allocation_id != candidate.allocation_id
265 });
266 }
267
268 fn charge_reported_model(
269 &mut self,
270 root: std::path::PathBuf,
271 model_id: &str,
272 allocation_id: String,
273 measured_weights_bytes: u64,
274 ) {
275 let root = car_inference::resource_policy::normalized_state_root_key(&root);
276 if self
277 .residents
278 .iter()
279 .any(|resident| resident.root == root && resident.logical_model_id == model_id)
280 {
281 return;
282 }
283 let Some(coordinator) = car_inference::resource_policy::local_admission_for_scope(&root)
284 else {
285 return;
286 };
287 coordinator.mark_teardown_pending_allocation_with_charge(
288 model_id,
289 &allocation_id,
290 measured_weights_bytes,
291 );
292 self.residents.push(WorkerResidentOwner {
293 root,
294 logical_model_id: model_id.to_string(),
295 allocation_id,
296 coordinator,
297 });
298 }
299
300 fn begin_teardown(&mut self) {
301 if self.teardown_started {
302 return;
303 }
304 self.teardown_started = true;
305 for resident in &self.residents {
306 resident.coordinator.mark_teardown_pending_allocation(
307 &resident.logical_model_id,
308 &resident.allocation_id,
309 );
310 }
311 }
312
313 fn finish_accounting(
314 residents: &[WorkerResidentOwner],
315 resident_models: &std::sync::Mutex<WorkerResidentMap>,
316 ) {
317 let mut tracked = resident_models
318 .lock()
319 .unwrap_or_else(std::sync::PoisonError::into_inner);
320 for resident in residents {
321 tracked.remove(&(resident.root.clone(), resident.logical_model_id.clone()));
322 resident
323 .coordinator
324 .finish_teardown_allocation(&resident.logical_model_id, &resident.allocation_id);
325 }
326 }
327
328 fn confirm_exited(mut self) {
329 self.worker.take();
330 Self::finish_accounting(&self.residents, &self.resident_models);
331 self.residents.clear();
332 }
333
334 fn return_to_slot(mut self, slot: &mut Option<WorkerProc>) {
335 *slot = self.worker.take();
336 self.residents.clear();
337 self.candidate = None;
338 }
339
340 async fn stop_and_confirm(mut self) -> Result<(), InferenceError> {
341 self.begin_teardown();
342 let worker = self.worker_mut();
343 match worker.child.try_wait().map_err(|error| {
344 InferenceError::InferenceFailed(format!(
345 "cannot inspect inference worker during teardown: {error}"
346 ))
347 })? {
348 Some(_) => {}
349 None => {
350 worker.child.kill().await.map_err(|error| {
351 InferenceError::InferenceFailed(format!(
352 "cannot stop inference worker during teardown: {error}"
353 ))
354 })?;
355 worker.child.wait().await.map_err(|error| {
356 InferenceError::InferenceFailed(format!(
357 "cannot reap inference worker during teardown: {error}"
358 ))
359 })?;
360 }
361 }
362 self.confirm_exited();
363 Ok(())
364 }
365}
366
367impl Drop for WorkerProcessGuard {
368 fn drop(&mut self) {
369 self.begin_teardown();
370 let Some(mut worker) = self.worker.take() else {
371 return;
372 };
373 let residents = std::mem::take(&mut self.residents);
374 let resident_models = self.resident_models.clone();
375 let _ = worker.child.start_kill();
376 if tokio::runtime::Handle::try_current().is_ok() {
377 tokio::spawn(async move {
378 if worker.child.wait().await.is_ok() {
379 Self::finish_accounting(&residents, &resident_models);
380 }
381 });
384 } else {
385 std::mem::forget(worker);
388 }
389 }
390}
391
392const IDLE_REAP_INTERVAL: Duration = Duration::from_secs(60);
397
398fn next_worker_allocation_scope() -> u64 {
399 static NEXT: AtomicU64 = AtomicU64::new(1);
400 NEXT.fetch_add(1, Ordering::Relaxed)
401}
402
403#[derive(Clone)]
406pub struct WorkerOffload {
407 inner: Arc<Mutex<Option<WorkerProc>>>,
408 program: OsString,
409 args: Arc<Vec<OsString>>,
410 policy_generation: Arc<AtomicU64>,
411 resident_models: Arc<std::sync::Mutex<WorkerResidentMap>>,
412 controls: Arc<std::sync::Mutex<HashMap<String, WorkerCancelSender>>>,
415 allocation_scope: u64,
416 #[cfg(test)]
417 release_delay_ms: Arc<AtomicU64>,
418}
419
420impl WorkerOffload {
421 fn allocation_id(&self, model_id: &str) -> String {
422 format!("worker:{}:{model_id}", self.allocation_scope)
423 }
424 pub fn new() -> std::io::Result<Self> {
427 let exe = std::env::current_exe()?;
428 Ok(Self::with_command(
429 exe,
430 vec![OsString::from("--mlx-worker")],
431 ))
432 }
433
434 pub fn with_command(program: impl Into<OsString>, args: Vec<OsString>) -> Self {
438 let me = Self {
439 inner: Arc::new(Mutex::new(None)),
440 program: program.into(),
441 args: Arc::new(args),
442 policy_generation: Arc::new(AtomicU64::new(1)),
443 resident_models: Arc::new(std::sync::Mutex::new(HashMap::new())),
444 controls: Arc::new(std::sync::Mutex::new(HashMap::new())),
445 allocation_scope: next_worker_allocation_scope(),
446 #[cfg(test)]
447 release_delay_ms: Arc::new(AtomicU64::new(0)),
448 };
449 me.spawn_idle_reaper();
450 me
451 }
452
453 fn spawn_idle_reaper(&self) {
464 if tokio::runtime::Handle::try_current().is_err() {
465 return;
468 }
469 let inner = Arc::clone(&self.inner);
470 let resident_models = Arc::clone(&self.resident_models);
471 tokio::spawn(async move {
472 let mut tick = tokio::time::interval(IDLE_REAP_INTERVAL);
473 tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
474 loop {
475 tick.tick().await;
476 if Arc::strong_count(&inner) == 1 {
480 let mut slot = inner.lock().await;
481 if let Some(worker) = slot.take() {
482 drop(WorkerProcessGuard::new(
483 worker,
484 resident_models.clone(),
485 None,
486 ));
487 }
488 return;
489 }
490 let mut slot = inner.lock().await;
491 let Some(p) = slot.as_mut() else { continue };
492 match p.child.try_wait() {
493 Ok(Some(_)) => {
494 let dead = slot.take().expect("idle worker checked");
495 WorkerProcessGuard::new(dead, resident_models.clone(), None)
496 .confirm_exited();
497 tracing::info!("reaped a dead idle on-device inference worker");
498 }
499 Ok(None) => {}
500 Err(error) => {
501 tracing::warn!(%error, "cannot inspect idle inference worker; retaining ownership and residency");
502 }
503 }
504 }
505 });
506 }
507
508 fn spawn(&self) -> Result<WorkerProc, InferenceError> {
509 let mut cmd = Command::new(&self.program);
510 cmd.args(self.args.iter())
511 .env(WORKER_ENV, "1")
512 .stdin(Stdio::piped())
513 .stdout(Stdio::piped())
514 .stderr(Stdio::inherit())
517 .kill_on_drop(true);
518 let mut child = cmd.spawn().map_err(|e| {
519 InferenceError::InferenceFailed(format!("failed to spawn inference worker: {e}"))
520 })?;
521 let stdin = child
522 .stdin
523 .take()
524 .ok_or_else(|| InferenceError::InferenceFailed("worker stdin unavailable".into()))?;
525 let stdout = child
526 .stdout
527 .take()
528 .ok_or_else(|| InferenceError::InferenceFailed("worker stdout unavailable".into()))?;
529 Ok(WorkerProc {
530 child,
531 stdin,
532 stdout: BufReader::new(stdout).lines(),
533 policy_generation: self.policy_generation.load(Ordering::Acquire),
534 state_root: None,
535 })
536 }
537
538 async fn take_or_spawn(
545 &self,
546 slot: &mut Option<WorkerProc>,
547 ) -> Result<WorkerProc, InferenceError> {
548 if let Some(mut p) = slot.take() {
549 match p.child.try_wait() {
555 Ok(None)
556 if p.policy_generation == self.policy_generation.load(Ordering::Acquire) =>
557 {
558 return Ok(p)
559 }
560 Ok(Some(_)) => {
561 WorkerProcessGuard::new(p, self.resident_models.clone(), None).confirm_exited();
562 }
563 Ok(None) => {
564 WorkerProcessGuard::new(p, self.resident_models.clone(), None)
565 .stop_and_confirm()
566 .await?;
567 }
568 Err(error) => {
569 drop(WorkerProcessGuard::new(
570 p,
571 self.resident_models.clone(),
572 None,
573 ));
574 return Err(InferenceError::InferenceFailed(format!(
575 "cannot inspect previous inference worker; teardown retained: {error}"
576 )));
577 }
578 }
579 }
580 self.spawn()
581 }
582
583 async fn take_or_spawn_for_scope(
584 &self,
585 slot: &mut Option<WorkerProc>,
586 state_root: &std::path::Path,
587 ) -> Result<WorkerProc, InferenceError> {
588 let state_root = car_inference::resource_policy::normalized_state_root_key(state_root);
589 let mut worker = self.take_or_spawn(slot).await?;
590 if worker
591 .state_root
592 .as_ref()
593 .is_some_and(|current| current != &state_root)
594 {
595 WorkerProcessGuard::new(worker, self.resident_models.clone(), None)
596 .stop_and_confirm()
597 .await?;
598 worker = self.spawn()?;
599 }
600 worker.state_root = Some(state_root);
601 Ok(worker)
602 }
603
604 fn remember_residency(
605 &self,
606 residency: &WorkerResidencyAck,
607 state_root: std::path::PathBuf,
608 last_used: Option<Instant>,
609 ) {
610 if residency.retention != car_inference::backend_cache::BackendRetention::Resident {
611 return;
612 }
613 let state_root = car_inference::resource_policy::normalized_state_root_key(&state_root);
614 let Some(coordinator) =
615 car_inference::resource_policy::local_admission_for_scope(&state_root)
616 else {
617 tracing::error!(root = %state_root.display(), model = %residency.model_id, "worker reported resident weights without a scoped admission owner");
618 return;
619 };
620 let allocation_id = self.allocation_id(&residency.model_id);
621 self.resident_models
622 .lock()
623 .unwrap_or_else(std::sync::PoisonError::into_inner)
624 .insert(
625 (state_root, residency.model_id.clone()),
626 WorkerResident {
627 allocation_id,
628 coordinator,
632 last_used,
633 },
634 );
635 }
636
637 fn mark_residency_used(
638 resident_models: &std::sync::Mutex<WorkerResidentMap>,
639 state_root: &std::path::Path,
640 model_id: &str,
641 ) {
642 let state_root = car_inference::resource_policy::normalized_state_root_key(state_root);
643 if let Some(resident) = resident_models
644 .lock()
645 .unwrap_or_else(std::sync::PoisonError::into_inner)
646 .get_mut(&(state_root, model_id.to_string()))
647 {
648 resident.last_used = Some(Instant::now());
649 }
650 }
651
652 async fn release_worker_locked(
655 &self,
656 slot: &mut Option<WorkerProc>,
657 model_id: &str,
658 ) -> Result<bool, InferenceError> {
659 let Some(worker) = slot.take() else {
660 return Err(InferenceError::InferenceFailed(format!(
661 "cannot confirm worker exit for resident model {model_id}: worker slot is empty"
662 )));
663 };
664 let mut release = WorkerProcessGuard::new(worker, self.resident_models.clone(), None);
665 release.begin_teardown();
666 #[cfg(test)]
667 tokio::time::sleep(Duration::from_millis(
668 self.release_delay_ms.load(Ordering::Acquire),
669 ))
670 .await;
671 let worker = release.worker_mut();
672 match worker.child.try_wait().map_err(|error| {
673 InferenceError::InferenceFailed(format!(
674 "cannot inspect worker before releasing {model_id}: {error}"
675 ))
676 })? {
677 Some(_) => {}
678 None => {
679 worker.child.kill().await.map_err(|error| {
680 InferenceError::InferenceFailed(format!(
681 "cannot stop worker before releasing {model_id}: {error}"
682 ))
683 })?;
684 worker.child.wait().await.map_err(|error| {
685 InferenceError::InferenceFailed(format!(
686 "cannot reap worker before releasing {model_id}: {error}"
687 ))
688 })?;
689 }
690 }
691 release.confirm_exited();
692 Ok(true)
693 }
694
695 async fn release_model_if_idle(
699 &self,
700 model_id: &str,
701 idle_for: Duration,
702 ) -> Result<Vec<String>, InferenceError> {
703 let mut slot = self.inner.lock().await;
704 let resident_ids = {
705 let residents = self
706 .resident_models
707 .lock()
708 .unwrap_or_else(std::sync::PoisonError::into_inner);
709 if !residents
710 .keys()
711 .any(|(_, resident_model)| resident_model == model_id)
712 {
713 return Ok(Vec::new());
714 }
715 let now = Instant::now();
716 if residents.values().any(|resident| {
717 resident.last_used.is_none_or(|last_used| {
718 now.checked_duration_since(last_used)
719 .is_none_or(|elapsed| elapsed < idle_for)
720 })
721 }) {
722 return Ok(Vec::new());
723 }
724 let mut ids = residents
725 .keys()
726 .map(|(_, resident_model)| resident_model.clone())
727 .collect::<Vec<_>>();
728 ids.sort();
729 ids.dedup();
730 ids
731 };
732 self.release_worker_locked(&mut slot, model_id).await?;
733 Ok(resident_ids)
734 }
735}
736
737async fn write_line<W, T>(w: &mut W, value: &T) -> std::io::Result<()>
738where
739 W: AsyncWrite + Unpin,
740 T: Serialize,
741{
742 let mut line = serde_json::to_string(value).map_err(std::io::Error::other)?;
743 line.push('\n');
744 w.write_all(line.as_bytes()).await?;
745 w.flush().await
746}
747
748async fn do_generate(
751 proc: &mut WorkerProc,
752 request: GenerateRequest,
753 admission: LocalWorkerAdmission,
754) -> Exchange<(InferenceResult, WorkerResidencyAck)> {
755 let req = WorkerRequest::Generate {
756 request: Box::new(request),
757 admission,
758 };
759 if let Err(e) = write_line(&mut proc.stdin, &req).await {
760 return Exchange::Dead(format!("write to inference worker failed: {e}"));
761 }
762 match proc.stdout.next_line().await {
763 Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
764 Ok(WorkerResponse::Result { result, residency }) => Exchange::Ok((*result, residency)),
765 Ok(WorkerResponse::Error(msg)) => {
766 Exchange::Reported(InferenceError::InferenceFailed(msg))
767 }
768 Ok(WorkerResponse::LocalResourceBlocked {
769 preflight,
770 recovery,
771 }) => Exchange::Reported(InferenceError::LocalResourceBlocked {
772 preflight,
773 recovery,
774 }),
775 Ok(_) => Exchange::Dead("inference worker sent an unexpected response frame".into()),
776 Err(e) => Exchange::Dead(format!("inference worker sent invalid JSON: {e}")),
777 },
778 Ok(None) => Exchange::Dead("inference worker exited mid-request (EOF)".into()),
779 Err(e) => Exchange::Dead(format!("read from inference worker failed: {e}")),
780 }
781}
782
783#[async_trait::async_trait]
784impl LocalGenerationOffload for WorkerOffload {
785 async fn generate(&self, _request: GenerateRequest) -> Result<InferenceResult, InferenceError> {
786 Err(InferenceError::InferenceFailed(
787 "WorkerOffload requires admission-aware dispatch".into(),
788 ))
789 }
790
791 async fn stream(
792 &self,
793 _request: GenerateRequest,
794 ) -> Result<tokio::sync::mpsc::Receiver<StreamEvent>, InferenceError> {
795 Err(InferenceError::InferenceFailed(
796 "WorkerOffload requires admission-aware dispatch".into(),
797 ))
798 }
799
800 async fn generate_admitted(
801 &self,
802 request: GenerateRequest,
803 admission: LocalWorkerAdmission,
804 ) -> Result<LocalOffloadResult, InferenceError> {
805 let mut slot = self.inner.lock().await;
811 let expected_model = request.model.clone().unwrap_or_default();
812 let candidate = (
813 admission.state_root.clone(),
814 expected_model.clone(),
815 self.allocation_id(&expected_model),
816 );
817 let proc = self
818 .take_or_spawn_for_scope(&mut slot, &admission.state_root)
819 .await?;
820 let mut ownership =
821 WorkerProcessGuard::new(proc, self.resident_models.clone(), Some(candidate));
822 ownership.charge_candidate(admission.measured_weights_bytes);
823 let state_root = admission.state_root.clone();
824 enum GenerateSelection {
825 Exchange(Exchange<(InferenceResult, WorkerResidencyAck)>),
826 Cancel(oneshot::Sender<bool>),
827 }
828 let inference_id = car_inference::current_inference_control_id();
829 let termination = car_inference::current_controlled_termination_token();
830 let (control_registration, mut cancel_rx) = inference_id
831 .map(|id| WorkerControlRegistration::install(self.controls.clone(), id))
832 .unzip();
833 let selected = if let Some(cancel_rx) = cancel_rx.as_mut() {
834 let worker = ownership.worker_mut();
835 tokio::select! {
836 exchange = do_generate(worker, request, admission) => GenerateSelection::Exchange(exchange),
837 request = cancel_rx.recv() => match request {
838 Some(ack) => GenerateSelection::Cancel(ack),
839 None => unreachable!("control registration owns the sender"),
840 },
841 }
842 } else {
843 GenerateSelection::Exchange(
844 do_generate(ownership.worker_mut(), request, admission).await,
845 )
846 };
847 let exchange = match selected {
848 GenerateSelection::Exchange(exchange) => exchange,
849 GenerateSelection::Cancel(ack) => {
850 let confirmed = ownership.stop_and_confirm().await.is_ok();
851 if confirmed {
852 if let Some(token) = termination.as_ref() {
853 token.confirm_exact_backend_termination();
854 }
855 }
856 let _ = ack.send(confirmed);
857 drop(control_registration);
858 return Err(if confirmed {
859 InferenceError::ControlledTermination
860 } else {
861 InferenceError::InferenceFailed(
862 "isolated inference worker termination was not confirmed".into(),
863 )
864 });
865 }
866 };
867 drop(control_registration);
868 match exchange {
869 Exchange::Ok((ir, residency)) => {
870 if residency.model_id != expected_model {
871 ownership.charge_reported_model(
872 state_root,
873 &residency.model_id,
874 self.allocation_id(&residency.model_id),
875 residency.measured_weights_bytes,
876 );
877 drop(ownership);
878 return Err(InferenceError::InferenceFailed(format!(
879 "local worker acknowledged model '{}' for requested '{}'",
880 residency.model_id, expected_model
881 )));
882 }
883 if residency.retention != car_inference::backend_cache::BackendRetention::Resident {
884 ownership.clear_candidate();
885 }
886 self.remember_residency(&residency, state_root, Some(Instant::now()));
887 ownership.return_to_slot(&mut slot);
888 Ok(LocalOffloadResult {
889 result: ir,
890 residency: LocalWorkerResidency {
891 model_id: residency.model_id,
892 measured_weights_bytes: residency.measured_weights_bytes,
893 },
894 retention: residency.retention,
895 })
896 }
897 Exchange::Reported(error) => {
898 if matches!(error, InferenceError::LocalResourceBlocked { .. }) {
899 ownership.clear_candidate();
901 ownership.return_to_slot(&mut slot);
902 } else {
903 drop(ownership);
907 }
908 Err(error)
909 }
910 Exchange::Dead(msg) => {
911 let exit = match tokio::time::timeout(
915 Duration::from_millis(100),
916 ownership.worker_mut().child.wait(),
917 )
918 .await
919 {
920 Ok(Ok(status)) => format!("; worker {status}"),
921 Ok(Err(error)) => format!("; exit status unavailable: {error}"),
922 Err(_) => "; worker exit not yet observed".to_string(),
923 };
924 drop(ownership);
928 Err(InferenceError::InferenceFailed(format!(
929 "on-device inference worker connection failed; \
930 this request failed but the daemon is up: {msg}{exit}. \
931 The next request starts a new worker."
932 )))
933 }
934 }
935 }
936
937 async fn stream_admitted(
938 &self,
939 request: GenerateRequest,
940 admission: LocalWorkerAdmission,
941 ) -> Result<LocalOffloadStream, InferenceError> {
942 let mut guard = self.inner.clone().lock_owned().await;
950 let expected_model = request.model.clone().unwrap_or_default();
951 let candidate = (
952 admission.state_root.clone(),
953 expected_model.clone(),
954 self.allocation_id(&expected_model),
955 );
956 let proc = self
957 .take_or_spawn_for_scope(&mut guard, &admission.state_root)
958 .await?;
959 let mut ownership =
960 WorkerProcessGuard::new(proc, self.resident_models.clone(), Some(candidate));
961 ownership.charge_candidate(admission.measured_weights_bytes);
962 let state_root = admission.state_root.clone();
963 let inference_id = car_inference::current_inference_control_id();
964 let termination = car_inference::current_controlled_termination_token();
965 let (control_registration, mut cancel_rx) = inference_id
966 .map(|id| WorkerControlRegistration::install(self.controls.clone(), id))
967 .unzip();
968 {
971 let req = WorkerRequest::Stream {
972 request: Box::new(request),
973 admission,
974 };
975 if let Err(e) = write_line(&mut ownership.worker_mut().stdin, &req).await {
976 return Err(InferenceError::InferenceFailed(format!(
977 "write to inference worker failed: {e}"
978 )));
979 }
980 }
981
982 enum StreamStartSelection {
983 Frame(std::io::Result<Option<String>>),
984 Cancel(oneshot::Sender<bool>),
985 }
986 let stream_start = if let Some(cancel_rx) = cancel_rx.as_mut() {
990 let worker = ownership.worker_mut();
991 tokio::select! {
992 frame = worker.stdout.next_line() => StreamStartSelection::Frame(frame),
993 request = cancel_rx.recv() => match request {
994 Some(ack) => StreamStartSelection::Cancel(ack),
995 None => unreachable!("control registration owns the sender"),
996 },
997 }
998 } else {
999 StreamStartSelection::Frame(ownership.worker_mut().stdout.next_line().await)
1000 };
1001 let stream_start = match stream_start {
1002 StreamStartSelection::Frame(frame) => frame,
1003 StreamStartSelection::Cancel(ack) => {
1004 let confirmed = ownership.stop_and_confirm().await.is_ok();
1005 if confirmed {
1006 if let Some(token) = termination.as_ref() {
1007 token.confirm_exact_backend_termination();
1008 }
1009 }
1010 let _ = ack.send(confirmed);
1011 return Err(if confirmed {
1012 InferenceError::ControlledTermination
1013 } else {
1014 InferenceError::InferenceFailed(
1015 "isolated inference worker termination before StreamStarted was not confirmed"
1016 .into(),
1017 )
1018 });
1019 }
1020 };
1021 let residency = match stream_start {
1022 Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
1023 Ok(WorkerResponse::StreamStarted { residency }) => residency,
1024 Ok(WorkerResponse::LocalResourceBlocked {
1025 preflight,
1026 recovery,
1027 }) => {
1028 ownership.clear_candidate();
1029 ownership.return_to_slot(&mut guard);
1030 return Err(InferenceError::LocalResourceBlocked {
1031 preflight,
1032 recovery,
1033 });
1034 }
1035 Ok(WorkerResponse::Error(message)) => {
1036 drop(ownership);
1038 return Err(InferenceError::InferenceFailed(message));
1039 }
1040 Ok(_) => {
1041 return Err(InferenceError::InferenceFailed(
1042 "inference worker streamed before a successful load acknowledgement".into(),
1043 ));
1044 }
1045 Err(error) => {
1046 return Err(InferenceError::InferenceFailed(format!(
1047 "inference worker sent invalid stream acknowledgement: {error}"
1048 )));
1049 }
1050 },
1051 Ok(None) => {
1052 return Err(InferenceError::InferenceFailed(
1053 "inference worker exited before loading the streaming model".into(),
1054 ));
1055 }
1056 Err(error) => {
1057 return Err(InferenceError::InferenceFailed(format!(
1058 "failed reading inference worker load acknowledgement: {error}"
1059 )));
1060 }
1061 };
1062 if residency.model_id != expected_model {
1063 ownership.charge_reported_model(
1064 state_root,
1065 &residency.model_id,
1066 self.allocation_id(&residency.model_id),
1067 residency.measured_weights_bytes,
1068 );
1069 drop(ownership);
1070 return Err(InferenceError::InferenceFailed(format!(
1071 "local worker acknowledged model '{}' for requested '{}'",
1072 residency.model_id, expected_model
1073 )));
1074 }
1075 if residency.retention != car_inference::backend_cache::BackendRetention::Resident {
1076 ownership.clear_candidate();
1077 }
1078 self.remember_residency(&residency, state_root.clone(), None);
1079 let resident_models = self.resident_models.clone();
1080 let resident_model_id = residency.model_id.clone();
1081
1082 let (tx, rx) = tokio::sync::mpsc::channel::<StreamEvent>(64);
1083 tokio::spawn(async move {
1084 let _control_registration = control_registration;
1085 let mut ownership = Some(ownership);
1086 let mut clean = false;
1087 loop {
1088 enum StreamSelection {
1089 Frame(std::io::Result<Option<String>>),
1090 Cancel(oneshot::Sender<bool>),
1091 }
1092 let selected = if let Some(cancel_rx) = cancel_rx.as_mut() {
1093 let worker = ownership
1094 .as_mut()
1095 .expect("worker ownership remains until terminal")
1096 .worker_mut();
1097 tokio::select! {
1098 frame = worker.stdout.next_line() => StreamSelection::Frame(frame),
1099 request = cancel_rx.recv() => match request {
1100 Some(ack) => StreamSelection::Cancel(ack),
1101 None => unreachable!("control registration owns the sender"),
1102 },
1103 }
1104 } else {
1105 StreamSelection::Frame(
1106 ownership
1107 .as_mut()
1108 .expect("worker ownership remains until terminal")
1109 .worker_mut()
1110 .stdout
1111 .next_line()
1112 .await,
1113 )
1114 };
1115 let frame = match selected {
1116 StreamSelection::Frame(frame) => frame,
1117 StreamSelection::Cancel(ack) => {
1118 let confirmed = ownership
1119 .take()
1120 .expect("worker ownership remains until terminal")
1121 .stop_and_confirm()
1122 .await
1123 .is_ok();
1124 if confirmed {
1125 if let Some(token) = termination.as_ref() {
1126 token.confirm_exact_backend_termination();
1127 }
1128 }
1129 let _ = ack.send(confirmed);
1130 return;
1131 }
1132 };
1133 match frame {
1134 Ok(Some(line)) => match serde_json::from_str::<WorkerResponse>(&line) {
1135 Ok(WorkerResponse::Event(ev)) => {
1136 if tx.send(*ev).await.is_err() {
1137 break; }
1139 }
1140 Ok(WorkerResponse::StreamEnd) => {
1141 clean = true;
1142 break;
1143 }
1144 Ok(WorkerResponse::Error(msg)) => {
1145 tracing::warn!(error = %msg, "inference worker stream error");
1146 let _ = tx.send(StreamEvent::StopReason("error".into())).await;
1149 break;
1150 }
1151 Ok(WorkerResponse::Result { .. })
1152 | Ok(WorkerResponse::StreamStarted { .. })
1153 | Ok(WorkerResponse::LocalResourceBlocked { .. })
1154 | Err(_) => {
1155 let _ = tx.send(StreamEvent::StopReason("error".into())).await;
1157 break;
1158 }
1159 },
1160 Ok(None) | Err(_) => {
1161 let _ = tx.send(StreamEvent::StopReason("error".into())).await;
1163 break;
1164 }
1165 }
1166 }
1167 if clean {
1168 Self::mark_residency_used(&resident_models, &state_root, &resident_model_id);
1171 ownership
1173 .take()
1174 .expect("worker ownership remains on clean end")
1175 .return_to_slot(&mut guard);
1176 } else {
1177 drop(ownership.take());
1180 }
1181 });
1182 Ok(LocalOffloadStream {
1183 events: rx,
1184 residency: LocalWorkerResidency {
1185 model_id: residency.model_id,
1186 measured_weights_bytes: residency.measured_weights_bytes,
1187 },
1188 retention: residency.retention,
1189 })
1190 }
1191
1192 fn refresh_resource_policy(&self, generation: u64) {
1193 self.policy_generation.store(generation, Ordering::Release);
1194 if let Ok(mut slot) = self.inner.try_lock() {
1195 if let Some(worker) = slot.take() {
1196 drop(WorkerProcessGuard::new(
1197 worker,
1198 self.resident_models.clone(),
1199 None,
1200 ));
1201 }
1202 }
1203 }
1204
1205 fn resident_allocation_id(&self, model_id: &str) -> Option<String> {
1206 Some(self.allocation_id(model_id))
1207 }
1208
1209 async fn resident_models(&self) -> Vec<String> {
1210 self.resident_models
1211 .lock()
1212 .unwrap_or_else(std::sync::PoisonError::into_inner)
1213 .keys()
1214 .map(|(_, model_id)| model_id.clone())
1215 .collect()
1216 }
1217
1218 async fn release_model(&self, model_id: &str) -> Result<bool, InferenceError> {
1219 let mut slot = self.inner.lock().await;
1220 let resident = self
1221 .resident_models
1222 .lock()
1223 .unwrap_or_else(std::sync::PoisonError::into_inner)
1224 .keys()
1225 .any(|(_, resident_model)| resident_model == model_id);
1226 if !resident {
1227 return Ok(false);
1228 }
1229 self.release_worker_locked(&mut slot, model_id).await
1230 }
1231
1232 async fn release_idle_models(&self, idle_for: Duration) -> Vec<String> {
1233 let candidate = self
1237 .resident_models
1238 .lock()
1239 .unwrap_or_else(std::sync::PoisonError::into_inner)
1240 .keys()
1241 .next()
1242 .map(|(_, model_id)| model_id.clone());
1243 let Some(candidate) = candidate else {
1244 return Vec::new();
1245 };
1246 match self.release_model_if_idle(&candidate, idle_for).await {
1247 Ok(released) => released,
1248 Err(error) => {
1249 tracing::warn!(%error, "failed to release idle on-device inference worker");
1250 Vec::new()
1251 }
1252 }
1253 }
1254
1255 async fn terminate_inference(&self, inference_id: &str) -> InferenceTerminationAck {
1256 let sender = self
1257 .controls
1258 .lock()
1259 .unwrap_or_else(std::sync::PoisonError::into_inner)
1260 .get(inference_id)
1261 .cloned();
1262 let Some(sender) = sender else {
1263 return InferenceTerminationAck::Unconfirmed;
1264 };
1265 let (ack_tx, ack_rx) = oneshot::channel();
1266 if sender.send(ack_tx).is_err() {
1267 return InferenceTerminationAck::Unconfirmed;
1268 }
1269 match ack_rx.await {
1270 Ok(true) => InferenceTerminationAck::Confirmed,
1271 Ok(false) | Err(_) => InferenceTerminationAck::Unconfirmed,
1272 }
1273 }
1274}
1275
1276pub async fn run_mlx_worker() {
1287 std::env::set_var(WORKER_ENV, "1");
1290
1291 let mut engine: Option<(u64, std::path::PathBuf, Arc<InferenceEngine>)> = None;
1292 let mut lines = BufReader::new(tokio::io::stdin()).lines();
1293 let mut stdout = tokio::io::stdout();
1294
1295 while let Ok(Some(line)) = lines.next_line().await {
1296 let line = line.trim();
1297 if line.is_empty() {
1298 continue;
1299 }
1300 let req: WorkerRequest = match serde_json::from_str(line) {
1301 Ok(r) => r,
1302 Err(e) => {
1303 let _ = write_line(
1304 &mut stdout,
1305 &WorkerResponse::Error(format!("malformed worker request: {e}")),
1306 )
1307 .await;
1308 continue;
1309 }
1310 };
1311 let admission = match &req {
1312 WorkerRequest::Generate { admission, .. } | WorkerRequest::Stream { admission, .. } => {
1313 admission
1314 }
1315 };
1316 let recreate = engine.as_ref().is_none_or(|(generation, root, _)| {
1317 *generation != admission.policy_generation || root != &admission.state_root
1318 });
1319 if recreate {
1320 let mut config = InferenceConfig::default();
1321 config.state_root = admission.state_root.clone();
1322 let candidate = Arc::new(InferenceEngine::new(config));
1323 candidate.apply_local_resource_policy(admission.policy.clone());
1324 engine = Some((
1325 admission.policy_generation,
1326 admission.state_root.clone(),
1327 candidate,
1328 ));
1329 }
1330 let active_engine = Arc::clone(&engine.as_ref().expect("worker engine initialized").2);
1331 active_engine.apply_local_resource_policy(admission.policy.clone());
1332
1333 match req {
1334 WorkerRequest::Generate {
1335 request,
1336 admission: _,
1337 } => {
1338 let model_id = request.model.clone().unwrap_or_default();
1344 let resp = match active_engine.generate_tracked(*request).await {
1345 Ok(ir) => {
1346 let measured = measured_worker_model_bytes(&active_engine, &model_id);
1347 let retention = worker_model_retention(&active_engine, &model_id);
1348 WorkerResponse::Result {
1349 result: Box::new(ir),
1350 residency: WorkerResidencyAck {
1351 model_id,
1352 measured_weights_bytes: measured,
1353 retention,
1354 },
1355 }
1356 }
1357 Err(InferenceError::LocalResourceBlocked {
1358 preflight,
1359 recovery,
1360 }) => WorkerResponse::LocalResourceBlocked {
1361 preflight,
1362 recovery,
1363 },
1364 Err(e) => WorkerResponse::Error(e.to_string()),
1365 };
1366 if write_line(&mut stdout, &resp).await.is_err() {
1367 break; }
1369 }
1370 WorkerRequest::Stream {
1371 request,
1372 admission: _,
1373 } => {
1374 let model_id = request.model.clone().unwrap_or_default();
1375 match active_engine.generate_tracked_stream(*request).await {
1376 Ok(mut tracked) => {
1377 let measured = measured_worker_model_bytes(&active_engine, &model_id);
1378 let retention = worker_model_retention(&active_engine, &model_id);
1379 if write_line(
1380 &mut stdout,
1381 &WorkerResponse::StreamStarted {
1382 residency: WorkerResidencyAck {
1383 model_id,
1384 measured_weights_bytes: measured,
1385 retention,
1386 },
1387 },
1388 )
1389 .await
1390 .is_err()
1391 {
1392 break;
1393 }
1394 let mut broke = false;
1395 while let Some(ev) = tracked.events.recv().await {
1396 if write_line(&mut stdout, &WorkerResponse::Event(Box::new(ev)))
1397 .await
1398 .is_err()
1399 {
1400 broke = true;
1401 break;
1402 }
1403 }
1404 if broke {
1405 break;
1406 }
1407 if write_line(&mut stdout, &WorkerResponse::StreamEnd)
1408 .await
1409 .is_err()
1410 {
1411 break;
1412 }
1413 }
1414 Err(InferenceError::LocalResourceBlocked {
1415 preflight,
1416 recovery,
1417 }) => {
1418 if write_line(
1419 &mut stdout,
1420 &WorkerResponse::LocalResourceBlocked {
1421 preflight,
1422 recovery,
1423 },
1424 )
1425 .await
1426 .is_err()
1427 {
1428 break;
1429 }
1430 }
1431 Err(e) => {
1432 if write_line(&mut stdout, &WorkerResponse::Error(e.to_string()))
1433 .await
1434 .is_err()
1435 {
1436 break;
1437 }
1438 }
1439 }
1440 }
1441 }
1442 }
1443}
1444
1445fn measured_worker_model_bytes(engine: &InferenceEngine, model_id: &str) -> u64 {
1446 let Some(schema) = engine
1447 .unified_registry
1448 .get(model_id)
1449 .or_else(|| engine.unified_registry.find_by_name(model_id))
1450 else {
1451 return 0;
1452 };
1453 car_inference::backend_cache::estimate_model_size(&engine.config.models_dir.join(&schema.name))
1454}
1455
1456fn worker_model_retention(
1457 engine: &InferenceEngine,
1458 model_id: &str,
1459) -> car_inference::backend_cache::BackendRetention {
1460 engine.local_model_retention(model_id)
1461}
1462
1463#[cfg(test)]
1464mod tests {
1465 use super::*;
1466
1467 #[test]
1468 fn local_model_preflight_worker_rechecks_before_allocation() {
1469 let source = include_str!("inference_worker.rs");
1470 assert!(source.contains("LOCAL_ADMISSION_BOUNDARY:worker-side-allocation"));
1471 assert!(source.contains("active_engine.generate_tracked(*request).await"));
1472 }
1473
1474 #[test]
1475 fn idle_backend_eviction_calls_worker_idle_release() {
1476 let source = include_str!("handler.rs");
1477 assert!(source.contains("release_idle_models"));
1478 }
1479
1480 const RESULT_LINE: &str = r#"{"Result":{"result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"stub","latency_ms":1,"usage":null},"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"resident"}}}"#;
1483 #[cfg(unix)]
1487 const TRANSIENT_RESULT_LINE: &str = r#"{"Result":{"result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"stub","latency_ms":1,"usage":null},"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"transient"}}}"#;
1488 #[cfg(unix)]
1489 const STREAM_STARTED_LINE: &str = r#"{"StreamStarted":{"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"resident"}}}"#;
1490
1491 #[cfg(unix)]
1492 fn test_admission() -> LocalWorkerAdmission {
1493 LocalWorkerAdmission {
1494 policy: car_inference::ResourcePolicy::custom_gb(8.0).unwrap(),
1495 policy_generation: 1,
1496 state_root: std::path::PathBuf::from("/tmp/car-worker-root"),
1497 measured_weights_bytes: 1,
1498 }
1499 }
1500
1501 #[test]
1502 fn envelope_serde_round_trips() {
1503 let admission = car_inference::LocalWorkerAdmission {
1505 policy: car_inference::ResourcePolicy::custom_gb(8.0).unwrap(),
1506 policy_generation: 7,
1507 state_root: std::path::PathBuf::from("/tmp/car-worker-root"),
1508 measured_weights_bytes: 3 * 1024 * 1024 * 1024,
1509 };
1510 let req = WorkerRequest::Generate {
1511 request: Box::new(GenerateRequest {
1512 prompt: "hi".into(),
1513 ..Default::default()
1514 }),
1515 admission: admission.clone(),
1516 };
1517 let json = serde_json::to_string(&req).unwrap();
1518 assert!(json.starts_with(r#"{"Generate":"#), "got {json}");
1519 let decoded = serde_json::from_str::<WorkerRequest>(&json).unwrap();
1520 let WorkerRequest::Generate {
1521 admission: decoded_admission,
1522 ..
1523 } = decoded
1524 else {
1525 panic!("wrong request variant")
1526 };
1527 assert_eq!(decoded_admission, admission);
1528
1529 assert!(matches!(
1531 serde_json::from_str::<WorkerResponse>(RESULT_LINE).unwrap(),
1532 WorkerResponse::Result { .. }
1533 ));
1534 let ev = WorkerResponse::Event(Box::new(StreamEvent::TextDelta("hi".into())));
1535 let ev_json = serde_json::to_string(&ev).unwrap();
1536 assert_eq!(ev_json, r#"{"Event":{"TextDelta":"hi"}}"#);
1537 assert_eq!(
1538 serde_json::to_string(&WorkerResponse::StreamEnd).unwrap(),
1539 r#""StreamEnd""#
1540 );
1541 let sr = StreamEvent::StopReason("length".into());
1543 let sr2: StreamEvent = serde_json::from_str(&serde_json::to_string(&sr).unwrap()).unwrap();
1544 assert!(matches!(sr2, StreamEvent::StopReason(s) if s == "length"));
1545 }
1546
1547 #[test]
1548 fn worker_protocol_preserves_structured_resource_rejection_and_residency_ack() {
1549 let preflight = car_inference::LocalLoadPreflight {
1550 model_id: "mlx/test".into(),
1551 estimate: car_inference::ModelMemoryEstimate {
1552 weights_mb: 6144,
1553 runtime_overhead_mb: 256,
1554 context_overhead_mb: 128,
1555 transient_margin_mb: 256,
1556 estimated_peak_mb: 6784,
1557 evidence: car_inference::ModelResourceEvidence::FileSystemMeasured,
1558 },
1559 configured_ceiling_mb: 4096,
1560 resident_model_mb: 0,
1561 active_reservations_mb: 0,
1562 estimated_incremental_mb: 6144,
1563 accelerator_total_mb: None,
1564 accelerator_resident_mb: None,
1565 accelerator_incremental_mb: None,
1566 live_available_mb: Some(8192),
1567 emergency_reserve_mb: 4096,
1568 verdict: car_inference::LocalLoadVerdict::ExceedsConfiguredCeiling,
1569 };
1570 let blocked = WorkerResponse::LocalResourceBlocked {
1571 preflight: preflight.clone(),
1572 recovery: "choose a smaller model".into(),
1573 };
1574 let decoded: WorkerResponse =
1575 serde_json::from_str(&serde_json::to_string(&blocked).unwrap()).unwrap();
1576 assert!(matches!(
1577 decoded,
1578 WorkerResponse::LocalResourceBlocked {
1579 preflight: actual,
1580 ..
1581 } if actual == preflight
1582 ));
1583
1584 let started = WorkerResponse::StreamStarted {
1585 residency: WorkerResidencyAck {
1586 model_id: "mlx/test".into(),
1587 measured_weights_bytes: 3 * 1024 * 1024 * 1024,
1588 retention: car_inference::backend_cache::BackendRetention::Resident,
1589 },
1590 };
1591 assert!(matches!(
1592 serde_json::from_str::<WorkerResponse>(&serde_json::to_string(&started).unwrap())
1593 .unwrap(),
1594 WorkerResponse::StreamStarted { .. }
1595 ));
1596 }
1597
1598 #[cfg(unix)]
1602 fn sh(script: &str) -> WorkerOffload {
1603 WorkerOffload::with_command("sh", vec![OsString::from("-c"), OsString::from(script)])
1604 }
1605
1606 #[cfg(unix)]
1607 fn stub_request() -> GenerateRequest {
1608 GenerateRequest {
1609 model: Some("stub".into()),
1610 ..Default::default()
1611 }
1612 }
1613
1614 #[cfg(unix)]
1615 #[tokio::test]
1616 async fn generate_round_trips_and_reuses_the_worker() {
1617 let off = sh(&format!(
1620 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1621 ));
1622 let r1 = off
1623 .generate_admitted(stub_request(), test_admission())
1624 .await
1625 .unwrap();
1626 assert_eq!(r1.result.text, "pong");
1627 let r2 = off
1628 .generate_admitted(stub_request(), test_admission())
1629 .await
1630 .unwrap();
1631 assert_eq!(r2.result.text, "pong");
1632 }
1633
1634 #[cfg(unix)]
1635 #[tokio::test]
1636 async fn exact_worker_cancel_confirms_only_after_kill_and_wait() {
1637 let off = sh("while IFS= read -r line; do sleep 60; done");
1638 let caller = off.clone();
1639 let generation = tokio::spawn(async move {
1640 car_inference::scope_inference_control_id("exact-worker-id".to_string(), async move {
1641 caller
1642 .generate_admitted(stub_request(), test_admission())
1643 .await
1644 })
1645 .await
1646 });
1647
1648 tokio::time::timeout(Duration::from_secs(2), async {
1649 loop {
1650 if off
1651 .controls
1652 .lock()
1653 .unwrap_or_else(std::sync::PoisonError::into_inner)
1654 .contains_key("exact-worker-id")
1655 {
1656 break;
1657 }
1658 tokio::task::yield_now().await;
1659 }
1660 })
1661 .await
1662 .expect("worker request must register its exact control ID");
1663
1664 assert_eq!(
1665 off.terminate_inference("exact-worker-id").await,
1666 InferenceTerminationAck::Confirmed
1667 );
1668 assert!(generation.await.unwrap().is_err());
1669 assert!(off.inner.lock().await.is_none());
1670 }
1671
1672 #[cfg(unix)]
1673 #[tokio::test]
1674 async fn transient_worker_ack_never_creates_parent_residency() {
1675 let off = sh(&format!(
1676 "while IFS= read -r line; do printf '%s\\n' '{TRANSIENT_RESULT_LINE}'; done"
1677 ));
1678 let result = off
1679 .generate_admitted(stub_request(), test_admission())
1680 .await
1681 .unwrap();
1682
1683 assert_eq!(
1684 result.retention,
1685 car_inference::backend_cache::BackendRetention::Transient
1686 );
1687 assert!(off.resident_models().await.is_empty());
1688 }
1689
1690 #[cfg(unix)]
1710 static LEDGER_GUARD: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
1711
1712 #[cfg(unix)]
1713 #[tokio::test]
1714 async fn generic_post_load_error_tears_down_before_releasing_pending_charge() {
1715 let _ledger = LEDGER_GUARD.lock().await;
1716 let fixture = tempfile::tempdir().unwrap();
1717 let admission = LocalWorkerAdmission {
1718 state_root: fixture.path().join("state"),
1719 measured_weights_bytes: 512 * 1024 * 1024,
1720 ..test_admission()
1721 };
1722 let coordinator = car_inference::resource_policy::scoped_local_admission(
1723 &admission.state_root,
1724 admission.policy.clone(),
1725 car_inference::hardware::HardwareInfo::detect(),
1726 );
1727 let resident_baseline_mb = coordinator.owned_resident_model_mb_for_testing();
1728 let off = sh(
1729 r#"while IFS= read -r line; do sleep 0.2; printf '%s\n' '{"Error":"failed after cache publication"}'; done"#,
1730 );
1731 let request = GenerateRequest {
1732 model: Some("stub".into()),
1733 ..Default::default()
1734 };
1735
1736 let baseline_mb = coordinator.resident_model_mb();
1743 let caller = off.clone();
1744 let exchange =
1745 tokio::spawn(async move { caller.generate_admitted(request, admission).await });
1746 tokio::time::sleep(Duration::from_millis(30)).await;
1747 assert!(coordinator.teardown_pending("stub"));
1748 assert_eq!(
1749 coordinator.resident_model_mb().saturating_sub(baseline_mb),
1750 512,
1751 "this worker's own 512 MB charge must still be held before release"
1752 );
1753 assert_eq!(
1754 coordinator.resident_allocation_ids("stub"),
1755 vec![off.allocation_id("stub")]
1756 );
1757 assert_eq!(
1758 coordinator
1759 .owned_resident_model_mb_for_testing()
1760 .checked_sub(resident_baseline_mb),
1761 Some(512)
1762 );
1763 assert!(exchange.await.unwrap().is_err());
1764 assert!(off.inner.lock().await.is_none());
1765
1766 tokio::time::timeout(Duration::from_secs(2), async {
1767 while coordinator.teardown_pending("stub") {
1768 tokio::time::sleep(Duration::from_millis(10)).await;
1769 }
1770 })
1771 .await
1772 .expect("worker must be reaped before its pending machine charge clears");
1773 assert!(coordinator.resident_allocation_ids("stub").is_empty());
1774 assert_eq!(
1775 coordinator.owned_resident_model_mb_for_testing(),
1776 resident_baseline_mb
1777 );
1778 }
1779
1780 #[cfg(unix)]
1781 #[tokio::test]
1782 async fn mismatched_worker_ack_tears_down_requested_and_reported_allocations() {
1783 let _ledger = LEDGER_GUARD.lock().await;
1784 let fixture = tempfile::tempdir().unwrap();
1785 let admission = LocalWorkerAdmission {
1786 state_root: fixture.path().join("state"),
1787 measured_weights_bytes: 64 * 1024 * 1024,
1788 ..test_admission()
1789 };
1790 let coordinator = car_inference::resource_policy::scoped_local_admission(
1791 &admission.state_root,
1792 admission.policy.clone(),
1793 car_inference::hardware::HardwareInfo::detect(),
1794 );
1795 let off = sh(&format!(
1796 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1797 ));
1798 let request = GenerateRequest {
1799 model: Some("requested-a".into()),
1800 ..Default::default()
1801 };
1802
1803 let error = off
1804 .generate_admitted(request, admission)
1805 .await
1806 .err()
1807 .expect("mismatched ack must fail");
1808 assert!(error.to_string().contains("acknowledged model 'stub'"));
1809 assert!(off.inner.lock().await.is_none());
1810 tokio::time::timeout(Duration::from_secs(2), async {
1811 while coordinator.teardown_pending("requested-a")
1812 || coordinator.teardown_pending("stub")
1813 {
1814 tokio::time::sleep(Duration::from_millis(10)).await;
1815 }
1816 })
1817 .await
1818 .expect("mismatch teardown must reap both exact allocation candidates");
1819 assert!(coordinator
1820 .resident_allocation_ids("requested-a")
1821 .is_empty());
1822 assert!(coordinator.resident_allocation_ids("stub").is_empty());
1823 }
1824
1825 #[cfg(unix)]
1826 #[tokio::test]
1827 async fn mismatched_stream_ack_never_returns_or_remembers_worker() {
1828 let _ledger = LEDGER_GUARD.lock().await;
1829 let fixture = tempfile::tempdir().unwrap();
1830 let admission = LocalWorkerAdmission {
1831 state_root: fixture.path().join("state"),
1832 measured_weights_bytes: 64 * 1024 * 1024,
1833 ..test_admission()
1834 };
1835 let coordinator = car_inference::resource_policy::scoped_local_admission(
1836 &admission.state_root,
1837 admission.policy.clone(),
1838 car_inference::hardware::HardwareInfo::detect(),
1839 );
1840 let off = sh(&format!(
1841 "while IFS= read -r line; do printf '%s\\n' '{STREAM_STARTED_LINE}'; done"
1842 ));
1843 let request = GenerateRequest {
1844 model: Some("requested-stream-a".into()),
1845 ..Default::default()
1846 };
1847
1848 let error = off
1849 .stream_admitted(request, admission)
1850 .await
1851 .err()
1852 .expect("mismatched stream ack must fail");
1853 assert!(error.to_string().contains("acknowledged model 'stub'"));
1854 assert!(off.inner.lock().await.is_none());
1855 tokio::time::timeout(Duration::from_secs(2), async {
1856 while coordinator.teardown_pending("requested-stream-a")
1857 || coordinator.teardown_pending("stub")
1858 {
1859 tokio::time::sleep(Duration::from_millis(10)).await;
1860 }
1861 })
1862 .await
1863 .expect("stream mismatch teardown must reap both allocation candidates");
1864 assert!(off.resident_models().await.is_empty());
1865 }
1866
1867 #[cfg(unix)]
1868 #[tokio::test]
1869 async fn replacement_worker_generation_is_charged_separately_until_old_exit() {
1870 let _ledger = LEDGER_GUARD.lock().await;
1871 let fixture = tempfile::tempdir().unwrap();
1872 let admission = LocalWorkerAdmission {
1873 state_root: fixture.path().join("state"),
1874 measured_weights_bytes: 64 * 1024 * 1024,
1875 ..test_admission()
1876 };
1877 let coordinator = car_inference::resource_policy::scoped_local_admission(
1878 &admission.state_root,
1879 admission.policy.clone(),
1880 car_inference::hardware::HardwareInfo::detect(),
1881 );
1882 let script = format!("while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done");
1883 let first = sh(&script);
1884 let second = sh(&script);
1885 assert_ne!(first.allocation_id("stub"), second.allocation_id("stub"));
1886
1887 let mut first_reservation = coordinator
1888 .reserve_measured_host("stub", 64 * 1024 * 1024, 0)
1889 .unwrap();
1890 let first_result = first
1891 .generate_admitted(
1892 GenerateRequest {
1893 model: Some("stub".into()),
1894 ..Default::default()
1895 },
1896 admission.clone(),
1897 )
1898 .await
1899 .unwrap();
1900 first_reservation.publish_resident_weights_as(
1901 &first.allocation_id("stub"),
1902 first_result.residency.measured_weights_bytes,
1903 );
1904 drop(first_reservation);
1905
1906 let mut second_reservation = coordinator
1907 .reserve_measured_host("stub", 64 * 1024 * 1024, 0)
1908 .unwrap();
1909 let second_result = second
1910 .generate_admitted(
1911 GenerateRequest {
1912 model: Some("stub".into()),
1913 ..Default::default()
1914 },
1915 admission,
1916 )
1917 .await
1918 .unwrap();
1919 second_reservation.publish_resident_weights_as(
1920 &second.allocation_id("stub"),
1921 second_result.residency.measured_weights_bytes,
1922 );
1923 drop(second_reservation);
1924 assert_eq!(coordinator.resident_allocation_ids("stub").len(), 2);
1925
1926 assert!(first.release_model("stub").await.unwrap());
1927 assert_eq!(coordinator.resident_allocation_ids("stub").len(), 1);
1928 assert!(second.release_model("stub").await.unwrap());
1929 }
1930
1931 #[cfg(unix)]
1932 #[tokio::test]
1933 async fn policy_generation_change_restarts_existing_worker() {
1934 let off = sh(r#"while IFS= read -r line; do
1935printf '{"Result":{"result":{"text":"pong","tool_calls":[],"trace_id":"t","model_used":"%s","latency_ms":1,"usage":null},"residency":{"model_id":"stub","measured_weights_bytes":1,"retention":"resident"}}}\n' "$$"
1936 done"#);
1937 let first = off
1938 .generate_admitted(stub_request(), test_admission())
1939 .await
1940 .unwrap()
1941 .result
1942 .model_used;
1943 off.refresh_resource_policy(2);
1944 let mut updated = test_admission();
1945 updated.policy_generation = 2;
1946 let second = off
1947 .generate_admitted(stub_request(), updated)
1948 .await
1949 .unwrap()
1950 .result
1951 .model_used;
1952 assert_ne!(first, second, "policy refresh must replace the old child");
1953 }
1954
1955 #[cfg(unix)]
1956 #[tokio::test]
1957 async fn targeted_release_waits_for_worker_exit_and_clears_residency() {
1958 let _ledger = LEDGER_GUARD.lock().await;
1959 let fixture = tempfile::tempdir().unwrap();
1960 let admission = LocalWorkerAdmission {
1961 state_root: fixture.path().join("state"),
1962 ..test_admission()
1963 };
1964 let coordinator = car_inference::resource_policy::scoped_local_admission(
1965 &admission.state_root,
1966 admission.policy.clone(),
1967 car_inference::hardware::HardwareInfo::detect(),
1968 );
1969 let off = sh(&format!(
1970 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1971 ));
1972 off.generate_admitted(stub_request(), admission)
1973 .await
1974 .unwrap();
1975 coordinator.mark_resident_allocation("stub", &off.allocation_id("stub"), 1);
1976 assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
1977 assert!(off.release_model("stub").await.unwrap());
1978 assert!(off.resident_models().await.is_empty());
1979 assert!(!coordinator.is_resident("stub"));
1980 assert!(off.inner.lock().await.is_none());
1981 }
1982
1983 #[cfg(unix)]
1984 #[tokio::test(start_paused = true)]
1985 async fn release_idle_models_releases_a_resident_older_than_ttl() {
1986 let _ledger = LEDGER_GUARD.lock().await;
1987 let fixture = tempfile::tempdir().unwrap();
1988 let admission = LocalWorkerAdmission {
1989 state_root: fixture.path().join("state"),
1990 ..test_admission()
1991 };
1992 let _coordinator = car_inference::resource_policy::scoped_local_admission(
1993 &admission.state_root,
1994 admission.policy.clone(),
1995 car_inference::hardware::HardwareInfo::detect(),
1996 );
1997 let off = sh(&format!(
1998 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
1999 ));
2000 off.generate_admitted(stub_request(), admission)
2001 .await
2002 .unwrap();
2003 assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2004
2005 tokio::time::advance(Duration::from_secs(11)).await;
2006 let released = off.release_idle_models(Duration::from_secs(10)).await;
2007
2008 assert_eq!(released, vec!["stub".to_string()]);
2009 assert!(off.resident_models().await.is_empty());
2010 assert!(off.inner.lock().await.is_none());
2011 }
2012
2013 #[cfg(unix)]
2014 #[tokio::test(start_paused = true)]
2015 async fn release_idle_models_keeps_a_recently_used_resident() {
2016 let _ledger = LEDGER_GUARD.lock().await;
2017 let fixture = tempfile::tempdir().unwrap();
2018 let admission = LocalWorkerAdmission {
2019 state_root: fixture.path().join("state"),
2020 ..test_admission()
2021 };
2022 let _coordinator = car_inference::resource_policy::scoped_local_admission(
2023 &admission.state_root,
2024 admission.policy.clone(),
2025 car_inference::hardware::HardwareInfo::detect(),
2026 );
2027 let off = sh(&format!(
2028 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2029 ));
2030 off.generate_admitted(stub_request(), admission)
2031 .await
2032 .unwrap();
2033
2034 tokio::time::advance(Duration::from_secs(9)).await;
2035 assert!(off
2036 .release_idle_models(Duration::from_secs(10))
2037 .await
2038 .is_empty());
2039 assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2040 assert!(off.release_model("stub").await.unwrap());
2041 }
2042
2043 #[cfg(unix)]
2044 #[tokio::test(start_paused = true)]
2045 async fn release_idle_models_keeps_a_never_completed_resident() {
2046 let _ledger = LEDGER_GUARD.lock().await;
2047 let fixture = tempfile::tempdir().unwrap();
2048 let admission = LocalWorkerAdmission {
2049 state_root: fixture.path().join("state"),
2050 ..test_admission()
2051 };
2052 let _coordinator = car_inference::resource_policy::scoped_local_admission(
2053 &admission.state_root,
2054 admission.policy.clone(),
2055 car_inference::hardware::HardwareInfo::detect(),
2056 );
2057 let off = sh(&format!(
2058 "while IFS= read -r line; do printf '%s\\n' '{TRANSIENT_RESULT_LINE}'; done"
2059 ));
2060 off.generate_admitted(stub_request(), admission.clone())
2061 .await
2062 .unwrap();
2063 off.remember_residency(
2064 &WorkerResidencyAck {
2065 model_id: "stub".into(),
2066 measured_weights_bytes: 1,
2067 retention: car_inference::backend_cache::BackendRetention::Resident,
2068 },
2069 admission.state_root,
2070 None,
2071 );
2072
2073 tokio::time::advance(Duration::from_secs(60)).await;
2074 assert!(off
2075 .release_idle_models(Duration::from_secs(10))
2076 .await
2077 .is_empty());
2078 assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2079 assert!(off.release_model("stub").await.unwrap());
2080 }
2081
2082 #[cfg(unix)]
2083 #[tokio::test(start_paused = true)]
2084 async fn release_idle_models_honors_completion_that_lands_while_sweep_waits() {
2085 let _ledger = LEDGER_GUARD.lock().await;
2086 let fixture = tempfile::tempdir().unwrap();
2087 let admission = LocalWorkerAdmission {
2088 state_root: fixture.path().join("state"),
2089 ..test_admission()
2090 };
2091 let _coordinator = car_inference::resource_policy::scoped_local_admission(
2092 &admission.state_root,
2093 admission.policy.clone(),
2094 car_inference::hardware::HardwareInfo::detect(),
2095 );
2096 let off = sh(&format!(
2097 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2098 ));
2099 off.generate_admitted(stub_request(), admission.clone())
2100 .await
2101 .unwrap();
2102 tokio::time::advance(Duration::from_secs(11)).await;
2103
2104 let inner = off.inner.clone();
2107 let worker_guard = inner.lock().await;
2108 let sweep_offload = off.clone();
2109 let mut sweep = Box::pin(async move {
2110 sweep_offload
2111 .release_idle_models(Duration::from_secs(10))
2112 .await
2113 });
2114 tokio::select! {
2115 released = &mut sweep => panic!("sweep released while worker lock was held: {released:?}"),
2116 _ = tokio::task::yield_now() => {}
2117 }
2118 WorkerOffload::mark_residency_used(&off.resident_models, &admission.state_root, "stub");
2119 drop(worker_guard);
2120
2121 assert!(sweep.await.is_empty());
2122 assert_eq!(off.resident_models().await, vec!["stub".to_string()]);
2123 assert!(off.release_model("stub").await.unwrap());
2124 }
2125
2126 #[cfg(unix)]
2127 #[tokio::test(start_paused = true)]
2128 async fn release_idle_models_keeps_worker_with_recent_co_resident() {
2129 let _ledger = LEDGER_GUARD.lock().await;
2130 let fixture = tempfile::tempdir().unwrap();
2131 let admission = LocalWorkerAdmission {
2132 state_root: fixture.path().join("state"),
2133 ..test_admission()
2134 };
2135 let _coordinator = car_inference::resource_policy::scoped_local_admission(
2136 &admission.state_root,
2137 admission.policy.clone(),
2138 car_inference::hardware::HardwareInfo::detect(),
2139 );
2140 let off = sh(&format!(
2141 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2142 ));
2143 off.generate_admitted(stub_request(), admission.clone())
2144 .await
2145 .unwrap();
2146 tokio::time::advance(Duration::from_secs(11)).await;
2147 off.remember_residency(
2148 &WorkerResidencyAck {
2149 model_id: "recent".into(),
2150 measured_weights_bytes: 1,
2151 retention: car_inference::backend_cache::BackendRetention::Resident,
2152 },
2153 admission.state_root,
2154 Some(Instant::now()),
2155 );
2156
2157 assert!(off
2158 .release_idle_models(Duration::from_secs(10))
2159 .await
2160 .is_empty());
2161 let mut residents = off.resident_models().await;
2162 residents.sort();
2163 assert_eq!(residents, vec!["recent".to_string(), "stub".to_string()]);
2164 assert!(off.release_model("stub").await.unwrap());
2165 }
2166
2167 #[cfg(unix)]
2168 #[tokio::test(start_paused = true)]
2169 async fn release_idle_models_releases_worker_once_when_all_residents_are_idle() {
2170 let _ledger = LEDGER_GUARD.lock().await;
2171 let fixture = tempfile::tempdir().unwrap();
2172 let admission = LocalWorkerAdmission {
2173 state_root: fixture.path().join("state"),
2174 ..test_admission()
2175 };
2176 let _coordinator = car_inference::resource_policy::scoped_local_admission(
2177 &admission.state_root,
2178 admission.policy.clone(),
2179 car_inference::hardware::HardwareInfo::detect(),
2180 );
2181 let off = sh(&format!(
2182 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2183 ));
2184 off.generate_admitted(stub_request(), admission.clone())
2185 .await
2186 .unwrap();
2187 off.remember_residency(
2188 &WorkerResidencyAck {
2189 model_id: "also-idle".into(),
2190 measured_weights_bytes: 1,
2191 retention: car_inference::backend_cache::BackendRetention::Resident,
2192 },
2193 admission.state_root,
2194 Some(Instant::now()),
2195 );
2196 tokio::time::advance(Duration::from_secs(11)).await;
2197
2198 let mut released = off.release_idle_models(Duration::from_secs(10)).await;
2199 released.sort();
2200
2201 assert_eq!(released, vec!["also-idle".to_string(), "stub".to_string()]);
2202 assert!(off.resident_models().await.is_empty());
2203 assert!(off.inner.lock().await.is_none());
2204 }
2205
2206 #[cfg(unix)]
2207 #[tokio::test]
2208 async fn cancelled_worker_release_keeps_accounting_until_background_reap() {
2209 let _ledger = LEDGER_GUARD.lock().await;
2210 let fixture = tempfile::tempdir().unwrap();
2211 let admission = LocalWorkerAdmission {
2212 state_root: fixture.path().join("state"),
2213 ..test_admission()
2214 };
2215 let coordinator = car_inference::resource_policy::scoped_local_admission(
2216 &admission.state_root,
2217 admission.policy.clone(),
2218 car_inference::hardware::HardwareInfo::detect(),
2219 );
2220 let off = sh(&format!(
2221 "while IFS= read -r line; do printf '%s\\n' '{RESULT_LINE}'; done"
2222 ));
2223 off.generate_admitted(stub_request(), admission)
2224 .await
2225 .unwrap();
2226 coordinator.mark_resident_allocation("stub", &off.allocation_id("stub"), 1);
2227 off.release_delay_ms.store(250, Ordering::Release);
2228
2229 let release_owner = off.clone();
2230 let release = tokio::spawn(async move { release_owner.release_model("stub").await });
2231 tokio::time::sleep(Duration::from_millis(20)).await;
2232 assert!(coordinator.teardown_pending("stub"));
2233 release.abort();
2234 let _ = release.await;
2235
2236 tokio::time::timeout(Duration::from_secs(2), async {
2237 while !off.resident_models().await.is_empty() || coordinator.is_resident("stub") {
2238 tokio::time::sleep(Duration::from_millis(10)).await;
2239 }
2240 })
2241 .await
2242 .expect("cancelled release guard must retain, kill, reap, and clear accounting");
2243 assert!(off.inner.lock().await.is_none());
2244 }
2245
2246 #[cfg(unix)]
2253 #[tokio::test]
2254 async fn a_dropped_worker_generation_is_killed_reaped_and_unaccounted() {
2255 let _ledger = LEDGER_GUARD.lock().await;
2256 let fixture = tempfile::tempdir().unwrap();
2257 let admission = LocalWorkerAdmission {
2258 state_root: fixture.path().join("state"),
2259 measured_weights_bytes: 256 * 1024 * 1024,
2260 ..test_admission()
2261 };
2262 let coordinator = car_inference::resource_policy::scoped_local_admission(
2263 &admission.state_root,
2264 admission.policy.clone(),
2265 car_inference::hardware::HardwareInfo::detect(),
2266 );
2267 let resident_baseline_mb = coordinator.owned_resident_model_mb_for_testing();
2268 let pid_file = fixture.path().join("worker.pid");
2272 let off = sh(&format!(
2273 "IFS= read -r line; echo $$ > '{tmp}'; mv '{tmp}' '{pid}'; exec sleep 60",
2274 tmp = pid_file.with_extension("tmp").display(),
2275 pid = pid_file.display(),
2276 ));
2277 let alive = |pid: u32| {
2278 std::process::Command::new("sh")
2279 .arg("-c")
2280 .arg(format!("kill -0 {pid} 2>/dev/null"))
2281 .status()
2282 .is_ok_and(|status| status.success())
2283 };
2284
2285 let mut generation = Box::pin(off.generate_admitted(stub_request(), admission));
2286 let pid = tokio::time::timeout(Duration::from_secs(5), async {
2287 tokio::select! {
2288 result = &mut generation => {
2289 panic!("the stub never answers, yet generation ended (ok: {})", result.is_ok())
2290 }
2291 pid = async {
2292 loop {
2293 if let Some(pid) = std::fs::read_to_string(&pid_file)
2294 .ok()
2295 .and_then(|text| text.trim().parse::<u32>().ok())
2296 {
2297 break pid;
2298 }
2299 tokio::time::sleep(Duration::from_millis(10)).await;
2300 }
2301 } => pid,
2302 }
2303 })
2304 .await
2305 .expect("the stub worker must receive the request");
2306 assert!(
2307 alive(pid),
2308 "positive control: the in-flight worker is running"
2309 );
2310 assert!(
2311 coordinator.teardown_pending("stub"),
2312 "positive control: the in-flight request holds its admission charge"
2313 );
2314
2315 assert!(
2318 tokio::time::timeout(Duration::from_millis(50), generation)
2319 .await
2320 .is_err(),
2321 "a stub that never answers must hit the deadline"
2322 );
2323
2324 tokio::time::timeout(Duration::from_secs(5), async {
2325 while alive(pid) || coordinator.teardown_pending("stub") {
2326 tokio::time::sleep(Duration::from_millis(10)).await;
2327 }
2328 })
2329 .await
2330 .expect("a dropped generation must kill and reap its worker and clear its charge");
2331 assert!(coordinator.resident_allocation_ids("stub").is_empty());
2332 assert_eq!(
2333 coordinator.owned_resident_model_mb_for_testing(),
2334 resident_baseline_mb
2335 );
2336 assert!(off.resident_models().await.is_empty());
2337 assert!(
2338 off.inner.lock().await.is_none(),
2339 "a killed worker must not return to the slot"
2340 );
2341 }
2342
2343 #[cfg(unix)]
2344 #[tokio::test]
2345 async fn worker_crash_fails_one_call_gracefully_then_respawns() {
2346 let off = sh("exit 1");
2353 let e1 = off
2354 .generate_admitted(GenerateRequest::default(), test_admission())
2355 .await
2356 .err()
2357 .expect("worker crash must fail");
2358 assert!(
2359 e1.to_string().contains("worker connection failed"),
2360 "unexpected error: {e1}"
2361 );
2362 let e2 = off
2363 .generate_admitted(GenerateRequest::default(), test_admission())
2364 .await
2365 .err()
2366 .expect("worker crash must fail");
2367 assert!(
2368 e2.to_string().contains("worker connection failed"),
2369 "unexpected error: {e2}"
2370 );
2371 }
2372
2373 #[cfg(unix)]
2374 #[tokio::test]
2375 async fn worker_failure_preserves_exit_status_and_does_not_claim_a_restart() {
2376 for (script, expected) in [
2377 ("IFS= read -r line; exit 17", "exit status: 17"),
2378 ("IFS= read -r line; kill -TERM $$", "signal: 15"),
2379 ] {
2380 let off = sh(script);
2381 let error = off
2382 .generate_admitted(GenerateRequest::default(), test_admission())
2383 .await
2384 .err()
2385 .expect("worker must fail")
2386 .to_string();
2387 assert!(error.contains(expected), "{error}");
2388 assert!(
2389 error.contains("next request starts a new worker"),
2390 "{error}"
2391 );
2392 assert!(!error.contains("was restarted"), "{error}");
2393 assert!(off.inner.lock().await.is_none());
2394 }
2395 }
2396
2397 #[cfg(unix)]
2398 #[tokio::test]
2399 async fn stream_forwards_events_then_closes() {
2400 let off = sh(&format!(
2402 "IFS= read -r line; \
2403 printf '%s\\n' '{STREAM_STARTED_LINE}'; \
2404 printf '%s\\n' '{{\"Event\":{{\"TextDelta\":\"he\"}}}}'; \
2405 printf '%s\\n' '{{\"Event\":{{\"TextDelta\":\"llo\"}}}}'; \
2406 printf '%s\\n' '\"StreamEnd\"'; \
2407 while IFS= read -r l; do :; done"
2408 ));
2409 let mut rx = off
2410 .stream_admitted(stub_request(), test_admission())
2411 .await
2412 .unwrap()
2413 .events;
2414 let mut got = String::new();
2415 while let Some(ev) = rx.recv().await {
2416 if let StreamEvent::TextDelta(t) = ev {
2417 got.push_str(&t);
2418 }
2419 }
2420 assert_eq!(got, "hello");
2421 }
2422
2423 #[cfg(unix)]
2424 #[tokio::test]
2425 async fn stream_worker_death_closes_with_error_end() {
2426 let off = sh(&format!("IFS= read -r line; printf '%s\\n' '{STREAM_STARTED_LINE}'; printf '%s\\n' '{{\"Event\":{{\"TextDelta\":\"partial\"}}}}'"));
2429 let mut rx = off
2430 .stream_admitted(stub_request(), test_admission())
2431 .await
2432 .unwrap()
2433 .events;
2434 let mut saw_text = false;
2435 let mut saw_error_end = false;
2436 while let Some(ev) = rx.recv().await {
2437 match ev {
2438 StreamEvent::TextDelta(t) if t == "partial" => saw_text = true,
2439 StreamEvent::StopReason(s) if s == "error" => saw_error_end = true,
2440 _ => {}
2441 }
2442 }
2443 assert!(saw_text, "should have forwarded the partial delta");
2444 assert!(
2445 saw_error_end,
2446 "mid-stream death should surface an error end"
2447 );
2448 }
2449}