lightning-sdk 2025.8.21__py3-none-any.whl → 2025.8.26__py3-none-any.whl

lightning_sdk/__init__.py

@@ -32,6 +32,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "2025.08.21"
+__version__ = "2025.08.26"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/studio_api.py

@@ -11,7 +11,7 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
     _create_app,
-    _download_studio_files,
+    _download_teamspace_files,
     _DummyBody,
     _DummyResponse,
     _FileUploader,
@@ -25,6 +25,7 @@ from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     CloudspaceIdRunsBody,
+    CloudspacesIdBody,
     Externalv1LightningappInstance,
     IdCodeconfigBody,
     IdExecuteBody1,
@@ -40,6 +41,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceState,
     V1ClusterAccelerator,
     V1EndpointType,
+    V1EnvVar,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
     V1LoginRequest,
@@ -677,7 +679,7 @@ class StudioApi:
         if prefix.endswith("/") is False:
             prefix = prefix + "/"
 
-        _download_studio_files(
+        _download_teamspace_files(
             client=self._client,
             teamspace_id=teamspace_id,
             cluster_id=cloud_account,
@@ -940,3 +942,68 @@ class StudioApi:
             plugin_type=plugin_type,
             **other_arguments,
         )
+
+    def _update_cloudspace(self, studio: V1CloudSpace, teamspace_id: str, key: str, value: Any) -> None:
+        body = CloudspacesIdBody(
+            code_url=studio.code_url,
+            data_connection_mounts=studio.data_connection_mounts,
+            description=studio.description,
+            display_name=studio.display_name,
+            env=studio.env,
+            featured=studio.featured,
+            hide_files=studio.hide_files,
+            is_cloudspace_private=studio.is_cloudspace_private,
+            is_code_private=studio.is_code_private,
+            is_favorite=studio.is_favorite,
+            is_published=studio.is_published,
+            license=studio.license,
+            license_url=studio.license_url,
+            message=studio.message,
+            multi_user_edit=studio.multi_user_edit,
+            operating_cost=studio.operating_cost,
+            paper_authors=studio.paper_authors,
+            paper_org=studio.paper_org,
+            paper_org_avatar_url=studio.paper_org_avatar_url,
+            paper_url=studio.paper_url,
+            switch_to_default_machine_on_idle=studio.switch_to_default_machine_on_idle,
+            tags=studio.tags,
+            thumbnail_file_type=studio.thumbnail_file_type,
+            user_metadata=studio.user_metadata,
+        )
+
+        setattr(body, key, value)
+
+        self._client.cloud_space_service_update_cloud_space(
+            id=studio.id,
+            project_id=teamspace_id,
+            body=body,
+        )
+
+    def set_env(
+        self,
+        studio: V1CloudSpace,
+        teamspace_id: str,
+        new_env: Dict[str, str],
+        partial: bool = True,
+    ) -> None:
+        """Set the environment variables for the Studio.
+
+        Args:
+            new_env: The new environment variables to set.
+            partial: Whether to only set the environment variables that are provided.
+                If False, existing environment variables that are not in new_env will be removed.
+                If True, existing environment variables that are not in new_env will be kept.
+        """
+        updated_env_dict = {}
+        if partial:
+            updated_env_dict = {env.name: env.value for env in studio.env}
+            updated_env_dict.update(new_env)
+        else:
+            updated_env_dict = new_env
+
+        updated_env = [V1EnvVar(name=key, value=value) for key, value in updated_env_dict.items()]
+
+        self._update_cloudspace(studio, teamspace_id, "env", updated_env)
+
+    def get_env(self, studio: V1CloudSpace) -> Dict[str, str]:
+        return {env.name: env.value for env in studio.env}

lightning_sdk/api/teamspace_api.py

@@ -1,6 +1,5 @@
 import os
-import tempfile
-import zipfile
+import re
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -9,6 +8,7 @@ from tqdm.auto import tqdm
 
 from lightning_sdk.api.utils import (
     _download_model_files,
+    _download_teamspace_files,
     _DummyBody,
     _FileUploader,
     _get_model_version,
@@ -22,6 +22,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     ModelsStoreApi,
     ProjectIdAgentsBody,
     ProjectIdModelsBody,
+    ProjectIdSecretsBody,
+    SecretsIdBody,
     V1Assistant,
     V1CloudSpace,
     V1ClusterAccelerator,
@@ -34,6 +36,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1Project,
     V1ProjectClusterBinding,
     V1PromptSuggestion,
+    V1Secret,
+    V1SecretType,
     V1UpstreamOpenAI,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -421,39 +425,65 @@ class TeamspaceApi:
         # TODO: Update this endpoint to permit basic auth
         auth = Auth()
         auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "clusterId": cloud_account,
-            "prefix": _resolve_teamspace_remote_path(path),
-            "token": token,
-        }
+        prefix = _resolve_teamspace_remote_path(path)
 
-        r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-            params=query_params,
-            stream=True,
+        # ensure we only download as a directory and not the entire prefix
+        if prefix.endswith("/") is False:
+            prefix = prefix + "/"
+
+        _download_teamspace_files(
+            client=self._client,
+            teamspace_id=teamspace_id,
+            cluster_id=cloud_account,
+            prefix=prefix,
+            download_dir=Path(target_path),
+            progress_bar=progress_bar,
         )
 
-        if progress_bar:
-            pbar = tqdm(
-                desc=f"Downloading {os.path.split(path)[1]}",
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1000,
-            )
+    def get_secrets(self, teamspace_id: str) -> Dict[str, str]:
+        """Get all secrets for a teamspace."""
+        secrets = self._get_secrets(teamspace_id)
+        # this returns encrypted values for security. It doesn't make sense to show them,
+        # so we just return a placeholder
+        # not a security issue to replace in the client as we get the encrypted values from the server.
+        return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED}
 
-            pbar_update = pbar.update
-        else:
-            pbar_update = lambda x: None
+    def set_secret(self, teamspace_id: str, key: str, value: str) -> None:
+        """Set a secret for a teamspace.
 
-        if target_path:
-            os.makedirs(target_path, exist_ok=True)
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets(teamspace_id)
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(teamspace_id, secret.id, value)
+        return self._create_secret(teamspace_id, key, value)
+
+    def _get_secrets(self, teamspace_id: str) -> List[V1Secret]:
+        return self._client.secret_service_list_secrets(project_id=teamspace_id).secrets
+
+    def _update_secret(self, teamspace_id: str, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_secret(
+            body=SecretsIdBody(value=value),
+            project_id=teamspace_id,
+            id=secret_id,
+        )
 
-        with tempfile.TemporaryFile() as f:
-            for chunk in r.iter_content(chunk_size=4096 * 8):
-                f.write(chunk)
-                pbar_update(len(chunk))
+    def _create_secret(
+        self,
+        teamspace_id: str,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_secret(
+            body=ProjectIdSecretsBody(name=key, value=value, type=V1SecretType.UNSPECIFIED), project_id=teamspace_id
+        )
 
-            with zipfile.ZipFile(f) as z:
-                z.extractall(target_path)
+    def verify_secret_name(self, name: str) -> bool:
+        """Verify if a secret name is valid.
+
+        A valid secret name starts with a letter or underscore, followed by letters, digits, or underscores.
+        """
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        return re.match(pattern, name) is not None

lightning_sdk/api/user_api.py

@@ -1,15 +1,20 @@
-from typing import List, Union
+import re
+from typing import Dict, List, Union
 
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
+    SecretsIdBody1,
     V1CloudSpace,
+    V1CreateUserSecretRequest,
     V1GetUserResponse,
     V1ListCloudSpacesResponse,
     V1Membership,
     V1Organization,
     V1SearchUser,
+    V1Secret,
     V1UserFeatures,
 )
+from lightning_sdk.lightning_cloud.openapi.models.v1_secret_type import V1SecretType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -69,3 +74,46 @@ class UserApi:
     def _get_feature_flags(self) -> V1UserFeatures:
         resp: V1GetUserResponse = self._client.auth_service_get_user()
         return resp.features
+
+    def get_secrets(self) -> Dict[str, str]:
+        """Get all secrets for the current user."""
+        secrets = self._get_secrets()
+        # this returns encrypted values for security. It doesn't make sense to show them,
+        # so we just return a placeholder
+        # not a security issue to replace in the client as we get the encrypted values from the server.
+        return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED}
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set a secret for the current user.
+
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets()
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(secret.id, value)
+        return self._create_secret(key, value)
+
+    def _get_secrets(self) -> List[V1Secret]:
+        return self._client.secret_service_list_user_secrets().secrets
+
+    def _update_secret(self, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_user_secret(
+            body=SecretsIdBody1(value=value),
+            id=secret_id,
+        )
+
+    def _create_secret(
+        self,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_user_secret(body=V1CreateUserSecretRequest(name=key, value=value))
+
+    def verify_secret_name(self, name: str) -> bool:
+        """Verify if a secret name is valid.
+
+        A valid secret name starts with a letter or underscore, followed by letters, digits, or underscores.
+        """
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        return re.match(pattern, name) is not None

lightning_sdk/api/utils.py

@@ -586,7 +586,7 @@ def _download_model_files(
         return response.filepaths
 
 
-def _download_studio_files(
+def _download_teamspace_files(
     client: LightningClient,
     teamspace_id: str,
     cluster_id: str,

lightning_sdk/cli/config/set.py

@@ -1,8 +1,8 @@
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+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.machine import CloudProvider
-from lightning_sdk.organization import Organization
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.config import Config, DefaultConfigKeys
 from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
@@ -59,23 +59,11 @@ def set_studio(studio_name: str) -> None:
 @click.argument("teamspace_name")
 def set_teamspace(teamspace_name: str) -> None:
     """Set the default teamspace name in the config."""
-    config = Config()
-
-    teamspace_resolved = resolve_teamspace_owner_name_format(teamspace_name)
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace=teamspace_name)
 
-    if teamspace_resolved is None:
-        # TODO: make this a generic CLI error
-        raise ValueError(
-            f"Could not resolve teamspace: '{teamspace_name}'. "
-            "Teamspace should be specified as 'owner/name'. Does the teamspace exist?"
-        )
-
-    setattr(config, DefaultConfigKeys.teamspace_name, teamspace_resolved.name)
-    setattr(config, DefaultConfigKeys.teamspace_owner, teamspace_resolved.owner.name)
-    if isinstance(teamspace_resolved.owner, Organization):
-        setattr(config, DefaultConfigKeys.teamspace_owner_type, "organization")
-    else:
-        setattr(config, DefaultConfigKeys.teamspace_owner_type, "user")
+    # explicit user action, so overwrite the config
+    save_teamspace_to_config(teamspace_resolved, overwrite=True)
 
 
 @set_value.command("cloud-account")

lightning_sdk/cli/legacy/create.py

@@ -7,7 +7,7 @@ import click
 from rich.console import Console
 
 from lightning_sdk import Machine, Studio
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import CloudProvider
 
 _MACHINE_VALUES = tuple(
@@ -78,8 +78,8 @@ def studio(
 
     NAME: the name of the studio to create. If already present within teamspace, will add a random suffix.
     """
-    menu = _TeamspacesMenu()
-    teamspace_resolved = menu._resolve_teamspace(teamspace)
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace)
 
     # default cloud account to current studios cloud account if run from studio
     # else it will fall back to teamspace default in the backend

lightning_sdk/cli/legacy/delete.py

@@ -5,7 +5,7 @@ from rich.console import Console
 
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lit_container import LitContainer
 from lightning_sdk.studio import Studio
@@ -30,8 +30,8 @@ def delete() -> None:
 def container(name: str, teamspace: Optional[str] = None) -> None:
     """Delete the docker container NAME."""
     api = LitContainer()
-    menu = _TeamspacesMenu()
-    resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
     try:
         api.delete_container(name, resolved_teamspace.name, resolved_teamspace.owner.name)
         Console().print(f"Container {name} deleted successfully.")

lightning_sdk/cli/legacy/deploy/_auth.py

@@ -10,7 +10,7 @@ from rich.prompt import Confirm
 
 from lightning_sdk import Teamspace
 from lightning_sdk.api import UserApi
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud import env
 from lightning_sdk.lightning_cloud.login import Auth, AuthServer
 from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
@@ -74,13 +74,13 @@ def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:
 def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
     if teamspace is None:
         user = _get_authed_user()
-        menu = _TeamspacesMenu()
+        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 menu(teamspace)
 
     return _resolve_teamspace(teamspace=teamspace, org=org, user=user)
 
@@ -180,7 +180,7 @@ class _Onboarding:
             return select_teamspace(teamspace, org, user)
 
         # Run only when user hasn't completed onboarding yet.
-        menu = _TeamspacesMenu()
+        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)

lightning_sdk/cli/legacy/download.py

@@ -11,7 +11,7 @@ from lightning_sdk.api.license_api import LicenseApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+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
@@ -110,8 +110,8 @@ def folder(
         path = _expand_remote_path(path)
         resolved_downloader = _resolve_studio(studio)
     elif teamspace:
-        menu = _TeamspacesMenu()
-        resolved_downloader = menu._resolve_teamspace(teamspace)
+        menu = TeamspacesMenu()
+        resolved_downloader = menu(teamspace)
     else:
         raise ValueError("Either --studio or --teamspace must be provided")
 
@@ -173,8 +173,8 @@ def file(path: str = "", studio: Optional[str] = None, teamspace: Optional[str]
     if studio:
         resolved_downloader = _resolve_studio(studio)
     elif teamspace:
-        menu = _TeamspacesMenu()
-        resolved_downloader = menu._resolve_teamspace(teamspace)
+        menu = TeamspacesMenu()
+        resolved_downloader = menu(teamspace)
     else:
         raise ValueError("Either --studio or --teamspace must be provided")
 
@@ -214,8 +214,8 @@ def download_container(
     CONTAINER: The name of the container to download.
     """
     console = Console()
-    menu = _TeamspacesMenu()
-    resolved_teamspace = menu._resolve_teamspace(teamspace)
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace)
     with console.status("Downloading container..."):
         api = LitContainerApi()
         api.download_container(container, resolved_teamspace, tag, cloud_account)

lightning_sdk/cli/legacy/job_and_mmt_action.py

@@ -2,12 +2,12 @@ from typing import Optional
 
 from lightning_sdk.cli.legacy.jobs_menu import _JobsMenu
 from lightning_sdk.cli.legacy.mmts_menu import _MMTsMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.job import Job
 from lightning_sdk.mmt import MMT
 
 
-class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
+class _JobAndMMTAction(TeamspacesMenu, _JobsMenu, _MMTsMenu):
     """Inspect resources of the Lightning AI platform to get additional details as JSON."""
 
     def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> Job:
@@ -20,7 +20,7 @@ class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
                 If not specified can be selected interactively.
 
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
+        resolved_teamspace = self(teamspace)
         return self._resolve_job(name, teamspace=resolved_teamspace)
 
     def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> MMT:
@@ -33,5 +33,5 @@ class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
                 If not specified can be selected interactively.
 
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
+        resolved_teamspace = self(teamspace)
         return self._resolve_mmt(name, teamspace=resolved_teamspace)

lightning_sdk/cli/legacy/list.py

@@ -8,7 +8,7 @@ from typing_extensions import Literal
 
 from lightning_sdk import Job, Machine, Studio, Teamspace
 from lightning_sdk.cli.legacy.clusters_menu import _ClustersMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi import V1MultiMachineJob
 from lightning_sdk.lit_container import LitContainer
 from lightning_sdk.utils.resolve import _get_authed_user
@@ -48,7 +48,7 @@ def studios(
 ) -> None:
     """List studios for a given teamspace."""
     studios = []
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     if all and not teamspace:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
@@ -56,7 +56,7 @@ def studios(
             teamspace = Teamspace(**ts)
             studios.extend(teamspace.studios)
     else:
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        resolved_teamspace = menu(teamspace=teamspace)
         studios = resolved_teamspace.studios
 
     table = Table(
@@ -111,7 +111,7 @@ def jobs(
 ) -> None:
     """List jobs for a given teamspace."""
     jobs = []
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     if all and not teamspace:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
@@ -119,7 +119,7 @@ def jobs(
             teamspace = Teamspace(**ts)
             jobs.extend(teamspace.jobs)
     else:
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        resolved_teamspace = menu(teamspace=teamspace)
         jobs = resolved_teamspace.jobs
 
     table = Table(pad_edge=True)
@@ -182,7 +182,7 @@ def mmts(
 ) -> None:
     """List multi-machine jobs for a given teamspace."""
     jobs = []
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     if all and not teamspace:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
@@ -190,7 +190,7 @@ def mmts(
             teamspace = Teamspace(**ts)
             jobs.extend(teamspace.multi_machine_jobs)
     else:
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        resolved_teamspace = menu(teamspace=teamspace)
         jobs = resolved_teamspace.multi_machine_jobs
 
     table = Table(pad_edge=True)
@@ -242,9 +242,9 @@ def mmts(
 def containers(teamspace: Optional[str] = None, cloud_account: Optional[str] = None) -> None:
     """Display the list of available containers."""
     api = LitContainer()
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     clusters_menu = _ClustersMenu()
-    resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+    resolved_teamspace = menu(teamspace=teamspace)
 
     if not cloud_account:
         cloud_account = clusters_menu._resolve_cluster(resolved_teamspace)

lightning_sdk/cli/legacy/open.py

@@ -6,8 +6,8 @@ from typing import Optional
 import click
 from rich.console import Console
 
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.cli.legacy.upload import _upload_folder
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.studio import Studio
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils.resolve import _get_studio_url
@@ -49,8 +49,8 @@ def open(path: str = ".", teamspace: Optional[str] = None, cloud_account: Option
     try:
         resolved_teamspace = Teamspace()
     except ValueError:
-        menu = _TeamspacesMenu()
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        menu = TeamspacesMenu()
+        resolved_teamspace = menu(teamspace=teamspace)
 
     # default cloud account to current studios cloud account if run from studio
     # else it will fall back to teamspace default in the backend

lightning_sdk/cli/legacy/upload.py

@@ -16,7 +16,7 @@ from lightning_sdk.api.lit_container_api import DockerNotRunningError, LCRAuthFa
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+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
@@ -137,8 +137,8 @@ def upload_container(
     platform: Optional[str] = "linux/amd64",
 ) -> None:
     """Upload a container to Lightning AI's container registry."""
-    menu = _TeamspacesMenu()
-    teamspace = menu._resolve_teamspace(teamspace)
+    menu = TeamspacesMenu()
+    teamspace = menu(teamspace)
     console = Console()
     with Progress(
         SpinnerColumn(),