SImulate Profile (#2)

* Refactor: move simulate profile to v_prof with name profile_sim

* add scalene implementation

* add extra scaelene args

* remove some comments

* remove unused imports

* bugfix

* add profiler arg

* add cli tests

* add run profile tests

* formatting

* make cprofile the default

* use mocker instead of unittest mock

* marginally improve compactness

* formatting

* add mprof back

* add to changelog

* add to help message

* fix help message

* fix help formatting

Author: patricktnast

CHANGELOG.rst

@@ -1,3 +1,8 @@
+**v0.1.0 - 10/24/2025**
+
+ - Convert bash script to python
+ - add profiling CLI with scalene backend
+
 **v0.0.1 - 10/20/2025**
 
  - Initial release

setup.py

@@ -53,6 +53,7 @@
         "memory_profiler",
         "matplotlib",
         "seaborn",
+        "scalene",
     ]
 
     setup_requires = ["setuptools_scm"]
@@ -96,5 +97,6 @@
             [console_scripts]
             make_artifacts=vivarium_profiling.tools.cli:make_artifacts
             run_benchmark=vivarium_profiling.tools.cli:run_benchmark
+            profile_sim=vivarium_profiling.tools.cli:profile_sim
         """,
     )

src/vivarium_profiling/tools/cli.py

@@ -1,15 +1,149 @@
 import glob
+import pstats
+import subprocess
+from datetime import datetime as dt
 from pathlib import Path
 
 import click
 from loguru import logger
+from vivarium.framework.logging import configure_logging_to_file
 from vivarium.framework.utilities import handle_exceptions
 
 from vivarium_profiling.constants import metadata, paths
 from vivarium_profiling.tools import build_artifacts, configure_logging_to_terminal
 from vivarium_profiling.tools.run_benchmark import run_benchmark_loop
 
 
+@click.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
+@click.argument(
+    "model_specification",
+    type=click.Path(exists=True, dir_okay=False, resolve_path=True),
+)
+@click.option(
+    "--results_directory",
+    "-o",
+    type=click.Path(resolve_path=True),
+    default=Path("~/vivarium_results/").expanduser(),
+    show_default=True,
+    help=(
+        "The directory to write results to. A folder will be created "
+        "in this directory with the same name as the configuration file."
+    ),
+)
+@click.option(
+    "--skip_writing",
+    is_flag=True,
+    help=(
+        "Skip writing the simulation results to the output directory; the time spent "
+        "normally writing simulation results to disk will not be included in the profiling "
+        "statistics."
+    ),
+)
+@click.option(
+    "--skip_processing",
+    is_flag=True,
+    help=(
+        "Skip processing the resulting binary file to a human-readable .txt file "
+        "sorted by cumulative runtime; the resulting .stats file can still be read "
+        "and processed later using the pstats module."
+    ),
+)
+@click.option(
+    "--profiler",
+    type=click.Choice(["cprofile", "scalene"]),
+    default="cprofile",
+    show_default=True,
+    help=(
+        "Profiling backend to use. cProfile provides the most detaile function-level"
+        "runtime information, while scalene provides detailed annotation of source"
+        "code that may be the source of bottlenecks."
+    ),
+)
+@click.pass_context
+def profile_sim(
+    ctx: click.Context,
+    model_specification: Path,
+    results_directory: Path,
+    skip_writing: bool,
+    skip_processing: bool,
+    profiler: str,
+) -> None:
+    """Run a simulation based on the provided MODEL_SPECIFICATION and profile the run."""
+    model_specification = Path(model_specification)
+    results_directory = Path(results_directory)
+    results_root = results_directory / f"{dt.now().strftime('%Y_%m_%d_%H_%M_%S')}"
+    configure_logging_to_file(output_directory=results_root)
+
+    if skip_writing:
+        configuration_override = {}
+    else:
+        output_data_root = results_root / "results"
+        output_data_root.mkdir(parents=True, exist_ok=False)
+        configuration_override = {
+            "output_data": {"results_directory": str(output_data_root)},
+        }
+
+    script_path = Path(__file__).parent / "run_profile.py"
+
+    config_str = repr(configuration_override)
+
+    # Get any extra arguments passed to the profiler
+    extra_args = ctx.args
+
+    if profiler == "scalene":
+        out_json_file = results_root / f"{model_specification.name}".replace("yaml", "json")
+        try:
+            cmd = [
+                "scalene",
+                "--json",
+                "--outfile",
+                str(out_json_file),
+                "--off",
+            ]
+            # Add any additional profiler arguments
+            cmd.extend(extra_args)
+            cmd.extend(
+                [
+                    str(script_path),
+                    str(model_specification),
+                    "--config-override",
+                    config_str,
+                ]
+            )
+            subprocess.run(cmd, check=True)
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Scalene profiling failed: {e}")
+            raise
+    elif profiler == "cprofile":
+        out_stats_file = results_root / f"{model_specification.name}".replace("yaml", "stats")
+        try:
+            subprocess.run(
+                [
+                    "python",
+                    str(script_path),
+                    str(model_specification),
+                    "--config-override",
+                    config_str,
+                    "--profiler",
+                    "cprofile",
+                    "--output",
+                    str(out_stats_file),
+                ],
+                check=True,
+            )
+
+            if not skip_processing:
+                out_txt_file = Path(str(out_stats_file) + ".txt")
+                with out_txt_file.open("w") as f:
+                    p = pstats.Stats(str(out_stats_file), stream=f)
+                    p.sort_stats("cumulative")
+                    p.print_stats()
+
+        except subprocess.CalledProcessError as e:
+            logger.error(f"cProfile profiling failed: {e}")
+            raise
+
+
 @click.command()
 @click.option(
     "-l",
@@ -66,7 +200,7 @@ def make_artifacts(
 @click.command()
 @click.option(
     "-m",
-    "--models",
+    "--model_specifications",
     multiple=True,
     required=True,
     help="Model specification files (supports glob patterns). Can be specified multiple times.",

src/vivarium_profiling/tools/run_benchmark.py

@@ -59,7 +59,7 @@ def initialize_results_file(results_dir: str) -> str:
 
 def run_memory_profiler(spec: str, output_dir: str) -> None:
     """Run memory profiler on a model specification."""
-    cmd = ["mprof", "run", "-CM", "simulate", "profile", spec, "-o", output_dir]
+    cmd = ["mprof", "run", "-CM", "profile_sim", spec, "-o", output_dir]
     result = subprocess.run(cmd, check=True, capture_output=False)
     if result.returncode != 0:
         raise click.ClickException(f"Memory profiler failed for {spec}")

src/vivarium_profiling/tools/run_profile.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Standalone script for profiling Vivarium simulations."""
+
+import argparse
+import ast
+import cProfile
+import sys
+
+from scalene import scalene_profiler
+from vivarium.framework.engine import SimulationContext
+
+
+def run_profile_scalene(sim: SimulationContext):
+    """Run a Vivarium simulation for scalene profiling purposes."""
+    with scalene_profiler.enable_profiling():
+        sim.run_simulation()
+
+
+def run_profile_cprofile(sim: SimulationContext, output_file: str):
+    """Run a Vivarium simulation for cProfile profiling purposes."""
+    with cProfile.Profile() as profiler:
+        sim.run_simulation()
+    profiler.dump_stats(output_file)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run Vivarium simulation for profiling")
+    parser.add_argument("model_specification", help="Path to model specification file")
+    parser.add_argument(
+        "--config-override", default="{}", help="Configuration override as JSON/dict string"
+    )
+    parser.add_argument(
+        "--profiler",
+        choices=["scalene", "cprofile"],
+        default="scalene",
+        help="Profiling backend to use",
+    )
+    parser.add_argument(
+        "--output", help="Output file for cProfile stats (required when using cprofile)"
+    )
+
+    args = parser.parse_args()
+
+    # Parse the configuration override
+    try:
+        configuration_override = ast.literal_eval(args.config_override)
+    except (ValueError, SyntaxError) as e:
+        print(f"Error parsing config override: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    sim = SimulationContext(args.model_specification, configuration=configuration_override)
+
+    if args.profiler == "scalene":
+        run_profile_scalene(sim)
+    elif args.profiler == "cprofile":
+        if not args.output:
+            print("Error: --output is required when using cprofile", file=sys.stderr)
+            sys.exit(1)
+        run_profile_cprofile(sim, args.output)
+
+
+if __name__ == "__main__":
+    main()