diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py index 8a465c02..62aa7362 100644 --- a/milatools/cli/commands.py +++ b/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,9 +930,13 @@ 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, @@ -946,24 +944,22 @@ def _add_allocation_options(parser: ArgumentParser): "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( diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index 98e3a40a..b833664a 100644 --- a/tests/cli/test_commands.py +++ b/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. ], diff --git a/tests/cli/test_commands/test_help_mila_code_.txt b/tests/cli/test_commands/test_help_mila_code_.txt index 65f2565a..13c572bc 100644 --- a/tests/cli/test_commands/test_help_mila_code_.txt +++ b/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. diff --git a/tests/cli/test_commands/test_help_mila_serve_aim_.txt b/tests/cli/test_commands/test_help_mila_serve_aim_.txt index b3bec903..0c6b2498 100644 --- a/tests/cli/test_commands/test_help_mila_serve_aim_.txt +++ b/tests/cli/test_commands/test_help_mila_serve_aim_.txt @@ -1,6 +1,6 @@ -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: @@ -8,10 +8,19 @@ positional arguments: 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. diff --git a/tests/cli/test_commands/test_help_mila_serve_lab_.txt b/tests/cli/test_commands/test_help_mila_serve_lab_.txt index 344f13dd..a9854dc4 100644 --- a/tests/cli/test_commands/test_help_mila_serve_lab_.txt +++ b/tests/cli/test_commands/test_help_mila_serve_lab_.txt @@ -1,6 +1,6 @@ -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: @@ -8,10 +8,19 @@ positional arguments: 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. diff --git a/tests/cli/test_commands/test_help_mila_serve_mlflow_.txt b/tests/cli/test_commands/test_help_mila_serve_mlflow_.txt index 6e6c5444..f7bd1f37 100644 --- a/tests/cli/test_commands/test_help_mila_serve_mlflow_.txt +++ b/tests/cli/test_commands/test_help_mila_serve_mlflow_.txt @@ -1,6 +1,6 @@ -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: @@ -8,10 +8,19 @@ positional arguments: 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. diff --git a/tests/cli/test_commands/test_help_mila_serve_notebook_.txt b/tests/cli/test_commands/test_help_mila_serve_notebook_.txt index b1067ded..e3bb24bc 100644 --- a/tests/cli/test_commands/test_help_mila_serve_notebook_.txt +++ b/tests/cli/test_commands/test_help_mila_serve_notebook_.txt @@ -1,6 +1,6 @@ -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: @@ -8,10 +8,19 @@ positional arguments: 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. diff --git a/tests/cli/test_commands/test_help_mila_serve_tensorboard_.txt b/tests/cli/test_commands/test_help_mila_serve_tensorboard_.txt index 77575c2f..ce5b5800 100644 --- a/tests/cli/test_commands/test_help_mila_serve_tensorboard_.txt +++ b/tests/cli/test_commands/test_help_mila_serve_tensorboard_.txt @@ -1,6 +1,7 @@ -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: @@ -8,10 +9,19 @@ positional arguments: 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.