lightning-sdk 0.1.29__py3-none-any.whl → 0.1.30__py3-none-any.whl

lightning_sdk/__init__.py

@@ -25,5 +25,5 @@ __all__ = [
     "Agent",
 ]
 
-__version__ = "0.1.29"
+__version__ = "0.1.30"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/cli/download.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Optional
 
 from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.models import _get_teamspace, _parse_model_name
 from lightning_sdk.cli.studios_menu import _StudiosMenu
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
@@ -12,6 +13,22 @@ from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 class _Downloads(_StudiosMenu):
     """Download files and folders from Lightning AI."""
 
+    def model(self, name: str, path: Optional[str] = None) -> None:
+        """Download a Model.
+
+        Args:
+          name: The name of the Model you want to download.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+          path: The path to the directory where the Model should be downloaded.
+        """
+        org_name, teamspace_name, model_name = _parse_model_name(name)
+        teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+        teamspace.download_model(
+            name=model_name,
+            download_dir=path or ".",
+            progress_bar=True,
+        )
+
     def _resolve_studio(self, studio: Optional[str]) -> Studio:
         user = _get_authed_user()
         # if no studio specify suggest/filter only user's studios
@@ -53,7 +70,7 @@ class _Downloads(_StudiosMenu):
             return Studio(**selected_studio)
 
     def folder(self, path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
-        """Download a folder from a studio.
+        """Download a folder from a Studio.
 
         Args:
           path: The relative path within the Studio you want to download.
@@ -86,7 +103,7 @@ class _Downloads(_StudiosMenu):
             ) from e
 
     def file(self, path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
-        """Download a file from a studio.
+        """Download a file from a Studio.
 
         Args:
           path: The relative path within the Studio you want to download.

lightning_sdk/cli/models.py

@@ -0,0 +1,38 @@
+from typing import Tuple
+
+from lightning_sdk.api import OrgApi, UserApi
+from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.utils.resolve import _get_authed_user
+
+
+def _parse_model_name(name: str) -> Tuple[str, str, str]:
+    """Parse the name argument into its components."""
+    try:
+        org_name, teamspace_name, model_name = name.split("/")
+    except ValueError as err:
+        raise StudioCliError(
+            f"Model name must be in the format 'organization/teamspace/model' but you provided '{name}'."
+        ) from err
+    return org_name, teamspace_name, model_name
+
+
+def _get_teamspace(name: str, organization: str) -> Teamspace:
+    """Get a Teamspace object from the SDK."""
+    org_api = OrgApi()
+    user = _get_authed_user()
+    teamspaces = {}
+    for ts in UserApi()._get_all_teamspace_memberships(""):
+        if ts.owner_type == "organization":
+            org = org_api._get_org_by_id(ts.owner_id)
+            teamspaces[f"{org.name}/{ts.name}"] = {"name": ts.name, "org": org.name}
+        elif ts.owner_type == "user":  # todo: check also the name
+            teamspaces[f"{user.name}/{ts.name}"] = {"name": ts.name, "user": user}
+        else:
+            raise StudioCliError(f"Unknown organization type {ts.owner_type}")
+
+    requested_teamspace = f"{organization}/{name}".lower()
+    if requested_teamspace not in teamspaces:
+        options = "\n\t".join(teamspaces.keys())
+        raise StudioCliError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: \n\t{options}")
+    return Teamspace(**teamspaces[requested_teamspace])

lightning_sdk/cli/upload.py

@@ -9,6 +9,7 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.models import _get_teamspace, _parse_model_name
 from lightning_sdk.cli.studios_menu import _StudiosMenu
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
@@ -19,6 +20,24 @@ class _Uploads(_StudiosMenu):
 
     _studio_upload_status_path = "~/.lightning/studios/uploads"
 
+    def model(self, name: str, path: Optional[str] = None, cloud_account: Optional[str] = None) -> None:
+        """Upload a Model.
+
+        Args:
+          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.
+        """
+        org_name, teamspace_name, model_name = _parse_model_name(name)
+        teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+        teamspace.upload_model(
+            path=path or ".",
+            name=model_name,
+            progress_bar=True,
+            cluster_id=cloud_account,
+        )
+
     def _resolve_studio(self, studio: Optional[str]) -> Studio:
         user = _get_authed_user()
         possible_studios = self._get_possible_studios(user)
@@ -43,7 +62,7 @@ class _Uploads(_StudiosMenu):
             return Studio(**selected_studio)
 
     def folder(self, path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
-        """Upload a file or folder to a studio.
+        """Upload a file or folder to a Studio.
 
         Args:
           path: The path to the file or directory you want to upload
@@ -97,7 +116,7 @@ class _Uploads(_StudiosMenu):
         print(f"See your files at {studio_url}")
 
     def file(self, path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
-        """Upload a file to a studio.
+        """Upload a file to a Studio.
 
         Args:
           path: The path to the file you want to upload

{lightning_sdk-0.1.29.dist-info → lightning_sdk-0.1.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.1.29
+Version: 0.1.30
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License

{lightning_sdk-0.1.29.dist-info → lightning_sdk-0.1.30.dist-info}/RECORD

@@ -1,5 +1,5 @@
 docs/source/conf.py,sha256=r8yX20eC-4mHhMTd0SbQb5TlSWHhO6wnJ0VJ_FBFpag,13249
-lightning_sdk/__init__.py,sha256=ZHv_QwJZQc1CtQzkcEr9KfO54uJJbiuHqwWmrd4yKZk,873
+lightning_sdk/__init__.py,sha256=3XBiirAYfvf8AM-l2cNOgRfPkmtJGyFdE9C7IFIjQFc,873
 lightning_sdk/agents.py,sha256=ly6Ma1j0ZgGPFyvPvMN28JWiB9dATIstFa5XM8pMi6I,1577
 lightning_sdk/constants.py,sha256=ztl1PTUBULnqTf3DyKUSJaV_O20hNtUYT6XvAYIrmIk,749
 lightning_sdk/helpers.py,sha256=RnQwUquc_YPotjh6YXOoJvZs8krX_QFhd7kGv4U_spQ,1844
@@ -21,12 +21,13 @@ lightning_sdk/api/teamspace_api.py,sha256=XXK9ApV7yya5tYUNr4uSVDQT0Z6KU-B2PDN1gO
 lightning_sdk/api/user_api.py,sha256=FcvJ7nVPrZIh94JGxtPscKMb-de-vV_FUOGEwbuyKcg,2312
 lightning_sdk/api/utils.py,sha256=Pnze_iy14cScQdOs8oELLXP-9WncpyigVEaohiT7uLM,20380
 lightning_sdk/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lightning_sdk/cli/download.py,sha256=LnG-cJftGRiysgz2GN4HV-Rz18fl20apn1vE4g17KiQ,5472
+lightning_sdk/cli/download.py,sha256=b6dpUMLhro6p0y6CleCHxggQnZlZMxCq0bjuf8B6_eA,6191
 lightning_sdk/cli/entrypoint.py,sha256=DD9rbmNqH6uE0IKMzfOenVymyrvfTMYi_zDPHr4Pk5E,1258
 lightning_sdk/cli/exceptions.py,sha256=QUF3OMAMZwBikvlusimSHSBjb6ywvHpfAumJBEaodSw,169
 lightning_sdk/cli/legacy.py,sha256=ocTVNwlsLRS5aMjbMkwFPjT3uEYvS8C40CJ0PeRRv8g,4707
+lightning_sdk/cli/models.py,sha256=_OT5sqk8lC40fKyutfx9XKvVhoj7QSrU7PhPJhBCSCQ,1647
 lightning_sdk/cli/studios_menu.py,sha256=0kQGqGel8gAbpdJtjOM1a6NEat_TnIqRNprNn8QiK58,3236
-lightning_sdk/cli/upload.py,sha256=kF567cxUkGWVQoiC-4ezxpPW7PtF2XdA4JEqxmSYpz0,8363
+lightning_sdk/cli/upload.py,sha256=05LVL5CnXuTefuU8dkyI5GLdSPoTkeLBCJIhznoKvVM,9246
 lightning_sdk/deployment/__init__.py,sha256=BLu7_cVLp97TYxe6qe-J1zKUSZXAVcvCjgcA7plV2k4,497
 lightning_sdk/deployment/deployment.py,sha256=47-aga7g8emJWrEkqr4FzfhMaNnsWhTSoinOl7FAMI0,15035
 lightning_sdk/job/__init__.py,sha256=1MxjQ6rHkyUHCypSW9RuXuVMVH11WiqhIXcU2LCFMwE,64
@@ -775,9 +776,9 @@ lightning_sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 lightning_sdk/utils/dynamic.py,sha256=glUTO1JC9APtQ6Gr9SO02a3zr56-sPAXM5C3NrTpgyQ,1959
 lightning_sdk/utils/enum.py,sha256=h2JRzqoBcSlUdanFHmkj_j5DleBHAu1esQYUsdNI-hU,4106
 lightning_sdk/utils/resolve.py,sha256=gU3MSko9Y7rE4jcnVwstNBaW83OFnSgvM-N44Ibrc_A,5148
-lightning_sdk-0.1.29.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
-lightning_sdk-0.1.29.dist-info/METADATA,sha256=y1Xcge9rW0vGORngiHSvJjWsleQ6wv0TFncwUuvKpJE,3941
-lightning_sdk-0.1.29.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-lightning_sdk-0.1.29.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
-lightning_sdk-0.1.29.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
-lightning_sdk-0.1.29.dist-info/RECORD,,
+lightning_sdk-0.1.30.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
+lightning_sdk-0.1.30.dist-info/METADATA,sha256=UaWQnogmihUzuy1qN_1MY9klBbNU4tHgA3x_wT666r8,3941
+lightning_sdk-0.1.30.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+lightning_sdk-0.1.30.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
+lightning_sdk-0.1.30.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
+lightning_sdk-0.1.30.dist-info/RECORD,,