sie_sdk/client/
meta.rs

1//! Model catalogue, health, and capacity waiting.
2
3use std::time::{Duration, Instant};
4
5use reqwest::Method;
6use serde_json::Value;
7
8use crate::client::Client;
9use crate::error::{Error, Result};
10use crate::http::{HttpResponse, headers};
11use crate::retry::{RetryPolicy, backoff};
12use crate::types::{CapacityInfo, HealthResponse, ModelInfo, WorkerInfo};
13
14/// Decode a successful JSON response, preserving the metering evidence on failure.
15pub(crate) fn parse_json<T: serde::de::DeserializeOwned>(
16    response: &HttpResponse,
17    owner: &str,
18) -> Result<T> {
19    serde_json::from_slice(&response.body)
20        .map_err(|err| Error::decode(format!("malformed {owner} response: {err}")))
21}
22
23impl Client {
24    /// Every model the server knows about.
25    pub async fn list_models(&self) -> Result<Vec<ModelInfo>> {
26        let request = self
27            .request(Method::GET, "/v1/models")?
28            .header("accept", headers::JSON_CONTENT_TYPE);
29        let response = self.send_once(request, RetryPolicy::NONE).await?;
30        let body: Value = parse_json(&response, "models")?;
31        let models = body
32            .get("models")
33            .cloned()
34            .ok_or_else(|| Error::decode("models response is missing its `models` array"))?;
35        serde_json::from_value(models)
36            .map_err(|err| Error::decode(format!("malformed models response: {err}")))
37    }
38
39    /// Details for one model.
40    pub async fn get_model(&self, model: &str) -> Result<ModelInfo> {
41        let request = self
42            .request(Method::GET, &format!("/v1/models/{model}"))?
43            .header("accept", headers::JSON_CONTENT_TYPE);
44        let response = self.send_once(request, RetryPolicy::NONE).await?;
45        parse_json(&response, "model")
46    }
47
48    /// Health, from either a gateway or a standalone worker.
49    ///
50    /// A gateway answers `/health` with the whole cluster's state. A worker running on its
51    /// own has no such view: it serves the Kubernetes-style `/healthz` probe, which is
52    /// plain text, so its reply is reported as a worker with that status and nothing else.
53    pub async fn health(&self) -> Result<HealthResponse> {
54        let request = self
55            .request(Method::GET, "/health")?
56            .header("accept", headers::JSON_CONTENT_TYPE);
57        match self.send_once(request, RetryPolicy::NONE).await {
58            Ok(response) => parse_json(&response, "health"),
59            Err(error) if error.status() == Some(404) => self.worker_health().await,
60            Err(error) => Err(error),
61        }
62    }
63
64    /// The liveness probe a standalone worker serves in place of `/health`.
65    async fn worker_health(&self) -> Result<HealthResponse> {
66        let request = self.request(Method::GET, "/healthz")?;
67        let response = self.send_once(request, RetryPolicy::NONE).await?;
68        let status = response.text().trim().to_string();
69        Ok(HealthResponse {
70            status: if status.is_empty() {
71                "ok".to_string()
72            } else {
73                status
74            },
75            kind: "worker".to_string(),
76            ..HealthResponse::default()
77        })
78    }
79
80    /// Cluster capacity, optionally narrowed to one GPU type.
81    ///
82    /// Requires a gateway: a worker's `/health` describes only itself.
83    pub async fn get_capacity(&self, gpu: Option<&str>) -> Result<CapacityInfo> {
84        let health = self.health().await?;
85        if health.kind != "gateway" {
86            return Err(Error::Request {
87                message: "get_capacity() requires a gateway endpoint. This appears to be a worker."
88                    .to_string(),
89                code: Some("not_gateway".to_string()),
90                status: 400,
91                request: None,
92            });
93        }
94
95        let workers: Vec<WorkerInfo> = match gpu {
96            Some(gpu) => {
97                let wanted = gpu.to_ascii_lowercase();
98                health
99                    .workers
100                    .into_iter()
101                    .filter(|worker| worker.gpu.eq_ignore_ascii_case(&wanted))
102                    .collect()
103            }
104            None => health.workers,
105        };
106
107        Ok(CapacityInfo {
108            status: health.status,
109            // A filtered view reports the matching workers, not the cluster total.
110            worker_count: if gpu.is_some() {
111                workers.len() as u32
112            } else {
113                health.cluster.worker_count
114            },
115            gpu_count: health.cluster.gpu_count,
116            models_loaded: health.cluster.models_loaded,
117            configured_gpu_types: health.configured_gpu_types,
118            live_gpu_types: health.live_gpu_types,
119            workers,
120        })
121    }
122
123    /// Block until a GPU type has capacity.
124    ///
125    /// With `model`, this warms the model as well: it issues one small encode with capacity
126    /// waiting enabled, which triggers both scale-up and model loading, then reports the
127    /// resulting capacity. Without it, `/health` is polled until a worker appears.
128    pub async fn wait_for_capacity(
129        &self,
130        gpu: &str,
131        model: Option<&str>,
132        timeout: Option<Duration>,
133        poll_interval: Duration,
134    ) -> Result<CapacityInfo> {
135        let budget = timeout.unwrap_or(backoff::DEFAULT_PROVISION_TIMEOUT);
136        let start = Instant::now();
137
138        if let Some(model) = model {
139            self.encode(model, [crate::types::Item::text("warmup")])
140                .gpu(gpu)
141                .wait_for_capacity(true)
142                .provision_timeout(budget)
143                .send()
144                .await?;
145            return self.get_capacity(Some(gpu)).await;
146        }
147
148        loop {
149            // A server that is still scaling up refuses connections and returns errors; both
150            // are expected here and only the deadline ends the wait.
151            if let Ok(capacity) = self.get_capacity(Some(gpu)).await
152                && capacity.worker_count > 0
153            {
154                return Ok(capacity);
155            }
156
157            let elapsed = start.elapsed();
158            if elapsed >= budget {
159                return Err(Error::Provisioning {
160                    message: format!(
161                        "Timeout after {:.1}s waiting for GPU '{gpu}' capacity",
162                        elapsed.as_secs_f64()
163                    ),
164                    gpu: Some(gpu.to_string()),
165                    retry_after: None,
166                });
167            }
168            tokio::time::sleep(poll_interval.min(budget.saturating_sub(elapsed))).await;
169        }
170    }
171}
sie_sdk/client/meta.rs

sie_sdk/client/
meta.rs
