Deduplicate top level lighting CLI command groups (#15761)

* unify remove and delete command groups & the add and delete command groups
* added changelog
* fix tests
* Apply suggestions from code review

Co-authored-by: Jirka Borovec <[email protected]>
(cherry picked from commit 7b2788e)

Author: rlizzo

src/lightning_app/CHANGELOG.md

@@ -13,7 +13,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
--
+- `lightning add ssh-key` CLI command has been transitioned to `lightning create ssh-key` with the same calling signature ([#15761](https://github.com/Lightning-AI/lightning/pull/15761))
+- `lightning remove ssh-key` CLI command has been transitioned to `lightning delete ssh-key` with the same calling signature ([#15761](https://github.com/Lightning-AI/lightning/pull/15761))
 
 
 ### Fixed

src/lightning_app/cli/cmd_ssh_keys.py

@@ -19,7 +19,7 @@ def as_json(self) -> str:
         return json.dumps(self.ssh_keys)
 
     def as_table(self) -> Table:
-        table = Table("id", "public_key", "created", show_header=True)
+        table = Table("id", "public_key", "created", show_header=True, header_style="bold green")
         for ssh_key in self.ssh_keys:
             table.add_row(
                 ssh_key.id,

src/lightning_app/cli/lightning_cli.py

@@ -24,11 +24,9 @@
     disconnect,
 )
 from lightning_app.cli.commands.logs import logs
-from lightning_app.cli.lightning_cli_add import cli_add
 from lightning_app.cli.lightning_cli_create import create
 from lightning_app.cli.lightning_cli_delete import delete
 from lightning_app.cli.lightning_cli_list import get_list
-from lightning_app.cli.lightning_cli_remove import cli_remove
 from lightning_app.core.constants import DEBUG, ENABLE_APP_COMMENT_COMMAND_EXECUTION, get_lightning_cloud_url
 from lightning_app.runners.runtime import dispatch
 from lightning_app.runners.runtime_type import RuntimeType
@@ -381,8 +379,6 @@ def stop() -> None:
 _main.add_command(get_list)
 _main.add_command(delete)
 _main.add_command(create)
-_main.add_command(cli_add)
-_main.add_command(cli_remove)
 
 
 @_main.command("ssh")

src/lightning_app/cli/lightning_cli_add.py

@@ -1,8 +1,12 @@
-from typing import Any
+import os
+from pathlib import Path
+from typing import Any, Optional, Union
 
 import click
+from lightning_cloud.openapi.rest import ApiException
 
 from lightning_app.cli.cmd_clusters import _check_cluster_name_is_valid, AWSClusterManager
+from lightning_app.cli.cmd_ssh_keys import _SSHKeyManager
 
 
 @click.group("create")
@@ -77,3 +81,29 @@ def create_cluster(
         cost_savings=not enable_performance,
         wait=wait,
     )
+
+
+@create.command("ssh-key")
+@click.option("--name", "key_name", default=None, help="name of ssh key")
+@click.option("--comment", "comment", default="", help="comment detailing your SSH key")
+@click.option(
+    "--public-key",
+    "public_key",
+    help="public key or path to public key file",
+    required=True,
+)
+def add_ssh_key(
+    public_key: Union[str, "os.PathLike[str]"], key_name: Optional[str] = None, comment: Optional[str] = None
+) -> None:
+    """Add a new Lightning AI ssh-key to your account."""
+    ssh_key_manager = _SSHKeyManager()
+
+    new_public_key = Path(str(public_key)).read_text() if os.path.isfile(str(public_key)) else public_key
+    try:
+        ssh_key_manager.add_key(name=key_name, comment=comment, public_key=str(new_public_key))
+    except ApiException as e:
+        # if we got an exception it might be the user passed the private key file
+        if os.path.isfile(str(public_key)) and os.path.isfile(f"{public_key}.pub"):
+            ssh_key_manager.add_key(name=key_name, comment=comment, public_key=Path(f"{public_key}.pub").read_text())
+        else:
+            raise e

src/lightning_app/cli/lightning_cli_create.py

@@ -1,6 +1,7 @@
 import click
 
 from lightning_app.cli.cmd_clusters import AWSClusterManager
+from lightning_app.cli.cmd_ssh_keys import _SSHKeyManager
 
 
 @click.group("delete")
@@ -49,3 +50,11 @@ def delete_cluster(cluster: str, force: bool = False, wait: bool = False) -> Non
     """
     cluster_manager = AWSClusterManager()
     cluster_manager.delete(cluster_id=cluster, force=force, wait=wait)
+
+
+@delete.command("ssh-key")
+@click.argument("key_id")
+def remove_ssh_key(key_id: str) -> None:
+    """Delete a ssh-key from your Lightning AI account."""
+    ssh_key_manager = _SSHKeyManager()
+    ssh_key_manager.remove_key(key_id=key_id)

src/lightning_app/cli/lightning_cli_delete.py

@@ -61,8 +61,6 @@ def test_main_lightning_cli_no_arguments():
     assert "create  " in res
     assert "show    " in res
     assert "ssh     " in res
-    assert "add     " in res
-    assert "remove  " in res
 
 
 def test_main_lightning_cli_help():
@@ -76,8 +74,6 @@ def test_main_lightning_cli_help():
     assert "create  " in res
     assert "show    " in res
     assert "ssh     " in res
-    assert "add     " in res
-    assert "remove  " in res
 
     res = os.popen("lightning run --help").read()
     assert "app  " in res
@@ -97,6 +93,16 @@ def test_main_lightning_cli_help():
     res = os.popen("lightning show cluster --help").read()
     assert "logs " in res
 
+    # inspect create group
+    res = os.popen("lightning create --help").read()
+    assert "cluster " in res
+    assert "ssh-key " in res
+
+    # inspect delete group
+    res = os.popen("lightning delete --help").read()
+    assert "cluster " in res
+    assert "ssh-key " in res
+
 
 @mock.patch("lightning_cloud.login.Auth.authenticate", MagicMock())
 @mock.patch("lightning_app.cli.cmd_clusters.AWSClusterManager.create")