qualia_core_db/solvers/qpu/
dispatcher.rs

1//! In-process QPU job queue and dispatcher.
2//!
3//! This module manages the lifecycle of QPU jobs in memory.  Actual HTTP
4//! egress to the 8 supported providers (IBM, D-Wave, IonQ, Rigetti, Azure,
5//! Braket, Google, Quantinuum) is handled by `qualia-client-core::qpu_dispatcher`.
6//!
7//! Ported from `qpu/src/dispatcher.rs`.
8
9use super::{JobStatus, QpuError, QpuJob, QpuResult};
10use std::cmp::Ordering;
11use std::collections::{BinaryHeap, HashMap};
12use std::sync::Arc;
13use tokio::sync::{Mutex, RwLock};
14
15// ── Job state ─────────────────────────────────────────────────────────────────
16
17#[derive(Debug, Clone)]
18struct JobState {
19    job: QpuJob,
20    enqueued_at_ms: u64,
21    retries: u32,
22    status: InternalStatus,
23}
24
25#[derive(Debug, Clone, PartialEq)]
26enum InternalStatus {
27    Queued,
28    Submitted,
29    Running,
30    Completed,
31    Failed,
32}
33
34fn now_ms() -> u64 {
35    std::time::SystemTime::now()
36        .duration_since(std::time::UNIX_EPOCH)
37        .map(|d| d.as_millis() as u64)
38        .unwrap_or(0)
39}
40
41// ── Queue statistics ──────────────────────────────────────────────────────────
42
43#[derive(Debug, Clone)]
44pub struct QueueStats {
45    pub pending_count: usize,
46    pub running_count: usize,
47    pub completed_count: usize,
48}
49
50// ── Job queue ─────────────────────────────────────────────────────────────────
51
52/// Prioritized job wrapper for QGroup heuristic
53#[derive(Debug, Clone)]
54pub struct PrioritizedJob(pub QpuJob);
55
56impl PartialEq for PrioritizedJob {
57    fn eq(&self, other: &Self) -> bool {
58        self.0.job_id == other.0.job_id
59    }
60}
61impl Eq for PrioritizedJob {}
62
63impl PartialOrd for PrioritizedJob {
64    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
65        Some(self.cmp(other))
66    }
67}
68
69impl Ord for PrioritizedJob {
70    fn cmp(&self, other: &Self) -> Ordering {
71        // QGroup heuristic: group by similar circuit depth and shot count.
72        // We sort descending so the BinaryHeap (max-heap) pops highest depth first,
73        // and jobs with similar depths will be contiguous.
74        let depth_ord = self
75            .0
76            .parameters
77            .circuit_depth
78            .cmp(&other.0.parameters.circuit_depth);
79        if depth_ord != Ordering::Equal {
80            return depth_ord;
81        }
82        self.0.parameters.shots.cmp(&other.0.parameters.shots)
83    }
84}
85
86/// In-process job queue.  HTTP dispatch is performed by the caller via
87/// `submit_fn` passed to `process_queue`.
88pub struct JobQueue {
89    pending: Arc<Mutex<BinaryHeap<PrioritizedJob>>>,
90    running: Arc<RwLock<HashMap<String, JobState>>>,
91    completed: Arc<Mutex<Vec<QpuResult>>>,
92}
93
94impl JobQueue {
95    pub fn new() -> Self {
96        Self {
97            pending: Arc::new(Mutex::new(BinaryHeap::new())),
98            running: Arc::new(RwLock::new(HashMap::new())),
99            completed: Arc::new(Mutex::new(Vec::new())),
100        }
101    }
102
103    /// Enqueue a job for dispatch.
104    pub async fn enqueue(&self, job: QpuJob) -> String {
105        let id = job.job_id.clone();
106        self.pending.lock().await.push(PrioritizedJob(job));
107        id
108    }
109
110    /// Drain the pending queue and call `dispatch` for each job.
111    ///
112    /// `dispatch` should return the provider-assigned job ID on success.
113    pub async fn process_queue<F>(&self, dispatch: F) -> Result<(), QpuError>
114    where
115        F: Fn(&QpuJob) -> Result<String, String>,
116    {
117        let jobs: Vec<QpuJob> = {
118            let mut pending = self.pending.lock().await;
119            let mut jobs = Vec::with_capacity(pending.len());
120            while let Some(job) = pending.pop() {
121                jobs.push(job.0);
122            }
123            jobs
124        };
125
126        for job in jobs {
127            let job_id = job.job_id.clone();
128
129            let mut state = JobState {
130                job: job.clone(),
131                enqueued_at_ms: now_ms(),
132                retries: 0,
133                status: InternalStatus::Queued,
134            };
135
136            let mut dispatch_result = dispatch(&job);
137
138            while dispatch_result.is_err() && state.retries < 3 {
139                state.retries += 1;
140                log::warn!(
141                    "Retrying QPU dispatch for {} (retry {})",
142                    job_id,
143                    state.retries
144                );
145                dispatch_result = dispatch(&job);
146            }
147
148            match dispatch_result {
149                Ok(provider_id) => {
150                    state.job.job_id = provider_id.clone();
151                    state.status = InternalStatus::Submitted;
152                    self.running
153                        .write()
154                        .await
155                        .insert(provider_id.clone(), state.clone());
156
157                    let mut running = self.running.write().await;
158                    if let Some(s) = running.get_mut(&provider_id) {
159                        s.status = InternalStatus::Running;
160                    }
161                }
162                Err(e) => {
163                    log::error!(
164                        "QPU dispatch failed for {} after {} retries: {}",
165                        job_id,
166                        state.retries,
167                        e
168                    );
169                    state.status = InternalStatus::Failed;
170                    let result = QpuResult::failed(job_id, e);
171                    self.completed.lock().await.push(result);
172                }
173            }
174        }
175        Ok(())
176    }
177
178    /// Mark a running job as completed or failed based on a polled status.
179    pub async fn record_result(&self, job_id: &str, result: QpuResult) {
180        let mut running = self.running.write().await;
181        if let Some(mut state) = running.remove(job_id) {
182            state.status = if result.error.is_some() {
183                InternalStatus::Failed
184            } else {
185                InternalStatus::Completed
186            };
187            let _duration = now_ms() - state.enqueued_at_ms;
188        }
189        self.completed.lock().await.push(result);
190    }
191
192    /// Drain and return all completed results.
193    pub async fn take_results(&self) -> Vec<QpuResult> {
194        std::mem::take(&mut *self.completed.lock().await)
195    }
196
197    pub async fn stats(&self) -> QueueStats {
198        QueueStats {
199            pending_count: self.pending.lock().await.len(),
200            running_count: self.running.read().await.len(),
201            completed_count: self.completed.lock().await.len(),
202        }
203    }
204}
205
206impl Default for JobQueue {
207    fn default() -> Self {
208        Self::new()
209    }
210}
211
212// ── Dispatcher ────────────────────────────────────────────────────────────────
213
214/// Wraps a `JobQueue` and exposes a simple submit/poll API.
215pub struct Dispatcher {
216    pub queue: Arc<JobQueue>,
217}
218
219impl Dispatcher {
220    pub fn new() -> Self {
221        Self {
222            queue: Arc::new(JobQueue::new()),
223        }
224    }
225
226    pub async fn submit(&self, job: QpuJob) -> String {
227        self.queue.enqueue(job).await
228    }
229
230    pub async fn flush<F>(&self, dispatch: F) -> Result<(), QpuError>
231    where
232        F: Fn(&QpuJob) -> Result<String, String>,
233    {
234        self.queue.process_queue(dispatch).await
235    }
236
237    pub async fn drain_results(&self) -> Vec<QpuResult> {
238        self.queue.take_results().await
239    }
240
241    pub async fn stats(&self) -> QueueStats {
242        self.queue.stats().await
243    }
244}
245
246impl Default for Dispatcher {
247    fn default() -> Self {
248        Self::new()
249    }
250}
251
252// ── Fallback handler ──────────────────────────────────────────────────────────
253
254/// Returns a classical-simulation result when no QPU provider is available.
255pub struct FallbackHandler {
256    pub enabled: bool,
257}
258
259impl FallbackHandler {
260    pub fn new(enabled: bool) -> Self {
261        Self { enabled }
262    }
263
264    pub fn simulate_classically(&self, job: &QpuJob) -> Result<QpuResult, QpuError> {
265        if !self.enabled {
266            return Err(QpuError::Api("Fallback is disabled".into()));
267        }
268        Ok(QpuResult {
269            job_id: job.job_id.clone(),
270            status: JobStatus::Completed,
271            result: Some(super::JobResultData {
272                measurements: vec![],
273                energies: Some(vec![0.0]),
274                metadata: serde_json::json!({"method": "classical_simulation"}),
275            }),
276            completed_at_ms: Some(
277                std::time::SystemTime::now()
278                    .duration_since(std::time::UNIX_EPOCH)
279                    .map(|d| d.as_millis() as u64)
280                    .unwrap_or(0),
281            ),
282            error: None,
283        })
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use crate::solvers::qpu::{JobParameters, ProblemType};
291
292    #[tokio::test]
293    async fn queue_enqueue_and_stats() {
294        let q = JobQueue::new();
295        let job = QpuJob::new(
296            "test-job-1".into(),
297            ProblemType::Annealing,
298            JobParameters::default(),
299        );
300        q.enqueue(job).await;
301        let stats = q.stats().await;
302        assert_eq!(stats.pending_count, 1);
303    }
304
305    #[test]
306    fn fallback_handler_enabled() {
307        let handler = FallbackHandler::new(true);
308        let job = QpuJob::default();
309        let result = handler.simulate_classically(&job).unwrap();
310        assert_eq!(result.status, JobStatus::Completed);
311    }
312
313    #[test]
314    fn fallback_handler_disabled() {
315        let handler = FallbackHandler::new(false);
316        let job = QpuJob::default();
317        assert!(handler.simulate_classically(&job).is_err());
318    }
319
320    #[tokio::test]
321    async fn test_qgroup_heuristic_sorting() {
322        let q = JobQueue::new();
323
324        let mut job1 = QpuJob::new("job1".into(), ProblemType::Vqe, JobParameters::default());
325        job1.parameters.circuit_depth = 10;
326        job1.parameters.shots = 1000;
327
328        let mut job2 = QpuJob::new("job2".into(), ProblemType::Vqe, JobParameters::default());
329        job2.parameters.circuit_depth = 50;
330        job2.parameters.shots = 1000;
331
332        let mut job3 = QpuJob::new("job3".into(), ProblemType::Vqe, JobParameters::default());
333        job3.parameters.circuit_depth = 10;
334        job3.parameters.shots = 2000;
335
336        let mut job4 = QpuJob::new("job4".into(), ProblemType::Vqe, JobParameters::default());
337        job4.parameters.circuit_depth = 50;
338        job4.parameters.shots = 500;
339
340        // Enqueue jobs in random order
341        q.enqueue(job1).await;
342        q.enqueue(job2).await;
343        q.enqueue(job3).await;
344        q.enqueue(job4).await;
345
346        let jobs = {
347            let mut pending = q.pending.lock().await;
348            let mut extracted = Vec::new();
349            while let Some(job) = pending.pop() {
350                extracted.push(job.0);
351            }
352            extracted
353        };
354
355        // BinaryHeap pops the largest first.
356        // Based on our Ord implementation: highest depth first, then highest shots.
357        assert_eq!(jobs[0].job_id, "job2"); // depth: 50, shots: 1000
358        assert_eq!(jobs[1].job_id, "job4"); // depth: 50, shots: 500
359        assert_eq!(jobs[2].job_id, "job3"); // depth: 10, shots: 2000
360        assert_eq!(jobs[3].job_id, "job1"); // depth: 10, shots: 1000
361    }
362}
qualia_core_db/solvers/qpu/dispatcher.rs

qualia_core_db/solvers/qpu/
dispatcher.rs
