Merge pull request #45 from gridfm/add_profiler

add lightning profiler cli argument

Author: romeokienzler

.gitignore

@@ -1,4 +1,8 @@
-venv
+*venv
+data_out
+logs
+site
+.julia
 venv_pp
 /data/
 __pycache__/
@@ -10,3 +14,4 @@ gridfm_graphkit.egg-info
 mlruns
 *.pt
 .DS_Store
+.venv

gridfm_graphkit/__main__.py

@@ -18,6 +18,13 @@ def main():
     train_parser.add_argument("--run_name", type=str, default="run")
     train_parser.add_argument("--log_dir", type=str, default="mlruns")
     train_parser.add_argument("--data_path", type=str, default="data")
+    train_parser.add_argument(
+        "--profiler",
+        type=str,
+        default=None,
+        choices=["simple", "advanced", "pytorch"],
+        help="Enable Lightning profiler: 'simple', 'advanced', or 'pytorch'.",
+    )
 
     # ---- FINETUNE SUBCOMMAND ----
     finetune_parser = subparsers.add_parser("finetune", help="Run fine-tuning")
@@ -27,6 +34,13 @@ def main():
     finetune_parser.add_argument("--run_name", type=str, default="run")
     finetune_parser.add_argument("--log_dir", type=str, default="mlruns")
     finetune_parser.add_argument("--data_path", type=str, default="data")
+    finetune_parser.add_argument(
+        "--profiler",
+        type=str,
+        default=None,
+        choices=["simple", "advanced", "pytorch"],
+        help="Enable Lightning profiler: 'simple', 'advanced', or 'pytorch'.",
+    )
 
     # ---- EVALUATE SUBCOMMAND ----
     evaluate_parser = subparsers.add_parser(
@@ -46,6 +60,13 @@ def main():
     evaluate_parser.add_argument("--run_name", type=str, default="run")
     evaluate_parser.add_argument("--log_dir", type=str, default="mlruns")
     evaluate_parser.add_argument("--data_path", type=str, default="data")
+    evaluate_parser.add_argument(
+        "--profiler",
+        type=str,
+        default=None,
+        choices=["simple", "advanced", "pytorch"],
+        help="Enable Lightning profiler: 'simple', 'advanced', or 'pytorch'.",
+    )
     evaluate_parser.add_argument(
         "--compute_dc_ac_metrics",
         action="store_true",
@@ -72,6 +93,13 @@ def main():
     predict_parser.add_argument("--log_dir", type=str, default="mlruns")
     predict_parser.add_argument("--data_path", type=str, default="data")
     predict_parser.add_argument("--output_path", type=str, default="data")
+    predict_parser.add_argument(
+        "--profiler",
+        type=str,
+        default=None,
+        choices=["simple", "advanced", "pytorch"],
+        help="Enable Lightning profiler: 'simple', 'advanced', or 'pytorch'.",
+    )
 
     args = parser.parse_args()
     main_cli(args)

gridfm_graphkit/cli.py

@@ -66,6 +66,8 @@ def main_cli(args):
         state_dict = torch.load(args.model_path, map_location="cpu")
         model.load_state_dict(state_dict)
 
+    profiler = getattr(args, "profiler", None)
+
     trainer = L.Trainer(
         logger=logger,
         accelerator=config_args.training.accelerator,
@@ -75,6 +77,7 @@ def main_cli(args):
         default_root_dir=args.log_dir,
         max_epochs=config_args.training.epochs,
         callbacks=get_training_callbacks(config_args),
+        profiler=profiler,
     )
     if args.command == "train" or args.command == "finetune":
         trainer.fit(model=model, datamodule=litGrid)
@@ -87,6 +90,7 @@ def main_cli(args):
             num_nodes=1,
             log_every_n_steps=1,
             default_root_dir=args.log_dir,
+            profiler=profiler,
         )
         test_trainer.test(model=model, datamodule=litGrid)
 
@@ -119,6 +123,7 @@ def main_cli(args):
             num_nodes=1,
             log_every_n_steps=1,
             default_root_dir=args.log_dir,
+            profiler=profiler,
         )
         predictions = predict_trainer.predict(model=model, datamodule=litGrid)