lightning-sdk 0.1.3__py3-none-any.whl → 0.1.47__py3-none-any.whl

lightning_sdk/cli/studios_menu.py

@@ -0,0 +1,78 @@
+from typing import Dict, List, Optional
+
+from simple_term_menu import TerminalMenu
+
+from lightning_sdk.api import OrgApi, TeamspaceApi
+from lightning_sdk.user import User
+
+
+class _StudiosMenu:
+    def _get_studio_from_interactive_menu(self, possible_studios: List[Dict[str, str]]) -> Dict[str, str]:
+        terminal_menu = self._prepare_terminal_menu_list_studios(possible_studios)
+        terminal_menu.show()
+        return possible_studios[terminal_menu.chosen_menu_index]
+
+    def _get_studio_from_name(self, studio: str, possible_studios: List[Dict[str, str]]) -> Dict[str, str]:
+        teamspace, name = studio.split("/", maxsplit=1)
+        for st in possible_studios:
+            if st["teamspace"] == teamspace and name == st["name"]:
+                return st
+
+        print("Could not find Studio {studio}, please select it from the list:")
+        return self._get_studio_from_interactive_menu(possible_studios)
+
+    @staticmethod
+    def _prepare_terminal_menu_list_studios(
+        possible_studios: List[Dict[str, str]], title: Optional[str] = None
+    ) -> TerminalMenu:
+        if title is None:
+            title = "Please select a Studio of the following studios:"
+
+        return TerminalMenu(
+            [f"{s['teamspace']}/{s['name']}" for s in possible_studios], title=title, clear_menu_on_exit=True
+        )
+
+    @staticmethod
+    def _get_possible_studios(user: User, is_owner: bool = True) -> List[Dict[str, str]]:
+        teamspace_api = TeamspaceApi()
+        org_api = OrgApi()
+        user_api = user._user_api
+        possible_studios = []
+
+        user_api._get_organizations_for_authed_user()
+        memberships = user_api._get_all_teamspace_memberships(user_id=user.id)
+
+        teamspaces = {}
+        # get all teamspace memberships
+        for membership in memberships:
+            teamspace_id = membership.project_id
+
+            if is_owner:
+                # get all studios for teamspace when user is owner
+                all_studios = user._user_api._get_cloudspaces_for_user(user_id=user.id, project_id=teamspace_id)
+            else:
+                all_studios = user._user_api._get_cloudspaces_for_user(project_id=teamspace_id)
+
+            for st in all_studios:
+                # populate teamspace info if necessary
+                if teamspace_id not in teamspaces:
+                    ts = teamspace_api._get_teamspace_by_id(teamspace_id)
+                    ts_name = ts.name
+
+                    # get organization if necessary
+                    if ts.owner_type == "organization":
+                        org_name = org_api._get_org_by_id(ts.owner_id).name
+                        user_name = None
+                    else:
+                        org_name = None
+
+                        # don't do a request if not necessary
+                        if ts.owner_id == user.id:
+                            user_name = user.name
+                        else:
+                            user_name = user_api._get_user_by_id(ts.owner_id).username
+
+                    teamspaces[teamspace_id] = {"user": user_name, "org": org_name, "teamspace": ts_name}
+                possible_studios.append({"name": st.name, **teamspaces[teamspace_id]})
+
+        return possible_studios

lightning_sdk/cli/teamspace_menu.py

@@ -0,0 +1,94 @@
+from typing import Dict, List, Optional
+
+from simple_term_menu import TerminalMenu
+
+from lightning_sdk.api import OrgApi
+from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _get_authed_user
+
+
+class _TeamspacesMenu:
+    def _get_teamspace_from_interactive_menu(self, possible_teamspaces: Dict[str, Dict[str, str]]) -> Dict[str, str]:
+        teamspace_ids = sorted(possible_teamspaces.keys())
+        terminal_menu = self._prepare_terminal_menu_teamspaces([possible_teamspaces[k] for k in teamspace_ids])
+        terminal_menu.show()
+
+        selected_id = teamspace_ids[terminal_menu.chosen_menu_index]
+        return possible_teamspaces[selected_id]
+
+    def _get_teamspace_from_name(
+        self, teamspace: str, possible_teamspaces: Dict[str, Dict[str, str]]
+    ) -> Dict[str, str]:
+        owner, name = teamspace.split("/", maxsplit=1)
+        for _, ts in possible_teamspaces.items():
+            if ts["name"] == name and (ts["user"] == owner or ts["org"] == owner):
+                return ts
+
+        print("Could not find Teamspace {teamspace}, please select it from the list:")
+        return self._get_teamspace_from_interactive_menu(possible_teamspaces)
+
+    @staticmethod
+    def _prepare_terminal_menu_teamspaces(
+        possible_teamspaces: List[Dict[str, str]], title: Optional[str] = None
+    ) -> TerminalMenu:
+        if title is None:
+            title = "Please select a Teamspace of the following:"
+
+        return TerminalMenu(
+            [f"{t['user'] or t['org']}/{t['name']}" for t in possible_teamspaces], title=title, clear_menu_on_exit=True
+        )
+
+    @staticmethod
+    def _get_possible_teamspaces(user: User, is_owner: bool = True) -> Dict[str, Dict[str, str]]:
+        org_api = OrgApi()
+        user_api = user._user_api
+
+        user_api._get_organizations_for_authed_user()
+        memberships = user_api._get_all_teamspace_memberships(user_id=user.id)
+
+        teamspaces = {}
+        # get all teamspace memberships
+        for membership in memberships:
+            teamspace_id = membership.project_id
+            teamspace_name = membership.name
+
+            # get organization if necessary
+            if membership.owner_type == "organization":
+                org_name = org_api._get_org_by_id(membership.owner_id).name
+                user_name = None
+            else:
+                org_name = None
+
+                # don't do a request if not necessary
+                if membership.owner_id == user.id:
+                    user_name = user.name
+                else:
+                    user_name = user_api._get_user_by_id(membership.owner_id).username
+
+            teamspaces[teamspace_id] = {"user": user_name, "org": org_name, "name": teamspace_name}
+
+        return teamspaces
+
+    def _resolve_teamspace(self, teamspace: Optional[str]) -> Teamspace:
+        try:
+            user = _get_authed_user()
+
+            possible_teamspaces = self._get_possible_teamspaces(user)
+            if teamspace is None:
+                teamspace_dict = self._get_teamspace_from_interactive_menu(possible_teamspaces=possible_teamspaces)
+            else:
+                teamspace_dict = self._get_teamspace_from_name(
+                    teamspace=teamspace, possible_teamspaces=possible_teamspaces
+                )
+
+            return Teamspace(**teamspace_dict)
+        except KeyboardInterrupt:
+            raise KeyboardInterrupt from None
+
+        except Exception as e:
+            raise StudioCliError(
+                f"Could not find the given Teamspace {teamspace}. "
+                "Please contact Lightning AI directly to resolve this issue."
+            ) from e

lightning_sdk/cli/upload.py

@@ -1,38 +1,37 @@
 import concurrent.futures
 import json
 import os
+from pathlib import Path
 from typing import Dict, List, Optional
 
 from simple_term_menu import TerminalMenu
 from tqdm import tqdm
 
-from lightning_sdk.api import OrgApi, TeamspaceApi
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.studios_menu import _StudiosMenu
+from lightning_sdk.models import upload_model
 from lightning_sdk.studio import Studio
-from lightning_sdk.user import User
-from lightning_sdk.utils import _get_authed_user, skip_studio_init
+from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 
 
-class _Uploads:
+class _Uploads(_StudiosMenu):
+    """Upload files and folders to Lightning AI."""
+
     _studio_upload_status_path = "~/.lightning/studios/uploads"
 
-    def upload(self, path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
-        """Upload a file or folder to a studio.
+    def model(self, name: str, path: str = ".", cloud_account: Optional[str] = None) -> None:
+        """Upload a Model.
 
         Args:
-          path: The path to the file or directory you want to upload
-          studio: The name of the studio to upload to. Will show a menu for selection if not specified.
-            If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME>
-          remote_path: The path where the uploaded file should appear on your Studio.
-            Has to be within your Studio's home directory and will be relative to that.
-            If not specified, will use the file or directory name of the path you want to upload
-            and place it in your home directory.
-
+          name: The name of the Model you want to upload.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+          path: The path to the file or directory you want to upload. Defaults to the current directory.
+          cloud_account: The name of the cloud account to store the Model in.
         """
-        if remote_path is None:
-            remote_path = os.path.basename(path)
+        upload_model(name, path, cloud_account=cloud_account)
 
+    def _resolve_studio(self, studio: Optional[str]) -> Studio:
         user = _get_authed_user()
         possible_studios = self._get_possible_studios(user)
 
@@ -52,19 +51,38 @@ class _Uploads:
                 "Please contact Lightning AI directly to resolve this issue."
             ) from e
 
-        print(f"Uploading to {selected_studio['teamspace']}/{selected_studio['name']}")
-        pairs = {}
-        if os.path.isdir(path):
-            for root, _, files in os.walk(path):
-                rel_root = os.path.relpath(root, path)
-                for f in files:
-                    pairs[os.path.join(root, f)] = os.path.join(remote_path, rel_root, f)
+        with skip_studio_init():
+            return Studio(**selected_studio)
 
-        else:
-            pairs[path] = remote_path
+    def folder(self, path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
+        """Upload a file or folder to a Studio.
 
-        with skip_studio_init():
-            selected_studio = Studio(**selected_studio)
+        Args:
+          path: The path to the file or directory you want to upload
+          studio: The name of the studio to upload to. Will show a menu for selection if not specified.
+            If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME>
+          remote_path: The path where the uploaded file should appear on your Studio.
+            Has to be within your Studio's home directory and will be relative to that.
+            If not specified, will use the file or directory name of the path you want to upload
+            and place it in your home directory.
+        """
+        if remote_path is None:
+            remote_path = os.path.basename(path)
+
+        if not Path(path).exists():
+            raise FileNotFoundError(f"The provided path does not exist: {path}.")
+        if not Path(path).is_dir():
+            raise StudioCliError(f"The provided path is not a folder: {path}. Use `lightning upload file` instead.")
+
+        selected_studio = self._resolve_studio(studio)
+
+        print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")
+
+        pairs = {}
+        for root, _, files in os.walk(path):
+            rel_root = os.path.relpath(root, path)
+            for f in files:
+                pairs[os.path.join(root, f)] = os.path.join(remote_path, rel_root, f)
 
         upload_state = self._resolve_previous_upload_state(selected_studio, remote_path, pairs)
 
@@ -91,19 +109,42 @@ class _Uploads:
         )
         print(f"See your files at {studio_url}")
 
-    def _get_studio_from_interactive_menu(self, possible_studios: List[Dict[str, str]]) -> Dict[str, str]:
-        terminal_menu = self._prepare_terminal_menu_all_studios(possible_studios)
-        terminal_menu.show()
-        return possible_studios[terminal_menu.chosen_menu_index]
+    def file(self, path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
+        """Upload a file to a Studio.
+
+        Args:
+          path: The path to the file you want to upload
+          studio: The name of the studio to upload to. Will show a menu for selection if not specified.
+            If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME>
+          remote_path: The path where the uploaded file should appear on your Studio.
+            Has to be within your Studio's home directory and will be relative to that.
+            If not specified, will use the name of the file you want to upload
+            and place it in your home directory.
+        """
+        if remote_path is None:
+            remote_path = os.path.basename(path)
+
+        if Path(path).is_dir():
+            raise StudioCliError(f"The provided path is a folder: {path}. Use `lightning upload folder` instead.")
+        if not Path(path).exists():
+            raise FileNotFoundError(f"The provided path does not exist: {path}.")
+
+        selected_studio = self._resolve_studio(studio)
 
-    def _get_studio_from_name(self, studio: str, possible_studios: List[Dict[str, str]]) -> Dict[str, str]:
-        teamspace, name = studio.split("/", maxsplit=1)
-        for st in possible_studios:
-            if st["teamspace"] == teamspace and name == st["name"]:
-                return st
+        print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")
 
-        print("Could not find Studio {studio}, please select it from the list:")
-        return self._get_studio_from_interactive_menu(possible_studios)
+        self._single_file_upload(selected_studio, path, remote_path, True)
+
+        studio_url = (
+            _get_cloud_url().replace(":443", "")
+            + "/"
+            + selected_studio.owner.name
+            + "/"
+            + selected_studio.teamspace.name
+            + "/studios/"
+            + selected_studio.name
+        )
+        print(f"See your file at {studio_url}")
 
     def _start_parallel_upload(
         self, executor: concurrent.futures.ThreadPoolExecutor, studio: Studio, upload_state: Dict[str, str]
@@ -125,57 +166,6 @@ class _Uploads:
         studio.upload_file(local_path, remote_path, progress_bar)
         return local_path
 
-    def _prepare_terminal_menu_all_studios(
-        self, possible_studios: List[Dict[str, str]], title: Optional[str] = None
-    ) -> TerminalMenu:
-        if title is None:
-            title = "Please select a Studio of the following studios:"
-
-        return TerminalMenu(
-            [f"{s['teamspace']}/{s['name']}" for s in possible_studios], title=title, clear_menu_on_exit=True
-        )
-
-    def _get_possible_studios(self, user: User) -> List[Dict[str, str]]:
-        teamspace_api = TeamspaceApi()
-        org_api = OrgApi()
-        user_api = user._user_api
-        possible_studios = []
-
-        user_api._get_organizations_for_authed_user()
-        memberships = user_api._get_all_teamspace_memberships(user_id=user.id)
-
-        teamspaces = {}
-        # get all teamspace memberships
-        for membership in memberships:
-            teamspace_id = membership.project_id
-
-            # get all studios for teamspace
-            all_studios = user._user_api._get_cloudspaces_for_user(user.id, teamspace_id)
-
-            for st in all_studios:
-                # populate teamspace info if necessary
-                if teamspace_id not in teamspaces:
-                    ts = teamspace_api._get_teamspace_by_id(teamspace_id)
-                    ts_name = ts.name
-
-                    # get organization if necessary
-                    if ts.owner_type == "organization":
-                        org_name = org_api._get_org_by_id(ts.owner_id).name
-                        user_name = None
-                    else:
-                        org_name = None
-
-                        # don't do a request if not necessary
-                        if ts.owner_id == user.id:
-                            user_name = user.name
-                        else:
-                            user_name = user_api._get_user_by_id(ts.owner_id).username
-
-                    teamspaces[teamspace_id] = {"user": user_name, "org": org_name, "teamspace": ts_name}
-                possible_studios.append({"name": st.name, **teamspaces[teamspace_id]})
-
-        return possible_studios
-
     def _dump_current_upload_state(self, studio: Studio, remote_path: str, state_dict: Dict[str, str]) -> None:
         """Dumps the current upload state so that we can safely resume later."""
         curr_path = os.path.abspath(

lightning_sdk/constants.py

@@ -1,3 +1,31 @@
 import os
+from uuid import uuid4
 
-_LIGHTNING_DEBUG = bool(int(os.getenv("LIGHTNING_DEBUG", "0")))
+_LIGHTNING_DEBUG = {
+    "": False,
+    "0": False,
+    "false": False,
+    "no": False,
+    "1": True,
+    "true": True,
+    "yes": True,
+}.get(os.getenv("LIGHTNING_DEBUG", "").lower(), False)
+
+
+class Store:
+    def __init__(self) -> None:
+        self._d = {}
+        # This is needed to ensure the ids are the same within created threads and processes.
+        self._list = [str(uuid4().hex) for _ in range(500)]
+
+    def __getitem__(self, key: str) -> str:
+        """Get item."""
+        if key in self._d:
+            return self._d[key]
+
+        value = self._list.pop(0) if self._list else str(uuid4().hex)
+        self._d[key] = value
+        return value
+
+
+__GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__ = Store()

lightning_sdk/deployment/__init__.py

@@ -0,0 +1,25 @@
+from lightning_sdk.api.deployment_api import (
+    AutoScaleConfig,
+    BasicAuth,
+    Env,
+    ExecHealthCheck,
+    HttpHealthCheck,
+    ReleaseStrategy,
+    RollingUpdateReleaseStrategy,
+    Secret,
+    TokenAuth,
+)
+from lightning_sdk.deployment.deployment import Deployment
+
+__all__ = [
+    "AutoScaleConfig",
+    "BasicAuth",
+    "Env",
+    "ExecHealthCheck",
+    "HttpHealthCheck",
+    "ReleaseStrategy",
+    "RollingUpdateReleaseStrategy",
+    "Secret",
+    "TokenAuth",
+    "Deployment",
+]