Skip to content

Commit 2188086

Browse files
LeechaelLeechael
authored
Merge pull request #288 from Dstack-TEE/imp-vmm-oneshot
feat(vmm): added one-shot
2 parents 2693bd0 + d6e0e9c commit 2188086

File tree

3 files changed

+512
-69
lines changed

3 files changed

+512
-69
lines changed

‎vmm/src/main.rs‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use std::{path::Path, time::Duration};
66

77
use anyhow::{anyhow, Context, Result};
88
use app::App;
9-
use clap::Parser;
9+
use clap::{Args as ClapArgs, Parser, Subcommand};
1010
use config::Config;
1111
use guest_api_service::GuestApiHandler;
1212
use host_api_service::HostApiHandler;
@@ -28,6 +28,7 @@ mod guest_api_service;
2828
mod host_api_service;
2929
mod main_routes;
3030
mod main_service;
31+
mod one_shot;
3132

3233
const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
3334
const GIT_REV: &str = git_version::git_version!(
@@ -46,6 +47,30 @@ struct Args {
4647
/// Path to the configuration file
4748
#[arg(short, long)]
4849
config: Option<String>,
50+
/// Subcommand to run
51+
#[command(subcommand)]
52+
command: Option<Command>,
53+
}
54+
55+
#[derive(Default, Subcommand)]
56+
enum Command {
57+
/// Start the VMM server (default mode)
58+
#[default]
59+
Serve,
60+
/// One-shot VM execution mode for debugging
61+
Run(RunArgs),
62+
}
63+
64+
#[derive(ClapArgs)]
65+
struct RunArgs {
66+
/// VM configuration file path
67+
vm_config: String,
68+
/// Working directory for one-shot mode (default: create in current directory)
69+
#[arg(long)]
70+
workdir: Option<String>,
71+
/// Dry run: only output QEMU command without executing
72+
#[arg(long)]
73+
dry_run: bool,
4974
}
5075

5176
async fn run_external_api(app: App, figment: Figment, api_auth: ApiToken) -> Result<()> {
@@ -134,6 +159,24 @@ async fn main() -> Result<()> {
134159
let args = Args::parse();
135160
let figment = config::load_config_figment(args.config.as_deref());
136161
let config = Config::extract_or_default(&figment)?.abs_path()?;
162+
163+
// Handle commands
164+
match args.command.unwrap_or_default() {
165+
Command::Run(run_args) => {
166+
// One-shot VM execution mode
167+
return one_shot::run_one_shot(
168+
&run_args.vm_config,
169+
config,
170+
run_args.workdir,
171+
run_args.dry_run,
172+
)
173+
.await;
174+
}
175+
Command::Serve => {
176+
// Default server mode - continue to main server logic
177+
}
178+
}
179+
137180
let api_auth = ApiToken::new(config.auth.tokens.clone(), config.auth.enabled);
138181
let supervisor = {
139182
let cfg = &config.supervisor;

‎vmm/src/main_service.rs‎

Lines changed: 86 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,21 @@ fn validate_label(label: &str) -> Result<()> {
6262
Ok(())
6363
}
6464

65-
fn resolve_gpus(gpu_cfg: &rpc::GpuConfig) -> Result<GpuConfig> {
65+
pub fn resolve_gpus_with_config(
66+
gpu_cfg: &rpc::GpuConfig,
67+
cvm_config: &crate::config::CvmConfig,
68+
) -> Result<GpuConfig> {
69+
if !cvm_config.gpu.enabled {
70+
bail!("GPU is not enabled");
71+
}
72+
let gpus = resolve_gpus(gpu_cfg)?;
73+
if !cvm_config.gpu.allow_attach_all && gpus.attach_mode.is_all() {
74+
bail!("Attaching all GPUs is not allowed");
75+
}
76+
Ok(gpus)
77+
}
78+
79+
pub fn resolve_gpus(gpu_cfg: &rpc::GpuConfig) -> Result<GpuConfig> {
6680
// Check the attach mode to determine how to handle GPUs
6781
match gpu_cfg.attach_mode.as_str() {
6882
"listed" => {
@@ -110,80 +124,84 @@ fn resolve_gpus(gpu_cfg: &rpc::GpuConfig) -> Result<GpuConfig> {
110124
}
111125
}
112126

127+
// Shared function to create manifest from VM configuration
128+
pub fn create_manifest_from_vm_config(
129+
request: VmConfiguration,
130+
cvm_config: &crate::config::CvmConfig,
131+
) -> Result<Manifest> {
132+
validate_label(&request.name)?;
133+
134+
let pm_cfg = &cvm_config.port_mapping;
135+
if !(request.ports.is_empty() || pm_cfg.enabled) {
136+
bail!("Port mapping is disabled");
137+
}
138+
let port_map = request
139+
.ports
140+
.iter()
141+
.map(|p| {
142+
let from = p.host_port.try_into().context("Invalid host port")?;
143+
let to = p.vm_port.try_into().context("Invalid vm port")?;
144+
if !pm_cfg.is_allowed(&p.protocol, from) {
145+
bail!("Port mapping is not allowed for {}:{}", p.protocol, from);
146+
}
147+
let protocol = p.protocol.parse().context("Invalid protocol")?;
148+
let address = if !p.host_address.is_empty() {
149+
p.host_address.parse().context("Invalid host address")?
150+
} else {
151+
pm_cfg.address
152+
};
153+
Ok(PortMapping {
154+
address,
155+
protocol,
156+
from,
157+
to,
158+
})
159+
})
160+
.collect::<Result<Vec<_>>>()?;
161+
162+
let app_id = match &request.app_id {
163+
Some(id) => id.strip_prefix("0x").unwrap_or(id).to_lowercase(),
164+
None => app_id_of(&request.compose_file),
165+
};
166+
let id = uuid::Uuid::new_v4().to_string();
167+
let now = SystemTime::now()
168+
.duration_since(UNIX_EPOCH)
169+
.unwrap_or_default()
170+
.as_millis() as u64;
171+
let gpus = match &request.gpus {
172+
Some(gpus) => resolve_gpus_with_config(gpus, cvm_config)?,
173+
None => GpuConfig::default(),
174+
};
175+
176+
Ok(Manifest::builder()
177+
.id(id)
178+
.name(request.name.clone())
179+
.app_id(app_id)
180+
.image(request.image.clone())
181+
.vcpu(request.vcpu)
182+
.memory(request.memory)
183+
.disk_size(request.disk_size)
184+
.port_map(port_map)
185+
.created_at_ms(now)
186+
.hugepages(request.hugepages)
187+
.pin_numa(request.pin_numa)
188+
.gpus(gpus)
189+
.kms_urls(request.kms_urls.clone())
190+
.gateway_urls(request.gateway_urls.clone())
191+
.build())
192+
}
193+
113194
impl RpcHandler {
114195
fn resolve_gpus(&self, gpu_cfg: &rpc::GpuConfig) -> Result<GpuConfig> {
115-
let gpus = resolve_gpus(gpu_cfg)?;
116-
if !self.app.config.cvm.gpu.enabled {
117-
bail!("GPU is not enabled");
118-
}
119-
if !self.app.config.cvm.gpu.allow_attach_all && gpus.attach_mode.is_all() {
120-
bail!("Attaching all GPUs is not allowed");
121-
}
122-
Ok(gpus)
196+
resolve_gpus_with_config(gpu_cfg, &self.app.config.cvm)
123197
}
124198
}
125199

126200
impl VmmRpc for RpcHandler {
127201
async fn create_vm(self, request: VmConfiguration) -> Result<Id> {
128-
validate_label(&request.name)?;
129-
130-
let pm_cfg = &self.app.config.cvm.port_mapping;
131-
if !(request.ports.is_empty() || pm_cfg.enabled) {
132-
bail!("Port mapping is disabled");
133-
}
134-
let port_map = request
135-
.ports
136-
.iter()
137-
.map(|p| {
138-
let from = p.host_port.try_into().context("Invalid host port")?;
139-
let to = p.vm_port.try_into().context("Invalid vm port")?;
140-
if !pm_cfg.is_allowed(&p.protocol, from) {
141-
bail!("Port mapping is not allowed for {}:{}", p.protocol, from);
142-
}
143-
let protocol = p.protocol.parse().context("Invalid protocol")?;
144-
let address = if !p.host_address.is_empty() {
145-
p.host_address.parse().context("Invalid host address")?
146-
} else {
147-
pm_cfg.address
148-
};
149-
Ok(PortMapping {
150-
address,
151-
protocol,
152-
from,
153-
to,
154-
})
155-
})
156-
.collect::<Result<Vec<_>>>()?;
157-
158-
let app_id = match &request.app_id {
159-
Some(id) => id.strip_prefix("0x").unwrap_or(id).to_lowercase(),
160-
None => app_id_of(&request.compose_file),
161-
};
162-
let id = uuid::Uuid::new_v4().to_string();
163-
let now = SystemTime::now()
164-
.duration_since(UNIX_EPOCH)
165-
.unwrap_or_default()
166-
.as_millis() as u64;
167-
let gpus = match &request.gpus {
168-
Some(gpus) => self.resolve_gpus(gpus)?,
169-
None => GpuConfig::default(),
170-
};
171-
let manifest = Manifest::builder()
172-
.id(id.clone())
173-
.name(request.name.clone())
174-
.app_id(app_id.clone())
175-
.image(request.image.clone())
176-
.vcpu(request.vcpu)
177-
.memory(request.memory)
178-
.disk_size(request.disk_size)
179-
.port_map(port_map)
180-
.created_at_ms(now)
181-
.hugepages(request.hugepages)
182-
.pin_numa(request.pin_numa)
183-
.gpus(gpus)
184-
.kms_urls(request.kms_urls.clone())
185-
.gateway_urls(request.gateway_urls.clone())
186-
.build();
202+
let manifest = create_manifest_from_vm_config(request.clone(), &self.app.config.cvm)?;
203+
let id = manifest.id.clone();
204+
let app_id = manifest.app_id.clone();
187205
let vm_work_dir = self.app.work_dir(&id);
188206
vm_work_dir
189207
.put_manifest(&manifest)

0 commit comments

Comments
 (0)