Put the --alloc/--salloc/--sbatch args last

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

milatools/cli/commands.py

@@ -204,7 +204,6 @@ def add_arguments(parser: argparse.ArgumentParser):
     code_parser = subparsers.add_parser(
         "code",
         help="Open a remote VSCode session on a compute node.",
-        formatter_class=SortingHelpFormatter,
     )
     code_parser.add_argument(
         "PATH",
@@ -359,7 +358,6 @@ def add_arguments(parser: argparse.ArgumentParser):
     serve_lab_parser = serve_subparsers.add_parser(
         "lab",
         help="Start a Jupyterlab server.",
-        formatter_class=SortingHelpFormatter,
     )
     serve_lab_parser.add_argument(
         "PATH",
@@ -375,7 +373,6 @@ def add_arguments(parser: argparse.ArgumentParser):
     serve_notebook_parser = serve_subparsers.add_parser(
         "notebook",
         help="Start a Jupyter Notebook server.",
-        formatter_class=SortingHelpFormatter,
     )
     serve_notebook_parser.add_argument(
         "PATH",
@@ -391,7 +388,6 @@ def add_arguments(parser: argparse.ArgumentParser):
     serve_tensorboard_parser = serve_subparsers.add_parser(
         "tensorboard",
         help="Start a Tensorboard server.",
-        formatter_class=SortingHelpFormatter,
     )
     serve_tensorboard_parser.add_argument(
         "LOGDIR", type=str, help="Path to the experiment logs"
@@ -404,7 +400,6 @@ def add_arguments(parser: argparse.ArgumentParser):
     serve_mlflow_parser = serve_subparsers.add_parser(
         "mlflow",
         help="Start an MLFlow server.",
-        formatter_class=SortingHelpFormatter,
     )
     serve_mlflow_parser.add_argument(
         "LOGDIR", type=str, help="Path to the experiment logs"
@@ -417,7 +412,6 @@ def add_arguments(parser: argparse.ArgumentParser):
     serve_aim_parser = serve_subparsers.add_parser(
         "aim",
         help="Start an AIM server.",
-        formatter_class=SortingHelpFormatter,
     )
     serve_aim_parser.add_argument(
         "LOGDIR", type=str, help="Path to the experiment logs"
@@ -936,34 +930,36 @@ def add_arguments(self, actions):
 
 
 def _add_allocation_options(parser: ArgumentParser):
+    # note: Ideally we'd like [--persist --alloc] | [--salloc] | [--sbatch] (i.e. a
+    # subgroup with alloc and persist within a mutually exclusive group with salloc and
+    # sbatch) but that doesn't seem possible with argparse as far as I can tell.
     arg_group = parser.add_argument_group(
         "Allocation options", description="Extra options to pass to slurm."
     )
+    alloc_group = arg_group.add_mutually_exclusive_group()
     common_kwargs = {
         "dest": "alloc",
         "nargs": argparse.REMAINDER,
         "action": AllocationFlagsAction,
         "metavar": "VALUE",
         "default": [],
     }
-    arg_group.add_argument(
-        "--alloc",
-        **common_kwargs,
-        help="Extra options to pass to salloc or to sbatch if --persist is set.",
-    )
-    alloc_group = arg_group.add_mutually_exclusive_group()
     alloc_group.add_argument(
         "--persist",
         action="store_true",
         help="Whether the server should persist or not when using --alloc",
     )
-
+    # --persist can be used with --alloc
+    arg_group.add_argument(
+        "--alloc",
+        **common_kwargs,
+        help="Extra options to pass to salloc or to sbatch if --persist is set.",
+    )
     # --persist cannot be used with --salloc or --sbatch.
     # Note: REMAINDER args like --alloc, --sbatch and --salloc are already mutually
     # exclusive in a sense, since it's only possible to use one correctly, the other
     # args are stored in the first one (e.g. mila code --alloc --salloc bob will have
     # alloc of ["--salloc", "bob"]).
-
     alloc_group.add_argument(
         "--salloc",
         **common_kwargs,
@@ -972,12 +968,11 @@ def _add_allocation_options(parser: ArgumentParser):
     alloc_group.add_argument(
         "--sbatch",
         **common_kwargs,
-        help="Extra options to pass to sbatch (equivalent to --persist --alloc [...])",
+        help="Extra options to pass to sbatch. Same as using --alloc with --persist.",
     )
 
 
 def _add_standard_server_args(parser: ArgumentParser):
-    _add_allocation_options(parser)
     parser.add_argument(
         "--job",
         type=int,
@@ -1013,6 +1008,8 @@ def _add_standard_server_args(parser: ArgumentParser):
         help="Name of the profile to use",
         metavar="VALUE",
     )
+    # Add these arguments last because we want them to show up last in the usage message
+    _add_allocation_options(parser)
 
 
 def _standard_server(

tests/cli/test_commands.py

@@ -63,7 +63,7 @@ def test_help(
     [
         "mila",  # Error: Missing a subcommand.
         "mila search conda",
-        "mila code",  # Error: Missing the required PATH argument.
+        "mila code --boo",  # Error: Unknown argument.
         "mila serve",  # Error: Missing the subcommand.
         "mila forward",  # Error: Missing the REMOTE argument.
     ],

tests/cli/test_commands/test_help_mila_code_.txt

@@ -1,18 +1,33 @@
 usage: mila code [-h] [--cluster {mila,cedar,narval,beluga,graham}]
-                 [--alloc ...] [--command VALUE] [--job VALUE] [--node VALUE]
-                 [--persist]
-                 PATH
+                 [--command VALUE] [--job JOB_ID] [--node NODE] [--persist]
+                 [--alloc ...] [--salloc ...] [--sbatch ...]
+                 [PATH]
 
 positional arguments:
-  PATH                  Path to open on the remote machine
+  PATH                  Path to open on the remote machine. Defaults to $HOME.
+                        Can be a relative or absolute path. When a relative
+                        path (that doesn't start with a '/', like foo/bar) is
+                        passed, the path is relative to the $HOME directory on
+                        the selected cluster. For example, foo/project will be
+                        interpreted as $HOME/foo/project.
 
 optional arguments:
   -h, --help            show this help message and exit
-  --alloc ...           Extra options to pass to slurm
   --cluster {mila,cedar,narval,beluga,graham}
                         Which cluster to connect to.
   --command VALUE       Command to use to start vscode (defaults to "code" or
                         the value of $MILATOOLS_CODE_COMMAND)
-  --job VALUE           Job ID to connect to
-  --node VALUE          Node to connect to
-  --persist             Whether the server should persist or not
+  --job JOB_ID          Job ID to connect to
+  --node NODE           Node to connect to
+
+Allocation optional arguments:
+  Extra options to pass to slurm.
+
+  --persist             Whether the server should persist or not when using
+                        --alloc
+  --alloc ...           Extra options to pass to salloc or to sbatch if
+                        --persist is set.
+  --salloc ...          Extra options to pass to salloc. Same as using --alloc
+                        without --persist.
+  --sbatch ...          Extra options to pass to sbatch. Same as using --alloc
+                        with --persist.

tests/cli/test_commands/test_help_mila_serve_aim_.txt

@@ -1,17 +1,26 @@
-usage: mila serve aim [-h] [--alloc ...] [--job VALUE] [--name VALUE]
-                      [--node VALUE] [--persist] [--port VALUE]
-                      [--profile VALUE]
+usage: mila serve aim [-h] [--job JOB_ID] [--name VALUE] [--node VALUE]
+                      [--port VALUE] [--profile VALUE] [--persist]
+                      [--alloc ...] [--salloc ...] [--sbatch ...]
                       LOGDIR
 
 positional arguments:
   LOGDIR           Path to the experiment logs
 
 optional arguments:
   -h, --help       show this help message and exit
-  --alloc ...      Extra options to pass to slurm
-  --job VALUE      Job ID to connect to
+  --job JOB_ID     Job ID to connect to
   --name VALUE     Name of the persistent server
   --node VALUE     Node to connect to
-  --persist        Whether the server should persist or not
   --port VALUE     Port to open on the local machine
   --profile VALUE  Name of the profile to use
+
+Allocation optional arguments:
+  Extra options to pass to slurm.
+
+  --persist        Whether the server should persist or not when using --alloc
+  --alloc ...      Extra options to pass to salloc or to sbatch if --persist
+                   is set.
+  --salloc ...     Extra options to pass to salloc. Same as using --alloc
+                   without --persist.
+  --sbatch ...     Extra options to pass to sbatch. Same as using --alloc with
+                   --persist.

tests/cli/test_commands/test_help_mila_serve_lab_.txt

@@ -1,17 +1,26 @@
-usage: mila serve lab [-h] [--alloc ...] [--job VALUE] [--name VALUE]
-                      [--node VALUE] [--persist] [--port VALUE]
-                      [--profile VALUE]
+usage: mila serve lab [-h] [--job JOB_ID] [--name VALUE] [--node VALUE]
+                      [--port VALUE] [--profile VALUE] [--persist]
+                      [--alloc ...] [--salloc ...] [--sbatch ...]
                       [PATH]
 
 positional arguments:
   PATH             Path to open on the remote machine
 
 optional arguments:
   -h, --help       show this help message and exit
-  --alloc ...      Extra options to pass to slurm
-  --job VALUE      Job ID to connect to
+  --job JOB_ID     Job ID to connect to
   --name VALUE     Name of the persistent server
   --node VALUE     Node to connect to
-  --persist        Whether the server should persist or not
   --port VALUE     Port to open on the local machine
   --profile VALUE  Name of the profile to use
+
+Allocation optional arguments:
+  Extra options to pass to slurm.
+
+  --persist        Whether the server should persist or not when using --alloc
+  --alloc ...      Extra options to pass to salloc or to sbatch if --persist
+                   is set.
+  --salloc ...     Extra options to pass to salloc. Same as using --alloc
+                   without --persist.
+  --sbatch ...     Extra options to pass to sbatch. Same as using --alloc with
+                   --persist.

tests/cli/test_commands/test_help_mila_serve_mlflow_.txt

@@ -1,17 +1,26 @@
-usage: mila serve mlflow [-h] [--alloc ...] [--job VALUE] [--name VALUE]
-                         [--node VALUE] [--persist] [--port VALUE]
-                         [--profile VALUE]
+usage: mila serve mlflow [-h] [--job JOB_ID] [--name VALUE] [--node VALUE]
+                         [--port VALUE] [--profile VALUE] [--persist]
+                         [--alloc ...] [--salloc ...] [--sbatch ...]
                          LOGDIR
 
 positional arguments:
   LOGDIR           Path to the experiment logs
 
 optional arguments:
   -h, --help       show this help message and exit
-  --alloc ...      Extra options to pass to slurm
-  --job VALUE      Job ID to connect to
+  --job JOB_ID     Job ID to connect to
   --name VALUE     Name of the persistent server
   --node VALUE     Node to connect to
-  --persist        Whether the server should persist or not
   --port VALUE     Port to open on the local machine
   --profile VALUE  Name of the profile to use
+
+Allocation optional arguments:
+  Extra options to pass to slurm.
+
+  --persist        Whether the server should persist or not when using --alloc
+  --alloc ...      Extra options to pass to salloc or to sbatch if --persist
+                   is set.
+  --salloc ...     Extra options to pass to salloc. Same as using --alloc
+                   without --persist.
+  --sbatch ...     Extra options to pass to sbatch. Same as using --alloc with
+                   --persist.

tests/cli/test_commands/test_help_mila_serve_notebook_.txt

@@ -1,17 +1,26 @@
-usage: mila serve notebook [-h] [--alloc ...] [--job VALUE] [--name VALUE]
-                           [--node VALUE] [--persist] [--port VALUE]
-                           [--profile VALUE]
+usage: mila serve notebook [-h] [--job JOB_ID] [--name VALUE] [--node VALUE]
+                           [--port VALUE] [--profile VALUE] [--persist]
+                           [--alloc ...] [--salloc ...] [--sbatch ...]
                            [PATH]
 
 positional arguments:
   PATH             Path to open on the remote machine
 
 optional arguments:
   -h, --help       show this help message and exit
-  --alloc ...      Extra options to pass to slurm
-  --job VALUE      Job ID to connect to
+  --job JOB_ID     Job ID to connect to
   --name VALUE     Name of the persistent server
   --node VALUE     Node to connect to
-  --persist        Whether the server should persist or not
   --port VALUE     Port to open on the local machine
   --profile VALUE  Name of the profile to use
+
+Allocation optional arguments:
+  Extra options to pass to slurm.
+
+  --persist        Whether the server should persist or not when using --alloc
+  --alloc ...      Extra options to pass to salloc or to sbatch if --persist
+                   is set.
+  --salloc ...     Extra options to pass to salloc. Same as using --alloc
+                   without --persist.
+  --sbatch ...     Extra options to pass to sbatch. Same as using --alloc with
+                   --persist.

tests/cli/test_commands/test_help_mila_serve_tensorboard_.txt

@@ -1,17 +1,27 @@
-usage: mila serve tensorboard [-h] [--alloc ...] [--job VALUE] [--name VALUE]
-                              [--node VALUE] [--persist] [--port VALUE]
-                              [--profile VALUE]
+usage: mila serve tensorboard [-h] [--job JOB_ID] [--name VALUE]
+                              [--node VALUE] [--port VALUE] [--profile VALUE]
+                              [--persist] [--alloc ...] [--salloc ...]
+                              [--sbatch ...]
                               LOGDIR
 
 positional arguments:
   LOGDIR           Path to the experiment logs
 
 optional arguments:
   -h, --help       show this help message and exit
-  --alloc ...      Extra options to pass to slurm
-  --job VALUE      Job ID to connect to
+  --job JOB_ID     Job ID to connect to
   --name VALUE     Name of the persistent server
   --node VALUE     Node to connect to
-  --persist        Whether the server should persist or not
   --port VALUE     Port to open on the local machine
   --profile VALUE  Name of the profile to use
+
+Allocation optional arguments:
+  Extra options to pass to slurm.
+
+  --persist        Whether the server should persist or not when using --alloc
+  --alloc ...      Extra options to pass to salloc or to sbatch if --persist
+                   is set.
+  --salloc ...     Extra options to pass to salloc. Same as using --alloc
+                   without --persist.
+  --sbatch ...     Extra options to pass to sbatch. Same as using --alloc with
+                   --persist.