Optional --output_path arg added to evaluate subcommand

README.md

@@ -243,6 +243,7 @@ gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.p
 | `--profiler` | `str` | Enable Lightning profiler (`simple`, `advanced`, `pytorch`). | `None` |
 | `--compute_dc_ac_metrics` | `flag` | Compute ground-truth AC/DC power balance metrics on the test split. | `False` |
 | `--save_output` | `flag` | Save predictions under MLflow artifacts (`.../artifacts/test`). For the PowerFlow task this writes `<grid_name>_predictions.parquet` (bus-level) and `<grid_name>_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `False` |
+| `--output_path` | `str` | Optional custom output directory to save evaluation predictions (when used with `--save_output`). Defaults to MLflow artifacts directory (`artifacts/test`). | `None` |
 | `--mp_context` | `str` | DataLoader multiprocessing start method (`spawn`, `fork`, `forkserver`). Defaults to PyTorch's automatic choice. On Linux, `spawn` is recommended for safety (CUDA + fork is unsafe); other choices emit a warning. | `None` |
 
 ### Example with saved normalizer stats

docs/quick_start/quick_start.md

@@ -109,6 +109,7 @@ gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.p
 | `--profiler` | `str` | Enable Lightning profiler (`simple`, `advanced`, `pytorch`). | `None` |
 | `--compute_dc_ac_metrics` | `flag` | Compute ground-truth AC/DC power balance metrics on the test split. | `False` |
 | `--save_output` | `flag` | Save predictions under MLflow artifacts (`.../artifacts/test`). For the PowerFlow task this writes `<grid_name>_predictions.parquet` (bus-level) and `<grid_name>_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `False` |
+| `--output_path` | `str` | Optional custom output directory to save evaluation predictions (when used with `--save_output`). Defaults to MLflow artifacts directory (`artifacts/test`). | `None` |
 | `--mp_context` | `str` | DataLoader multiprocessing start method (`spawn`, `fork`, `forkserver`). Defaults to PyTorch's automatic choice. On Linux, `spawn` is recommended for safety (CUDA + fork is unsafe); other choices emit a warning. | `None` |
 
 ### Example with saved normalizer stats

gridfm_graphkit/__main__.py

@@ -376,6 +376,12 @@ def main():
         "--converted_config_path",
         **_converted_config_path_kwargs,
     )
+    evaluate_parser.add_argument(
+        "--output_path",
+        type=str,
+        default=None,
+        help="Optional custom output directory to save evaluation predictions. Defaults to MLflow artifacts/test directory.",
+    )
     evaluate_parser.add_argument("--stream-partitions", **_stream_partitions_kwargs)
 
     # ---- PREDICT SUBCOMMAND ----

gridfm_graphkit/cli.py

@@ -438,7 +438,10 @@ def main_cli(args):
         if args.command == "predict":
             output_dir = args.output_path
         else:
-            output_dir = os.path.join(artifacts_dir, "test")
+            output_dir = getattr(args, "output_path", None) or os.path.join(
+                artifacts_dir,
+                "test",
+            )
         os.makedirs(output_dir, exist_ok=True)
         first_prediction = predictions[0]
         if any(isinstance(value, dict) for value in first_prediction.values()):