Merge pull request #636 from bennybp/cprof

Enable ability to profile fractal server

qcfractal/cli/qcfractal_server.py

@@ -78,7 +78,7 @@ def parse_args():
 
     # Allow some config settings to be altered via the command line
     fractal_args = start.add_argument_group("Server Settings")
-    for field in ["port", "logfile", "loglevel"]:
+    for field in ["port", "logfile", "loglevel", "cprofile"]:
         cli_name = "--" + field.replace("_", "-")
         fractal_args.add_argument(cli_name, **FractalServerSettings.help_info(field))
 
@@ -572,6 +572,16 @@ def main(args=None):
 
     config = FractalConfig(**config_args)
 
+    # If desired, enable profiling
+    if config.fractal.cprofile is not None:
+        print("!" * 80)
+        print(f"! Enabling profiling via cProfile. Outputting data file to {config.fractal.cprofile}")
+        print("!" * 80)
+        import cProfile
+
+        pr = cProfile.Profile()
+        pr.enable()
+
     # Merge files
     if command != "init":
         if not config.base_path.exists():
@@ -609,6 +619,14 @@ def main(args=None):
     elif command == "restore":
         server_restore(args, config)
 
+    # Everything finished. If profiling is enabled, write out the
+    # data file
+    if config.fractal.cprofile is not None:
+        print(f"! Writing profiling data to {config.fractal.cprofile}")
+        print("! Read using the Stats class of the pstats package")
+        pr.disable()
+        pr.dump_stats(config.fractal.cprofile)
+
 
 if __name__ == "__main__":
     main()

qcfractal/config.py

@@ -110,6 +110,9 @@ class FractalServerSettings(ConfigSettings):
     query_limit: int = Field(1000, description="The maximum number of records to return per query.")
     logfile: Optional[str] = Field("qcfractal_server.log", description="The logfile to write server logs.")
     loglevel: str = Field("info", description="Level of logging to enable (debug, info, warning, error, critical)")
+    cprofile: Optional[str] = Field(
+        None, description="Enable profiling via cProfile, and output cprofile data to this path"
+    )
     service_frequency: int = Field(60, description="The frequency to update the QCFractal services.")
     max_active_services: int = Field(20, description="The maximum number of concurrent active services.")
     heartbeat_frequency: int = Field(1800, description="The frequency (in seconds) to check the heartbeat of workers.")