Merge pull request #935 from linsword13/line-profile

Support per-line profiling for phase functions

Author: douglasjacobsen

lib/ramble/ramble/application.py

@@ -120,10 +120,19 @@ def _run_phase_hook(obj, workspace, pipeline, hook):
 
     hook_func_name = f"_{hook}"
     if hasattr(obj, hook_func_name):
-        phase_func = getattr(obj, hook_func_name)
+        phase_func = _get_phase_func_wrapper(workspace, getattr(obj, hook_func_name), hook)
         phase_func(workspace)
 
 
+def _get_phase_func_wrapper(workspace, phase_func, phase_name):
+    if workspace.profile_config is None:
+        return phase_func
+    (profiler, profile_phases) = workspace.profile_config
+    if phase_name not in profile_phases:
+        return phase_func
+    return profiler(phase_func)
+
+
 class ApplicationBase(metaclass=ApplicationMeta):
     name = None
     _builtin_name = NS_SEPARATOR.join(("builtin", "{name}"))
@@ -648,7 +657,7 @@ def run_phase(self, pipeline, phase, workspace):
 
         for _, obj in self._objects(exclude_types=[ramble.repository.ObjectTypes.applications]):
             _run_phase_hook(obj, workspace, pipeline, phase)
-        phase_func = phase_node.attribute
+        phase_func = _get_phase_func_wrapper(workspace, phase_node.attribute, phase)
         phase_func(workspace, app_inst=self)
         self._phase_times[phase] = time.time() - start_time
 

lib/ramble/ramble/cmd/common/arguments.py

@@ -148,6 +148,19 @@ def include_phase_dependencies():
     )
 
 
+@arg
+def profile_phases():
+    return Args(
+        "--profile-phase",
+        nargs="+",
+        action="append",
+        default=None,
+        dest="profile_phases",
+        help="phases to be profiled by line_profiler",
+        required=False,
+    )
+
+
 @arg
 def where():
     return Args(

lib/ramble/ramble/cmd/workspace.py

@@ -10,6 +10,7 @@
 import sys
 import tempfile
 import argparse
+import itertools
 
 import llnl.util.tty as tty
 import llnl.util.tty.color as color
@@ -443,12 +444,25 @@ def workspace_concretize(args):
 
 
 def workspace_run_pipeline(args, pipeline):
-    include_phase_dependencies = getattr(args, "include_phase_dependencies", None)
-    if include_phase_dependencies:
-        with ramble.config.override("config:include_phase_dependencies", True):
+    profile_phases = getattr(args, "profile_phases", None)
+    if profile_phases:
+        import line_profiler
+
+        profiler = line_profiler.LineProfiler()
+        profiler.enable()
+        p_phases = set(itertools.chain.from_iterable(profile_phases))
+        pipeline.workspace.profile_config = (profiler, p_phases)
+    try:
+        include_phase_dependencies = getattr(args, "include_phase_dependencies", None)
+        if include_phase_dependencies:
+            with ramble.config.override("config:include_phase_dependencies", True):
+                pipeline.run()
+        else:
             pipeline.run()
-    else:
-        pipeline.run()
+    finally:
+        if profile_phases:
+            profiler.disable()
+            profiler.print_stats()
 
 
 def workspace_setup_setup_parser(subparser):
@@ -464,7 +478,14 @@ def workspace_setup_setup_parser(subparser):
 
     arguments.add_common_arguments(
         subparser,
-        ["phases", "include_phase_dependencies", "where", "exclude_where", "filter_tags"],
+        [
+            "phases",
+            "include_phase_dependencies",
+            "where",
+            "exclude_where",
+            "filter_tags",
+            "profile_phases",
+        ],
     )
 
 
@@ -530,7 +551,14 @@ def workspace_analyze_setup_parser(subparser):
 
     arguments.add_common_arguments(
         subparser,
-        ["phases", "include_phase_dependencies", "where", "exclude_where", "filter_tags"],
+        [
+            "phases",
+            "include_phase_dependencies",
+            "where",
+            "exclude_where",
+            "filter_tags",
+            "profile_phases",
+        ],
     )
 
 
@@ -587,7 +615,9 @@ def workspace_push_to_cache_setup_parser(subparser):
         "-d", dest="cache_path", default=None, required=True, help="Path to cache."
     )
 
-    arguments.add_common_arguments(subparser, ["where", "exclude_where", "filter_tags"])
+    arguments.add_common_arguments(
+        subparser, ["where", "exclude_where", "filter_tags", "profile_phases"]
+    )
 
 
 def workspace_info_setup_parser(subparser):
@@ -981,7 +1011,8 @@ def workspace_archive_setup_parser(subparser):
     )
 
     arguments.add_common_arguments(
-        subparser, ["phases", "include_phase_dependencies", "where", "exclude_where"]
+        subparser,
+        ["phases", "include_phase_dependencies", "where", "exclude_where", "profile_phases"],
     )
 
 
@@ -1024,7 +1055,8 @@ def workspace_mirror_setup_parser(subparser):
     )
 
     arguments.add_common_arguments(
-        subparser, ["phases", "include_phase_dependencies", "where", "exclude_where"]
+        subparser,
+        ["phases", "include_phase_dependencies", "where", "exclude_where", "profile_phases"],
     )
 
 

lib/ramble/ramble/workspace/workspace.py

@@ -487,6 +487,9 @@ def __init__(self, root, dry_run=False, read_default_template=True):
 
         self.deployment_name = self.name
 
+        # Used by profiling phases
+        self.profile_config = None
+
     def _re_read(self):
         """Reinitialize the workspace object if it has been written (this
         may not be true if the workspace was just created in this running

requirements-dev.txt

@@ -6,3 +6,4 @@ coverage
 pre-commit
 pytest-xdist
 pytest-cov
+line_profiler

share/ramble/ramble-completion.bash

@@ -646,7 +646,7 @@ _ramble_workspace_activate() {
 }
 
 _ramble_workspace_archive() {
-    RAMBLE_COMPREPLY="-h --help --tar-archive -t --prefix -p --upload-url -u --include-secrets --phases --include-phase-dependencies --where --exclude-where"
+    RAMBLE_COMPREPLY="-h --help --tar-archive -t --prefix -p --upload-url -u --include-secrets --phases --include-phase-dependencies --where --exclude-where --profile-phase"
 }
 
 _ramble_workspace_deactivate() {
@@ -667,15 +667,15 @@ _ramble_workspace_concretize() {
 }
 
 _ramble_workspace_setup() {
-    RAMBLE_COMPREPLY="-h --help --dry-run --phases --include-phase-dependencies --where --exclude-where --filter-tags"
+    RAMBLE_COMPREPLY="-h --help --dry-run --phases --include-phase-dependencies --where --exclude-where --filter-tags --profile-phase"
 }
 
 _ramble_workspace_analyze() {
-    RAMBLE_COMPREPLY="-h --help -f --formats -u --upload -p --print-results -s --summary-only --phases --include-phase-dependencies --where --exclude-where --filter-tags"
+    RAMBLE_COMPREPLY="-h --help -f --formats -u --upload -p --print-results -s --summary-only --phases --include-phase-dependencies --where --exclude-where --filter-tags --profile-phase"
 }
 
 _ramble_workspace_push_to_cache() {
-    RAMBLE_COMPREPLY="-h --help -d --where --exclude-where --filter-tags"
+    RAMBLE_COMPREPLY="-h --help -d --where --exclude-where --filter-tags --profile-phase"
 }
 
 _ramble_workspace_info() {
@@ -692,7 +692,7 @@ _ramble_workspace_edit() {
 }
 
 _ramble_workspace_mirror() {
-    RAMBLE_COMPREPLY="-h --help -d --dry-run --phases --include-phase-dependencies --where --exclude-where"
+    RAMBLE_COMPREPLY="-h --help -d --dry-run --phases --include-phase-dependencies --where --exclude-where --profile-phase"
 }
 
 _ramble_workspace_experiment_logs() {