Skip to main content

mold_core/
runpod.rs

1//! RunPod REST API client.
2//!
3//! Wraps `https://rest.runpod.io/v1/` for pod lifecycle management from within
4//! the mold CLI. Uses Bearer-token auth via `RUNPOD_API_KEY` env var or an
5//! explicit key. All methods are async.
6
7use crate::error::MoldError;
8use anyhow::Result;
9use reqwest::{Client, StatusCode};
10use serde::{Deserialize, Serialize};
11use std::collections::HashSet;
12use std::fmt;
13use std::time::Duration;
14
15/// Default REST base URL.
16pub const DEFAULT_ENDPOINT: &str = "https://rest.runpod.io/v1";
17
18/// GraphQL endpoint (used for /user since REST doesn't expose it).
19pub const GRAPHQL_ENDPOINT: &str = "https://api.runpod.io/graphql";
20
21/// Environment variable that holds the RunPod API key.
22pub const API_KEY_ENV: &str = "RUNPOD_API_KEY";
23
24/// Live RunPod REST limits. The generated API schema currently advertises a
25/// wider range, but the production service rejects sizes below 10 GB and 4 TB.
26pub const NETWORK_VOLUME_MIN_GB: u32 = 10;
27pub const NETWORK_VOLUME_MAX_GB: u32 = 3999;
28
29pub fn valid_network_volume_size(size: u32) -> bool {
30    (NETWORK_VOLUME_MIN_GB..=NETWORK_VOLUME_MAX_GB).contains(&size)
31}
32
33/// Persisted configuration under `[runpod]` in `config.toml`.
34#[derive(Debug, Clone, Deserialize, Serialize, Default)]
35pub struct RunPodSettings {
36    /// API key stored in config. Env var `RUNPOD_API_KEY` takes precedence.
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub api_key: Option<String>,
39
40    /// Preferred GPU, e.g. `"NVIDIA GeForce RTX 5090"`.
41    #[serde(default, skip_serializing_if = "Option::is_none")]
42    pub default_gpu: Option<String>,
43
44    /// Preferred datacenter id, e.g. `"EUR-IS-2"`.
45    #[serde(default, skip_serializing_if = "Option::is_none")]
46    pub default_datacenter: Option<String>,
47
48    /// Attach this network volume to new pods (id from RunPod console).
49    #[serde(default, skip_serializing_if = "Option::is_none")]
50    pub default_network_volume_id: Option<String>,
51
52    /// When `true`, `mold runpod run` deletes the pod after generating.
53    /// When `false`, the pod is left running for reuse. Default `false`.
54    #[serde(default)]
55    pub auto_teardown: bool,
56
57    /// After this many minutes of idle time, background reap deletes the pod.
58    /// `0` disables the idle reaper. Default `20`.
59    #[serde(default = "default_auto_teardown_idle_mins")]
60    pub auto_teardown_idle_mins: u32,
61
62    /// Fail UAT or `run` if cumulative pod spend for the session exceeds this
63    /// many USD. `0.0` disables the guard. Default `0.0`.
64    #[serde(default)]
65    pub cost_alert_usd: f64,
66
67    /// Override the REST endpoint (mostly for testing).
68    #[serde(default, skip_serializing_if = "Option::is_none")]
69    pub endpoint: Option<String>,
70}
71
72fn default_auto_teardown_idle_mins() -> u32 {
73    20
74}
75
76/// Redact `api_key` when logging.
77impl RunPodSettings {
78    pub fn redacted_debug(&self) -> String {
79        format!(
80            "RunPodSettings {{ api_key: {}, default_gpu: {:?}, default_datacenter: {:?}, \
81             default_network_volume_id: {:?}, auto_teardown: {}, auto_teardown_idle_mins: {}, \
82             cost_alert_usd: {}, endpoint: {:?} }}",
83            if self.api_key.is_some() {
84                "Some(\"<redacted>\")"
85            } else {
86                "None"
87            },
88            self.default_gpu,
89            self.default_datacenter,
90            self.default_network_volume_id,
91            self.auto_teardown,
92            self.auto_teardown_idle_mins,
93            self.cost_alert_usd,
94            self.endpoint,
95        )
96    }
97}
98
99// ─── API response types ────────────────────────────────────────────────────
100
101/// `GET /user` response.
102#[derive(Debug, Clone, Deserialize, Serialize)]
103pub struct UserInfo {
104    pub id: String,
105    pub email: String,
106    #[serde(default)]
107    pub client_balance: f64,
108    #[serde(default)]
109    pub current_spend_per_hr: f64,
110    #[serde(default)]
111    pub spend_limit: Option<f64>,
112}
113
114/// One entry from `GET /gputypes`.
115#[derive(Debug, Clone, Deserialize, Serialize)]
116pub struct GpuType {
117    #[serde(default)]
118    pub id: Option<String>,
119    #[serde(rename = "displayName", default)]
120    pub display_name: String,
121    #[serde(rename = "gpuId", default)]
122    pub gpu_id: String,
123    #[serde(rename = "memoryInGb", default)]
124    pub memory_in_gb: Option<u32>,
125    #[serde(rename = "secureCloud", default)]
126    pub secure_cloud: bool,
127    #[serde(rename = "communityCloud", default)]
128    pub community_cloud: bool,
129    #[serde(rename = "stockStatus", default)]
130    pub stock_status: Option<String>,
131    #[serde(default)]
132    pub available: bool,
133}
134
135impl GpuType {
136    /// Provider identity accepted by the Pod API. GraphQL normally supplies
137    /// `id`; older/alternate inventory shapes may supply `gpuId` instead.
138    pub fn authoritative_type_id(&self) -> Option<&str> {
139        self.id
140            .as_deref()
141            .map(str::trim)
142            .filter(|id| !id.is_empty())
143            .or_else(|| {
144                let gpu_id = self.gpu_id.trim();
145                (!gpu_id.is_empty()).then_some(gpu_id)
146            })
147    }
148
149    /// Allocation identity accepted by the Pod API. Provider IDs remain
150    /// authoritative; display-name inference exists only for legacy inventory
151    /// shapes that omitted both provider identity fields.
152    pub fn allocation_type_id(&self) -> Option<&str> {
153        self.authoritative_type_id().or_else(|| {
154            let legacy_id = legacy_gpu_type_id_from_display_name(&self.display_name);
155            (!legacy_id.is_empty()).then_some(legacy_id)
156        })
157    }
158}
159
160/// Legacy RunPod inventory exposed a human display label without an allocation
161/// ID. Keep this mapping centralized so every launcher makes the same fallback.
162pub fn legacy_gpu_type_id_from_display_name(display_name: &str) -> &str {
163    match display_name.trim() {
164        "RTX 4090" => "NVIDIA GeForce RTX 4090",
165        "RTX 5090" => "NVIDIA GeForce RTX 5090",
166        "RTX 3090" => "NVIDIA GeForce RTX 3090",
167        "L40S" => "NVIDIA L40S",
168        "L40" => "NVIDIA L40",
169        "A100 PCIe" => "NVIDIA A100 80GB PCIe",
170        "A100 SXM" => "NVIDIA A100-SXM4-80GB",
171        "H100 SXM" => "NVIDIA H100 80GB HBM3",
172        "H100 NVL" => "NVIDIA H100 NVL",
173        "RTX A6000" => "NVIDIA RTX A6000",
174        other => other,
175    }
176}
177
178pub fn normalized_gpu_type_identity(value: &str) -> String {
179    value
180        .split_whitespace()
181        .collect::<Vec<_>>()
182        .join(" ")
183        .to_ascii_lowercase()
184}
185
186pub fn canonical_supported_gpu_type_id<'a>(
187    supported_gpu_ids: &'a HashSet<String>,
188    candidate_id: &str,
189) -> Option<&'a str> {
190    let candidate = normalized_gpu_type_identity(candidate_id);
191    if candidate.is_empty() {
192        return None;
193    }
194    supported_gpu_ids
195        .iter()
196        .filter(|supported| normalized_gpu_type_identity(supported) == candidate)
197        .map(String::as_str)
198        .min()
199}
200
201/// One entry from `GET /datacenters`.
202#[derive(Debug, Clone, Deserialize, Serialize)]
203pub struct Datacenter {
204    pub id: String,
205    #[serde(default)]
206    pub name: String,
207    #[serde(default)]
208    pub location: Option<String>,
209    #[serde(rename = "gpuAvailability", default)]
210    pub gpu_availability: Vec<GpuAvailability>,
211}
212
213#[derive(Debug, Clone, Deserialize, Serialize)]
214pub struct GpuAvailability {
215    #[serde(rename = "displayName", default)]
216    pub display_name: String,
217    #[serde(rename = "gpuId", default)]
218    pub gpu_id: String,
219    #[serde(rename = "stockStatus", default)]
220    pub stock_status: Option<String>,
221}
222
223/// `GET /pods` / `GET /pods/{id}` response.
224///
225/// Only fields we actually use are deserialized — anything else is allowed via
226/// `#[serde(default)]` on the struct to avoid breaking on RunPod API drift.
227#[derive(Debug, Clone, Deserialize, Serialize)]
228pub struct Pod {
229    pub id: String,
230    #[serde(default)]
231    pub name: Option<String>,
232    #[serde(rename = "desiredStatus", default)]
233    pub desired_status: String,
234    #[serde(rename = "imageName", default)]
235    pub image_name: Option<String>,
236    #[serde(rename = "gpuCount", default)]
237    pub gpu_count: u32,
238    #[serde(rename = "costPerHr", default)]
239    pub cost_per_hr: f64,
240    #[serde(rename = "uptimeSeconds", default)]
241    pub uptime_seconds: u64,
242    #[serde(rename = "lastStatusChange", default)]
243    pub last_status_change: Option<String>,
244    #[serde(rename = "memoryInGb", default)]
245    pub memory_in_gb: u32,
246    #[serde(rename = "vcpuCount", default)]
247    pub vcpu_count: u32,
248    #[serde(rename = "volumeInGb", default)]
249    pub volume_in_gb: u32,
250    #[serde(rename = "volumeMountPath", default)]
251    pub volume_mount_path: Option<String>,
252    #[serde(default)]
253    pub ports: serde_json::Value,
254    #[serde(default)]
255    pub env: serde_json::Value,
256    #[serde(default)]
257    pub machine: Option<PodMachine>,
258    /// Current REST responses expose the assigned GPU here even when the
259    /// optional machine expansion omits `gpuDisplayName`.
260    #[serde(default)]
261    pub gpu: Option<PodGpu>,
262    #[serde(default)]
263    pub runtime: Option<serde_json::Value>,
264    #[serde(rename = "networkVolume", default)]
265    pub network_volume: Option<NetworkVolume>,
266    /// Current REST list/get responses expose only this id; create responses
267    /// may additionally include the expanded `networkVolume` object above.
268    #[serde(rename = "networkVolumeId", default)]
269    pub network_volume_id: Option<String>,
270}
271
272impl Pod {
273    pub fn attached_network_volume_id(&self) -> Option<&str> {
274        self.network_volume
275            .as_ref()
276            .map(|volume| volume.id.as_str())
277            .or(self.network_volume_id.as_deref())
278    }
279
280    pub fn gpu_name(&self) -> Option<&str> {
281        self.gpu
282            .as_ref()
283            .and_then(|gpu| gpu.display_name.as_deref())
284            .or_else(|| {
285                self.machine
286                    .as_ref()
287                    .and_then(|machine| machine.gpu_display_name.as_deref())
288            })
289            .or_else(|| {
290                self.machine
291                    .as_ref()
292                    .and_then(|machine| machine.gpu_type_id.as_deref())
293            })
294            .or_else(|| self.gpu.as_ref().and_then(|gpu| gpu.id.as_deref()))
295    }
296
297    pub fn datacenter_id(&self) -> Option<&str> {
298        self.machine
299            .as_ref()
300            .and_then(|machine| machine.data_center_id.as_deref())
301    }
302}
303
304#[derive(Debug, Clone, Deserialize, Serialize)]
305pub struct PodMachine {
306    #[serde(rename = "gpuDisplayName", default)]
307    pub gpu_display_name: Option<String>,
308    #[serde(rename = "gpuTypeId", default)]
309    pub gpu_type_id: Option<String>,
310    #[serde(rename = "dataCenterId", default)]
311    pub data_center_id: Option<String>,
312    #[serde(default)]
313    pub location: Option<String>,
314}
315
316#[derive(Debug, Clone, Deserialize, Serialize)]
317pub struct PodGpu {
318    #[serde(default)]
319    pub id: Option<String>,
320    #[serde(rename = "displayName", default)]
321    pub display_name: Option<String>,
322    #[serde(default)]
323    pub count: Option<u32>,
324}
325
326/// Body for `POST /pods`.
327#[derive(Debug, Clone, Serialize, Default)]
328pub struct CreatePodRequest {
329    pub name: String,
330    #[serde(rename = "imageName")]
331    pub image_name: String,
332    #[serde(rename = "gpuTypeIds")]
333    pub gpu_type_ids: Vec<String>,
334    #[serde(rename = "cloudType")]
335    pub cloud_type: String,
336    #[serde(rename = "dataCenterIds", skip_serializing_if = "Option::is_none")]
337    pub data_center_ids: Option<Vec<String>>,
338    #[serde(rename = "gpuCount")]
339    pub gpu_count: u32,
340    #[serde(rename = "containerDiskInGb")]
341    pub container_disk_in_gb: u32,
342    #[serde(rename = "volumeInGb")]
343    pub volume_in_gb: u32,
344    #[serde(rename = "volumeMountPath")]
345    pub volume_mount_path: String,
346    pub ports: Vec<String>,
347    pub env: serde_json::Map<String, serde_json::Value>,
348    #[serde(rename = "networkVolumeId", skip_serializing_if = "Option::is_none")]
349    pub network_volume_id: Option<String>,
350}
351
352/// One entry from `GET /networkvolumes`.
353#[derive(Debug, Clone, Deserialize, Serialize)]
354pub struct NetworkVolume {
355    pub id: String,
356    pub name: String,
357    #[serde(rename = "dataCenterId", default)]
358    pub data_center_id: String,
359    pub size: u32,
360}
361
362/// Body for `POST /networkvolumes`.
363#[derive(Debug, Clone, Serialize)]
364pub struct CreateNetworkVolumeRequest {
365    pub name: String,
366    pub size: u32,
367    #[serde(rename = "dataCenterId")]
368    pub data_center_id: String,
369}
370
371/// Body for `PATCH /networkvolumes/{id}`. Sizes may only increase.
372#[derive(Debug, Clone, Serialize)]
373pub struct UpdateNetworkVolumeRequest {
374    #[serde(skip_serializing_if = "Option::is_none")]
375    pub name: Option<String>,
376    #[serde(skip_serializing_if = "Option::is_none")]
377    pub size: Option<u32>,
378}
379
380// ─── Client ────────────────────────────────────────────────────────────────
381
382#[derive(Clone)]
383pub struct RunPodClient {
384    endpoint: String,
385    graphql_endpoint: String,
386    api_key: String,
387    http: Client,
388}
389
390impl fmt::Debug for RunPodClient {
391    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
392        f.debug_struct("RunPodClient")
393            .field("endpoint", &self.endpoint)
394            .field("api_key", &"<redacted>")
395            .finish()
396    }
397}
398
399impl RunPodClient {
400    /// Construct with explicit endpoint + key. The GraphQL endpoint
401    /// defaults to `GRAPHQL_ENDPOINT` when the REST endpoint is production,
402    /// and falls back to the same URL when the REST endpoint is overridden
403    /// (so tests pointing at a mock server route GraphQL calls there too).
404    pub fn new(endpoint: impl Into<String>, api_key: impl Into<String>) -> Self {
405        let rest = endpoint.into();
406        let graphql = if rest.starts_with(DEFAULT_ENDPOINT) {
407            GRAPHQL_ENDPOINT.to_string()
408        } else {
409            rest.clone()
410        };
411        Self::new_with_graphql(rest, graphql, api_key)
412    }
413
414    /// Construct with explicit REST + GraphQL endpoints.
415    pub fn new_with_graphql(
416        endpoint: impl Into<String>,
417        graphql_endpoint: impl Into<String>,
418        api_key: impl Into<String>,
419    ) -> Self {
420        let http = Client::builder()
421            .timeout(Duration::from_secs(30))
422            .build()
423            .unwrap_or_default();
424        Self {
425            endpoint: endpoint.into(),
426            graphql_endpoint: graphql_endpoint.into(),
427            api_key: api_key.into(),
428            http,
429        }
430    }
431
432    /// Construct from config + environment. `RUNPOD_API_KEY` overrides
433    /// `settings.api_key`. Returns `RunPodAuth` error if no key is available.
434    pub fn from_settings(settings: &RunPodSettings) -> std::result::Result<Self, MoldError> {
435        let key = std::env::var(API_KEY_ENV)
436            .ok()
437            .filter(|k| !k.is_empty())
438            .or_else(|| settings.api_key.clone())
439            .ok_or_else(|| {
440                MoldError::RunPodAuth(format!(
441                    "RunPod API key not set — export {API_KEY_ENV} or run \
442                     `mold config set runpod.api_key <key>`"
443                ))
444            })?;
445        let endpoint = settings
446            .endpoint
447            .clone()
448            .unwrap_or_else(|| DEFAULT_ENDPOINT.to_string());
449        Ok(Self::new(endpoint, key))
450    }
451
452    fn url(&self, path: &str) -> String {
453        format!("{}{}", self.endpoint.trim_end_matches('/'), path)
454    }
455
456    async fn get_json<T: for<'de> Deserialize<'de>>(&self, path: &str) -> Result<T> {
457        let resp = self
458            .http
459            .get(self.url(path))
460            .bearer_auth(&self.api_key)
461            .send()
462            .await
463            .map_err(|e| MoldError::RunPod(format!("RunPod {path}: {e}")))?;
464        let status = resp.status();
465        if status.is_success() {
466            let body = resp
467                .text()
468                .await
469                .map_err(|e| MoldError::RunPod(format!("RunPod {path} body: {e}")))?;
470            serde_json::from_str(&body).map_err(|e| {
471                MoldError::RunPod(format!(
472                    "RunPod {path}: failed to parse response: {e} — body: {}",
473                    truncate_for_error(&body)
474                ))
475                .into()
476            })
477        } else {
478            Err(http_error(path, status, resp).await.into())
479        }
480    }
481
482    async fn post_json<B: Serialize, T: for<'de> Deserialize<'de>>(
483        &self,
484        path: &str,
485        body: &B,
486    ) -> Result<T> {
487        let resp = self
488            .http
489            .post(self.url(path))
490            .bearer_auth(&self.api_key)
491            .json(body)
492            .send()
493            .await
494            .map_err(|e| MoldError::RunPod(format!("RunPod {path}: {e}")))?;
495        let status = resp.status();
496        if status.is_success() {
497            let text = resp
498                .text()
499                .await
500                .map_err(|e| MoldError::RunPod(format!("RunPod {path} body: {e}")))?;
501            serde_json::from_str(&text).map_err(|e| {
502                MoldError::RunPod(format!(
503                    "RunPod {path}: failed to parse response: {e} — body: {}",
504                    truncate_for_error(&text)
505                ))
506                .into()
507            })
508        } else {
509            Err(http_error(path, status, resp).await.into())
510        }
511    }
512
513    async fn post_empty(&self, path: &str) -> Result<()> {
514        let resp = self
515            .http
516            .post(self.url(path))
517            .bearer_auth(&self.api_key)
518            .send()
519            .await
520            .map_err(|e| MoldError::RunPod(format!("RunPod {path}: {e}")))?;
521        let status = resp.status();
522        if status.is_success() {
523            Ok(())
524        } else {
525            Err(http_error(path, status, resp).await.into())
526        }
527    }
528
529    async fn patch_json<B: Serialize, T: for<'de> Deserialize<'de>>(
530        &self,
531        path: &str,
532        body: &B,
533    ) -> Result<T> {
534        let resp = self
535            .http
536            .patch(self.url(path))
537            .bearer_auth(&self.api_key)
538            .json(body)
539            .send()
540            .await
541            .map_err(|e| MoldError::RunPod(format!("RunPod {path}: {e}")))?;
542        let status = resp.status();
543        if status.is_success() {
544            let text = resp
545                .text()
546                .await
547                .map_err(|e| MoldError::RunPod(format!("RunPod {path} body: {e}")))?;
548            serde_json::from_str(&text).map_err(|e| {
549                MoldError::RunPod(format!(
550                    "RunPod {path}: failed to parse response: {e} — body: {}",
551                    truncate_for_error(&text)
552                ))
553                .into()
554            })
555        } else {
556            Err(http_error(path, status, resp).await.into())
557        }
558    }
559
560    async fn delete(&self, path: &str) -> Result<()> {
561        let resp = self
562            .http
563            .delete(self.url(path))
564            .bearer_auth(&self.api_key)
565            .send()
566            .await
567            .map_err(|e| MoldError::RunPod(format!("RunPod {path}: {e}")))?;
568        let status = resp.status();
569        if status.is_success() {
570            Ok(())
571        } else {
572            Err(http_error(path, status, resp).await.into())
573        }
574    }
575
576    // ─── Typed endpoints ────────────────────────────────────────────
577
578    /// User/account info isn't exposed by the REST API, so we fall back to
579    /// the GraphQL endpoint (same API key works for both).
580    pub async fn user(&self) -> Result<UserInfo> {
581        let query = serde_json::json!({
582            "query": "query { myself { id email clientBalance currentSpendPerHr spendLimit } }"
583        });
584        let resp = self
585            .http
586            .post(&self.graphql_endpoint)
587            .bearer_auth(&self.api_key)
588            .json(&query)
589            .send()
590            .await
591            .map_err(|e| MoldError::RunPod(format!("RunPod graphql /user: {e}")))?;
592        let status = resp.status();
593        if !status.is_success() {
594            return Err(http_error("graphql /user", status, resp).await.into());
595        }
596        let body: serde_json::Value = resp
597            .json()
598            .await
599            .map_err(|e| MoldError::RunPod(format!("RunPod graphql /user json: {e}")))?;
600        if let Some(errs) = body.get("errors") {
601            return Err(MoldError::RunPod(format!("RunPod graphql errors: {errs}")).into());
602        }
603        let myself = body
604            .get("data")
605            .and_then(|d| d.get("myself"))
606            .ok_or_else(|| MoldError::RunPod("graphql: missing data.myself".into()))?;
607        let info = UserInfo {
608            id: myself
609                .get("id")
610                .and_then(|v| v.as_str())
611                .unwrap_or("")
612                .to_string(),
613            email: myself
614                .get("email")
615                .and_then(|v| v.as_str())
616                .unwrap_or("")
617                .to_string(),
618            client_balance: myself
619                .get("clientBalance")
620                .and_then(|v| v.as_f64())
621                .unwrap_or(0.0),
622            current_spend_per_hr: myself
623                .get("currentSpendPerHr")
624                .and_then(|v| v.as_f64())
625                .unwrap_or(0.0),
626            spend_limit: myself.get("spendLimit").and_then(|v| v.as_f64()),
627        };
628        Ok(info)
629    }
630
631    /// Query GPU types via GraphQL (not exposed in REST v1).
632    /// Stock status is aggregated: the highest stock level across all DCs.
633    pub async fn gpu_types(&self) -> Result<Vec<GpuType>> {
634        let query = serde_json::json!({
635            "query": "query { gpuTypes { id displayName memoryInGb secureCloud communityCloud } dataCenters { gpuAvailability { displayName stockStatus } } }"
636        });
637        let body = self.graphql(&query).await?;
638        let data = body
639            .get("data")
640            .ok_or_else(|| MoldError::RunPod("graphql: missing data".into()))?;
641        let types: Vec<GpuType> = serde_json::from_value(
642            data.get("gpuTypes")
643                .cloned()
644                .unwrap_or(serde_json::Value::Array(vec![])),
645        )
646        .map_err(|e| MoldError::RunPod(format!("parse gpuTypes: {e}")))?;
647        // Aggregate stock across datacenters.
648        let mut best_stock: std::collections::HashMap<String, String> =
649            std::collections::HashMap::new();
650        if let Some(dcs) = data.get("dataCenters").and_then(|v| v.as_array()) {
651            for dc in dcs {
652                if let Some(avail) = dc.get("gpuAvailability").and_then(|v| v.as_array()) {
653                    for a in avail {
654                        if let (Some(name), Some(stock)) = (
655                            a.get("displayName").and_then(|v| v.as_str()),
656                            a.get("stockStatus").and_then(|v| v.as_str()),
657                        ) {
658                            let current = best_stock.get(name).cloned().unwrap_or_default();
659                            if stock_rank(stock) > stock_rank(&current) {
660                                best_stock.insert(name.to_string(), stock.to_string());
661                            }
662                        }
663                    }
664                }
665            }
666        }
667        let mut out = types;
668        for g in out.iter_mut() {
669            if let Some(s) = best_stock.get(&g.display_name) {
670                if !s.is_empty() {
671                    g.stock_status = Some(s.clone());
672                }
673            }
674            g.available = g.stock_status.as_deref().is_some_and(|s| s != "None");
675        }
676        Ok(out)
677    }
678
679    /// Query datacenters with per-GPU availability via GraphQL.
680    pub async fn datacenters(&self) -> Result<Vec<Datacenter>> {
681        let query = serde_json::json!({
682            "query": "query { dataCenters { id name listed gpuAvailability { id displayName stockStatus } } }"
683        });
684        let body = self.graphql(&query).await?;
685        let arr = body
686            .get("data")
687            .and_then(|d| d.get("dataCenters"))
688            .cloned()
689            .unwrap_or(serde_json::Value::Array(vec![]));
690        // Map GraphQL `id` → `gpuId` so we can reuse the same Datacenter type.
691        let arr = match arr {
692            serde_json::Value::Array(mut dcs) => {
693                for dc in dcs.iter_mut() {
694                    if let Some(avail) =
695                        dc.get_mut("gpuAvailability").and_then(|v| v.as_array_mut())
696                    {
697                        for a in avail.iter_mut() {
698                            if let Some(id) = a.get("id").and_then(|v| v.as_str()) {
699                                let id = id.to_string();
700                                if let Some(obj) = a.as_object_mut() {
701                                    obj.insert("gpuId".into(), serde_json::Value::String(id));
702                                }
703                            }
704                        }
705                    }
706                }
707                serde_json::Value::Array(dcs)
708            }
709            other => other,
710        };
711        let dcs: Vec<Datacenter> = serde_json::from_value(arr)
712            .map_err(|e| MoldError::RunPod(format!("parse dataCenters: {e}")))?;
713        Ok(dcs)
714    }
715
716    async fn graphql(&self, query: &serde_json::Value) -> Result<serde_json::Value> {
717        let resp = self
718            .http
719            .post(&self.graphql_endpoint)
720            .bearer_auth(&self.api_key)
721            .json(query)
722            .send()
723            .await
724            .map_err(|e| MoldError::RunPod(format!("RunPod graphql: {e}")))?;
725        let status = resp.status();
726        if !status.is_success() {
727            return Err(http_error("graphql", status, resp).await.into());
728        }
729        let body: serde_json::Value = resp
730            .json()
731            .await
732            .map_err(|e| MoldError::RunPod(format!("graphql body: {e}")))?;
733        if let Some(errs) = body
734            .get("errors")
735            .filter(|e| !e.as_array().map(|a| a.is_empty()).unwrap_or(true))
736        {
737            return Err(MoldError::RunPod(format!("graphql errors: {errs}")).into());
738        }
739        Ok(body)
740    }
741
742    pub async fn list_pods(&self) -> Result<Vec<Pod>> {
743        self.get_json("/pods?includeMachine=true").await
744    }
745
746    pub async fn get_pod(&self, id: &str) -> Result<Pod> {
747        self.get_json(&format!("/pods/{id}?includeMachine=true"))
748            .await
749    }
750
751    pub async fn create_pod(&self, req: &CreatePodRequest) -> Result<Pod> {
752        self.post_json("/pods", req).await
753    }
754
755    /// GPU IDs accepted by the REST Pod create endpoint. RunPod's GraphQL
756    /// inventory can advertise GPUs before the REST create schema accepts
757    /// them, so launchers must intersect inventory with this live contract.
758    pub async fn supported_pod_gpu_type_ids(&self) -> Result<HashSet<String>> {
759        let spec: serde_json::Value = self.get_json("/openapi.json").await?;
760        parse_pod_gpu_type_ids(&spec)
761    }
762
763    pub async fn stop_pod(&self, id: &str) -> Result<()> {
764        self.post_empty(&format!("/pods/{id}/stop")).await
765    }
766
767    pub async fn start_pod(&self, id: &str) -> Result<()> {
768        self.post_empty(&format!("/pods/{id}/start")).await
769    }
770
771    pub async fn delete_pod(&self, id: &str) -> Result<()> {
772        self.delete(&format!("/pods/{id}")).await
773    }
774
775    pub async fn network_volumes(&self) -> Result<Vec<NetworkVolume>> {
776        self.get_json("/networkvolumes").await
777    }
778
779    pub async fn get_network_volume(&self, id: &str) -> Result<NetworkVolume> {
780        self.get_json(&format!("/networkvolumes/{id}")).await
781    }
782
783    pub async fn create_network_volume(
784        &self,
785        req: &CreateNetworkVolumeRequest,
786    ) -> Result<NetworkVolume> {
787        self.post_json("/networkvolumes", req).await
788    }
789
790    pub async fn update_network_volume(
791        &self,
792        id: &str,
793        req: &UpdateNetworkVolumeRequest,
794    ) -> Result<NetworkVolume> {
795        self.patch_json(&format!("/networkvolumes/{id}"), req).await
796    }
797
798    pub async fn delete_network_volume(&self, id: &str) -> Result<()> {
799        self.delete(&format!("/networkvolumes/{id}")).await
800    }
801
802    /// Permanently delete a network volume only after proving that no Pod is
803    /// attached. The attachment check intentionally fails closed: an auth,
804    /// transport, or API error must never fall through to destructive delete.
805    pub async fn delete_network_volume_if_detached(&self, id: &str) -> Result<()> {
806        let pods = self.list_pods().await?;
807        if let Some(pod) = pods
808            .iter()
809            .find(|pod| pod.attached_network_volume_id() == Some(id))
810        {
811            return Err(MoldError::RunPod(format!(
812                "delete pod {} before deleting its attached network volume",
813                pod.id
814            ))
815            .into());
816        }
817        self.delete_network_volume(id).await
818    }
819}
820
821// ─── Helpers ────────────────────────────────────────────────────────────────
822
823async fn http_error(path: &str, status: StatusCode, resp: reqwest::Response) -> MoldError {
824    let body = resp.text().await.unwrap_or_default();
825    let msg = truncate_for_error(&body);
826    match status {
827        StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
828            MoldError::RunPodAuth(format!("RunPod {path} {status}: {msg}"))
829        }
830        StatusCode::NOT_FOUND => {
831            MoldError::RunPodNotFound(format!("RunPod {path} {status}: {msg}"))
832        }
833        StatusCode::CONFLICT
834        | StatusCode::SERVICE_UNAVAILABLE
835        | StatusCode::INTERNAL_SERVER_ERROR
836            if {
837                let lower = msg.to_lowercase();
838                lower.contains("does not have the resources")
839                    || lower.contains("no instances currently available")
840            } =>
841        {
842            MoldError::RunPodNoStock(format!("RunPod {path} {status}: {msg}"))
843        }
844        _ => MoldError::RunPod(format!("RunPod {path} {status}: {msg}")),
845    }
846}
847
848fn stock_rank(s: &str) -> u8 {
849    match s {
850        "High" => 3,
851        "Medium" => 2,
852        "Low" => 1,
853        _ => 0,
854    }
855}
856
857fn parse_pod_gpu_type_ids(spec: &serde_json::Value) -> Result<HashSet<String>> {
858    let values = spec
859        .pointer("/components/schemas/PodCreateInput/properties/gpuTypeIds/items/enum")
860        .and_then(serde_json::Value::as_array)
861        .ok_or_else(|| MoldError::RunPod("OpenAPI schema is missing Pod GPU types".into()))?;
862    Ok(values
863        .iter()
864        .filter_map(serde_json::Value::as_str)
865        .map(str::to_string)
866        .collect())
867}
868
869fn truncate_for_error(s: &str) -> String {
870    const MAX: usize = 400;
871    let s = s.trim();
872    if s.len() <= MAX {
873        s.to_string()
874    } else {
875        format!("{}…", &s[..MAX])
876    }
877}
878
879/// Map a RunPod GPU `displayName` (e.g. `"RTX 4090"`) to the matching
880/// `ghcr.io/utensils/mold` image tag.
881pub fn image_tag_for_gpu(
882    display_name: &str,
883    version: &str,
884) -> Result<String, crate::cuda_distribution::UnsupportedPublishedImagePlatform> {
885    crate::cuda_distribution::image_tag_for_gpu_name(display_name, version)
886}
887
888/// Ranked preference when auto-picking GPUs. Higher index = more preferred.
889pub const GPU_PREFERENCE: &[&str] = &[
890    "A100 PCIe",
891    "L40",
892    "L40S",
893    "RTX A6000",
894    "RTX 5090",
895    "RTX 4090",
896];
897
898#[cfg(test)]
899mod tests {
900    use super::*;
901
902    #[test]
903    fn gpu_type_authority_prefers_id_then_gpu_id_and_rejects_blank_values() {
904        let gpu: GpuType = serde_json::from_value(serde_json::json!({
905            "id": "  primary-id  ",
906            "gpuId": "  alternate-id  ",
907            "displayName": "Display label"
908        }))
909        .unwrap();
910        assert_eq!(gpu.authoritative_type_id(), Some("primary-id"));
911
912        let gpu: GpuType = serde_json::from_value(serde_json::json!({
913            "id": "  ",
914            "gpuId": "  alternate-id  ",
915            "displayName": "Display label"
916        }))
917        .unwrap();
918        assert_eq!(gpu.authoritative_type_id(), Some("alternate-id"));
919
920        let gpu: GpuType = serde_json::from_value(serde_json::json!({
921            "id": "  ",
922            "gpuId": "\t",
923            "displayName": "Display label"
924        }))
925        .unwrap();
926        assert_eq!(gpu.authoritative_type_id(), None);
927    }
928
929    #[test]
930    fn gpu_type_allocation_identity_uses_display_only_as_legacy_fallback() {
931        let gpu: GpuType = serde_json::from_value(serde_json::json!({
932            "id": "  ",
933            "gpuId": "\t",
934            "displayName": "  RTX 5090 "
935        }))
936        .unwrap();
937        assert_eq!(gpu.allocation_type_id(), Some("NVIDIA GeForce RTX 5090"));
938
939        let gpu: GpuType = serde_json::from_value(serde_json::json!({
940            "id": " provider-id ",
941            "gpuId": "alternate-id",
942            "displayName": "RTX 5090"
943        }))
944        .unwrap();
945        assert_eq!(gpu.allocation_type_id(), Some("provider-id"));
946    }
947
948    #[test]
949    fn gpu_type_memory_preserves_missing_and_present_zero() {
950        let missing: GpuType = serde_json::from_value(serde_json::json!({
951            "displayName": "RTX 5090"
952        }))
953        .unwrap();
954        let zero: GpuType = serde_json::from_value(serde_json::json!({
955            "displayName": "RTX 5090",
956            "memoryInGb": 0
957        }))
958        .unwrap();
959        assert_eq!(missing.memory_in_gb, None);
960        assert_eq!(zero.memory_in_gb, Some(0));
961    }
962
963    #[test]
964    fn supported_gpu_identity_matching_normalizes_case_and_whitespace() {
965        let supported = ["NVIDIA A100-SXM4-80GB".to_string()].into_iter().collect();
966        assert_eq!(
967            canonical_supported_gpu_type_id(&supported, "  nvidia   a100-sxm4-80gb "),
968            Some("NVIDIA A100-SXM4-80GB")
969        );
970        assert!(canonical_supported_gpu_type_id(&supported, " \t ").is_none());
971    }
972
973    #[test]
974    fn image_tag_mapping() {
975        assert_eq!(image_tag_for_gpu("RTX 4090", "latest").unwrap(), "latest");
976        assert_eq!(
977            image_tag_for_gpu("NVIDIA GeForce RTX 4090", "0.10.0").unwrap(),
978            "0.10.0"
979        );
980        assert_eq!(image_tag_for_gpu("L40S", "latest").unwrap(), "latest");
981        assert_eq!(
982            image_tag_for_gpu("RTX 5090", "latest").unwrap(),
983            "latest-sm120"
984        );
985        assert_eq!(
986            image_tag_for_gpu("NVIDIA GeForce RTX 5090", "0.10.0").unwrap(),
987            "0.10.0-sm120"
988        );
989        assert_eq!(
990            image_tag_for_gpu("RTX PRO 4500", "latest").unwrap(),
991            "latest-sm120"
992        );
993        assert_eq!(
994            image_tag_for_gpu("NVIDIA B200", "latest").unwrap(),
995            "latest-sm100"
996        );
997        assert!(image_tag_for_gpu("NVIDIA GB200", "0.10.0").is_err());
998        assert_eq!(
999            image_tag_for_gpu("A100 80GB", "latest").unwrap(),
1000            "latest-sm80"
1001        );
1002        assert_eq!(
1003            image_tag_for_gpu("A100 PCIe", "latest").unwrap(),
1004            "latest-sm80"
1005        );
1006        assert_eq!(
1007            image_tag_for_gpu("RTX 3090", "latest").unwrap(),
1008            "latest-sm86"
1009        );
1010        assert_eq!(
1011            image_tag_for_gpu("NVIDIA A40", "latest").unwrap(),
1012            "latest-sm86"
1013        );
1014        assert_eq!(
1015            image_tag_for_gpu("H100 SXM", "latest").unwrap(),
1016            "latest-sm90"
1017        );
1018        assert_eq!(
1019            image_tag_for_gpu("H200 SXM", "latest").unwrap(),
1020            "latest-sm90"
1021        );
1022        assert_eq!(
1023            image_tag_for_gpu("NVIDIA A10", "latest").unwrap(),
1024            "latest-sm86"
1025        );
1026        assert_eq!(
1027            image_tag_for_gpu("NVIDIA RTX A6000", "latest").unwrap(),
1028            "latest-sm86"
1029        );
1030        assert_eq!(
1031            image_tag_for_gpu("NVIDIA A16", "latest").unwrap(),
1032            "latest-sm86"
1033        );
1034        assert_eq!(
1035            image_tag_for_gpu("NVIDIA A2", "latest").unwrap(),
1036            "latest-sm86"
1037        );
1038        assert_eq!(
1039            image_tag_for_gpu("NVIDIA A30", "latest").unwrap(),
1040            "latest-sm80"
1041        );
1042        assert_eq!(
1043            image_tag_for_gpu("NVIDIA B300", "latest").unwrap(),
1044            "latest-sm100"
1045        );
1046        assert!(image_tag_for_gpu("NVIDIA GB300", "latest").is_err());
1047        assert_eq!(
1048            image_tag_for_gpu("NVIDIA Blackwell", "latest").unwrap(),
1049            "latest",
1050            "generic Blackwell must not guess between incompatible targets"
1051        );
1052    }
1053
1054    #[test]
1055    fn live_network_volume_size_bounds_are_enforced() {
1056        assert!(!valid_network_volume_size(9));
1057        assert!(valid_network_volume_size(10));
1058        assert!(valid_network_volume_size(3999));
1059        assert!(!valid_network_volume_size(4000));
1060    }
1061
1062    #[test]
1063    fn pod_reads_gpu_from_top_level_rest_shape() {
1064        let pod: Pod = serde_json::from_value(serde_json::json!({
1065            "id": "pod-1",
1066            "desiredStatus": "RUNNING",
1067            "gpu": { "id": "NVIDIA RTX 6000", "displayName": "RTX PRO 6000 Blackwell" }
1068        }))
1069        .unwrap();
1070        assert_eq!(
1071            pod.gpu.and_then(|gpu| gpu.display_name).as_deref(),
1072            Some("RTX PRO 6000 Blackwell")
1073        );
1074    }
1075
1076    #[test]
1077    fn pod_reads_current_machine_gpu_and_network_volume_id_shape() {
1078        let pod: Pod = serde_json::from_value(serde_json::json!({
1079            "id": "pod-1",
1080            "desiredStatus": "RUNNING",
1081            "networkVolumeId": "nv-1",
1082            "machine": {
1083                "gpuTypeId": "NVIDIA GeForce RTX 4090",
1084                "dataCenterId": "EU-RO-1",
1085                "location": "RO"
1086            }
1087        }))
1088        .unwrap();
1089        assert_eq!(pod.attached_network_volume_id(), Some("nv-1"));
1090        assert_eq!(pod.gpu_name(), Some("NVIDIA GeForce RTX 4090"));
1091        assert_eq!(pod.datacenter_id(), Some("EU-RO-1"));
1092        assert_eq!(
1093            pod.machine
1094                .as_ref()
1095                .and_then(|machine| machine.gpu_type_id.as_deref()),
1096            Some("NVIDIA GeForce RTX 4090")
1097        );
1098    }
1099
1100    #[test]
1101    fn pod_location_is_not_treated_as_an_exact_datacenter_id() {
1102        let pod: Pod = serde_json::from_value(serde_json::json!({
1103            "id": "pod-1",
1104            "machine": { "location": "RO" }
1105        }))
1106        .unwrap();
1107
1108        assert_eq!(pod.datacenter_id(), None);
1109    }
1110
1111    #[test]
1112    fn parses_rest_pod_gpu_ids_from_openapi() {
1113        let spec = serde_json::json!({
1114            "components": { "schemas": { "PodCreateInput": { "properties": {
1115                "gpuTypeIds": { "items": { "enum": ["NVIDIA GeForce RTX 5090", "NVIDIA L40S"] } }
1116            } } } }
1117        });
1118        assert_eq!(
1119            parse_pod_gpu_type_ids(&spec).unwrap(),
1120            ["NVIDIA GeForce RTX 5090", "NVIDIA L40S"]
1121                .into_iter()
1122                .map(str::to_string)
1123                .collect()
1124        );
1125    }
1126
1127    #[test]
1128    fn redacted_debug_hides_api_key() {
1129        let s = RunPodSettings {
1130            api_key: Some("secret-key".to_string()),
1131            ..Default::default()
1132        };
1133        let out = s.redacted_debug();
1134        assert!(!out.contains("secret-key"));
1135        assert!(out.contains("<redacted>"));
1136    }
1137
1138    #[test]
1139    fn from_settings_requires_key() {
1140        std::env::remove_var(API_KEY_ENV);
1141        let err = RunPodClient::from_settings(&RunPodSettings::default()).unwrap_err();
1142        assert!(matches!(err, MoldError::RunPodAuth(_)));
1143    }
1144
1145    #[test]
1146    fn truncate_for_error_boundary() {
1147        let short = "short";
1148        assert_eq!(truncate_for_error(short), "short");
1149        let long = "x".repeat(500);
1150        let truncated = truncate_for_error(&long);
1151        assert!(truncated.ends_with('…'));
1152        assert!(truncated.chars().count() <= 401);
1153    }
1154
1155    #[test]
1156    fn runpod_settings_toml_roundtrip() {
1157        let original = RunPodSettings {
1158            api_key: Some("k".to_string()),
1159            default_gpu: Some("RTX 5090".to_string()),
1160            default_datacenter: Some("EUR-IS-2".to_string()),
1161            default_network_volume_id: Some("nv-123".to_string()),
1162            auto_teardown: true,
1163            auto_teardown_idle_mins: 30,
1164            cost_alert_usd: 3.5,
1165            endpoint: None,
1166        };
1167        let toml_s = toml::to_string(&original).unwrap();
1168        let round: RunPodSettings = toml::from_str(&toml_s).unwrap();
1169        assert_eq!(round.api_key, original.api_key);
1170        assert_eq!(round.default_gpu, original.default_gpu);
1171        assert_eq!(round.default_datacenter, original.default_datacenter);
1172        assert_eq!(
1173            round.default_network_volume_id,
1174            original.default_network_volume_id
1175        );
1176        assert_eq!(round.auto_teardown, original.auto_teardown);
1177        assert_eq!(
1178            round.auto_teardown_idle_mins,
1179            original.auto_teardown_idle_mins
1180        );
1181        assert_eq!(round.cost_alert_usd, original.cost_alert_usd);
1182    }
1183
1184    #[test]
1185    fn default_auto_teardown_idle_mins_is_20() {
1186        let s: RunPodSettings = toml::from_str("").unwrap();
1187        assert_eq!(s.auto_teardown_idle_mins, 20);
1188    }
1189}