Introduce MachineLocator for Validator

.github/workflows/_go-tests.yml

@@ -167,7 +167,7 @@ jobs:
 
             # 1. Start server in background
             echo "Starting server ($MODE_NAME)..."
-            target/bin/validator --module-root-path ./target/machines/latest/module-root.txt $SERVER_ARGS &
+            target/bin/validator $SERVER_ARGS &
             SERVER_PID=$!
 
             # 2. Wait for server to respond (up to 5 seconds)

Cargo.toml

@@ -72,6 +72,7 @@ paste = { version = "1.0.15" }
 rand = { version = "0.8.4", default-features = false }
 rand_pcg = { version = "0.3.1", default-features = false }
 rayon = { version = "1.5.1" }
+reqwest = { version = "0.13.1" }
 ruint2 = { version = "1.9.0" }
 rustc-demangle = { version = "0.1.21" }
 serde = { version = "1.0.130" }

changelog/bragaigor-nit-4346.md

@@ -0,0 +1,2 @@
+### Internal
+- Introduce MachineLocator for Validator

crates/validator/Cargo.toml

@@ -29,4 +29,5 @@ tracing-subscriber = { workspace = true, features = [
 validation = { workspace = true }
 
 [dev-dependencies]
-reqwest = "0.13.1"
+rand = { workspace = true, default-features = true }
+reqwest = { workspace = true }

crates/validator/src/config.rs

@@ -8,42 +8,36 @@
 //! into strongly-typed configuration objects used throughout the application.
 
 use anyhow::Result;
-use clap::{Args, Parser, ValueEnum};
-use std::fs::read_to_string;
+use clap::{Parser, ValueEnum};
 use std::net::SocketAddr;
 use std::path::PathBuf;
 
-use crate::engine::config::{JitManagerConfig, ModuleRoot};
 use crate::engine::machine::JitProcessManager;
+use crate::engine::machine_locator::MachineLocator;
 
 #[derive(Debug)]
 pub struct ServerState {
     pub mode: InputMode,
-    pub binary: PathBuf,
-    pub module_root: ModuleRoot,
-    /// Jit manager responsible for computing next GlobalState. Not wrapped
-    /// in Arc<> since the caller of ServerState is wrapped in Arc<>. This field
-    /// is optional because it's only available in continuous InputMode
+    /// Machine locator is responsible for locating replay.wasm binary and building
+    /// a map of module roots to their respective location + binary
+    pub locator: MachineLocator,
+    /// Jit manager is responsible for computing next GlobalState. Not wrapped
+    /// in Arc<> since the caller of ServerState is wrapped in Arc<>.
     pub jit_manager: JitProcessManager,
     pub available_workers: usize,
 }
 
 impl ServerState {
     pub fn new(config: &ServerConfig, available_workers: usize) -> Result<Self> {
-        // TODO: Load multiple module roots via MachineLocator (NIT-4346)
-        let module_root = config.get_module_root()?;
-        let manager_config = JitManagerConfig {
-            prover_bin_path: config.binary.clone(),
-            ..Default::default()
-        };
+        let locator = MachineLocator::new(&config.root_path)?;
+
         let jit_manager = match config.mode {
-            InputMode::Continuous => JitProcessManager::new(&manager_config, module_root)?,
-            InputMode::Native => JitProcessManager::new_empty(&manager_config),
+            InputMode::Continuous => JitProcessManager::new(&locator)?,
+            InputMode::Native => JitProcessManager::new_empty(),
         };
         Ok(ServerState {
             mode: config.mode,
-            binary: config.binary.clone(),
-            module_root,
+            locator,
             jit_manager,
             available_workers,
         })
@@ -72,10 +66,6 @@ pub struct ServerConfig {
     #[clap(long, default_value = "0.0.0.0:4141")]
     pub address: SocketAddr,
 
-    /// Path to the `replay.wasm` binary.
-    #[clap(long, default_value = "./target/machines/latest/replay.wasm")]
-    pub binary: PathBuf,
-
     /// Logging format configuration.
     #[clap(long, value_enum, default_value_t = LoggingFormat::Text)]
     pub logging_format: LoggingFormat,
@@ -84,46 +74,15 @@ pub struct ServerConfig {
     #[clap(long, value_enum, default_value_t = InputMode::Native)]
     pub mode: InputMode,
 
-    #[clap(flatten)]
-    module_root_config: ModuleRootConfig,
-
     #[clap(long)]
     workers: Option<usize>,
-}
-
-#[derive(Clone, Debug, Args)]
-#[group(required = true, multiple = false)]
-struct ModuleRootConfig {
-    /// Supported module root.
-    #[clap(long)]
-    module_root: Option<ModuleRoot>,
 
-    /// Path to the file containing the module root.
+    /// Root path to where 0x1234.../replay.wasm machines are located
     #[clap(long)]
-    module_root_path: Option<PathBuf>,
+    pub root_path: Option<PathBuf>,
 }
 
 impl ServerConfig {
-    pub fn get_module_root(&self) -> anyhow::Result<ModuleRoot> {
-        match (
-            self.module_root_config.module_root,
-            &self.module_root_config.module_root_path,
-        ) {
-            (Some(root), None) => Ok(root),
-            (None, Some(ref path)) => {
-                let content = read_to_string(path)?;
-                let root = content
-                    .trim()
-                    .parse::<ModuleRoot>()
-                    .map_err(|e| anyhow::anyhow!(e))?;
-                Ok(root)
-            }
-            _ => Err(anyhow::anyhow!(
-                "Either module_root or module_root_path must be specified"
-            )),
-        }
-    }
-
     pub fn get_workers(&self) -> Result<usize> {
         if let Some(workers) = self.workers {
             Ok(workers)
@@ -153,69 +112,9 @@ mod tests {
         ServerConfig::command().debug_assert()
     }
 
-    #[test]
-    fn module_root_parsing() {
-        assert!(
-            ServerConfig::try_parse_from([
-                "server",
-                "--module-root",
-                "0x0000000000000000000000000000000000000000000000000000000000000000"
-            ])
-            .is_ok(),
-            "Valid module root should parse correctly"
-        );
-
-        assert!(
-            ServerConfig::try_parse_from([
-                "server",
-                "--module-root",
-                "0000000000000000000000000000000000000000000000000000000000000000"
-            ])
-            .is_ok(),
-            "Valid module root (without 0x prefix) should parse correctly"
-        );
-
-        assert!(
-            ServerConfig::try_parse_from(["server", "--module-root", "0xinvalidhex"]).is_err(),
-            "Invalid module root should fail to parse"
-        );
-
-        assert!(
-            ServerConfig::try_parse_from([
-                "server",
-                "--module-root-path",
-                "/some/path/to/module/root.txt"
-            ])
-            .is_ok(),
-            "Valid module root path should parse correctly"
-        );
-
-        assert!(
-            ServerConfig::try_parse_from([
-                "server",
-                "--module-root",
-                "0x0000000000000000000000000000000000000000000000000000000000000000",
-                "--module-root-path",
-                "/some/path/to/module/root.txt"
-            ])
-            .is_err(),
-            "Specifying both module root and module root path should fail"
-        );
-
-        assert!(
-            ServerConfig::try_parse_from(["server"]).is_err(),
-            "Not specifying either module root or module root path should fail"
-        );
-    }
-
     #[test]
     fn capacity_parsing() {
-        let server_config = ServerConfig::try_parse_from([
-            "server",
-            "--module-root",
-            "0x0000000000000000000000000000000000000000000000000000000000000000",
-        ])
-        .unwrap();
+        let server_config = ServerConfig::try_parse_from(["server"]).unwrap();
 
         assert!(server_config.workers.is_none());
         let workers = server_config.get_workers().unwrap();

crates/validator/src/engine/execution.rs

@@ -17,38 +17,52 @@
 //!    binary version targeting.
 
 use axum::Json;
-use validation::{local_target, BatchInfo, GoGlobalState, ValidationInput};
+use tracing::info;
+use validation::{local_target, BatchInfo, GoGlobalState};
 
 use crate::{
-    config::ServerState, engine::config::DEFAULT_JIT_CRANELIFT,
+    config::ServerState,
+    engine::{replay_binary, DEFAULT_JIT_CRANELIFT},
     spawner_endpoints::ValidationRequest,
 };
 
 pub async fn validate_native(
     server_state: &ServerState,
-    request: ValidationInput,
+    request: ValidationRequest,
 ) -> Result<Json<GoGlobalState>, String> {
-    let delayed_inbox = match request.has_delayed_msg {
+    let delayed_inbox = match request.validation_input.has_delayed_msg {
         true => vec![BatchInfo {
-            number: request.delayed_msg_nr,
-            data: request.delayed_msg,
+            number: request.validation_input.delayed_msg_nr,
+            data: request.validation_input.delayed_msg,
         }],
         false => vec![],
     };
 
+    let binary_path = if let Some(module_root) = request.module_root {
+        server_state.locator.get_machine_path(module_root)?
+    } else {
+        server_state
+            .locator
+            .latest_wasm_module_root()
+            .path
+            .to_path_buf()
+    };
+    let binary = replay_binary(binary_path);
+    info!("validate native serving request with module root at {binary:?}");
+
     let opts = jit::Opts {
         validator: jit::ValidatorOpts {
-            binary: server_state.binary.clone(), // wasm binary
+            binary,
             cranelift: DEFAULT_JIT_CRANELIFT,
             debug: false, // JIT's debug messages are using printlns, which would clutter the server logs
             require_success: false, // Relevant for JIT binary only.
         },
         input_mode: jit::InputMode::Native(jit::NativeInput {
-            old_state: request.start_state.into(),
-            inbox: request.batch_info,
+            old_state: request.validation_input.start_state.into(),
+            inbox: request.validation_input.batch_info,
             delayed_inbox,
-            preimages: request.preimages,
-            programs: request.user_wasms[local_target()].clone(),
+            preimages: request.validation_input.preimages,
+            programs: request.validation_input.user_wasms[local_target()].clone(),
         }),
     };
 
@@ -64,7 +78,11 @@ pub async fn validate_continuous(
     server_state: &ServerState,
     request: ValidationRequest,
 ) -> Result<Json<GoGlobalState>, String> {
-    let module_root = request.module_root.unwrap_or(server_state.module_root);
+    let module_root = request
+        .module_root
+        .unwrap_or_else(|| server_state.locator.latest_wasm_module_root().module_root);
+
+    info!("validate continuous serving request with module_root 0x{module_root}");
 
     let new_state = server_state
         .jit_manager