parallel execution using rayon

Author: msabvid

Cargo.toml

@@ -17,3 +17,6 @@ funty = "2.0.0"
 clap = { version = "4.5.1", features = ["derive"] }
 serde = "1.0.198"
 serde_json = "1.0.116"
+rayon = "1.10.0"
+features = "0.10.0"
+kdam = { version = "0.5.1", features = ["rayon"] }

src/aave/agents/borrow_agent.rs

@@ -61,7 +61,7 @@ impl Agent for BorrowAgent {
                     self.address,
                     self.pool_address,
                     self.supply_token_address,
-                    U256::from(10_u128.pow(20)),
+                    U256::from(10_u128.pow(18)),
                 );
                 self.has_supplied = true;
 
@@ -79,7 +79,7 @@ impl Agent for BorrowAgent {
 
                 let exp = U256::from(10u128).pow(self.borrow_token_decimals - U256::from(4u128));
                 // Agent borrows a random fraction of the available to borrow amount
-                let u = U256::from(rng.gen_range(9000..10000));
+                let u = U256::from(rng.gen_range(8000..9000));
                 let available_borrow = exp * available_borrow_base * u / borrow_asset_price;
 
                 if available_borrow > U256::ZERO {

src/aave/agents/mod.rs

@@ -25,7 +25,8 @@ pub struct AgentStates {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct AgentData {
+pub struct SimData {
+    pub seed: u64,
     pub borrow_agents: Vec<Vec<U256>>,
     pub liquidation_agents: Vec<Vec<UserData>>,
     pub uniswap_price_agent: Vec<(i128, i128)>,

src/aave/mod.rs

@@ -6,20 +6,18 @@ mod initialisation;
 mod protocol;
 pub mod types;
 
-use agents::AgentData;
+pub use agents::SimData;
 use verbs_rs::env::GasPriorityValidator;
 use verbs_rs::{agent::RecordedAgentSet, sim_runner};
 
 use self::initialisation::initialise_sim;
 
-use serde_json::json;
-
 pub fn aave_sim_from_fork(
     seed: u64,
     n_steps: usize,
     params: types::ForkedSimParameters,
     alchemy_key: String,
-) -> serde_json::Value {
+) -> SimData {
     println!("Initialising Simulation");
 
     let validator = GasPriorityValidator {};
@@ -30,33 +28,25 @@ pub fn aave_sim_from_fork(
     println!("Running");
 
     sim_runner::run(&mut env, &mut agent_sets, seed, n_steps);
-    let sim_data = AgentData {
+    SimData {
+        seed,
         borrow_agents: agent_sets.borrow_agents.take_records(),
         liquidation_agents: agent_sets.liquidation_agents.take_records(),
         uniswap_price_agent: agent_sets.uniswap_price_agent.take_records(),
         uniswap_noise_agents: agent_sets.uniswap_noise_agents.take_records(),
-    };
-    let sim_data = json!({
-        "seed": seed,
-        "sim_data": sim_data
-    });
-    sim_data
+    }
 }
 
-pub fn aave_sim(seed: u64, n_steps: usize, params: types::SimParameters) -> serde_json::Value {
+pub fn aave_sim(seed: u64, n_steps: usize, params: types::SimParameters) -> SimData {
     let validator = GasPriorityValidator {};
     let (mut env, mut agent_sets, _, _, _) = initialise_sim(params, validator);
 
     sim_runner::run(&mut env, &mut agent_sets, seed, n_steps);
-    let sim_data = AgentData {
+    SimData {
+        seed,
         borrow_agents: agent_sets.borrow_agents.take_records(),
         liquidation_agents: agent_sets.liquidation_agents.take_records(),
         uniswap_price_agent: agent_sets.uniswap_price_agent.take_records(),
         uniswap_noise_agents: agent_sets.uniswap_noise_agents.take_records(),
-    };
-    let sim_data = json!({
-        "seed": seed,
-        "sim_data": sim_data
-    });
-    sim_data
+    }
 }

src/main.rs

@@ -1,13 +1,18 @@
-use aave::aave_sim;
+use aave::SimData;
 use clap::Parser;
+use kdam::TqdmParallelIterator;
+use rayon::prelude::*;
+
 mod aave;
 
+use std::fs;
+
 #[derive(Parser, Debug)]
 #[command(about, long_about = None)]
 struct Args {
     /// Random seed
     #[arg(long)]
-    seed: u64,
+    n_seeds: u64,
     /// Number of simulation steps
     #[arg(long)]
     n_steps: usize,
@@ -22,10 +27,10 @@ struct Args {
 fn main() {
     let args = Args::parse();
 
-    let seed = args.seed;
+    let seeds = Vec::from_iter(10..10 + args.n_seeds);
     let n_steps = args.n_steps;
 
-    let results = match args.fork {
+    let results: Vec<SimData> = match args.fork {
         true => match args.key {
             Some(k) => {
                 let params = aave::types::ForkedSimParameters {
@@ -40,7 +45,13 @@ fn main() {
                     block_number: 18564279u64,
                 };
 
-                aave::aave_sim_from_fork(seed, n_steps, params, k)
+                seeds
+                    .par_iter()
+                    .map(|i| {
+                        let k = k.clone();
+                        aave::aave_sim_from_fork(*i, n_steps, params, k)
+                    })
+                    .collect()
             }
             None => panic!("Alchemy key argument required for forked simulation"),
         },
@@ -50,7 +61,7 @@ fn main() {
                 n_liquidators: 1,
                 prices_mu: 0f64,
                 prices_dt: 0.01f64,
-                prices_sigma: 0.4f64,
+                prices_sigma: 0.3f64,
                 borrow_activation_rate: 0.1f64,
                 token_a_initial_price: 100000000000i128,
                 token_b_initial_price: 100000000i128,
@@ -62,7 +73,13 @@ fn main() {
                 adversarial: false,
                 uniswap_fee: 500u32,
             };
-            aave_sim(seed, n_steps, params)
+            seeds
+                .par_iter()
+                .tqdm()
+                .map(|i| aave::aave_sim(*i, n_steps, params))
+                .collect()
         }
     };
+    let json = serde_json::to_string(&results).expect("Could not serialise to json string");
+    let _ = fs::write("sim_dat.txt", json);
 }