lightning-sdk 2025.9.23__py3-none-any.whl → 2025.9.30__py3-none-any.whl

lightning_sdk/__init__.py

@@ -9,7 +9,7 @@ from lightning_sdk.mmt import MMT
 from lightning_sdk.organization import Organization
 from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
 from lightning_sdk.status import Status
-from lightning_sdk.studio import Studio
+from lightning_sdk.studio import VM, Studio
 from lightning_sdk.teamspace import ConnectionType, FolderLocation, Teamspace
 from lightning_sdk.user import User
 
@@ -32,8 +32,9 @@ __all__ = [
     "Studio",
     "Teamspace",
     "User",
+    "VM",
 ]
 
-__version__ = "2025.09.23"
+__version__ = "2025.09.30"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/cli/entrypoint.py

@@ -21,6 +21,7 @@ from lightning_sdk.cli.groups import (
     # job,
     # mmt,
     studio,
+    vm,
 )
 from lightning_sdk.cli.utils import CustomHelpFormatter, rich_to_str
 from lightning_sdk.constants import _LIGHTNING_DEBUG
@@ -40,7 +41,7 @@ def _notify_exception(exception_type: Type[BaseException], value: BaseException,
     if _LIGHTNING_DEBUG:
         tb_text = "".join(traceback.format_exception(exception_type, value, tb))
         renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
-        renderables.append(Syntax(tb_text, "python", theme="monokai", line_numbers=False, word_wrap=True))
+        renderables.append(Syntax(tb_text, "python", theme="monokai light", line_numbers=False, word_wrap=True))
     else:
         renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
 
@@ -83,6 +84,7 @@ main_cli.add_command(config)
 # main_cli.add_command(job)
 # main_cli.add_command(mmt)
 main_cli.add_command(studio)
+main_cli.add_command(vm)
 if os.environ.get("LIGHTNING_EXPERIMENTAL_CLI_ONLY", "0") != "1":
     #### LEGACY COMMANDS ####
     # these commands are currently supported for backwards compatibility, but will potentially be removed in the future.

lightning_sdk/cli/groups.py

@@ -6,6 +6,7 @@ from lightning_sdk.cli.config import register_commands as register_config_comman
 from lightning_sdk.cli.job import register_commands as register_job_commands
 from lightning_sdk.cli.mmt import register_commands as register_mmt_commands
 from lightning_sdk.cli.studio import register_commands as register_studio_commands
+from lightning_sdk.cli.vm import register_commands as register_vm_commands
 
 
 @click.group(name="studio")
@@ -25,7 +26,12 @@ def mmt() -> None:
 
 @click.group(name="config")
 def config() -> None:
-    """Manage Lightning SDK and CLIconfiguration."""
+    """Manage Lightning SDK and CLI configuration."""
+
+
+@click.group(name="vm")
+def vm() -> None:
+    """Manage Lightning AI VMs."""
 
 
 # Register config commands with the main config group
@@ -33,3 +39,4 @@ register_job_commands(job)
 register_mmt_commands(mmt)
 register_studio_commands(studio)
 register_config_commands(config)
+register_vm_commands(vm)

lightning_sdk/cli/legacy/download.py

@@ -14,7 +14,7 @@ from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.models import download_model
 from lightning_sdk.studio import Studio
-from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
+from lightning_sdk.utils.resolve import _get_authed_user
 
 
 def _expand_remote_path(path: str) -> str:
@@ -242,7 +242,7 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
                     and (re.match(studio_name, st["name"]) or studio_name in st["name"]),
                     possible_studios,
                 )
-                if not possible_studios:
+                if not len(possible_studios):
                     raise ValueError(
                         f"Could not find Studio like '{studio}', please consider update your filtering pattern."
                     )
@@ -260,8 +260,7 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
             "Please contact Lightning AI directly to resolve this issue."
         ) from e
 
-    with skip_studio_init():
-        return Studio(**selected_studio)
+    return Studio(**selected_studio)
 
 
 @download.command(name="licenses")

lightning_sdk/cli/legacy/entrypoint.py

@@ -48,7 +48,7 @@ def _notify_exception(exception_type: Type[BaseException], value: BaseException,
     if _LIGHTNING_DEBUG:
         tb_text = "".join(traceback.format_exception(exception_type, value, tb))
         renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
-        renderables.append(Syntax(tb_text, "python", theme="monokai", line_numbers=False, word_wrap=True))
+        renderables.append(Syntax(tb_text, "python", theme="monokai light", line_numbers=False, word_wrap=True))
     else:
         renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
 

lightning_sdk/cli/legacy/upload.py

@@ -20,7 +20,7 @@ from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.models import upload_model as _upload_model
 from lightning_sdk.studio import Studio
-from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
+from lightning_sdk.utils.resolve import _get_authed_user
 
 _STUDIO_UPLOAD_STATUS_PATH = "~/.lightning/studios/uploads"
 
@@ -285,8 +285,7 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
             "Please contact Lightning AI directly to resolve this issue."
         ) from e
 
-    with skip_studio_init():
-        return Studio(**selected_studio)
+    return Studio(**selected_studio)
 
 
 def _print_docker_push(lines: Generator, console: Console, progress: Progress, push_task: rich.progress.TaskID) -> None:

lightning_sdk/cli/studio/__init__.py

@@ -5,6 +5,7 @@ import click
 
 def register_commands(group: click.Group) -> None:
     """Register studio commands with the given group."""
+    from lightning_sdk.cli.studio.connect import connect_studio
     from lightning_sdk.cli.studio.create import create_studio
     from lightning_sdk.cli.studio.delete import delete_studio
     from lightning_sdk.cli.studio.list import list_studios
@@ -20,3 +21,4 @@ def register_commands(group: click.Group) -> None:
     group.add_command(start_studio)
     group.add_command(stop_studio)
     group.add_command(switch_studio)
+    group.add_command(connect_studio)

lightning_sdk/cli/studio/connect.py

@@ -0,0 +1,94 @@
+"""Studio connect command."""
+
+import subprocess
+import sys
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config, save_teamspace_to_config
+from lightning_sdk.cli.utils.ssh_connection import download_ssh_keys
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.names import random_unique_name
+
+
+@click.command("connect")
+@click.argument("name", required=False)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--cloud-provider",
+    help="The cloud provider to start the studio on. Defaults to teamspace default.",
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+@click.option(
+    "--cloud-account",
+    help="The cloud account to create the studio on. Defaults to teamspace default.",
+    type=click.STRING,
+)
+@click.option("--gpus", help="The number of GPUs to start the studio on. ", type=click.INT)
+def connect_studio(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    gpus: Optional[int] = None,
+) -> None:
+    """Connect to a Studio.
+
+    Example:
+        lightning studio connect
+    """
+    menu = TeamspacesMenu()
+
+    resolved_teamspace = menu(teamspace)
+    save_teamspace_to_config(resolved_teamspace, overwrite=False)
+
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
+
+    name = name or random_unique_name()
+
+    try:
+        studio = Studio(
+            name=name,
+            teamspace=resolved_teamspace,
+            create_ok=True,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+        )
+    except (RuntimeError, ValueError, ApiException):
+        raise ValueError(f"Could not create Studio: '{name}'") from None
+
+    click.echo(f"Connecting to Studio '{studio_name_link(studio)}' ...")
+
+    machine = "CPU"
+    Studio.show_progress = True
+    if gpus:
+        # TODO: handle something like gpus=4:L4
+        pass
+
+    save_studio_to_config(studio)
+    studio.start(machine, interruptible=True)
+
+    ssh_private_key_path = _configure_ssh_internal()
+
+    try:
+        ssh_command = (
+            f"ssh -i {ssh_private_key_path} -o UserKnownHostsFile=/dev/null s_{studio._studio.id}@ssh.lightning.ai"
+        )
+        subprocess.run(ssh_command.split())
+    except Exception as ex:
+        print(f"Failed to establish SSH connection: {ex}")
+        sys.exit(1)
+
+
+def _configure_ssh_internal() -> str:
+    """Internal function to configure SSH without Click decorators."""
+    auth = Auth()
+    auth.authenticate()
+    return download_ssh_keys(auth.api_key, force_download=False)

lightning_sdk/cli/studio/create.py

@@ -9,7 +9,7 @@ from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.machine import CloudProvider
-from lightning_sdk.studio import Studio
+from lightning_sdk.studio import VM, Studio
 
 
 @click.command("create")
@@ -36,6 +36,16 @@ def create_studio(
     Example:
         lightning studio create
     """
+    create_impl(name=name, teamspace=teamspace, cloud_provider=cloud_provider, cloud_account=cloud_account, vm=False)
+
+
+def create_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+    cloud_provider: Optional[str],
+    cloud_account: Optional[str],
+    vm: bool,
+) -> None:
     menu = TeamspacesMenu()
 
     resolved_teamspace = menu(teamspace)
@@ -44,8 +54,12 @@ def create_studio(
     if cloud_provider is not None:
         cloud_provider = CloudProvider(cloud_provider)
 
+    create_cls = VM if vm else Studio
+    cls_name = create_cls.__qualname__
+
     try:
-        studio = Studio(
+        create_cls = VM if vm else Studio
+        studio = create_cls(
             name=name,
             teamspace=resolved_teamspace,
             create_ok=True,
@@ -54,7 +68,7 @@ def create_studio(
         )
     except (RuntimeError, ValueError, ApiException):
         if name:
-            raise ValueError(f"Could not create Studio: '{name}'. Does the Studio exist?") from None
-        raise ValueError(f"Could not create Studio: '{name}'. Please provide a Studio name") from None
+            raise ValueError(f"Could not create {cls_name}: '{name}'. Does the {cls_name} exist?") from None
+        raise ValueError(f"Could not create {cls_name}: '{name}'. Please provide a {cls_name} name") from None
 
-    click.echo(f"Studio {studio_name_link(studio)} created successfully")
+    click.echo(f"{cls_name} {studio_name_link(studio)} created successfully")

lightning_sdk/cli/studio/delete.py

@@ -12,7 +12,7 @@ from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 @click.option(
     "--name",
     help=(
-        "The name of the studio to start. "
+        "The name of the studio to delete. "
         "If not provided, will try to infer from environment, "
         "use the default value from the config or prompt for interactive selection."
     ),
@@ -25,21 +25,25 @@ def delete_studio(name: Optional[str] = None, teamspace: Optional[str] = None) -
       lightning studio delete --name my-studio
 
     """
+    return delete_impl(name=name, teamspace=teamspace, vm=False)
+
+
+def delete_impl(name: Optional[str], teamspace: Optional[str], vm: bool) -> None:
     menu = TeamspacesMenu()
     resolved_teamspace = menu(teamspace=teamspace)
 
-    menu = StudiosMenu(resolved_teamspace)
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
     studio = menu(studio=name)
 
     studio_name = f"{studio.teamspace.owner.name}/{studio.teamspace.name}/{studio.name}"
     confirmed = click.confirm(
-        f"Are you sure you want to delete studio '{studio_name}'?",
+        f"Are you sure you want to delete {studio._cls_name} '{studio_name}'?",
         abort=True,
     )
     if not confirmed:
-        click.echo("Studio deletion cancelled")
+        click.echo(f"{studio._cls_name} deletion cancelled")
         return
 
     studio.delete()
 
-    click.echo(f"Studio '{studio.name}' deleted successfully")
+    click.echo(f"{studio._cls_name} '{studio.name}' deleted successfully")

lightning_sdk/cli/studio/list.py

@@ -35,13 +35,17 @@ def list_studios(teamspace: Optional[str] = None, all: bool = False, sort_by: Op
         lightning studio list --teamspace owner/teamspace
 
     """
+    return list_impl(teamspace=teamspace, all=all, sort_by=sort_by, vm=False)
+
+
+def list_impl(teamspace: Optional[str], all: bool, sort_by: Optional[str], vm: bool) -> None:  # noqa: A002
     menu = TeamspacesMenu()
     teamspace_resolved = menu(teamspace=teamspace)
     save_teamspace_to_config(teamspace_resolved, overwrite=False)
 
     user = _get_authed_user()
 
-    studios = teamspace_resolved.studios
+    studios = teamspace_resolved.vms if vm else teamspace_resolved.studios
 
     table = Table(
         pad_edge=True,

lightning_sdk/cli/studio/ssh.py

@@ -1,26 +1,22 @@
 """Studio SSH command."""
 
-import os
-import platform
 import subprocess
-import uuid
-from pathlib import Path
 from typing import List, Optional
 
 import click
 
 from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.ssh_connection import download_ssh_keys
 from lightning_sdk.cli.utils.studio_selection import StudiosMenu
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
 
 
 @click.command("ssh")
 @click.option(
     "--name",
     help=(
-        "The name of the studio to start. "
+        "The name of the studio to ssh into. "
         "If not provided, will try to infer from environment, "
         "use the default value from the config or prompt for interactive selection."
     ),
@@ -39,15 +35,21 @@ def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, opti
     Example:
         lightning studio ssh --name my-studio
     """
+    return ssh_impl(name=name, teamspace=teamspace, option=option, vm=False)
+
+
+def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[List[str]], vm: bool) -> None:
     auth = Auth()
     auth.authenticate()
-    ssh_private_key_path = _download_ssh_keys(auth.api_key, force_download=False)
+    ssh_private_key_path = download_ssh_keys(auth.api_key, force_download=False)
 
     menu = TeamspacesMenu()
     resolved_teamspace = menu(teamspace=teamspace)
 
-    menu = StudiosMenu(resolved_teamspace)
-    studio = menu(studio=name)
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
+    studio = menu(
+        studio=name,
+    )
     save_studio_to_config(studio)
 
     ssh_options = " -o " + " -o ".join(option) if option else ""
@@ -57,53 +59,9 @@ def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, opti
         subprocess.run(ssh_command.split())
     except Exception:
         # redownload the keys to be sure they are up to date
-        _download_ssh_keys(auth.api_key, force_download=True)
+        download_ssh_keys(auth.api_key, force_download=True)
         try:
             subprocess.run(ssh_command.split())
         except Exception:
             # TODO: make this a generic CLI error
             raise RuntimeError("Failed to establish SSH connection") from None
-
-
-def _download_ssh_keys(
-    api_key: str,
-    force_download: bool = False,
-    ssh_key_name: str = "lightning_rsa",
-) -> None:
-    """Download the SSH key for a User."""
-    ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name)
-
-    os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True)
-
-    if not os.path.isfile(ssh_private_key_path) or force_download:
-        key_id = str(uuid.uuid4())
-        _download_file(
-            f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
-            ssh_private_key_path,
-            overwrite=True,
-            chmod=0o600,
-        )
-        _download_file(
-            f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}",
-            ssh_private_key_path + ".pub",
-            overwrite=True,
-        )
-
-    return ssh_private_key_path
-
-
-def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
-    """Download a file from a URL."""
-    import requests
-
-    if os.path.isfile(local_path) and not overwrite:
-        raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
-
-    response = requests.get(url, stream=True)
-    response.raise_for_status()
-
-    with open(local_path, "wb") as file:
-        for chunk in response.iter_content(chunk_size=8192):
-            file.write(chunk)
-    if chmod is not None:
-        os.chmod(local_path, 0o600)

lightning_sdk/cli/studio/start.py

@@ -9,7 +9,7 @@ from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
 from lightning_sdk.cli.utils.studio_selection import StudiosMenu
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import CloudProvider, Machine
-from lightning_sdk.studio import Studio
+from lightning_sdk.studio import VM, Studio
 
 
 @click.command("start")
@@ -57,6 +57,28 @@ def start_studio(
         lightning studio start --name my-studio
 
     """
+    return start_impl(
+        name=name,
+        teamspace=teamspace,
+        create=create,
+        machine=machine,
+        interruptible=interruptible,
+        cloud_provider=cloud_provider,
+        cloud_account=cloud_account,
+        vm=False,
+    )
+
+
+def start_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+    create: bool,
+    machine: str,
+    interruptible: bool,
+    cloud_provider: Optional[str],
+    cloud_account: Optional[str],
+    vm: bool,
+) -> None:
     menu = TeamspacesMenu()
     resolved_teamspace = menu(teamspace=teamspace)
 
@@ -64,10 +86,11 @@ def start_studio(
         cloud_provider = CloudProvider(cloud_provider)
 
     if not create:
-        menu = StudiosMenu(resolved_teamspace)
+        menu = StudiosMenu(resolved_teamspace, vm=vm)
         studio = menu(studio=name)
     else:
-        studio = Studio(
+        create_cls = VM if vm else Studio
+        studio = create_cls(
             name=name,
             teamspace=resolved_teamspace,
             create_ok=create,

lightning_sdk/cli/studio/stop.py

@@ -14,7 +14,7 @@ from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 @click.option(
     "--name",
     help=(
-        "The name of the studio to start. "
+        "The name of the studio to stop. "
         "If not provided, will try to infer from environment, "
         "use the default value from the config or prompt for interactive selection."
     ),
@@ -27,15 +27,19 @@ def stop_studio(name: Optional[str] = None, teamspace: Optional[str] = None) ->
         lightning studio stop --name my-studio
 
     """
+    return stop_impl(name=name, teamspace=teamspace, vm=False)
+
+
+def stop_impl(name: Optional[str], teamspace: Optional[str], vm: bool) -> None:
     # missing studio_name and teamspace are handled by the studio class
     menu = TeamspacesMenu()
     resolved_teamspace = menu(teamspace=teamspace)
 
-    menu = StudiosMenu(resolved_teamspace)
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
     studio = menu(studio=name)
 
     studio.stop()
 
     save_studio_to_config(studio)
 
-    click.echo(f"Studio {studio_name_link(studio)} stopped successfully")
+    click.echo(f"{studio._cls_name} {studio_name_link(studio)} stopped successfully")

lightning_sdk/cli/studio/switch.py

@@ -9,14 +9,13 @@ from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
 from lightning_sdk.cli.utils.studio_selection import StudiosMenu
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import Machine
-from lightning_sdk.studio import Studio
 
 
 @click.command("switch")
 @click.option(
     "--name",
     help=(
-        "The name of the studio to start. "
+        "The name of the studio to switch to a different machine. "
         "If not provided, will try to infer from environment, "
         "use the default value from the config or prompt for interactive selection."
     ),
@@ -35,16 +34,33 @@ def switch_studio(
     interruptible: bool = False,
 ) -> None:
     """Switch a Studio to a different machine type."""
+    return switch_impl(
+        name=name,
+        teamspace=teamspace,
+        machine=machine,
+        interruptible=interruptible,
+        vm=False,
+    )
+
+
+def switch_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+    machine: Optional[str],
+    interruptible: bool,
+    vm: bool,
+) -> None:
     menu = TeamspacesMenu()
     resolved_teamspace = menu(teamspace=teamspace)
 
-    menu = StudiosMenu(resolved_teamspace)
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
     studio = menu(studio=name)
 
     resolved_machine = Machine.from_str(machine)
-    Studio.show_progress = True
+
+    studio.__class__.show_progress = True
     studio.switch_machine(resolved_machine, interruptible=interruptible)
 
     save_studio_to_config(studio)
 
-    click.echo(f"Studio {studio_name_link(studio)} switched to machine '{resolved_machine}' successfully")
+    click.echo(f"{studio._cls_name} {studio_name_link(studio)} switched to machine '{resolved_machine}' successfully")

lightning_sdk/cli/utils/ssh_connection.py

@@ -0,0 +1,51 @@
+import os
+import platform
+import uuid
+from pathlib import Path
+from typing import Optional
+
+from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
+
+
+def download_ssh_keys(
+    api_key: str | None,
+    force_download: bool = False,
+    ssh_key_name: str = "lightning_rsa",
+) -> str:
+    """Download the SSH key for a User."""
+    ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name)
+
+    os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True)
+
+    if not os.path.isfile(ssh_private_key_path) or force_download:
+        key_id = str(uuid.uuid4())
+        download_file(
+            f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
+            Path(ssh_private_key_path),
+            overwrite=True,
+            chmod=0o600,
+        )
+        download_file(
+            f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}",
+            Path(ssh_private_key_path + ".pub"),
+            overwrite=True,
+        )
+
+    return ssh_private_key_path
+
+
+def download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
+    """Download a file from a URL."""
+    import requests
+
+    if os.path.isfile(local_path) and not overwrite:
+        raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
+
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    with open(local_path, "wb") as file:
+        for chunk in response.iter_content(chunk_size=8192):
+            file.write(chunk)
+    if chmod is not None:
+        os.chmod(local_path, 0o600)