Skip to content

Commit bb8371c

Browse files
fffonionfffonion
committed
pd-edge: jit on/off switch
1 parent 9ca8fe3 commit bb8371c

6 files changed

Lines changed: 56 additions & 2 deletions

File tree

‎pd-edge/src/bin/pd-edge-console.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
5252
let vm_execution = VmExecutionConfig {
5353
interrupt: cli.vm_interrupt_config()?,
5454
execution_mode: VmExecutionMode::Async,
55+
jit_enabled: true,
5556
};
5657
let store_limits = cli.runtime_store_limits();
5758
if cli.disable_metrics {
@@ -1023,6 +1024,7 @@ mod tests {
10231024
let state = SharedState::new(1024 * 1024).with_vm_execution_config(VmExecutionConfig {
10241025
interrupt: VmInterruptConfig::None,
10251026
execution_mode: VmExecutionMode::Async,
1027+
jit_enabled: true,
10261028
});
10271029
load_program_from_path(&state, &program_path)
10281030
.await
@@ -1054,6 +1056,7 @@ mod tests {
10541056
let state = SharedState::new(1024 * 1024).with_vm_execution_config(VmExecutionConfig {
10551057
interrupt: VmInterruptConfig::None,
10561058
execution_mode: VmExecutionMode::Async,
1059+
jit_enabled: true,
10571060
});
10581061
load_program_from_path(&state, &program_path)
10591062
.await

‎pd-edge/src/bin/pd-edge-http-proxy.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
7171
let vm_execution = VmExecutionConfig {
7272
interrupt: cli.vm_interrupt_config()?,
7373
execution_mode: cli.vm_execution_mode.unwrap_or_default(),
74+
jit_enabled: cli.vm_jit,
7475
};
7576
let store_limits = cli.runtime_store_limits();
7677
if cli.disable_metrics {
@@ -80,6 +81,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
8081
.with_metrics_collection_enabled(!cli.disable_metrics)
8182
.with_vm_execution_config(vm_execution);
8283
info!("vm execution mode={}", vm_execution.execution_mode.as_str());
84+
info!("vm jit enabled={}", vm_execution.jit_enabled);
8385
match vm_execution.interrupt {
8486
VmInterruptConfig::None => {}
8587
VmInterruptConfig::Fuel {
@@ -212,6 +214,7 @@ struct CliArgs {
212214
vm_epoch_deadline: Option<u64>,
213215
vm_epoch_check_interval: Option<u32>,
214216
vm_execution_mode: Option<VmExecutionMode>,
217+
vm_jit: bool,
215218
control_plane_url: Option<String>,
216219
edge_id: Option<String>,
217220
edge_name: Option<String>,
@@ -378,6 +381,9 @@ where
378381
let value = next_arg_value("--vm-execution-mode", &mut args)?;
379382
cli.vm_execution_mode = Some(parse_vm_execution_mode(&value)?);
380383
}
384+
"--vm-jit" => {
385+
cli.vm_jit = true;
386+
}
381387
"--edge-id" => {
382388
cli.edge_id = Some(next_arg_value("--edge-id", &mut args)?);
383389
}
@@ -471,6 +477,7 @@ fn print_cli_help() {
471477
" --vm-epoch-deadline <TICKS> Enable cooperative VM epoch slices per request (1 tick = 1ms wall clock)\n",
472478
" --vm-epoch-check-interval <OPS> Epoch check interval when --vm-epoch-deadline is enabled (default: 1)\n",
473479
" --vm-execution-mode <MODE> VM execution mode: async|threading (default: async)\n",
480+
" --vm-jit Enable VM JIT/trace execution (default: off)\n",
474481
" --control-plane-url <URL> Enable active control-plane RPC client\n",
475482
" --edge-id <UUID> Explicit edge UUID used by active control-plane client\n",
476483
" --edge-name <NAME> Friendly edge name (default: hostname)\n",
@@ -667,6 +674,7 @@ mod tests {
667674
vm_epoch_deadline: None,
668675
vm_epoch_check_interval: None,
669676
vm_execution_mode: None,
677+
vm_jit: false,
670678
control_plane_url: Some("http://127.0.0.1:9100".to_string()),
671679
edge_id: Some("123e4567-e89b-12d3-a456-426614174000".to_string()),
672680
edge_name: Some("test-edge".to_string()),
@@ -781,6 +789,16 @@ mod tests {
781789
assert_eq!(cli.vm_execution_mode, Some(VmExecutionMode::Threading));
782790
}
783791

792+
#[test]
793+
fn parse_cli_args_from_parses_vm_jit_flag() {
794+
let action = parse_cli_args_from(["--vm-jit".to_string()]).expect("parse should succeed");
795+
796+
let CliAction::Run(cli) = action else {
797+
panic!("expected run action");
798+
};
799+
assert!(cli.vm_jit);
800+
}
801+
784802
#[test]
785803
fn runtime_store_limits_uses_defaults_and_overrides() {
786804
let cli = CliArgs {

‎pd-edge/src/bin/pd-edge-transport-proxy.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
6969
let vm_execution = VmExecutionConfig {
7070
interrupt: cli.vm_interrupt_config()?,
7171
execution_mode: cli.vm_execution_mode.unwrap_or_default(),
72+
jit_enabled: cli.vm_jit,
7273
};
7374
let store_limits = cli.runtime_store_limits();
7475
if cli.disable_metrics {
@@ -78,6 +79,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
7879
.with_metrics_collection_enabled(!cli.disable_metrics)
7980
.with_vm_execution_config(vm_execution);
8081
info!("vm execution mode={}", vm_execution.execution_mode.as_str());
82+
info!("vm jit enabled={}", vm_execution.jit_enabled);
8183
match vm_execution.interrupt {
8284
VmInterruptConfig::None => {}
8385
VmInterruptConfig::Fuel {
@@ -154,6 +156,7 @@ struct CliArgs {
154156
vm_epoch_deadline: Option<u64>,
155157
vm_epoch_check_interval: Option<u32>,
156158
vm_execution_mode: Option<VmExecutionMode>,
159+
vm_jit: bool,
157160
control_plane_url: Option<String>,
158161
edge_id: Option<String>,
159162
edge_name: Option<String>,
@@ -297,6 +300,9 @@ where
297300
let value = next_arg_value("--vm-execution-mode", &mut args)?;
298301
cli.vm_execution_mode = Some(parse_vm_execution_mode(&value)?);
299302
}
303+
"--vm-jit" => {
304+
cli.vm_jit = true;
305+
}
300306
"--edge-id" => {
301307
cli.edge_id = Some(next_arg_value("--edge-id", &mut args)?);
302308
}
@@ -387,6 +393,7 @@ fn print_cli_help() {
387393
" --vm-epoch-deadline <TICKS> Enable cooperative VM epoch slices per connection (1 tick = 1ms wall clock)\n",
388394
" --vm-epoch-check-interval <OPS> Epoch check interval when --vm-epoch-deadline is enabled (default: 1)\n",
389395
" --vm-execution-mode <MODE> VM execution mode: async|threading (default: async)\n",
396+
" --vm-jit Enable VM JIT/trace execution (default: off)\n",
390397
" --control-plane-url <URL> Enable active control-plane RPC client\n",
391398
" --edge-id <UUID> Explicit edge UUID used by active control-plane client\n",
392399
" --edge-name <NAME> Friendly edge name (default: hostname)\n",
@@ -568,6 +575,7 @@ mod tests {
568575
vm_epoch_deadline: None,
569576
vm_epoch_check_interval: None,
570577
vm_execution_mode: None,
578+
vm_jit: false,
571579
control_plane_url: Some("http://127.0.0.1:9100".to_string()),
572580
edge_id: Some("123e4567-e89b-12d3-a456-426614174000".to_string()),
573581
edge_name: Some("test-edge".to_string()),
@@ -593,6 +601,16 @@ mod tests {
593601
assert!(cli.disable_metrics);
594602
}
595603

604+
#[test]
605+
fn parse_cli_args_from_parses_vm_jit_flag() {
606+
let action = parse_cli_args_from(["--vm-jit".to_string()]).expect("parse should succeed");
607+
608+
let CliAction::Run(cli) = action else {
609+
panic!("expected run action");
610+
};
611+
assert!(cli.vm_jit);
612+
}
613+
596614
#[test]
597615
fn runtime_store_limits_uses_defaults_and_overrides() {
598616
let cli = CliArgs {

‎pd-edge/src/runtime.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,13 +107,15 @@ impl VmInterruptConfig {
107107
pub struct VmExecutionConfig {
108108
pub interrupt: VmInterruptConfig,
109109
pub execution_mode: VmExecutionMode,
110+
pub jit_enabled: bool,
110111
}
111112

112113
impl Default for VmExecutionConfig {
113114
fn default() -> Self {
114115
Self {
115116
interrupt: VmInterruptConfig::None,
116117
execution_mode: VmExecutionMode::default(),
118+
jit_enabled: true,
117119
}
118120
}
119121
}

‎pd-edge/src/runtime/http_plane/proxy_path.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1192,6 +1192,7 @@ mod tests {
11921192
let state = SharedState::new(1024 * 1024).with_vm_execution_config(VmExecutionConfig {
11931193
interrupt: crate::runtime::VmInterruptConfig::None,
11941194
execution_mode: VmExecutionMode::Threading,
1195+
jit_enabled: true,
11951196
});
11961197
let source = r#"
11971198
use http;
@@ -1339,6 +1340,7 @@ mod tests {
13391340
let state = SharedState::new(1024 * 1024).with_vm_execution_config(VmExecutionConfig {
13401341
interrupt: crate::runtime::VmInterruptConfig::None,
13411342
execution_mode: VmExecutionMode::Threading,
1343+
jit_enabled: true,
13421344
});
13431345
let source = format!(
13441346
r#"
@@ -1458,6 +1460,7 @@ mod tests {
14581460
let state = SharedState::new(1024 * 1024).with_vm_execution_config(VmExecutionConfig {
14591461
interrupt: crate::runtime::VmInterruptConfig::None,
14601462
execution_mode: VmExecutionMode::Threading,
1463+
jit_enabled: true,
14611464
});
14621465
let source = format!(
14631466
r#"

‎pd-edge/src/runtime/vm_runner.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ pub async fn execute_vm_with_context(
103103
vm_context,
104104
async_ops,
105105
!debug.attach_debugger && matches!(vm_execution.interrupt, VmInterruptConfig::None),
106+
vm_execution.jit_enabled,
106107
);
107108
let execution_mode = if debug.force_threading {
108109
VmExecutionMode::Threading
@@ -303,8 +304,11 @@ fn new_vm_runner_store(
303304
vm_context: SharedProxyVmContext,
304305
async_ops: SharedVmAsyncOps,
305306
prefer_aot: bool,
307+
jit_enabled: bool,
306308
) -> VmRunnerStore {
307-
let prefer_aot = prefer_aot && std::env::var_os("PD_EDGE_DISABLE_NO_INTERRUPT_AOT").is_none();
309+
let prefer_aot = prefer_aot
310+
&& jit_enabled
311+
&& std::env::var_os("PD_EDGE_DISABLE_NO_INTERRUPT_AOT").is_none();
308312
let mut vm = if prefer_aot {
309313
if let Some(bundle) = program.no_interrupt_aot_bundle.as_ref() {
310314
match Vm::from_aot_bundle_bytes(bundle.as_ref().as_slice()) {
@@ -321,6 +325,11 @@ fn new_vm_runner_store(
321325
} else {
322326
Vm::new_shared(program.program.clone())
323327
};
328+
if !jit_enabled {
329+
let mut jit_config = *vm.jit_config();
330+
jit_config.enabled = false;
331+
vm.set_jit_config(jit_config);
332+
}
324333
vm.set_async_bridge(Box::new(VmAsyncOpBridge::new(async_ops.clone())));
325334
Store::new(vm, VmRunnerStoreData::new(vm_context, async_ops))
326335
}
@@ -530,7 +539,7 @@ mod tests {
530539
Arc::new(RateLimiterStore::new()),
531540
));
532541
let async_ops = new_shared_vm_async_ops();
533-
let store = new_vm_runner_store(&loaded_program, context, async_ops, false);
542+
let store = new_vm_runner_store(&loaded_program, context, async_ops, false, true);
534543
let debug = VmDebugInvocation {
535544
attach_debugger: false,
536545
force_threading: false,
@@ -542,6 +551,7 @@ mod tests {
542551
check_interval: 1,
543552
},
544553
execution_mode: VmExecutionMode::Threading,
554+
jit_enabled: true,
545555
};
546556

547557
let result = tokio::task::spawn_blocking(move || {

0 commit comments

Comments
 (0)