lightning-sdk 2025.10.14__py3-none-any.whl → 2025.10.22__py3-none-any.whl

lightning_sdk/__init__.py

@@ -2,7 +2,7 @@ from lightning_sdk.agents import Agent
 from lightning_sdk.ai_hub import AIHub
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noqa: F401
 from lightning_sdk.deployment import Deployment
-from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
+from lightning_sdk.helpers import VersionChecker, _set_tqdm_envvars_noninteractive
 from lightning_sdk.job import Job
 from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.mmt import MMT
@@ -35,6 +35,9 @@ __all__ = [
     "VM",
 ]
 
-__version__ = "2025.10.14"
-_check_version_and_prompt_upgrade(__version__)
+__version__ = "2025.10.22"
+
+_version_checker = VersionChecker()
+_version_checker.check_and_prompt_upgrade(__version__)
+
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/base_studio_api.py

@@ -15,24 +15,28 @@ class BaseStudioApi:
     def __init__(self) -> None:
         self._client = LightningClient(retry=False, max_tries=0)
 
-    def get_base_studio(self, base_studio_id: str, org_id: str) -> V1CloudSpaceEnvironmentTemplate:
+    def get_base_studio(self, base_studio_id: str, org_id: Optional[str] = None) -> V1CloudSpaceEnvironmentTemplate:
         """Retrieve the base studio by its ID."""
         try:
             return self._client.cloud_space_environment_template_service_get_cloud_space_environment_template(
-                base_studio_id, org_id=org_id
+                base_studio_id, org_id=org_id or ""
             )
         except ValueError as e:
             raise ValueError(f"Base studio {base_studio_id} does not exist") from e
 
-    def get_all_base_studios(self, org_id: str, managed: bool = True) -> V1ListCloudSpaceEnvironmentTemplatesResponse:
+    def get_all_base_studios(self, org_id: Optional[str]) -> V1ListCloudSpaceEnvironmentTemplatesResponse:
         """Retrieve all base studios for a given organization."""
-        if managed:
-            return self._client.cloud_space_environment_template_service_list_managed_cloud_space_environment_templates(
-                org_id=org_id
-            )
-        return self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
-            org_id=org_id
+        result = self._client.cloud_space_environment_template_service_list_managed_cloud_space_environment_templates(
+            org_id=org_id or ""
         )
+        if org_id is not None:
+            org_templates = (
+                self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
+                    org_id=org_id
+                )
+            )
+            result.templates = result.templates + org_templates.templates
+        return result
 
     def update_base_studio(
         self,

lightning_sdk/api/license_api.py

@@ -1,70 +1,37 @@
-import os
-from typing import Optional
-from urllib.parse import urlencode
-
-from lightning_sdk.lightning_cloud import env
-from lightning_sdk.lightning_cloud.rest_client import LightningClient
-
-LICENSE_CODE = os.environ.get("LICENSE_CODE", "d9s79g79ss")
-# https://lightning.ai/home?settings=licenses
-LICENSE_SIGNING_URL = f"{env.LIGHTNING_CLOUD_URL}?settings=licenses"
-
-
-def generate_url_user_settings(name: str, redirect_to: str = LICENSE_SIGNING_URL) -> str:
-    params = urlencode({"redirectTo": redirect_to, "okbhrt": LICENSE_CODE, "licenseName": name})
-    return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.lightning_cloud.openapi import LicenseKeyValidateBody, ProductLicenseServiceApi
 
 
 class LicenseApi:
-    _client_authenticated: LightningClient = None
-    _client_public: LightningClient = None
-
-    @property
-    def client_public(self) -> LightningClient:
-        if not self._client_public:
-            self._client_public = LightningClient(retry=False, max_tries=0, with_auth=False)
-        return self._client_public
-
-    @property
-    def client_authenticated(self) -> LightningClient:
-        if not self._client_authenticated:
-            self._client_authenticated = LightningClient(retry=True, max_tries=3, with_auth=True)
-        return self._client_authenticated
+    def __init__(self, login_token: str) -> None:
+        self._cloud_url = _cloud_url()
+        self._auth = Auth()
+        self._auth.token_login(login_token, save_token=True)
+        self._client = self._auth.create_api_client()
+        self._api = ProductLicenseServiceApi(self._client)
 
-    def valid_license(
-        self,
-        license_key: str,
-        product_name: str,
-        product_version: Optional[str] = None,
-        product_type: str = "package",
-    ) -> bool:
-        """Check if the license key is valid.
+    def validate_license(self, license_key: str, product_id: str) -> bool:
+        """Validate a license key for a specific product.
 
         Args:
-            license_key: The license key to check.
-            product_name: The name of the product.
-            product_version: The version of the product.
-            product_type: The type of the product. Default is "package".
+            license_key: The license key to validate
+            product_id: The product ID
 
         Returns:
-            True if the license key is valid, False otherwise.
-        """
-        response = self.client_public.product_license_service_validate_product_license(
-            license_key=license_key,
-            product_name=product_name,
-            product_version=product_version,
-            product_type=product_type,
-        )
-        return response.valid
+            bool: True if license is valid, False otherwise
 
-    def list_user_licenses(self, user_id: str) -> list:
-        """List all licenses for a user.
+        Raises:
+            Exception: If license validation fails
+        """
+        try:
+            response = self._api.product_license_service_validate_license(
+                body=LicenseKeyValidateBody(product_id=product_id), license_key=license_key
+            )
+            return response.is_valid
+        except Exception:
+            raise InvalidLicenseError(f"Invalid license key {license_key} for product {product_id}") from None
 
-        Args:
-            user_id: The ID of the user.
 
-        Returns:
-            A list of licenses for the user.
-        """
-        response = self.client_authenticated.product_license_service_list_user_licenses(user_id=user_id)
-        return response.licenses
+class InvalidLicenseError(Exception):
+    pass

lightning_sdk/api/studio_api.py

@@ -589,7 +589,12 @@ class StudioApi:
         )
 
     def duplicate_studio(
-        self, studio_id: str, teamspace_id: str, target_teamspace_id: str, machine: Machine = Machine.CPU
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        target_teamspace_id: str,
+        machine: Machine = Machine.CPU,
+        new_name: Optional[str] = None,
     ) -> Dict[str, Any]:
         """Duplicates the given Studio from a given Teamspace into a given target Teamspace."""
         target_teamspace = self._client.projects_service_get_project(target_teamspace_id)
@@ -604,7 +609,7 @@ class StudioApi:
             init_kwargs["org"] = OrgApi()._get_org_by_id(target_teamspace.owner_id).name
 
         new_cloudspace = self._client.cloud_space_service_fork_cloud_space(
-            IdForkBody1(target_project_id=target_teamspace_id), project_id=teamspace_id, id=studio_id
+            IdForkBody1(target_project_id=target_teamspace_id, new_name=new_name), project_id=teamspace_id, id=studio_id
         )
 
         while self.get_studio_by_id(new_cloudspace.id, target_teamspace_id).state != V1CloudSpaceState.READY:

lightning_sdk/base_studio.py

@@ -3,11 +3,11 @@ from typing import List, Optional, Union
 
 from lightning_sdk.api.base_studio_api import BaseStudioApi
 from lightning_sdk.api.user_api import UserApi
-from lightning_sdk.lightning_cloud import login
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_type import V1CloudSpaceEnvironmentType
 from lightning_sdk.organization import Organization
+from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
+from lightning_sdk.utils.resolve import _resolve_teamspace
 
 
 @dataclass
@@ -24,6 +24,7 @@ class BaseStudio:
     def __init__(
         self,
         name: Optional[str] = None,
+        teamspace: Optional[Union[str, Teamspace]] = None,
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
     ) -> None:
@@ -38,26 +39,35 @@ class BaseStudio:
         Raises:
             ConnectionError: If there is an issue with the authentication process.
         """
-        self._auth = login.Auth()
-        self._user = None
+        self._teamspace = None
 
-        try:
-            self._auth.authenticate()
-            if user is None:
-                self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
-        except ConnectionError as e:
-            raise e
+        _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        if _teamspace is None:
+            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
 
-        self._user = _resolve_user(self._user or user)
-        self._org = _resolve_org(org)
+        self._teamspace = _teamspace
+
+        # self._auth = login.Auth()
+        # self._user = None
+
+        # try:
+        #     self._auth.authenticate()
+        #     if user is None:
+        #         self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
+        # except ConnectionError as e:
+        #     raise e
+
+        # self._user = _resolve_user(self._user or user)
+        # self._org = _resolve_org(org)
 
         self._base_studio_api = BaseStudioApi()
 
         if name is not None:
-            base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+            org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+            base_studio = self._base_studio_api.get_base_studio(name, org_id)
 
             if base_studio is None:
-                raise ValueError(f"Base studio with name {name} does not exist in organization {self._org.name}")
+                raise ValueError(f"Base studio with name {name} does not exist")
             self._base_studio = base_studio
 
     def update(
@@ -70,9 +80,11 @@ class BaseStudio:
         machine_image_version: Optional[str] = None,
         setup_script_text: Optional[str] = None,
     ) -> None:
+        org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+        # TODO: if not in an org, can't update them
         self._base_studio = self._base_studio_api.update_base_studio(
             self._base_studio.id,
-            self._org.id,
+            org_id,
             name=name,
             allowed_machines=allowed_machines,
             default_machine=default_machine,
@@ -82,7 +94,7 @@ class BaseStudio:
             disabled=disabled,
         )
 
-    def list(self, managed: bool = True, include_disabled: bool = False) -> List[BaseStudioInfo]:
+    def list(self, include_disabled: bool = False) -> List[BaseStudioInfo]:
         """List all base studios in the organization.
 
         Args:
@@ -92,7 +104,8 @@ class BaseStudio:
         Returns:
             List[BaseStudioInfo]: A list of base studio templates.
         """
-        templates = self._base_studio_api.get_all_base_studios(self._org.id, managed).templates
+        org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+        templates = self._base_studio_api.get_all_base_studios(org_id).templates
 
         return [
             BaseStudioInfo(

lightning_sdk/cli/base_studio/list.py

@@ -21,9 +21,7 @@ def list_base_studios(include_disabled: bool) -> None:
 
 def list_impl(include_disabled: bool) -> None:
     base_studio_cls = BaseStudio()
-    base_studios = base_studio_cls.list(include_disabled=include_disabled) + base_studio_cls.list(
-        managed=False, include_disabled=include_disabled
-    )
+    base_studios = base_studio_cls.list(include_disabled=include_disabled)
 
     table = Table(
         pad_edge=True,

lightning_sdk/cli/entrypoint.py

@@ -2,15 +2,8 @@
 
 import os
 import sys
-import traceback
-from types import TracebackType
-from typing import Type
 
 import click
-from rich.console import Group
-from rich.panel import Panel
-from rich.syntax import Syntax
-from rich.text import Text
 
 from lightning_sdk import __version__
 from lightning_sdk.api.studio_api import _cloud_url
@@ -24,38 +17,19 @@ from lightning_sdk.cli.groups import (
     studio,
     vm,
 )
-from lightning_sdk.cli.utils import CustomHelpFormatter, rich_to_str
-from lightning_sdk.constants import _LIGHTNING_DEBUG
+from lightning_sdk.cli.utils import CustomHelpFormatter
+from lightning_sdk.cli.utils.logging import CommandLoggingGroup, logging_excepthook
 from lightning_sdk.lightning_cloud.login import Auth
 
 
-def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:
-    """CLI won't show tracebacks, just print the exception message."""
-    message = str(value.args[0]) if value.args else str(value) or "An unknown error occurred"
-
-    error_text = Text()
-    error_text.append(f"{exception_type.__name__}: ", style="bold red")
-    error_text.append(message, style="white")
-
-    renderables = [error_text]
-
-    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 light", line_numbers=False, word_wrap=True))
-    else:
-        renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
-
-    renderables.append(Text("\n📘 Need help? Run: lightning <command> --help", style="cyan"))
-
-    text = rich_to_str(Panel(Group(*renderables), title="⚡ Lightning CLI Error", border_style="red"))
-    click.echo(text, color=True)
-
-
-@click.group(name="lightning", help="Command line interface (CLI) to interact with/manage Lightning AI Studios.")
+@click.group(
+    name="lightning",
+    help="Command line interface (CLI) to interact with/manage Lightning AI Studios.",
+    cls=CommandLoggingGroup,
+)
 @click.version_option(__version__, message="Lightning CLI version %(version)s")
 def main_cli() -> None:
-    sys.excepthook = _notify_exception
+    sys.excepthook = logging_excepthook
 
 
 main_cli.context_class.formatter_class = CustomHelpFormatter

lightning_sdk/cli/studio/connect.py

@@ -2,11 +2,13 @@
 
 import subprocess
 import sys
-from typing import Dict, Optional, Set
+from contextlib import suppress
+from typing import Optional
 
 import click
 
-from lightning_sdk.base_studio import BaseStudio
+from lightning_sdk.cli.utils.get_base_studio import get_base_studio_id
+from lightning_sdk.cli.utils.handle_machine_and_gpus_args import handle_machine_and_gpus_args
 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 configure_ssh_internal
@@ -16,81 +18,42 @@ from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.names import random_unique_name
 
-DEFAULT_MACHINE = "CPU"
 
+def _parse_args_or_get_from_current_studio(
+    teamspace: Optional[str],
+    cloud_account: Optional[str],
+    studio_type: Optional[str],
+    machine: Optional[str],
+    gpus: Optional[str],
+    cloud_provider: Optional[str],
+    name: Optional[str],
+) -> tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+    # Parse args provided by user
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace)
+    save_teamspace_to_config(resolved_teamspace, overwrite=False)
 
-def _split_gpus_spec(gpus: str) -> tuple[str, int]:
-    machine_name, machine_val = gpus.split(":", 1)
-    machine_name = machine_name.strip()
-    machine_val = machine_val.strip()
-
-    if not machine_val.isdigit() or int(machine_val) <= 0:
-        raise ValueError(f"Invalid GPU count '{machine_val}'. Must be a positive integer.")
-
-    machine_num = int(machine_val)
-    return machine_name, machine_num
-
-
-def _construct_available_gpus(machine_options: Dict[str, str]) -> Set[str]:
-    # returns available gpus:count
-    available_gpus = set()
-    for v in machine_options.values():
-        if "_X_" in v:
-            gpu_type_num = v.replace("_X_", ":")
-            available_gpus.add(gpu_type_num)
-        else:
-            available_gpus.add(v)
-    return available_gpus
-
-
-def _get_machine_from_gpus(gpus: str) -> Machine:
-    machine_name = gpus
-    machine_num = 1
-
-    if ":" in gpus:
-        machine_name, machine_num = _split_gpus_spec(gpus)
-
-    machine_options = {
-        m.name.lower(): m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli
-    }
-
-    if machine_num == 1:
-        # e.g. gpus=L4 or gpus=L4:1
-        gpu_key = machine_name.lower()
-        try:
-            return machine_options[gpu_key]
-        except KeyError:
-            available = ", ".join(_construct_available_gpus(machine_options))
-            raise ValueError(f"Invalid GPU type '{machine_name}'. Available options: {available}") from None
-
-    # Else: e.g. gpus=L4:4
-    gpu_key = f"{machine_name.lower()}_x_{machine_num}"
-    try:
-        return machine_options[gpu_key]
-    except KeyError:
-        available = ", ".join(_construct_available_gpus(machine_options))
-        raise ValueError(f"Invalid GPU configuration '{gpus}'. Available options: {available}") from None
+    template_id = get_base_studio_id(studio_type)
 
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
 
-def _get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
-    base_studios = BaseStudio()
-    base_studios = base_studios.list()
-    template_id = None
+    name = name or random_unique_name()
 
-    if base_studios and len(base_studios):
-        # if not specified by user, use the first existing template studio
-        template_id = base_studios[0].id
-        # else, try to match the provided studio_type to base studio name
-        if studio_type:
-            normalized_studio_type = studio_type.lower().replace(" ", "-")
-            match = next(
-                (s for s in base_studios if s.name.lower().replace(" ", "-") == normalized_studio_type),
-                None,
-            )
-            if match:
-                template_id = match.id
+    with suppress(ValueError):
+        # Gets current studio context to use its parameters as defaults
+        s = Studio()
+        if not teamspace:
+            resolved_teamspace = s.teamspace
+            save_teamspace_to_config(resolved_teamspace, overwrite=False)
+        if not cloud_account:
+            cloud_account = s.cloud_account
+        if not template_id:
+            template_id = s._studio.environment_template_id
+        if not machine and not gpus:
+            machine = s.machine
 
-    return template_id
+    return resolved_teamspace, cloud_account, template_id, machine, cloud_provider, name
 
 
 @click.command("connect")
@@ -124,6 +87,7 @@ def _get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
     "Defaults to the first available template.",
     type=click.STRING,
 )
+@click.option("--interruptible", is_flag=True, help="Start the studio on an interruptible instance.")
 def connect_studio(
     name: Optional[str] = None,
     teamspace: Optional[str] = None,
@@ -132,29 +96,21 @@ def connect_studio(
     machine: Optional[str] = None,
     gpus: Optional[str] = None,
     studio_type: Optional[str] = None,
+    interruptible: bool = False,
 ) -> 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()
-
-    # check for available base studios
-    template_id = _get_base_studio_id(studio_type)
+    teamspace, cloud_account, template_id, machine, cloud_provider, name = _parse_args_or_get_from_current_studio(
+        teamspace, cloud_account, studio_type, machine, gpus, cloud_provider, name
+    )
 
     try:
         studio = Studio(
             name=name,
-            teamspace=resolved_teamspace,
+            teamspace=teamspace,
             create_ok=True,
             cloud_provider=cloud_provider,
             cloud_account=cloud_account,
@@ -167,16 +123,10 @@ def connect_studio(
 
     Studio.show_progress = True
 
-    if machine and gpus:
-        raise click.UsageError("Options --machine and --gpu are mutually exclusive. Provide only one.")
-    elif gpus:
-        machine = _get_machine_from_gpus(gpus.strip())
-    elif not machine:
-        machine = DEFAULT_MACHINE
+    machine = handle_machine_and_gpus_args(machine, gpus)
 
     save_studio_to_config(studio)
-    # by default, interruptible is False
-    studio.start(machine=machine, interruptible=False)
+    studio.start(machine=machine, interruptible=interruptible)
 
     ssh_private_key_path = configure_ssh_internal()
 

lightning_sdk/cli/studio/create.py

@@ -4,6 +4,7 @@ from typing import Optional
 
 import click
 
+from lightning_sdk.cli.utils.get_base_studio import get_base_studio_id
 from lightning_sdk.cli.utils.richt_print import studio_name_link
 from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
@@ -25,18 +26,34 @@ from lightning_sdk.studio import VM, Studio
     help="The cloud account to create the studio on. Defaults to teamspace default.",
     type=click.STRING,
 )
+@click.option(
+    "--studio-type",
+    help="The base studio template name to use for creating the studio. "
+    "Must be lowercase and hyphenated (use '-' instead of spaces). "
+    "Run 'lightning base-studio list' to see all available templates. "
+    "Defaults to the first available template.",
+    type=click.STRING,
+)
 def create_studio(
     name: Optional[str] = None,
     teamspace: Optional[str] = None,
     cloud_provider: Optional[str] = None,
     cloud_account: Optional[str] = None,
+    studio_type: Optional[str] = None,
 ) -> None:
     """Create a new Studio.
 
     Example:
         lightning studio create
     """
-    create_impl(name=name, teamspace=teamspace, cloud_provider=cloud_provider, cloud_account=cloud_account, vm=False)
+    create_impl(
+        name=name,
+        teamspace=teamspace,
+        cloud_provider=cloud_provider,
+        cloud_account=cloud_account,
+        vm=False,
+        studio_type=studio_type,
+    )
 
 
 def create_impl(
@@ -45,6 +62,7 @@ def create_impl(
     cloud_provider: Optional[str],
     cloud_account: Optional[str],
     vm: bool,
+    studio_type: Optional[str],
 ) -> None:
     menu = TeamspacesMenu()
 
@@ -57,6 +75,9 @@ def create_impl(
     create_cls = VM if vm else Studio
     cls_name = create_cls.__qualname__
 
+    # check for available base studios
+    template_id = get_base_studio_id(studio_type)
+
     try:
         create_cls = VM if vm else Studio
         studio = create_cls(
@@ -65,6 +86,7 @@ def create_impl(
             create_ok=True,
             cloud_provider=cloud_provider,
             cloud_account=cloud_account,
+            template_id=template_id,
         )
     except (RuntimeError, ValueError, ApiException):
         if name:

lightning_sdk/cli/studio/start.py

@@ -4,6 +4,7 @@ from typing import Optional
 
 import click
 
+from lightning_sdk.cli.utils.handle_machine_and_gpus_args import handle_machine_and_gpus_args
 from lightning_sdk.cli.utils.richt_print import studio_name_link
 from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
 from lightning_sdk.cli.utils.studio_selection import StudiosMenu
@@ -32,8 +33,7 @@ from lightning_sdk.studio import VM, Studio
 @click.option(
     "--cloud-provider",
     help=(
-        "The cloud provider to start the studio on. Defaults to teamspace default. "
-        "Only used if --create is specified."
+        "The cloud provider to start the studio on. Defaults to teamspace default. Only used if --create is specified."
     ),
     type=click.Choice(m.name for m in list(CloudProvider)),
 )
@@ -42,11 +42,17 @@ from lightning_sdk.studio import VM, Studio
     help="The cloud account to start the studio on. Defaults to teamspace default. Only used if --create is specified.",
     type=click.STRING,
 )
+@click.option(
+    "--gpus",
+    help="The number and type of GPUs to start the studio on (format: TYPE:COUNT, e.g. L4:4)",
+    type=click.STRING,
+)
 def start_studio(
     name: Optional[str] = None,
     teamspace: Optional[str] = None,
     create: bool = False,
     machine: str = "CPU",
+    gpus: Optional[str] = None,
     interruptible: bool = False,
     cloud_provider: Optional[str] = None,
     cloud_account: Optional[str] = None,
@@ -62,6 +68,7 @@ def start_studio(
         teamspace=teamspace,
         create=create,
         machine=machine,
+        gpus=gpus,
         interruptible=interruptible,
         cloud_provider=cloud_provider,
         cloud_account=cloud_account,
@@ -74,6 +81,7 @@ def start_impl(
     teamspace: Optional[str],
     create: bool,
     machine: str,
+    gpus: Optional[str],
     interruptible: bool,
     cloud_provider: Optional[str],
     cloud_account: Optional[str],
@@ -98,6 +106,8 @@ def start_impl(
             cloud_account=cloud_account,
         )
 
+    machine = handle_machine_and_gpus_args(machine, gpus)
+
     save_studio_to_config(studio)
 
     Studio.show_progress = True