1use 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
15pub const DEFAULT_ENDPOINT: &str = "https://rest.runpod.io/v1";
17
18pub const GRAPHQL_ENDPOINT: &str = "https://api.runpod.io/graphql";
20
21pub const API_KEY_ENV: &str = "RUNPOD_API_KEY";
23
24pub 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#[derive(Debug, Clone, Deserialize, Serialize, Default)]
35pub struct RunPodSettings {
36 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub api_key: Option<String>,
39
40 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub default_gpu: Option<String>,
43
44 #[serde(default, skip_serializing_if = "Option::is_none")]
46 pub default_datacenter: Option<String>,
47
48 #[serde(default, skip_serializing_if = "Option::is_none")]
50 pub default_network_volume_id: Option<String>,
51
52 #[serde(default)]
55 pub auto_teardown: bool,
56
57 #[serde(default = "default_auto_teardown_idle_mins")]
60 pub auto_teardown_idle_mins: u32,
61
62 #[serde(default)]
65 pub cost_alert_usd: f64,
66
67 #[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
76impl 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#[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#[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 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 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
160pub 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#[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#[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 #[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 #[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#[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#[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#[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#[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#[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 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 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 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 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 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 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(¤t) {
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 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 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 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 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
821async 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
879pub 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
888pub 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}