lightning-sdk 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

lightning_sdk/cli/{serve.py → deploy/serve.py}

@@ -1,15 +1,11 @@
-import concurrent.futures
 import os
 import socket
 import subprocess
-import time
 import webbrowser
 from datetime import datetime
-from enum import Enum
 from pathlib import Path
 from threading import Thread
-from typing import Dict, List, Optional, TypedDict, Union
-from urllib.parse import urlencode
+from typing import Optional, Union
 
 import click
 from rich.console import Console
@@ -17,26 +13,19 @@ from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.prompt import Confirm
 
 from lightning_sdk import Machine, Teamspace
-from lightning_sdk.api import UserApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.api.utils import _get_registry_url
-from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
-from lightning_sdk.cli.upload import (
-    _dump_current_upload_state,
-    _resolve_previous_upload_state,
-    _start_parallel_upload,
+from lightning_sdk.cli.deploy._auth import (
+    _AuthMode,
+    _Onboarding,
+    authenticate,
+    poll_verified_status,
+    select_teamspace,
 )
-from lightning_sdk.lightning_cloud import env
-from lightning_sdk.lightning_cloud.login import Auth, AuthServer
-from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
-from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.cli.deploy.devbox import _handle_devbox
 from lightning_sdk.serve import _LitServeDeployer
-from lightning_sdk.studio import Studio
-from lightning_sdk.utils.resolve import _get_authed_user, _get_studio_url, _resolve_teamspace
 
 _MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
-_POLL_TIMEOUT = 600
-LITSERVE_CODE = os.environ.get("LITSERVE_CODE", "j39bzk903h")
 
 
 class _ServeGroup(click.Group):
@@ -53,12 +42,12 @@ def deploy() -> None:
     """Deploy a LitServe model.
 
     Example:
-        lightning deploy server.py  # deploy to the cloud
+        lightning deploy server.py --cloud # deploy to the cloud
 
     Example:
-        lightning deploy server.py --local  # run locally
+        lightning deploy server.py  # run locally
 
-    You can deploy the API to the cloud by running `lightning deploy server.py`.
+    You can deploy the API to the cloud by running `lightning deploy server.py --cloud`.
     This will build a docker container for the server.py script and deploy it to the Lightning AI platform.
     """
 
@@ -73,11 +62,11 @@ def deploy() -> None:
     help="Generate a client for the model",
 )
 @click.option(
-    "--local",
+    "--cloud",
     is_flag=True,
     default=False,
     flag_value=True,
-    help="Run the model locally",
+    help="Run the model on cloud",
 )
 @click.option("--name", default=None, help="Name of the deployed API (e.g., 'classification-api', 'Llama-api')")
 @click.option(
@@ -133,7 +122,7 @@ def deploy() -> None:
 @click.option("--port", default=8000, help="The port to expose the API on.")
 @click.option("--min_replica", "--min-replica", default=0, help="Number of replicas to start with.")
 @click.option("--max_replica", "--max-replica", default=1, help="Number of replicas to scale up to.")
-@click.option("--replicas", "--replicas", default=1, help="Deployment will start with this many replicas.")
+@click.option("--replicas", default=1, help="Deployment will start with this many replicas.")
 @click.option(
     "--no_credentials",
     "--no-credentials",
@@ -145,7 +134,7 @@ def deploy() -> None:
 def api(
     script_path: str,
     easy: bool,
-    local: bool,
+    cloud: bool,
     name: Optional[str],
     non_interactive: bool,
     machine: Optional[str],
@@ -165,7 +154,7 @@ def api(
     return api_impl(
         script_path=script_path,
         easy=easy,
-        local=local,
+        cloud=cloud,
         name=name,
         non_interactive=non_interactive,
         machine=machine,
@@ -186,7 +175,7 @@ def api(
 def api_impl(
     script_path: Union[str, Path],
     easy: bool = False,
-    local: bool = False,
+    cloud: bool = False,
     name: Optional[str] = None,
     tag: Optional[str] = None,
     non_interactive: bool = False,
@@ -217,7 +206,7 @@ def api_impl(
         timestr = datetime.now().strftime("%b-%d-%H_%M")
         name = f"litserve-{timestr}".lower()
 
-    if local:
+    if not cloud and not devbox:
         try:
             subprocess.run(
                 ["python", str(script_path)],
@@ -230,7 +219,8 @@ def api_impl(
             raise RuntimeError(error_msg) from None
 
     if devbox:
-        return _handle_devbox(name, script_path, console, non_interactive, devbox, interruptible, teamspace, org, user)
+        machine = Machine.from_str(devbox)
+        return _handle_devbox(name, script_path, console, non_interactive, machine, interruptible, teamspace, org, user)
 
     machine = Machine.from_str(machine)
     return _handle_cloud(
@@ -253,165 +243,6 @@ def api_impl(
     )
 
 
-class _AuthServer(AuthServer):
-    def get_auth_url(self, port: int) -> str:
-        redirect_uri = f"http://localhost:{port}/login-complete"
-        params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITSERVE_CODE})
-        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
-
-
-class _Auth(Auth):
-    def __init__(self, shall_confirm: bool = False) -> None:
-        super().__init__()
-        self._shall_confirm = shall_confirm
-
-    def _run_server(self) -> None:
-        if self._shall_confirm:
-            proceed = Confirm.ask(
-                "Authenticating with Lightning AI. This will open a browser window. Continue?", default=True
-            )
-            if not proceed:
-                raise RuntimeError(
-                    "Login cancelled. Please login to Lightning AI to deploy the API."
-                    " Run `lightning login` to login."
-                ) from None
-        print("Opening browser for authentication...")
-        print("Please come back to the terminal after logging in.")
-        time.sleep(3)
-        _AuthServer().login_with_browser(self)
-
-
-def authenticate(shall_confirm: bool = True) -> None:
-    auth = _Auth(shall_confirm)
-    auth.authenticate()
-
-
-def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
-    if teamspace is None:
-        user = _get_authed_user()
-        menu = _TeamspacesMenu()
-        possible_teamspaces = menu._get_possible_teamspaces(user)
-        if len(possible_teamspaces) == 1:
-            name = next(iter(possible_teamspaces.values()))["name"]
-            return Teamspace(name=name, org=org, user=user)
-
-        return menu._resolve_teamspace(teamspace)
-
-    return _resolve_teamspace(teamspace=teamspace, org=org, user=user)
-
-
-class _UserStatus(TypedDict):
-    verified: bool
-    onboarded: bool
-
-
-def poll_verified_status(timeout: int = _POLL_TIMEOUT) -> _UserStatus:
-    """Polls the verified status of the user until it is True or a timeout occurs."""
-    user_api = UserApi()
-    user = _get_authed_user()
-    start_time = datetime.now()
-    result = {"onboarded": False, "verified": False}
-    while True:
-        user_resp = user_api.get_user(name=user.name)
-        result["onboarded"] = user_resp.status.completed_project_onboarding
-        result["verified"] = user_resp.status.verified
-        if user_resp.status.verified:
-            return result
-        if (datetime.now() - start_time).total_seconds() > timeout:
-            break
-        time.sleep(5)
-    return result
-
-
-class _OnboardingStatus(Enum):
-    NOT_VERIFIED = "not_verified"
-    ONBOARDING = "onboarding"
-    ONBOARDED = "onboarded"
-
-
-class _Onboarding:
-    def __init__(self, console: Console) -> None:
-        self.console = console
-        self.user = _get_authed_user()
-        self.user_api = UserApi()
-        self.client = LightningClient(max_tries=7)
-
-    @property
-    def verified(self) -> bool:
-        return self.user_api.get_user(name=self.user.name).status.verified
-
-    @property
-    def is_onboarded(self) -> bool:
-        return self.user_api.get_user(name=self.user.name).status.completed_project_onboarding
-
-    @property
-    def can_join_org(self) -> bool:
-        return len(self.client.organizations_service_list_joinable_organizations().joinable_organizations) > 0
-
-    @property
-    def status(self) -> _OnboardingStatus:
-        if not self.verified:
-            return _OnboardingStatus.NOT_VERIFIED
-        if self.is_onboarded:
-            return _OnboardingStatus.ONBOARDED
-        return _OnboardingStatus.ONBOARDING
-
-    def _wait_user_onboarding(self, timeout: int = _POLL_TIMEOUT) -> None:
-        """Wait for user onboarding if they can join the teamspace otherwise move to select a teamspace."""
-        status = self.status
-        if status == _OnboardingStatus.ONBOARDED:
-            return
-
-        self.console.print("Waiting for account setup. Visit lightning.ai")
-        start_time = datetime.now()
-        while self.status != _OnboardingStatus.ONBOARDED:
-            time.sleep(5)
-            if self.is_onboarded:
-                return
-            if (datetime.now() - start_time).total_seconds() > timeout:
-                break
-
-        raise RuntimeError("Timed out waiting for onboarding status")
-
-    def get_cloudspace_id(self, teamspace: Teamspace) -> Optional[str]:
-        cloudspaces: List[V1CloudSpace] = self.client.cloud_space_service_list_cloud_spaces(teamspace.id).cloudspaces
-        cloudspaces = sorted(cloudspaces, key=lambda cloudspace: cloudspace.created_at, reverse=True)
-        if len(cloudspaces) == 0:
-            raise RuntimeError("Error creating deployment! Finish account setup at lightning.ai first.")
-        # get the first cloudspace
-        cloudspace = cloudspaces[0]
-        if "scratch-studio" in cloudspace.name or "scratch-studio" in cloudspace.display_name:
-            return cloudspace.id
-        return None
-
-    def select_teamspace(self, teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
-        """Select a teamspace while onboarding.
-
-        If user is being onboarded and can't join any org, the teamspace it will be resolved to the default
-         personal teamspace.
-        If user is being onboarded and can join an org then it will select default teamspace from the org.
-        """
-        if self.is_onboarded:
-            return select_teamspace(teamspace, org, user)
-
-        # Run only when user hasn't completed onboarding yet.
-        menu = _TeamspacesMenu()
-        self._wait_user_onboarding()
-        # Onboarding has been completed - user already selected organization if they could
-        possible_teamspaces = menu._get_possible_teamspaces(self.user)
-        if len(possible_teamspaces) == 1:
-            # User didn't select any org
-            value = next(iter(possible_teamspaces.values()))
-            return Teamspace(name=value["name"], org=value["org"], user=value["user"])
-
-        for _, value in possible_teamspaces.items():
-            # User select an org
-            # Onboarding teamspace will be the default teamspace in the selected org
-            if value["org"]:
-                return Teamspace(name=value["name"], org=value["org"], user=value["user"])
-        raise RuntimeError("Unable to select teamspace. Visit lightning.ai")
-
-
 def is_connected(host: str = "8.8.8.8", port: int = 53, timeout: int = 10) -> bool:
     try:
         socket.setdefaulttimeout(timeout)
@@ -453,96 +284,6 @@ def _upload_container(
     return True
 
 
-# TODO: Move the rest of the devbox logic here
-class _LitServeDevbox:
-    """Build LitServe API in a Studio."""
-
-    def resolve_previous_upload(self, studio: Studio, folder: str) -> Dict[str, str]:
-        remote_path = "."
-        pairs = {}
-        for root, _, files in os.walk(folder):
-            rel_root = os.path.relpath(root, folder)
-            for f in files:
-                pairs[os.path.join(root, f)] = os.path.join(remote_path, rel_root, f)
-        return _resolve_previous_upload_state(studio, remote_path, pairs)
-
-    def upload_folder(self, studio: Studio, folder: str, upload_state: Dict[str, str]) -> None:
-        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
-            futures = _start_parallel_upload(executor, studio, upload_state)
-            total_files = len(upload_state)
-
-            with Progress(
-                SpinnerColumn(),
-                TextColumn("[progress.description]{task.description}"),
-                TimeElapsedColumn(),
-                console=Console(),
-                transient=True,
-            ) as progress:
-                upload_task = progress.add_task(f"[cyan]Uploading {total_files} files to Studio...", total=total_files)
-                for f in concurrent.futures.as_completed(futures):
-                    upload_state.pop(f.result())
-                    _dump_current_upload_state(studio, ".", upload_state)
-                    progress.update(upload_task, advance=1)
-
-
-def _handle_devbox(
-    name: str,
-    script_path: Path,
-    console: Console,
-    non_interactive: bool = False,
-    devbox: Machine = "CPU",
-    interruptible: bool = False,
-    teamspace: Optional[str] = None,
-    org: Optional[str] = None,
-    user: Optional[str] = None,
-) -> None:
-    if script_path.suffix != ".py":
-        console.print("❌ Error: Only Python files (.py) are supported for development servers", style="red")
-        return
-
-    resolved_teamspace = select_teamspace(teamspace, org, user)
-    studio = Studio(name=name, teamspace=resolved_teamspace)
-    lit_devbox = _LitServeDevbox()
-
-    studio_url = _get_studio_url(studio, turn_on=True)
-    pathlib_path = Path(script_path).resolve()
-    ok = False
-    studio_path = f"{studio.owner.name}/{studio.teamspace.name}/{studio.name}"
-
-    console.print("\n=== Lightning Studio Setup ===")
-    console.print(f"🔧 [bold]Setting up Studio:[/bold] {studio_path}")
-    console.print(f"📁 [bold]Local project:[/bold] {pathlib_path.parent}")
-
-    upload_state = lit_devbox.resolve_previous_upload(studio, pathlib_path.parent)
-    if non_interactive:
-        console.print(f"🌐 [bold]Opening Studio:[/bold] [link={studio_url}]{studio_url}[/link]")
-        ok = webbrowser.open(studio_url)
-    else:
-        if Confirm.ask("Would you like to open your Studio in the browser?", default=True):
-            console.print(f"🌐 [bold]Opening Studio:[/bold] [link={studio_url}]{studio_url}[/link]")
-            ok = webbrowser.open(studio_url)
-
-    if not ok:
-        console.print(f"🔗 [bold]Access Studio:[/bold] [link={studio_url}]{studio_url}[/link]")
-
-    console.print("\n⚡ Initializing Studio (this typically takes 1-2 minutes)...")
-    studio.start(machine=devbox, interruptible=interruptible)
-    studio.install_plugin("custom-port")
-    console.print("🔌 Configuring server port...")
-    studio.run_plugin("custom-port", port=8000)  # TODO: Remove hardcoded port and fetch from LitServe
-
-    console.print("📤 Syncing project files to Studio...")
-    lit_devbox.upload_folder(studio, pathlib_path.parent, upload_state)
-
-    # Add completion message with next steps
-    console.print("\n✅ Studio ready!")
-    console.print("\n📋 [bold]Next steps:[/bold]")
-    console.print("  [bold]1.[/bold] Server code will be available in the Studio")
-    console.print("  [bold]2.[/bold] The Studio is now running with the specified configuration")
-    console.print("  [bold]3.[/bold] Modify and run your server directly in the Studio")
-    # TODO: Once server running is implemented
-
-
 def _handle_cloud(
     script_path: Union[str, Path],
     console: Console,
@@ -563,7 +304,7 @@ def _handle_cloud(
 ) -> None:
     if not is_connected():
         console.print("❌ Internet connection required to deploy to the cloud.", style="red")
-        console.print("To run locally instead, use: `lightning serve [SCRIPT | server.py] --local`")
+        console.print("To run locally instead, use: `lightning serve [SCRIPT | server.py]`")
         return
 
     deployment_name = os.path.basename(repository)
@@ -605,7 +346,7 @@ def _handle_cloud(
     # Push the container to the registry
     console.print("\nPushing container to registry. It may take a while...", style="bold")
     # Authenticate with LitServe affiliate
-    authenticate(shall_confirm=not non_interactive)
+    authenticate(_AuthMode.DEPLOY, shall_confirm=not non_interactive)
     user_status = poll_verified_status()
     cloudspace_id: Optional[str] = None
     from_onboarding = False

lightning_sdk/cli/download.py

@@ -1,3 +1,4 @@
+import json
 import os
 import re
 from pathlib import Path
@@ -6,6 +7,7 @@ from typing import Optional
 import click
 from rich.console import Console
 
+from lightning_sdk.api.license_api import LicenseApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.studios_menu import _StudiosMenu
@@ -46,13 +48,18 @@ def model(name: str, download_dir: str = ".") -> None:
     "--studio",
     default=None,
     help=(
-        "The name of the studio to upload to. "
+        "The name of the studio to download from. "
         "Will show a menu with user's owned studios for selection if not specified. "
         "If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME> where the names are case-sensitive. "
         "The teamspace and studio names can be regular expressions to match, "
         "a menu filtered studios will be shown for final selection."
     ),
 )
+@click.option(
+    "--teamspace",
+    default=None,
+    help="The teamspace the drive folder is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
+)
 @click.option(
     "--local-path",
     "--local_path",
@@ -60,32 +67,41 @@ def model(name: str, download_dir: str = ".") -> None:
     type=click.Path(file_okay=False, dir_okay=True),
     help="The path to the directory you want to download the folder to.",
 )
-def folder(path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
-    """Download a folder from a Studio.
+def folder(
+    path: str = "", studio: Optional[str] = None, teamspace: Optional[str] = None, local_path: str = "."
+) -> None:
+    """Download a folder from a Studio or a Teamspace drive folder.
 
     Example:
       lightning download folder PATH
 
-    PATH: The relative path within the Studio you want to download.
-    Defaults to the entire studio.
+    PATH: The relative path within the Studio or drive folder you want to download.
+    Defaults to the entire Studio or drive folder.
     """
     local_path = Path(local_path)
     if not local_path.is_dir():
         raise NotADirectoryError(f"'{local_path}' is not a directory")
 
-    resolved_studio = _resolve_studio(studio)
+    if studio and teamspace:
+        raise ValueError("Either --studio or --teamspace must be provided, not both")
+
+    if studio:
+        resolved_downloader = _resolve_studio(studio)
+    elif teamspace:
+        menu = _TeamspacesMenu()
+        resolved_downloader = menu._resolve_teamspace(teamspace)
 
     if not path:
-        local_path /= resolved_studio.name
+        local_path /= resolved_downloader.name
         path = ""
 
     try:
         if not path:
             raise FileNotFoundError()
-        resolved_studio.download_folder(remote_path=path, target_path=str(local_path))
+        resolved_downloader.download_folder(remote_path=path, target_path=str(local_path))
     except Exception as e:
         raise StudioCliError(
-            f"Could not download the folder from the given Studio {studio}. "
+            f"Could not download the folder from the given Studio {studio} or Teamspace {teamspace}. "
             "Please contact Lightning AI directly to resolve this issue."
         ) from e
 
@@ -103,6 +119,11 @@ def folder(path: str = "", studio: Optional[str] = None, local_path: str = ".")
         "a menu filtered studios will be shown for final selection."
     ),
 )
+@click.option(
+    "--teamspace",
+    default=None,
+    help="The teamspace the file is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
+)
 @click.option(
     "--local-path",
     "--local_path",
@@ -110,31 +131,40 @@ def folder(path: str = "", studio: Optional[str] = None, local_path: str = ".")
     type=click.Path(file_okay=False, dir_okay=True),
     help="The path to the directory you want to download the file to.",
 )
-def file(path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
-    """Download a file from a Studio.
+def file(path: str = "", studio: Optional[str] = None, teamspace: Optional[str] = None, local_path: str = ".") -> None:
+    """Download a file from a Studio or Teamspace drive file.
 
     Example:
       lightning download file PATH
 
-    PATH: The relative path to the file within the Studio you want to download.
+    PATH: The relative path to the file within the Studio or Teamspace drive file you want to download.
     """
     local_path = Path(local_path)
     if not local_path.is_dir():
         raise NotADirectoryError(f"'{local_path}' is not a directory")
 
-    resolved_studio = _resolve_studio(studio)
+    if studio and teamspace:
+        raise ValueError("Either --studio or --teamspace must be provided, not both")
+
+    if studio:
+        resolved_downloader = _resolve_studio(studio)
+    elif teamspace:
+        menu = _TeamspacesMenu()
+        resolved_downloader = menu._resolve_teamspace(teamspace)
+    else:
+        raise ValueError("Either --studio or --teamspace must be provided")
 
     if not path:
-        local_path /= resolved_studio.name
+        local_path /= resolved_downloader.name
         path = ""
 
     try:
         if not path:
             raise FileNotFoundError()
-        resolved_studio.download_file(remote_path=path, file_path=str(local_path / os.path.basename(path)))
+        resolved_downloader.download_file(remote_path=path, file_path=str(local_path / os.path.basename(path)))
     except Exception as e:
         raise StudioCliError(
-            f"Could not download the file from the given Studio {studio}. "
+            f"Could not download the file from the given Studio {studio} or Teamspace {teamspace}. "
             "Please contact Lightning AI directly to resolve this issue."
         ) from e
 
@@ -208,3 +238,26 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
 
     with skip_studio_init():
         return Studio(**selected_studio)
+
+
+@download.command(name="licenses")
+def licenses() -> None:
+    """Download licenses for all products/packages.
+
+    Example:
+      lightning download licenses
+
+    """
+    user = _get_authed_user()
+    api = LicenseApi()
+    licenses = api.list_user_licenses(user.id)
+
+    user_home = Path.home()
+    lit_dir = user_home / ".lightning"
+    lit_dir.mkdir(parents=True, exist_ok=True)
+    licenses_file = lit_dir / "licenses.json"
+
+    licenses_short = {ll.product_name: ll.license_key for ll in licenses if ll.is_valid}
+    with licenses_file.open("w") as fp:
+        json.dump(licenses_short, fp, indent=4)
+    Console().print(f"Licenses downloaded to {licenses_file}", style="green")

lightning_sdk/cli/entrypoint.py

@@ -14,6 +14,7 @@ from lightning_sdk.cli.configure import configure
 from lightning_sdk.cli.connect import connect
 from lightning_sdk.cli.create import create
 from lightning_sdk.cli.delete import delete
+from lightning_sdk.cli.deploy.serve import deploy
 from lightning_sdk.cli.docker_cli import dockerize
 from lightning_sdk.cli.download import download
 from lightning_sdk.cli.generate import generate
@@ -21,7 +22,6 @@ from lightning_sdk.cli.inspect import inspect
 from lightning_sdk.cli.list import list_cli
 from lightning_sdk.cli.open import open
 from lightning_sdk.cli.run import run
-from lightning_sdk.cli.serve import deploy
 from lightning_sdk.cli.start import start
 from lightning_sdk.cli.stop import stop
 from lightning_sdk.cli.switch import switch

lightning_sdk/cli/open.py

@@ -1,4 +1,5 @@
 import webbrowser
+from contextlib import suppress
 from pathlib import Path
 from typing import Optional
 
@@ -23,7 +24,17 @@ from lightning_sdk.utils.resolve import _get_studio_url
         "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
     ),
 )
-def open(path: str = ".", teamspace: Optional[str] = None) -> None:  # noqa: A001
+@click.option(
+    "--cloud-account",
+    "--cloud_account",
+    default=None,
+    help=(
+        "The cloud account to create the studio on. "
+        "If not specified, will try to infer from the environment (e.g. when run from within a Studio.) "
+        "or fall back to the teamspace default."
+    ),
+)
+def open(path: str = ".", teamspace: Optional[str] = None, cloud_account: Optional[str] = None) -> None:  # noqa: A001
     """Open a local file or folder in a Lightning Studio.
 
     Example:
@@ -41,7 +52,15 @@ def open(path: str = ".", teamspace: Optional[str] = None) -> None:  # noqa: A00
         menu = _TeamspacesMenu()
         resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
 
-    new_studio = Studio(name=pathlib_path.stem, teamspace=resolved_teamspace)
+    # default cloud account to current studios cloud account if run from studio
+    # else it will fall back to teamspace default in the backend
+    if cloud_account is None:
+        with suppress(ValueError):
+            s = Studio()
+            if s.teamspace.name == resolved_teamspace.name and s.teamspace.owner.name == resolved_teamspace.owner.name:
+                cloud_account = s.cloud_account
+
+    new_studio = Studio(name=pathlib_path.stem, teamspace=resolved_teamspace, cloud_account=cloud_account)
 
     console.print(
         f"[bold]Uploading {path} to {new_studio.owner.name}/{new_studio.teamspace.name}/{new_studio.name}[/bold]"

lightning_sdk/cli/start.py

@@ -4,8 +4,10 @@ import click
 
 from lightning_sdk import Machine, Studio
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.studio import Provider
 
 _MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+_PROVIDER_VALUES = tuple([provider.value for provider in Provider])
 
 
 @click.group("start")
@@ -33,7 +35,14 @@ def start() -> None:
     type=click.Choice(_MACHINE_VALUES),
     help="The machine type to start the studio on.",
 )
-def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
+@click.option(
+    "--provider",
+    default=None,
+    show_default=True,
+    type=click.Choice(_PROVIDER_VALUES),
+    help="The provider to start the studio on.",
+)
+def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU", provider: Optional[str] = None) -> None:
     """Start a studio on a given machine.
 
     Example:
@@ -50,9 +59,9 @@ def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU") ->
         owner, teamspace = None, None
 
     try:
-        studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+        studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False, provider=provider)
     except (RuntimeError, ValueError, ApiException):
-        studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+        studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False, provider=provider)
 
     try:
         resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))