From e824630d0b6352cd3385ab34c303eec9cd45b423 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Tue, 16 Jan 2024 23:22:45 +0000
Subject: [PATCH] anyhow!() is expensive; don't create unless needed

---
 aicirt/src/main.rs           | 38 +++++++++++++++++++++++++++++-------
 aicirt/src/moduleinstance.rs |  2 +-
 aicirt/src/worker.rs         |  4 ++--
 rllm/src/iface.rs            | 22 +++++++++++----------
 4 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/aicirt/src/main.rs b/aicirt/src/main.rs
index 976802e2..5bb2c87f 100644
--- a/aicirt/src/main.rs
+++ b/aicirt/src/main.rs
@@ -398,7 +398,7 @@ impl ModuleRegistry {
 
     fn run_main(&self, req_id: &String) -> Result<()> {
         let req_instances = self.req_instances.lock().unwrap();
-        let inst = req_instances.get(req_id).ok_or(anyhow!("invalid req_id"))?;
+        let inst = req_instances.get(req_id).ok_or_else(|| anyhow!("invalid req_id"))?;
         inst.run_main()
     }
 
@@ -442,7 +442,7 @@ impl Stepper {
         Ok(self
             .instances
             .get(&id)
-            .ok_or(anyhow!("invalid id {}", id))?)
+            .ok_or_else(|| anyhow!("invalid id {}", id))?)
     }
 
     fn token_name(&self, idx: usize) -> String {
@@ -727,7 +727,12 @@ impl Stepper {
         let mut used_ids = Vec::new();
         let mut outputs = HashMap::default();
 
-        log::debug!("post_process0: {:?}", t0.elapsed());
+        log::warn!(
+            "post_process0: {:?} {} {}",
+            t0.elapsed(),
+            self.instances.len(),
+            req.ops.len()
+        );
 
         for op in req.ops.into_iter() {
             let instid = op.id;
@@ -768,7 +773,7 @@ impl Stepper {
             prev = Instant::now();
         }
 
-        log::debug!("post_process: {:?}", all_dur);
+        log::warn!("post_process: {:?}", all_dur);
 
         Ok(AiciPostProcessResp { seqs: outputs })
     }
@@ -949,6 +954,22 @@ impl CmdRespChannel {
     }
 }
 
+fn bench_hashmap() {
+    let mut h =  HashMap::<u64, u64>::default();
+    for x in 10..50 {
+        h.insert(x, x * x);
+    }
+    for _ in 0..10 {
+        let t0 = Instant::now();
+        let mut sum = 0;
+        for x in 10..50 {
+            let v = h.get(&x).unwrap();
+            sum += v;
+        }
+        println!("hashmap: {:?} {}", t0.elapsed(), sum);
+    }
+}
+
 fn bench_cmd_resp_busy(cli: &Cli, limits: &AiciLimits) {
     match fork_child::<u8, u8>(limits).unwrap() {
         worker::ForkResult::Parent { handle } => {
@@ -1105,9 +1126,12 @@ fn main() -> () {
     };
 
     if cli.bench {
-        bench_ipc(&limits);
-        bench_cmd_resp_busy(&cli, &limits);
-        bench_cmd_resp(&cli, &limits);
+        bench_hashmap();
+        if false {
+            bench_ipc(&limits);
+            bench_cmd_resp_busy(&cli, &limits);
+            bench_cmd_resp(&cli, &limits);
+        }
         return ();
     }
 
diff --git a/aicirt/src/moduleinstance.rs b/aicirt/src/moduleinstance.rs
index f452e6f0..b8905113 100644
--- a/aicirt/src/moduleinstance.rs
+++ b/aicirt/src/moduleinstance.rs
@@ -206,7 +206,7 @@ impl ModuleInstance {
         let instance = ctx.linker.instantiate(&mut store, &module)?;
         let memory = instance
             .get_memory(&mut store, "memory")
-            .ok_or(anyhow!("memory missing"))?;
+            .ok_or_else(|| anyhow!("memory missing"))?;
         store.data_mut().instance = Some(instance);
         store.data_mut().memory = Some(memory);
 
diff --git a/aicirt/src/worker.rs b/aicirt/src/worker.rs
index 3cba1f43..18b57e6c 100644
--- a/aicirt/src/worker.rs
+++ b/aicirt/src/worker.rs
@@ -875,11 +875,11 @@ impl WorkerForker {
                 Some(
                     req.prompt
                         .as_array()
-                        .ok_or(anyhow!("expecting string or int array as prompt"))?
+                        .ok_or_else(|| anyhow!("expecting string or int array as prompt"))?
                         .iter()
                         .map(|x| -> Result<u32> {
                             x.as_u64()
-                                .ok_or(anyhow!("expecting number as token"))?
+                                .ok_or_else(|| anyhow!("expecting number as token"))?
                                 .try_into()
                                 .map_err(|e: std::num::TryFromIntError| anyhow!(e))
                         })
diff --git a/rllm/src/iface.rs b/rllm/src/iface.rs
index af0fab0d..e8ebd3d8 100644
--- a/rllm/src/iface.rs
+++ b/rllm/src/iface.rs
@@ -93,10 +93,12 @@ impl CmdChannel {
             .as_object_mut()
             .unwrap()
             .remove("data")
-            .ok_or(anyhow::anyhow!(
-                "Bad response ({ctx}) - no 'data': {}",
-                limit_bytes(&bytes, 500)
-            ))?;
+            .ok_or_else(|| {
+                anyhow::anyhow!(
+                    "Bad response ({ctx}) - no 'data': {}",
+                    limit_bytes(&bytes, 500)
+                )
+            })?;
         let resp = serde_json::from_value(data).map_err(|e| {
             anyhow::anyhow!("Bad response ({ctx}): {e} {}", limit_bytes(&bytes, 500))
         })?;
@@ -305,14 +307,14 @@ impl AsyncCmdChannel {
 
         match resp["type"].as_str() {
             Some("ok") => {
-                let data = resp
-                    .as_object_mut()
-                    .unwrap()
-                    .remove("data")
-                    .ok_or(anyhow::anyhow!(
+                let data = resp.as_object_mut().unwrap().remove("data");
+                if data.is_none() {
+                    anyhow::bail!(
                         "Bad response ({op}) - no 'data': {}",
                         limit_bytes(&serde_json::to_vec(&resp)?, 500)
-                    ))?;
+                    );
+                }
+                let data = data.unwrap();
                 let data_copy = limit_bytes(&serde_json::to_vec(&data).unwrap(), 500);
                 let resp = serde_json::from_value(data)
                     .map_err(|e| anyhow::anyhow!("Bad response ({op}): {e} {}", data_copy))?;