Merge pull request #105 from imbue-ai/danver/cost-estimate

Implement Cost Estimate function

Author: DanverImbue

.beads/issues.jsonl

@@ -1 +1 @@
-code-102
+code-109

.beads/last-touched

@@ -436,8 +436,17 @@ def run(command: str):
     multiple=True,
     help="Environment variable (format: KEY=VALUE)",
 )
+@click.option(
+    "--cpu",
+    type=float,
+    default=None,
+    help="CPU cores per sandbox",
+)
 def create_from_image(
-    image_id: str, copy_dirs: tuple[str, ...] = (), env_vars: tuple[str, ...] = ()
+    image_id: str,
+    copy_dirs: tuple[str, ...] = (),
+    env_vars: tuple[str, ...] = (),
+    cpu: float | None = None,
 ):
     """Create sandbox using existing image_id.
 
@@ -483,13 +492,16 @@ def create_from_image(
 
     logger.debug("[%.2fs] Creating sandbox...", time.time() - t0)
     try:
-        sandbox = modal.Sandbox.create(
+        create_kwargs = dict(
             app=app,
             image=image,
             workdir="/app",
             timeout=3600,
             secrets=secrets,
         )
+        if cpu is not None:
+            create_kwargs["cpu"] = cpu
+        sandbox = modal.Sandbox.create(**create_kwargs)
     except Exception as e:
         logger.error("Failed to create sandbox with image %s: %s", image_id, e)
         logger.error(

scripts/modal_sandbox.py

@@ -178,6 +178,10 @@ pub struct ModalProviderConfig {
     /// making sandbox creation faster.
     #[serde(default)]
     pub copy_dirs: Vec<String>,
+
+    /// CPU cores per sandbox (default: 0.125).
+    #[serde(default = "default_modal_cpu_cores")]
+    pub cpu_cores: f64,
 }
 
 /// Configuration for custom remote execution provider.
@@ -272,12 +276,24 @@ pub struct DefaultProviderConfig {
     /// These are merged with (and override) the current environment.
     #[serde(default)]
     pub env: HashMap<String, String>,
+
+    /// CPU cores per sandbox (default: 1.0).
+    #[serde(default = "default_cpu_cores")]
+    pub cpu_cores: f64,
 }
 
 fn default_remote_timeout() -> u64 {
     3600 // 1 hour
 }
 
+fn default_cpu_cores() -> f64 {
+    1.0
+}
+
+fn default_modal_cpu_cores() -> f64 {
+    0.125
+}
+
 /// Configuration for a test group.
 ///
 /// Groups allow segmenting tests for different retry behaviors or filtering.

src/config/schema.rs

@@ -72,6 +72,13 @@ enum Commands {
         /// Emit a Perfetto trace to {output_dir}/trace.json
         #[arg(long)]
         trace: bool,
+
+        /// Show estimated sandbox cost after run.
+        ///
+        /// Note: This is calculated client-side using simple formulas and
+        /// may not reflect actual billing, discounts, or pricing adjustments.
+        #[arg(long)]
+        show_estimated_cost: bool,
     },
 
     /// Discover tests without running them
@@ -139,6 +146,7 @@ async fn main() -> Result<()> {
             env_vars,
             no_cache,
             trace,
+            show_estimated_cost,
         } => {
             run_tests(
                 &cli.config,
@@ -149,6 +157,7 @@ async fn main() -> Result<()> {
                 no_cache,
                 cli.verbose,
                 trace,
+                show_estimated_cost,
             )
             .await
         }
@@ -247,6 +256,7 @@ async fn dispatch_framework<P: offload::provider::SandboxProvider>(
     copy_dirs: &[CopyDir],
     verbose: bool,
     tracer: &offload::trace::Tracer,
+    show_estimated_cost: bool,
 ) -> Result<i32> {
     match &config.framework {
         FrameworkConfig::Pytest(f_cfg) => {
@@ -258,6 +268,7 @@ async fn dispatch_framework<P: offload::provider::SandboxProvider>(
                 copy_dirs,
                 verbose,
                 tracer,
+                show_estimated_cost,
             )
             .await
         }
@@ -270,6 +281,7 @@ async fn dispatch_framework<P: offload::provider::SandboxProvider>(
                 copy_dirs,
                 verbose,
                 tracer,
+                show_estimated_cost,
             )
             .await
         }
@@ -282,6 +294,7 @@ async fn dispatch_framework<P: offload::provider::SandboxProvider>(
                 copy_dirs,
                 verbose,
                 tracer,
+                show_estimated_cost,
             )
             .await
         }
@@ -294,6 +307,7 @@ async fn dispatch_framework<P: offload::provider::SandboxProvider>(
                 copy_dirs,
                 verbose,
                 tracer,
+                show_estimated_cost,
             )
             .await
         }
@@ -310,6 +324,7 @@ async fn run_tests(
     no_cache: bool,
     verbose: bool,
     trace: bool,
+    show_estimated_cost: bool,
 ) -> Result<()> {
     let tracer = if trace {
         offload::trace::Tracer::new()
@@ -435,6 +450,7 @@ async fn run_tests(
                 &copy_dirs,
                 verbose,
                 &tracer,
+                show_estimated_cost,
             )
             .await?
         }
@@ -465,7 +481,16 @@ async fn run_tests(
                 info!("No tests to run");
                 return Ok(());
             }
-            dispatch_framework(&config, &all_tests, provider, &copy_dirs, verbose, &tracer).await?
+            dispatch_framework(
+                &config,
+                &all_tests,
+                provider,
+                &copy_dirs,
+                verbose,
+                &tracer,
+                show_estimated_cost,
+            )
+            .await?
         }
         ProviderConfig::Modal(p_cfg) => {
             // Run discovery and image preparation concurrently
@@ -494,7 +519,16 @@ async fn run_tests(
                 info!("No tests to run");
                 return Ok(());
             }
-            dispatch_framework(&config, &all_tests, provider, &copy_dirs, verbose, &tracer).await?
+            dispatch_framework(
+                &config,
+                &all_tests,
+                provider,
+                &copy_dirs,
+                verbose,
+                &tracer,
+                show_estimated_cost,
+            )
+            .await?
         }
     };
 
@@ -515,6 +549,7 @@ async fn run_tests(
 
 /// Run all tests with a single orchestrator call.
 /// Returns the exit code (0 = success, 1 = failures/not run, 2 = flaky only).
+#[allow(clippy::too_many_arguments)]
 async fn run_all_tests<P, D>(
     config: &config::Config,
     tests: &[TestRecord],
@@ -523,6 +558,7 @@ async fn run_all_tests<P, D>(
     copy_dirs: &[CopyDir],
     verbose: bool,
     tracer: &offload::trace::Tracer,
+    show_estimated_cost: bool,
 ) -> Result<i32>
 where
     P: offload::provider::SandboxProvider,
@@ -565,7 +601,13 @@ where
         .context("Failed to create sandboxes")?;
     drop(_pool_span);
 
-    let orchestrator = Orchestrator::new(config.clone(), framework, verbose, tracer.clone());
+    let orchestrator = Orchestrator::new(
+        config.clone(),
+        framework,
+        verbose,
+        tracer.clone(),
+        show_estimated_cost,
+    );
 
     let result = orchestrator.run_with_tests(tests, sandbox_pool).await?;
 
@@ -663,6 +705,7 @@ fn init_config(provider: &str, framework: &str) -> Result<()> {
             timeout_secs: 3600,
             copy_dirs: vec![],
             env: HashMap::new(),
+            cpu_cores: 1.0,
         }),
         _ => {
             eprintln!("Unknown provider: {}. Use: local, default", provider);

src/main.rs

@@ -14,7 +14,7 @@ use tracing::{debug, error, info, warn};
 
 use crate::config::Config;
 use crate::framework::{TestFramework, TestInstance, TestRecord};
-use crate::provider::Sandbox;
+use crate::provider::{CostEstimate, Sandbox};
 use crate::report::{MasterJunitReport, load_test_durations, print_summary};
 
 pub use pool::SandboxPool;
@@ -66,6 +66,9 @@ pub struct RunResult {
 
     /// Wall-clock duration of the entire test run.
     pub duration: Duration,
+
+    /// Estimated cost of the test run (aggregated from all sandboxes).
+    pub estimated_cost: CostEstimate,
 }
 
 impl RunResult {
@@ -109,6 +112,7 @@ pub struct Orchestrator<S, D> {
     framework: D,
     verbose: bool,
     tracer: crate::trace::Tracer,
+    show_cost: bool,
     _sandbox: std::marker::PhantomData<S>,
 }
 
@@ -125,12 +129,20 @@ where
     /// * `framework` - Test framework for running tests
     /// * `verbose` - Whether to show verbose output (streaming test output)
     /// * `tracer` - Performance tracer for emitting trace events
-    pub fn new(config: Config, framework: D, verbose: bool, tracer: crate::trace::Tracer) -> Self {
+    /// * `show_cost` - Whether to display cost estimate in summary
+    pub fn new(
+        config: Config,
+        framework: D,
+        verbose: bool,
+        tracer: crate::trace::Tracer,
+        show_cost: bool,
+    ) -> Self {
         Self {
             config,
             framework,
             verbose,
             tracer,
+            show_cost,
             _sandbox: std::marker::PhantomData,
         }
     }
@@ -199,6 +211,7 @@ where
                 flaky: 0,
                 not_run: 0,
                 duration: start.elapsed(),
+                estimated_cost: CostEstimate::default(),
             });
         }
 
@@ -413,20 +426,22 @@ where
 
         // Use the JUnit total as the authoritative count (passed + failed + flaky = total)
         // This ensures passed can never exceed total
+        // Note: estimated_cost is set to default here and updated after sandbox cleanup
         let run_result = RunResult {
             total_tests: total_in_junit,
             passed: passed + flaky_count, // Flaky tests count as passed
             failed,
             flaky: flaky_count,
             not_run,
             duration: start.elapsed(),
+            estimated_cost: CostEstimate::default(),
         };
         drop(_agg_span);
 
         progress.finish_and_clear();
-        print_summary(&run_result);
 
         // Terminate all sandboxes in parallel (after printing results)
+        // Aggregate cost estimates BEFORE terminating (cost_estimate uses elapsed time)
         let _cleanup_span = self.tracer.span(
             "sandbox_cleanup",
             "orchestrator",
@@ -440,6 +455,17 @@ where
                 Vec::new()
             }
         };
+
+        // Aggregate cost estimates before terminating sandboxes
+        let estimated_cost = sandboxes
+            .iter()
+            .fold(CostEstimate::default(), |mut acc, sb| {
+                let cost = sb.cost_estimate();
+                acc.cpu_seconds += cost.cpu_seconds;
+                acc.estimated_cost_usd += cost.estimated_cost_usd;
+                acc
+            });
+
         let terminate_futures = sandboxes.into_iter().map(|sandbox| async move {
             if let Err(e) = sandbox.terminate().await {
                 warn!("Failed to terminate sandbox {}: {}", sandbox.id(), e);
@@ -448,6 +474,14 @@ where
         futures::future::join_all(terminate_futures).await;
         drop(_cleanup_span);
 
+        // Update run_result with estimated_cost
+        let run_result = RunResult {
+            estimated_cost,
+            ..run_result
+        };
+
+        print_summary(&run_result, self.show_cost);
+
         Ok(run_result)
     }
 }

src/orchestrator.rs

@@ -70,7 +70,7 @@ impl<S: Sandbox> Default for SandboxPool<S> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::provider::{OutputStream, ProviderResult};
+    use crate::provider::{CostEstimate, OutputStream, ProviderResult};
     use async_trait::async_trait;
     use std::path::Path;
 
@@ -95,6 +95,9 @@ mod tests {
         async fn terminate(&self) -> ProviderResult<()> {
             Ok(())
         }
+        fn cost_estimate(&self) -> CostEstimate {
+            CostEstimate::default()
+        }
     }
 
     struct FakeProvider;