Skip to content

Commit e9b9db8

Browse files
committed
feat(vmm): add OCI registry image discovery and pull support
Add ability for VMM to discover available guest images from an OCI registry and pull them on-demand through the web UI. Images are pulled in the background with status tracked server-side, surviving page refreshes. The UI auto-refreshes every 3s while the registry panel is open. - New `image_registry` config field (e.g., "cr.kvin.wang/dstack/guest-image") - New RPC: ListRegistryImages, PullRegistryImage - Registry module: list tags via Docker Registry HTTP API v2, pull and extract via `docker export` - Background pull with pulling state in App memory - UI: Image Registry button + dialog with pull/status per tag
1 parent d46ba28 commit e9b9db8

12 files changed

Lines changed: 564 additions & 16 deletions

File tree

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

vmm/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ fatfs.workspace = true
5555
fscommon.workspace = true
5656
or-panic.workspace = true
5757
url.workspace = true
58+
reqwest.workspace = true
5859

5960
[dev-dependencies]
6061
insta.workspace = true

vmm/rpc/proto/vmm_rpc.proto

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,11 @@ service Vmm {
350350
rpc SvStop(Id) returns (google.protobuf.Empty);
351351
// Remove a stopped supervisor process by ID.
352352
rpc SvRemove(Id) returns (google.protobuf.Empty);
353+
354+
// List images available in the configured OCI registry.
355+
rpc ListRegistryImages(google.protobuf.Empty) returns (RegistryImageListResponse);
356+
// Pull an image from the OCI registry to local storage.
357+
rpc PullRegistryImage(PullRegistryImageRequest) returns (google.protobuf.Empty);
353358
}
354359

355360
// DHCP lease event reported by the host DHCP server.
@@ -365,6 +370,27 @@ message SvListResponse {
365370
repeated SvProcessInfo processes = 1;
366371
}
367372

373+
// Available images discovered from the OCI registry.
374+
message RegistryImageListResponse {
375+
repeated RegistryImageInfo images = 1;
376+
}
377+
378+
// Metadata for an image tag in the OCI registry.
379+
message RegistryImageInfo {
380+
// Tag name (e.g., "0.5.8", "nvidia-0.5.8")
381+
string tag = 1;
382+
// Whether this image is already downloaded locally
383+
bool local = 2;
384+
// Whether this image is currently being pulled
385+
bool pulling = 3;
386+
}
387+
388+
// Request to pull an image from the OCI registry.
389+
message PullRegistryImageRequest {
390+
// Tag to pull (e.g., "0.5.8")
391+
string tag = 1;
392+
}
393+
368394
// Information about a single supervisor process.
369395
message SvProcessInfo {
370396
string id = 1;

vmm/src/app.rs

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ pub use qemu::{VmConfig, VmWorkDir};
3535
mod id_pool;
3636
mod image;
3737
mod qemu;
38+
pub(crate) mod registry;
3839

3940
#[derive(Deserialize, Serialize, Debug, Clone)]
4041
pub struct PortMapping {
@@ -124,6 +125,8 @@ pub struct App {
124125
pub supervisor: SupervisorClient,
125126
state: Arc<Mutex<AppState>>,
126127
forward_service: Arc<tokio::sync::Mutex<ForwardService>>,
128+
/// Tags currently being pulled from the image registry.
129+
pub(crate) pulling_tags: Arc<Mutex<std::collections::HashSet<String>>>,
127130
}
128131

129132
impl App {
@@ -152,6 +155,7 @@ impl App {
152155
})),
153156
config: Arc::new(config),
154157
forward_service: Arc::new(tokio::sync::Mutex::new(ForwardService::new())),
158+
pulling_tags: Arc::new(Mutex::new(std::collections::HashSet::new())),
155159
}
156160
}
157161

@@ -172,7 +176,7 @@ impl App {
172176
{
173177
bail!("Invalid image name");
174178
}
175-
let image_path = self.config.image_path.join(&manifest.image);
179+
let image_path = self.config.image.path.join(&manifest.image);
176180
let image = Image::load(&image_path).context("Failed to load image")?;
177181
let vm_id = manifest.id.clone();
178182
let app_compose = vm_work_dir
@@ -739,7 +743,7 @@ impl App {
739743
{
740744
bail!("Invalid image name");
741745
}
742-
let image_path = self.config.image_path.join(&manifest.image);
746+
let image_path = self.config.image.path.join(&manifest.image);
743747
let image = Image::load(&image_path).context("Failed to load image")?;
744748
let vm_id = manifest.id.clone();
745749
let already_running = cids_assigned.contains_key(&vm_id);
@@ -854,7 +858,7 @@ impl App {
854858
}
855859

856860
pub fn list_images(&self) -> Result<Vec<(String, ImageInfo)>> {
857-
let image_path = self.config.image_path.clone();
861+
let image_path = self.config.image.path.clone();
858862
let images = fs::read_dir(image_path).context("Failed to read image directory")?;
859863
Ok(images
860864
.flat_map(|entry| {
@@ -1115,7 +1119,7 @@ fn rotate_serial_log(work_dir: &VmWorkDir, max_bytes: u64) {
11151119
}
11161120

11171121
pub(crate) fn make_sys_config(cfg: &Config, manifest: &Manifest) -> Result<String> {
1118-
let image_path = cfg.image_path.join(&manifest.image);
1122+
let image_path = cfg.image.path.join(&manifest.image);
11191123
let image = Image::load(image_path).context("Failed to load image info")?;
11201124
let img_ver = image.info.version_tuple().unwrap_or((0, 0, 0));
11211125
let kms_urls = if manifest.kms_urls.is_empty() {

0 commit comments

Comments
 (0)