1use crate::error::MoldError;
4use anyhow::Result;
5use reqwest::{Client, StatusCode};
6use serde::{Deserialize, Serialize};
7use std::time::Duration;
8
9pub const DEFAULT_ENDPOINT: &str = "https://cloud.lambda.ai/api/v1";
10pub const API_KEY_ENV: &str = "LAMBDA_API_KEY";
11pub const DEFAULT_IMAGE_REPOSITORY: &str = "ghcr.io/utensils/mold";
12
13#[derive(Debug, Clone, Deserialize, Serialize)]
14pub struct LambdaSettings {
15 #[serde(default, skip_serializing_if = "Option::is_none")]
16 pub api_key: Option<String>,
17 #[serde(
18 default = "default_endpoint_opt",
19 skip_serializing_if = "Option::is_none"
20 )]
21 pub endpoint: Option<String>,
22 #[serde(
23 default = "default_image_repository_opt",
24 skip_serializing_if = "Option::is_none"
25 )]
26 pub image_repository: Option<String>,
27 #[serde(default, skip_serializing_if = "Option::is_none")]
28 pub ssh_key_name: Option<String>,
29 #[serde(default, skip_serializing_if = "Option::is_none")]
30 pub ssh_private_key_path: Option<String>,
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub filesystem_prefix: Option<String>,
33 #[serde(default = "default_filesystem_mount_path")]
34 pub filesystem_mount_path: String,
35 #[serde(default = "default_confirm_hourly_usd")]
36 pub confirm_hourly_usd: f64,
37 #[serde(default = "default_local_port")]
38 pub local_port: u16,
39}
40
41impl Default for LambdaSettings {
42 fn default() -> Self {
43 Self {
44 api_key: None,
45 endpoint: default_endpoint_opt(),
46 image_repository: default_image_repository_opt(),
47 ssh_key_name: None,
48 ssh_private_key_path: None,
49 filesystem_prefix: None,
50 filesystem_mount_path: default_filesystem_mount_path(),
51 confirm_hourly_usd: default_confirm_hourly_usd(),
52 local_port: default_local_port(),
53 }
54 }
55}
56
57fn default_endpoint_opt() -> Option<String> {
58 Some(DEFAULT_ENDPOINT.to_string())
59}
60
61fn default_image_repository_opt() -> Option<String> {
62 Some(DEFAULT_IMAGE_REPOSITORY.to_string())
63}
64
65fn default_filesystem_mount_path() -> String {
66 "/data/mold".to_string()
67}
68
69fn default_confirm_hourly_usd() -> f64 {
70 5.0
71}
72
73fn default_local_port() -> u16 {
74 7680
75}
76
77impl LambdaSettings {
78 pub fn resolved_api_key(&self) -> Option<String> {
79 std::env::var(API_KEY_ENV)
80 .ok()
81 .filter(|s| !s.is_empty())
82 .or_else(|| self.api_key.clone())
83 }
84
85 pub fn endpoint(&self) -> &str {
86 self.endpoint.as_deref().unwrap_or(DEFAULT_ENDPOINT)
87 }
88
89 pub fn image_repository(&self) -> &str {
90 self.image_repository
91 .as_deref()
92 .unwrap_or(DEFAULT_IMAGE_REPOSITORY)
93 }
94
95 pub fn redacted_debug(&self) -> String {
96 format!(
97 "LambdaSettings {{ api_key: {}, endpoint: {:?}, image_repository: {:?}, \
98 ssh_key_name: {:?}, ssh_private_key_path: {:?}, filesystem_prefix: {:?}, \
99 filesystem_mount_path: {:?}, confirm_hourly_usd: {}, local_port: {} }}",
100 if self.api_key.is_some() {
101 "Some(\"<redacted>\")"
102 } else {
103 "None"
104 },
105 self.endpoint,
106 self.image_repository,
107 self.ssh_key_name,
108 self.ssh_private_key_path,
109 self.filesystem_prefix,
110 self.filesystem_mount_path,
111 self.confirm_hourly_usd,
112 self.local_port,
113 )
114 }
115}
116
117#[derive(Debug, Clone, Deserialize, Serialize)]
118pub struct ApiList<T> {
119 #[serde(default)]
120 pub data: Vec<T>,
121}
122
123#[derive(Debug, Clone, Deserialize, Serialize)]
124pub struct ApiItem<T> {
125 pub data: T,
126}
127
128#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
129pub struct Region {
130 pub name: String,
131 #[serde(default)]
132 pub description: String,
133}
134
135#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
136pub struct InstanceTypeSpecs {
137 #[serde(default)]
138 pub gpus: u32,
139 #[serde(default)]
140 pub gpu_description: String,
141 #[serde(default)]
142 pub memory_gib: u32,
143 #[serde(default)]
144 pub storage_gib: u32,
145 #[serde(default)]
146 pub vcpus: u32,
147}
148
149#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
150pub struct InstanceType {
151 pub name: String,
152 #[serde(default)]
153 pub description: String,
154 #[serde(default)]
155 pub gpu_description: String,
156 #[serde(default)]
157 pub price_cents_per_hour: u32,
158 #[serde(default)]
159 pub specs: InstanceTypeSpecs,
160 #[serde(default)]
161 pub regions_with_capacity_available: Vec<Region>,
162}
163
164#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
165pub struct SshKey {
166 pub id: String,
167 pub name: String,
168 pub public_key: String,
169}
170
171#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
172pub struct Filesystem {
173 pub id: String,
174 pub name: String,
175 #[serde(default)]
176 pub mount_point: String,
177 #[serde(default)]
178 pub region: Option<Region>,
179 #[serde(default)]
180 pub bytes_used: Option<u64>,
181}
182
183#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
184pub struct Tag {
185 pub key: String,
186 pub value: String,
187}
188
189#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
190pub struct Instance {
191 pub id: String,
192 #[serde(default)]
193 pub name: Option<String>,
194 #[serde(default)]
195 pub status: String,
196 #[serde(default)]
197 pub ip: Option<String>,
198 #[serde(default)]
199 pub private_ip: Option<String>,
200 #[serde(default)]
201 pub instance_type: Option<InstanceType>,
202 #[serde(default)]
203 pub region: Option<Region>,
204 #[serde(default)]
205 pub ssh_key_names: Vec<String>,
206 #[serde(default)]
207 pub file_system_names: Vec<String>,
208 #[serde(default)]
209 pub tags: Vec<Tag>,
210}
211
212#[derive(Debug, Clone, Default, Deserialize, Serialize)]
213pub struct CreateSshKeyRequest {
214 pub name: String,
215 pub public_key: String,
216}
217
218#[derive(Debug, Clone, Default, Deserialize, Serialize)]
219pub struct CreateFilesystemRequest {
220 pub name: String,
221 pub region: String,
222}
223
224#[derive(Debug, Clone, Default, Deserialize, Serialize)]
225pub struct LaunchInstancesRequest {
226 pub region_name: String,
227 pub instance_type_name: String,
228 pub ssh_key_names: Vec<String>,
229 #[serde(skip_serializing_if = "Vec::is_empty", default)]
230 pub file_system_names: Vec<String>,
231 #[serde(skip_serializing_if = "Vec::is_empty", default)]
232 pub file_system_mounts: Vec<FilesystemMount>,
233 #[serde(skip_serializing_if = "Option::is_none")]
234 pub hostname: Option<String>,
235 pub name: String,
236 #[serde(skip_serializing_if = "Option::is_none")]
237 pub image: Option<LaunchImage>,
238 #[serde(skip_serializing_if = "Option::is_none")]
239 pub user_data: Option<String>,
240 #[serde(skip_serializing_if = "Vec::is_empty", default)]
241 pub tags: Vec<Tag>,
242}
243
244#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)]
245pub struct InstanceLaunchResponse {
246 #[serde(default)]
247 pub instance_ids: Vec<String>,
248}
249
250#[derive(Debug, Clone, Deserialize)]
251struct InstanceTypesResponse {
252 #[serde(default)]
253 data: std::collections::BTreeMap<String, InstanceTypeOffering>,
254}
255
256#[derive(Debug, Clone, Deserialize)]
257struct InstanceTypeOffering {
258 instance_type: InstanceType,
259 #[serde(default)]
260 regions_with_capacity_available: Vec<Region>,
261}
262
263#[derive(Debug, Clone, Deserialize, Serialize)]
264pub struct FilesystemMount {
265 pub mount_point: String,
266 #[serde(skip_serializing_if = "Option::is_none")]
267 pub file_system_name: Option<String>,
268 #[serde(skip_serializing_if = "Option::is_none")]
269 pub file_system_id: Option<String>,
270}
271
272#[derive(Debug, Clone, Deserialize, Serialize)]
273pub struct LaunchImage {
274 pub id: String,
275}
276
277pub struct LaunchRequestInput<'a> {
278 pub region_name: &'a str,
279 pub instance_type_name: &'a str,
280 pub ssh_key_name: &'a str,
281 pub filesystem_name: &'a str,
282 pub filesystem_id: Option<&'a str>,
283 pub filesystem_mount_path: &'a str,
284 pub instance_name: &'a str,
285 pub image_id: Option<&'a str>,
286 pub user_data: &'a str,
287}
288
289pub fn build_launch_request(input: LaunchRequestInput<'_>) -> LaunchInstancesRequest {
290 LaunchInstancesRequest {
291 region_name: input.region_name.to_string(),
292 instance_type_name: input.instance_type_name.to_string(),
293 ssh_key_names: vec![input.ssh_key_name.to_string()],
294 file_system_names: vec![input.filesystem_name.to_string()],
295 file_system_mounts: vec![FilesystemMount {
296 mount_point: input.filesystem_mount_path.to_string(),
297 file_system_name: input
298 .filesystem_id
299 .is_none()
300 .then(|| input.filesystem_name.to_string()),
301 file_system_id: input.filesystem_id.map(str::to_string),
302 }],
303 hostname: None,
304 name: input.instance_name.to_string(),
305 image: input.image_id.map(|id| LaunchImage { id: id.to_string() }),
306 user_data: Some(input.user_data.to_string()),
307 tags: vec![Tag {
308 key: "managed-by".to_string(),
309 value: "mold".to_string(),
310 }],
311 }
312}
313
314#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
315pub struct AvailabilityRow {
316 pub instance_type: String,
317 pub region: String,
318 pub gpu_description: String,
319 pub gpu_count: u32,
320 pub generation_slots: u32,
321 pub price_per_hour_usd: f64,
322 pub memory_gib: u32,
323 pub storage_gib: u32,
324 pub image: String,
325}
326
327impl AvailabilityRow {
328 pub fn from_instance_type(
329 instance_type: &InstanceType,
330 image_repository: &str,
331 version: &str,
332 ) -> Self {
333 let region = instance_type
334 .regions_with_capacity_available
335 .first()
336 .map(|r| r.name.clone())
337 .unwrap_or_default();
338 let image = match image_tag_for_gpu(&instance_type.specs.gpu_description, version) {
339 Ok(tag) if !gpu_uses_unsupported_linux_arm64(&instance_type.name) => {
340 format!("{image_repository}:{tag}")
341 }
342 Ok(_) | Err(_) => "unsupported: linux/arm64 host".to_string(),
343 };
344 Self {
345 instance_type: instance_type.name.clone(),
346 region,
347 gpu_description: instance_type.specs.gpu_description.clone(),
348 gpu_count: instance_type.specs.gpus,
349 generation_slots: instance_type.specs.gpus,
350 price_per_hour_usd: instance_type.price_cents_per_hour as f64 / 100.0,
351 memory_gib: instance_type.specs.memory_gib,
352 storage_gib: instance_type.specs.storage_gib,
353 image,
354 }
355 }
356}
357
358pub fn image_tag_for_gpu(
359 gpu_description: &str,
360 version: &str,
361) -> Result<String, crate::cuda_distribution::UnsupportedPublishedImagePlatform> {
362 crate::cuda_distribution::image_tag_for_gpu_name(gpu_description, version)
363}
364
365pub fn gpu_uses_unsupported_linux_arm64(gpu_description: &str) -> bool {
366 crate::cuda_distribution::gpu_name_uses_unsupported_grace_platform(gpu_description)
367}
368
369pub fn filesystem_name(settings: &LambdaSettings, region: &str) -> String {
370 let prefix = settings.filesystem_prefix.as_deref().unwrap_or("mold");
371 format!("{prefix}-{region}")
372}
373
374#[derive(Debug, Clone)]
375pub struct CloudInitOptions {
376 pub image: String,
377 pub mount_path: String,
378 pub env_file: String,
379}
380
381pub fn render_cloud_init(opts: &CloudInitOptions) -> String {
382 format!(
383 r#"#cloud-config
384write_files:
385 - path: /etc/systemd/system/mold-lambda.service
386 permissions: '0644'
387 content: |
388 [Unit]
389 Description=mold Lambda container
390 After=docker.service network-online.target
391 Wants=network-online.target
392
393 [Service]
394 Restart=always
395 RestartSec=10
396 ExecStartPre=-/usr/bin/docker rm -f mold
397 ExecStartPre=/usr/bin/docker pull {image}
398 ExecStart=/usr/bin/docker run --name mold --gpus all --restart unless-stopped --env-file {env_file} -e MOLD_PORT=7680 -p 127.0.0.1:7680:7680 -v {mount_path}:/workspace {image}
399 ExecStop=/usr/bin/docker stop mold
400
401 [Install]
402 WantedBy=multi-user.target
403runcmd:
404 - [ mkdir, -p, /etc/mold ]
405 - [ sh, -c, "touch {env_file} && chmod 600 {env_file}" ]
406 - [ systemctl, daemon-reload ]
407 - [ systemctl, enable, --now, mold-lambda.service ]
408"#,
409 image = opts.image,
410 mount_path = opts.mount_path,
411 env_file = opts.env_file,
412 )
413}
414
415#[derive(Clone)]
416pub struct LambdaClient {
417 client: Client,
418 endpoint: String,
419 api_key: String,
420}
421
422impl LambdaClient {
423 pub fn from_settings(settings: &LambdaSettings) -> Result<Self> {
424 let api_key = settings.resolved_api_key().ok_or_else(|| {
425 MoldError::Config("missing Lambda API key; set LAMBDA_API_KEY or lambda.api_key".into())
426 })?;
427 Ok(Self {
428 client: Client::builder().timeout(Duration::from_secs(60)).build()?,
429 endpoint: settings.endpoint().trim_end_matches('/').to_string(),
430 api_key,
431 })
432 }
433
434 pub fn new(endpoint: impl Into<String>, api_key: impl Into<String>) -> Self {
435 Self {
436 client: Client::new(),
437 endpoint: endpoint.into().trim_end_matches('/').to_string(),
438 api_key: api_key.into(),
439 }
440 }
441
442 async fn get_list<T: for<'de> Deserialize<'de> + Default>(&self, path: &str) -> Result<Vec<T>> {
443 let resp = self
444 .client
445 .get(format!("{}{}", self.endpoint, path))
446 .basic_auth(&self.api_key, Some(""))
447 .send()
448 .await?;
449 decode_list(resp).await
450 }
451
452 async fn post_item<B: Serialize, T: for<'de> Deserialize<'de>>(
453 &self,
454 path: &str,
455 body: &B,
456 ) -> Result<T> {
457 let resp = self
458 .client
459 .post(format!("{}{}", self.endpoint, path))
460 .basic_auth(&self.api_key, Some(""))
461 .json(body)
462 .send()
463 .await?;
464 decode_item(resp).await
465 }
466
467 pub async fn list_instance_types(&self) -> Result<Vec<InstanceType>> {
468 let resp = self
469 .client
470 .get(format!("{}/instance-types", self.endpoint))
471 .basic_auth(&self.api_key, Some(""))
472 .send()
473 .await?;
474 if !resp.status().is_success() {
475 return Err(lambda_error(resp).await.into());
476 }
477 decode_instance_types_body(&resp.text().await?)
478 }
479
480 pub async fn list_instances(&self) -> Result<Vec<Instance>> {
481 self.get_list("/instances").await
482 }
483
484 pub async fn get_instance(&self, id: &str) -> Result<Instance> {
485 let resp = self
486 .client
487 .get(format!("{}/instances/{id}", self.endpoint))
488 .basic_auth(&self.api_key, Some(""))
489 .send()
490 .await?;
491 decode_item(resp).await
492 }
493
494 pub async fn launch_instance(
495 &self,
496 req: &LaunchInstancesRequest,
497 ) -> Result<InstanceLaunchResponse> {
498 self.post_item("/instance-operations/launch", req).await
499 }
500
501 pub async fn terminate_instance(&self, id: &str) -> Result<()> {
502 let body = serde_json::json!({ "instance_ids": [id] });
503 let resp = self
504 .client
505 .post(format!("{}/instance-operations/terminate", self.endpoint))
506 .basic_auth(&self.api_key, Some(""))
507 .json(&body)
508 .send()
509 .await?;
510 ensure_success(resp).await
511 }
512
513 pub async fn list_ssh_keys(&self) -> Result<Vec<SshKey>> {
514 self.get_list("/ssh-keys").await
515 }
516
517 pub async fn create_ssh_key(&self, req: &CreateSshKeyRequest) -> Result<SshKey> {
518 self.post_item("/ssh-keys", req).await
519 }
520
521 pub async fn list_filesystems(&self) -> Result<Vec<Filesystem>> {
522 self.get_list("/file-systems").await
523 }
524
525 pub async fn create_filesystem(&self, req: &CreateFilesystemRequest) -> Result<Filesystem> {
526 self.post_item("/filesystems", req).await
527 }
528
529 pub async fn delete_filesystem(&self, id: &str) -> Result<()> {
530 let resp = self
531 .client
532 .delete(format!("{}/filesystems/{id}", self.endpoint))
533 .basic_auth(&self.api_key, Some(""))
534 .send()
535 .await?;
536 ensure_success(resp).await
537 }
538}
539
540async fn decode_list<T: for<'de> Deserialize<'de> + Default>(
541 resp: reqwest::Response,
542) -> Result<Vec<T>> {
543 if !resp.status().is_success() {
544 return Err(lambda_error(resp).await.into());
545 }
546 Ok(resp.json::<ApiList<T>>().await?.data)
547}
548
549async fn decode_item<T: for<'de> Deserialize<'de>>(resp: reqwest::Response) -> Result<T> {
550 if !resp.status().is_success() {
551 return Err(lambda_error(resp).await.into());
552 }
553 Ok(resp.json::<ApiItem<T>>().await?.data)
554}
555
556async fn ensure_success(resp: reqwest::Response) -> Result<()> {
557 if !resp.status().is_success() {
558 return Err(lambda_error(resp).await.into());
559 }
560 Ok(())
561}
562
563async fn lambda_error(resp: reqwest::Response) -> MoldError {
564 let status = resp.status();
565 let body = resp.text().await.unwrap_or_default();
566 let message = if status == StatusCode::UNAUTHORIZED {
567 "Lambda API authentication failed".to_string()
568 } else {
569 format!(
570 "Lambda API request failed with {status}: {}",
571 truncate(&body)
572 )
573 };
574 MoldError::Config(message)
575}
576
577fn truncate(s: &str) -> String {
578 const MAX: usize = 400;
579 if s.chars().count() <= MAX {
580 return s.to_string();
581 }
582 let mut out = s.chars().take(MAX).collect::<String>();
583 out.push('…');
584 out
585}
586
587pub fn decode_instance_types_body(body: &str) -> Result<Vec<InstanceType>> {
588 let response: InstanceTypesResponse = serde_json::from_str(body)?;
589 Ok(response
590 .data
591 .into_values()
592 .map(|offering| {
593 let mut instance_type = offering.instance_type;
594 if instance_type.specs.gpu_description.is_empty() {
595 instance_type.specs.gpu_description = instance_type.gpu_description.clone();
596 }
597 instance_type.regions_with_capacity_available =
598 offering.regions_with_capacity_available;
599 instance_type
600 })
601 .collect())
602}
603
604#[cfg(test)]
605mod tests {
606 use super::*;
607
608 #[test]
609 fn instance_types_decode_lambda_map_shape() {
610 let body = r#"{
611 "data": {
612 "gpu_1x_a10": {
613 "instance_type": {
614 "name": "gpu_1x_a10",
615 "description": "1x A10",
616 "gpu_description": "A10",
617 "price_cents_per_hour": 75,
618 "specs": {
619 "vcpus": 30,
620 "memory_gib": 200,
621 "storage_gib": 1400,
622 "gpus": 1
623 }
624 },
625 "regions_with_capacity_available": [
626 {"name": "us-west-1", "description": "California"}
627 ]
628 }
629 }
630 }"#;
631
632 let decoded = decode_instance_types_body(body).unwrap();
633 assert_eq!(decoded.len(), 1);
634 assert_eq!(decoded[0].name, "gpu_1x_a10");
635 assert_eq!(decoded[0].specs.gpu_description, "A10");
636 assert_eq!(
637 decoded[0].regions_with_capacity_available[0].name,
638 "us-west-1"
639 );
640 }
641
642 #[test]
643 fn lambda_settings_toml_roundtrip_and_defaults() {
644 let settings: LambdaSettings = toml::from_str("").unwrap();
645 assert_eq!(settings.endpoint.as_deref(), Some(DEFAULT_ENDPOINT));
646 assert_eq!(
647 settings.image_repository.as_deref(),
648 Some(DEFAULT_IMAGE_REPOSITORY)
649 );
650 assert_eq!(settings.filesystem_mount_path, "/data/mold");
651 assert_eq!(settings.confirm_hourly_usd, 5.0);
652 assert_eq!(settings.local_port, 7680);
653
654 let original = LambdaSettings {
655 api_key: Some("secret".into()),
656 endpoint: Some("http://localhost:9999".into()),
657 image_repository: Some("ghcr.io/example/mold".into()),
658 ssh_key_name: Some("mold-key".into()),
659 ssh_private_key_path: Some("~/.ssh/mold_lambda_ed25519".into()),
660 filesystem_prefix: Some("mold".into()),
661 filesystem_mount_path: "/mnt/mold".into(),
662 confirm_hourly_usd: 9.5,
663 local_port: 7777,
664 };
665 let encoded = toml::to_string(&original).unwrap();
666 let decoded: LambdaSettings = toml::from_str(&encoded).unwrap();
667 assert_eq!(decoded.api_key, original.api_key);
668 assert_eq!(decoded.filesystem_mount_path, "/mnt/mold");
669 assert_eq!(decoded.local_port, 7777);
670 }
671
672 #[test]
673 fn auth_prefers_lambda_api_key_env_over_config() {
674 let _guard = crate::test_support::ENV_LOCK.lock().unwrap();
675 std::env::set_var(API_KEY_ENV, "from-env");
676 let settings = LambdaSettings {
677 api_key: Some("from-config".into()),
678 ..Default::default()
679 };
680 assert_eq!(settings.resolved_api_key().as_deref(), Some("from-env"));
681 std::env::remove_var(API_KEY_ENV);
682 }
683
684 #[test]
685 fn image_tag_maps_gpu_generations() {
686 for name in ["NVIDIA A30", "NVIDIA A100-SXM4-80GB", "Ampere"] {
687 assert_eq!(
688 image_tag_for_gpu(name, "0.10.0").unwrap(),
689 "0.10.0-sm80",
690 "{name}"
691 );
692 }
693 for name in [
694 "NVIDIA A10",
695 "NVIDIA A40",
696 "NVIDIA RTX A6000",
697 "NVIDIA A16",
698 "NVIDIA A2",
699 "NVIDIA RTX 3090",
700 ] {
701 assert_eq!(
702 image_tag_for_gpu(name, "0.10.0").unwrap(),
703 "0.10.0-sm86",
704 "{name}"
705 );
706 }
707 assert_eq!(
708 image_tag_for_gpu("NVIDIA L40S", "0.10.0").unwrap(),
709 "0.10.0"
710 );
711 for name in ["NVIDIA H100 PCIe", "NVIDIA H200 SXM"] {
712 assert_eq!(
713 image_tag_for_gpu(name, "0.10.0").unwrap(),
714 "0.10.0-sm90",
715 "{name}"
716 );
717 }
718 for name in ["NVIDIA B200", "NVIDIA B300"] {
719 assert_eq!(
720 image_tag_for_gpu(name, "0.10.0").unwrap(),
721 "0.10.0-sm100",
722 "{name}"
723 );
724 }
725 for name in ["NVIDIA GH200", "NVIDIA GB200", "NVIDIA GB300"] {
726 assert!(image_tag_for_gpu(name, "0.10.0").is_err(), "{name}");
727 }
728 for name in ["NVIDIA RTX PRO 6000", "NVIDIA GeForce RTX 5090"] {
729 assert_eq!(
730 image_tag_for_gpu(name, "v0.10.0").unwrap(),
731 "0.10.0-sm120",
732 "{name}"
733 );
734 }
735 assert_eq!(
736 image_tag_for_gpu("NVIDIA Blackwell", "latest").unwrap(),
737 "latest",
738 "generic Blackwell must not guess between incompatible sm100 and sm120"
739 );
740 }
741
742 #[test]
743 fn grace_gpus_are_not_supported_by_published_linux_images() {
744 for name in [
745 "GH200 (96 GB)",
746 "gpu_1x_gh200",
747 "NVIDIA GB200",
748 "gpu_1x_gb300",
749 "Grace Hopper Superchip",
750 "Grace Blackwell Superchip",
751 ] {
752 assert!(gpu_uses_unsupported_linux_arm64(name), "{name}");
753 }
754 assert!(!gpu_uses_unsupported_linux_arm64("NVIDIA H100 PCIe"));
755 assert!(!gpu_uses_unsupported_linux_arm64("NVIDIA A100-SXM4-80GB"));
756 assert!(!gpu_uses_unsupported_linux_arm64("NVIDIA B200"));
757 assert!(!gpu_uses_unsupported_linux_arm64("NVIDIA B300"));
758 }
759
760 #[test]
761 fn availability_marks_every_grace_family_as_unsupported() {
762 for (instance_name, gpu_description) in [
763 ("gpu_1x_gh200", "GH200 (96 GB)"),
764 ("gpu_1x_gb200", "NVIDIA GB200"),
765 ("gpu_1x_gb300", "NVIDIA GB300"),
766 ] {
767 let ty = InstanceType {
768 name: instance_name.into(),
769 description: format!("1x {gpu_description}"),
770 gpu_description: gpu_description.into(),
771 price_cents_per_hour: 229,
772 specs: InstanceTypeSpecs {
773 gpus: 1,
774 gpu_description: gpu_description.into(),
775 memory_gib: 432,
776 storage_gib: 4096,
777 ..Default::default()
778 },
779 regions_with_capacity_available: vec![Region {
780 name: "us-east-3".into(),
781 description: "Austin".into(),
782 }],
783 };
784 let row = AvailabilityRow::from_instance_type(&ty, "ghcr.io/utensils/mold", "0.10.0");
785 assert_eq!(
786 row.image, "unsupported: linux/arm64 host",
787 "{gpu_description}"
788 );
789 }
790 }
791
792 #[test]
793 fn availability_row_uses_gpu_count_as_generation_slots() {
794 let ty = InstanceType {
795 name: "gpu_8x_h100".into(),
796 description: "8x H100".into(),
797 gpu_description: "NVIDIA H100".into(),
798 price_cents_per_hour: 15920,
799 specs: InstanceTypeSpecs {
800 gpus: 8,
801 gpu_description: "NVIDIA H100".into(),
802 memory_gib: 1800,
803 storage_gib: 200,
804 ..Default::default()
805 },
806 regions_with_capacity_available: vec![Region {
807 name: "us-east-1".into(),
808 description: "Virginia".into(),
809 }],
810 };
811 let row = AvailabilityRow::from_instance_type(&ty, "ghcr.io/utensils/mold", "0.10.0");
812 assert_eq!(row.generation_slots, 8);
813 assert_eq!(row.image, "ghcr.io/utensils/mold:0.10.0-sm90");
814 assert_eq!(row.price_per_hour_usd, 159.20);
815 }
816
817 #[test]
818 fn filesystem_name_defaults_to_prefix_region() {
819 let settings = LambdaSettings::default();
820 assert_eq!(filesystem_name(&settings, "us-west-1"), "mold-us-west-1");
821 let custom = LambdaSettings {
822 filesystem_prefix: Some("team-mold".into()),
823 ..Default::default()
824 };
825 assert_eq!(filesystem_name(&custom, "us-east-1"), "team-mold-us-east-1");
826 }
827
828 #[test]
829 fn launch_request_contains_expected_shape() {
830 let req = build_launch_request(LaunchRequestInput {
831 region_name: "us-west-1",
832 instance_type_name: "gpu_1x_a10",
833 ssh_key_name: "mold-laptop",
834 filesystem_name: "mold-us-west-1",
835 filesystem_id: None,
836 filesystem_mount_path: "/data/mold",
837 instance_name: "mold-us-west-1",
838 image_id: None,
839 user_data: "#cloud-config\n",
840 });
841 let json = serde_json::to_value(req).unwrap();
842 assert_eq!(json["region_name"], "us-west-1");
843 assert_eq!(json["ssh_key_names"], serde_json::json!(["mold-laptop"]));
844 assert_eq!(
845 json["file_system_names"],
846 serde_json::json!(["mold-us-west-1"])
847 );
848 assert_eq!(json["file_system_mounts"][0]["mount_point"], "/data/mold");
849 assert_eq!(json["tags"][0]["key"], "managed-by");
850 assert_eq!(json["tags"][0]["value"], "mold");
851 }
852
853 #[test]
854 fn launch_request_uses_filesystem_id_when_available() {
855 let req = build_launch_request(LaunchRequestInput {
856 region_name: "us-west-1",
857 instance_type_name: "gpu_1x_a10",
858 ssh_key_name: "mold-laptop",
859 filesystem_name: "mold-us-west-1",
860 filesystem_id: Some("fs-123"),
861 filesystem_mount_path: "/data/mold",
862 instance_name: "mold-us-west-1",
863 image_id: None,
864 user_data: "#cloud-config\n",
865 });
866 let mount = &req.file_system_mounts[0];
867 assert_eq!(mount.file_system_id.as_deref(), Some("fs-123"));
868 assert!(mount.file_system_name.is_none());
869 }
870
871 #[test]
872 fn create_filesystem_request_uses_lambda_region_field() {
873 let req = CreateFilesystemRequest {
874 name: "mold-us-east-1".into(),
875 region: "us-east-1".into(),
876 };
877 let json = serde_json::to_value(req).unwrap();
878 assert_eq!(json["name"], "mold-us-east-1");
879 assert_eq!(json["region"], "us-east-1");
880 assert!(json.get("region_name").is_none());
881 }
882
883 #[test]
884 fn cloud_init_keeps_service_private_and_omits_secrets_by_default() {
885 let rendered = render_cloud_init(&CloudInitOptions {
886 image: "ghcr.io/utensils/mold:0.10.0-sm90".into(),
887 mount_path: "/data/mold".into(),
888 env_file: "/etc/mold/lambda.env".into(),
889 });
890 assert!(rendered.contains("-p 127.0.0.1:7680:7680"));
891 assert!(rendered.contains("-v /data/mold:/workspace"));
892 assert!(rendered.contains("--gpus all"));
893 assert!(rendered.contains("ghcr.io/utensils/mold:0.10.0-sm90"));
894 assert!(!rendered.contains("HF_TOKEN"));
895 assert!(!rendered.contains("CIVITAI_TOKEN"));
896 }
897}