Filter PIDs to profile

gprofiler/gprofiler_types.py

@@ -6,7 +6,7 @@
 import re
 from collections import Counter
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, MutableMapping, Optional, Union
+from typing import Any, Callable, Dict, List, MutableMapping, Optional, Union
 
 import configargparse
 
@@ -76,6 +76,16 @@ def nonnegative_integer(value_str: str) -> int:
     return value
 
 
+def integers_list(value_str: str) -> List[int]:
+    try:
+        pids = [int(pid) for pid in value_str.split(",")]
+    except ValueError:
+        raise configargparse.ArgumentTypeError(
+            "IntegerList should be a single integer, or comma separated list of integers f.e. 13,452,2388"
+        )
+    return pids
+
+
 def integer_range(min_range: int, max_range: int) -> Callable[[str], int]:
     def integer_range_check(value_str: str) -> int:
         value = int(value_str)

gprofiler/main.py

@@ -15,7 +15,7 @@
 from pathlib import Path
 from threading import Event
 from types import FrameType, TracebackType
-from typing import Iterable, Optional, Type, cast
+from typing import Any, Iterable, Optional, Type, cast
 
 import configargparse
 import humanfriendly
@@ -38,7 +38,7 @@
 from gprofiler.databricks_client import DatabricksClient
 from gprofiler.diagnostics import log_diagnostics, set_diagnostics
 from gprofiler.exceptions import APIError, NoProfilersEnabledError
-from gprofiler.gprofiler_types import ProcessToProfileData, UserArgs, positive_integer
+from gprofiler.gprofiler_types import ProcessToProfileData, UserArgs, integers_list, positive_integer
 from gprofiler.log import RemoteLogsHandler, initial_root_logger_setup
 from gprofiler.merge import concatenate_from_external_file, concatenate_profiles, merge_profiles
 from gprofiler.metadata.application_identifiers import ApplicationIdentifiers
@@ -102,71 +102,57 @@ def sigint_handler(sig: int, frame: Optional[FrameType]) -> None:
 
 
 class GProfiler:
-    def __init__(
-        self,
-        output_dir: str,
-        flamegraph: bool,
-        rotating_output: bool,
-        profiler_api_client: Optional[ProfilerAPIClient],
-        collect_metrics: bool,
-        collect_metadata: bool,
-        enrichment_options: EnrichmentOptions,
-        state: State,
-        usage_logger: UsageLoggerInterface,
-        user_args: UserArgs,
-        duration: int,
-        profile_api_version: str,
-        profiling_mode: str,
-        profile_spawned_processes: bool = True,
-        remote_logs_handler: Optional[RemoteLogsHandler] = None,
-        controller_process: Optional[Process] = None,
-        spark_sampler: Optional[SparkSampler] = None,
-    ):
-        self._output_dir = output_dir
-        self._flamegraph = flamegraph
-        self._rotating_output = rotating_output
-        self._profiler_api_client = profiler_api_client
-        self._state = state
-        self._remote_logs_handler = remote_logs_handler
-        self._profile_api_version = profile_api_version
-        self._collect_metrics = collect_metrics
-        self._collect_metadata = collect_metadata
-        self._enrichment_options = enrichment_options
+    def __init__(self, **kwargs: Any):
+        self._output_dir: str = kwargs.pop("output_dir")
+        self._flamegraph: bool = kwargs.pop("flamegraph")
+        self._rotating_output: bool = kwargs.pop("rotating_output")
+        self._profiler_api_client: Optional[ProfilerAPIClient] = kwargs.pop("profiler_api_client")
+        self._state: State = kwargs.pop("state")
+        self._remote_logs_handler: Optional[RemoteLogsHandler] = kwargs.get("remote_logs_handler", None)
+        self._profile_api_version: str = kwargs.pop("profile_api_version")
+        self._collect_metrics: bool = kwargs.pop("collect_metrics")
+        self._collect_metadata: bool = kwargs.pop("collect_metadata")
+        self._enrichment_options: EnrichmentOptions = kwargs.pop("enrichment_options")
+        self._user_args: UserArgs = kwargs.pop("user_args")
         self._static_metadata: Optional[Metadata] = None
         self._spawn_time = time.time()
         self._last_diagnostics = 0.0
         self._gpid = ""
-        self._controller_process = controller_process
-        self._duration = duration
-        if collect_metadata:
-            self._static_metadata = get_static_metadata(spawn_time=self._spawn_time, run_args=user_args)
+        self._controller_process: Optional[Process] = kwargs.get("controller_process", None)
+        self._duration: int = kwargs.pop("duration")
+        if self._collect_metadata:
+            self._static_metadata = get_static_metadata(spawn_time=self._spawn_time, run_args=self._user_args)
         self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=10)
         # TODO: we actually need 2 types of temporary directories.
         # 1. accessible by everyone - for profilers that run code in target processes, like async-profiler
         # 2. accessible only by us.
         # the latter can be root only. the former can not. we should do this separation so we don't expose
         # files unnecessarily.
         container_names_client = ContainerNamesClient() if self._enrichment_options.container_names else None
-        self._profiler_state = ProfilerState(
-            Event(),
-            TEMPORARY_STORAGE_PATH,
-            profile_spawned_processes,
-            bool(user_args.get("insert_dso_name")),
-            profiling_mode,
-            container_names_client,
+        profiler_state_kwargs = {
+            "stop_event": Event(),
+            "profile_spawned_processes": kwargs.get("profile_spawned_processes", True),
+            "insert_dso_name": bool(self._user_args.get("insert_dso_name")),
+            "profiling_mode": kwargs.pop("profiling_mode"),
+            "container_names_client": container_names_client,
+            "pids_to_profile": kwargs.pop("pids_to_profile"),
+            "storage_dir": TEMPORARY_STORAGE_PATH,
+        }
+        self._profiler_state = ProfilerState(**profiler_state_kwargs)
+        self.system_profiler, self.process_profilers = get_profilers(
+            self._user_args, profiler_state=self._profiler_state
         )
-        self.system_profiler, self.process_profilers = get_profilers(user_args, profiler_state=self._profiler_state)
-        self._usage_logger = usage_logger
-        if collect_metrics:
+        self._usage_logger: UsageLoggerInterface = kwargs.pop("usage_logger")
+        if self._collect_metrics:
             self._system_metrics_monitor: SystemMetricsMonitorBase = SystemMetricsMonitor(
                 self._profiler_state.stop_event
             )
         else:
             self._system_metrics_monitor = NoopSystemMetricsMonitor()
 
-        self._spark_sampler = spark_sampler
+        self._spark_sampler: Optional[SparkSampler] = kwargs.get("spark_sampler", None)
 
-        if isinstance(self.system_profiler, NoopProfiler) and not self.process_profilers and not spark_sampler:
+        if isinstance(self.system_profiler, NoopProfiler) and not self.process_profilers and not self._spark_sampler:
             raise NoProfilersEnabledError()
 
     @property
@@ -521,6 +507,15 @@ def parse_cmd_args() -> configargparse.Namespace:
     parser.add_argument(
         "--rotating-output", action="store_true", default=False, help="Keep only the last profile result"
     )
+    parser.add_argument(
+        "--pids",
+        dest="pids_to_profile",
+        action="extend",
+        default=None,
+        type=integers_list,
+        help="Comma separated list of processes that will be filtered to profile,"
+        " given multiple times will append pids to one list",
+    )
 
     _add_profilers_arguments(parser)
 
@@ -791,6 +786,9 @@ def parse_cmd_args() -> configargparse.Namespace:
     if args.nodejs_mode in ("perf", "attach-maps") and args.perf_mode not in ("fp", "smart"):
         parser.error("--nodejs-mode perf or attach-maps requires --perf-mode 'fp' or 'smart'")
 
+    if args.profile_spawned_processes and args.pids_to_profile is not None:
+        parser.error("--pids is not allowed when profiling spawned processes")
+
     return args
 
 
@@ -851,6 +849,19 @@ def verify_preconditions(args: configargparse.Namespace) -> None:
         print("--log-usage is available only when run as a container!", file=sys.stderr)
         sys.exit(1)
 
+    if args.pids_to_profile is not None:
+        pids_to_profile = []
+        for pid in args.pids_to_profile:
+            try:
+                Process(pid)
+                pids_to_profile.append(pid)
+            except NoSuchProcess:
+                continue
+        setattr(args, "pids_to_profile", pids_to_profile)
+        if len(pids_to_profile) == 0:
+            print("There aren't any alive processes from provided via --pid PID list")
+            sys.exit(1)
+
 
 def setup_signals() -> None:
     # When we run under staticx & PyInstaller, both of them forward (some of the) signals to gProfiler.
@@ -1033,26 +1044,27 @@ def main() -> None:
 
         ApplicationIdentifiers.init(enrichment_options)
         set_diagnostics(args.diagnostics)
-
-        gprofiler = GProfiler(
-            args.output_dir,
-            args.flamegraph,
-            args.rotating_output,
-            profiler_api_client,
-            args.collect_metrics,
-            args.collect_metadata,
-            enrichment_options,
-            state,
-            usage_logger,
-            args.__dict__,
-            args.duration,
-            args.profile_api_version,
-            args.profiling_mode,
-            args.profile_spawned_processes,
-            remote_logs_handler,
-            controller_process,
-            spark_sampler,
-        )
+        gprofiler_kwargs = {
+            "output_dir": args.output_dir,
+            "flamegraph": args.flamegraph,
+            "rotating_output": args.rotating_output,
+            "profiler_api_client": profiler_api_client,
+            "collect_metrics": args.collect_metrics,
+            "collect_metadata": args.collect_metadata,
+            "enrichment_options": enrichment_options,
+            "state": state,
+            "usage_logger": usage_logger,
+            "user_args": args.__dict__,
+            "duration": args.duration,
+            "profile_api_version": args.profile_api_version,
+            "profiling_mode": args.profiling_mode,
+            "pids_to_profile": args.pids_to_profile,
+            "profile_spawned_processes": args.profile_spawned_processes,
+            "remote_logs_handler": remote_logs_handler,
+            "controller_process": controller_process,
+            "spark_sampler": spark_sampler,
+        }
+        gprofiler = GProfiler(**gprofiler_kwargs)
         logger.info("gProfiler initialized and ready to start profiling")
         if args.continuous:
             gprofiler.run_continuous()

gprofiler/profiler_state.py

@@ -1,5 +1,5 @@
 from threading import Event
-from typing import Optional
+from typing import Any, List, Optional
 
 from gprofiler.containers_client import ContainerNamesClient
 from gprofiler.utils import TemporaryDirectoryWithMode
@@ -8,22 +8,15 @@
 class ProfilerState:
     # Class for storing generic state parameters. These parameters are the same for each profiler.
     # Thanks to that class adding new state parameters to profilers won't result changing code in every profiler.
-    def __init__(
-        self,
-        stop_event: Event,
-        storage_dir: str,
-        profile_spawned_processes: bool,
-        insert_dso_name: bool,
-        profiling_mode: str,
-        container_names_client: Optional[ContainerNamesClient],
-    ) -> None:
-        self._stop_event = stop_event
-        self._profile_spawned_processes = profile_spawned_processes
-        self._insert_dso_name = insert_dso_name
-        self._profiling_mode = profiling_mode
-        self._temporary_dir = TemporaryDirectoryWithMode(dir=storage_dir, mode=0o755)
+    def __init__(self, **kwargs: Any) -> None:
+        self._stop_event: Event = kwargs.pop("stop_event")
+        self._profile_spawned_processes: bool = kwargs.pop("profile_spawned_processes")
+        self._insert_dso_name: bool = kwargs.pop("insert_dso_name")
+        self._profiling_mode: str = kwargs.pop("profiling_mode")
+        self._temporary_dir = TemporaryDirectoryWithMode(dir=kwargs.pop("storage_dir"), mode=0o755)
         self._storage_dir = self._temporary_dir.name
-        self._container_names_client = container_names_client
+        self._container_names_client: Optional[ContainerNamesClient] = kwargs.get("container_names_client", None)
+        self._pids_to_profile: Optional[List[int]] = kwargs.pop("pids_to_profile")
 
     @property
     def stop_event(self) -> Event:
@@ -49,6 +42,10 @@ def profiling_mode(self) -> str:
     def container_names_client(self) -> Optional[ContainerNamesClient]:
         return self._container_names_client
 
+    @property
+    def pids_to_profile(self) -> Optional[List[int]]:
+        return self._pids_to_profile
+
     def get_container_name(self, pid: int) -> str:
         if self._container_names_client is not None:
             return self._container_names_client.get_container_name(pid)

gprofiler/profilers/perf.py

@@ -160,7 +160,10 @@ def merge_global_perfs(
 
     total_fp_samples = sum([sum(stacks.values()) for stacks in fp_perf.values()])
     total_dwarf_samples = sum([sum(stacks.values()) for stacks in dwarf_perf.values()])
-    fp_to_dwarf_sample_ratio = total_fp_samples / total_dwarf_samples
+    if total_dwarf_samples == 0:
+        fp_to_dwarf_sample_ratio = 0.0  # ratio can be 0 because if total_dwarf_samples is 0 then it will be never used
+    else:
+        fp_to_dwarf_sample_ratio = total_fp_samples / total_dwarf_samples
 
     # The FP perf is used here as the "main" perf, to which the DWARF perf is scaled.
     merged_pid_to_stacks_counters: ProcessToStackSampleCounters = defaultdict(Counter)
@@ -389,6 +392,10 @@ def __init__(
         self._node_processes: List[Process] = []
         self._node_processes_attached: List[Process] = []
         self._perf_memory_restart = perf_memory_restart
+        pid_args: List[str] = []
+        if self._profiler_state.pids_to_profile is not None:
+            pid_args.append("--pid")
+            pid_args.append(",".join([str(pid) for pid in self._profiler_state.pids_to_profile]))
 
         if perf_mode in ("fp", "smart"):
             self._perf_fp: Optional[PerfProcess] = PerfProcess(
@@ -397,7 +404,7 @@ def __init__(
                 os.path.join(self._profiler_state.storage_dir, "perf.fp"),
                 False,
                 perf_inject,
-                [],
+                pid_args,
             )
             self._perfs.append(self._perf_fp)
         else:
@@ -410,7 +417,7 @@ def __init__(
                 os.path.join(self._profiler_state.storage_dir, "perf.dwarf"),
                 True,
                 False,  # no inject in dwarf mode, yet
-                ["--call-graph", f"dwarf,{perf_dwarf_stack_size}"],
+                pid_args + ["--call-graph", f"dwarf,{perf_dwarf_stack_size}"],
             )
             self._perfs.append(self._perf_dwarf)
         else:

gprofiler/profilers/php.py

@@ -194,6 +194,11 @@ def extract_metadata_section(re_expr: Pattern, metadata_line: str) -> str:
 
         profiles: ProcessToProfileData = {}
         for pid in results:
+            # Because of https://github.com/Granulate/gprofiler/issues/763,
+            # for now we only filter output of phpspy to return only profiles from chosen pids
+            if profiler_state.pids_to_profile is not None:
+                if pid not in profiler_state.pids_to_profile:
+                    continue
             # TODO: appid & app metadata for php!
             appid = None
             app_metadata = None

gprofiler/profilers/profiler_base.py

@@ -171,6 +171,11 @@ def _profiling_error_stack(
     def snapshot(self) -> ProcessToProfileData:
         processes_to_profile = self._select_processes_to_profile()
         logger.debug(f"{self.__class__.__name__}: selected {len(processes_to_profile)} processes to profile")
+        if self._profiler_state.pids_to_profile is not None:
+            processes_to_profile = [
+                process for process in processes_to_profile if process.pid in self._profiler_state.pids_to_profile
+            ]
+            logger.debug(f"{self.__class__.__name__}: processes left after filtering: {len(processes_to_profile)}")
         self._notify_selected_processes(processes_to_profile)
 
         if not processes_to_profile:

gprofiler/profilers/python_ebpf.py

@@ -238,6 +238,11 @@ def snapshot(self) -> ProcessToProfileData:
             parsed = python._add_versions_to_stacks(parsed)
         profiles = {}
         for pid in parsed:
+            # Because of https://github.com/Granulate/gprofiler/issues/764,
+            # for now we only filter output of pyperf to return only profiles from chosen pids
+            if self._profiler_state.pids_to_profile is not None:
+                if pid not in self._profiler_state.pids_to_profile:
+                    continue
             try:
                 process = Process(pid)
                 appid = application_identifiers.get_python_app_id(process)

tests/conftest.py

@@ -90,7 +90,16 @@ def java_args() -> Tuple[str]:
 
 @fixture()
 def profiler_state(tmp_path: Path, insert_dso_name: bool) -> ProfilerState:
-    return ProfilerState(Event(), str(tmp_path), False, insert_dso_name, CPU_PROFILING_MODE, ContainerNamesClient())
+    profiler_state_kwargs = {
+        "stop_event": Event(),
+        "profile_spawned_processes": False,
+        "insert_dso_name": insert_dso_name,
+        "profiling_mode": CPU_PROFILING_MODE,
+        "container_names_client": ContainerNamesClient(),
+        "pids_to_profile": None,
+        "storage_dir": str(tmp_path),
+    }
+    return ProfilerState(**profiler_state_kwargs)
 
 
 def make_path_world_accessible(path: Path) -> None: