lightning-sdk 0.1.37__py3-none-any.whl → 0.1.39__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/validate.py

@@ -44,8 +44,7 @@ class Validate(object):
         'aws': 'V1AwsDataConnection',
         'check_is_public': 'bool',
         'cluster_ids': 'list[str]',
-        'efs': 'V1EfsDataConnection',
-        'efs_folder': 'V1EFSFolderDataConnection',
+        'efs': 'V1EfsConfig',
         'gcp': 'V1GcpDataConnection',
         'gcs_folder': 'V1GCSFolderDataConnection',
         's3_folder': 'V1S3FolderDataConnection'
@@ -56,19 +55,17 @@ class Validate(object):
         'check_is_public': 'checkIsPublic',
         'cluster_ids': 'clusterIds',
         'efs': 'efs',
-        'efs_folder': 'efsFolder',
         'gcp': 'gcp',
         'gcs_folder': 'gcsFolder',
         's3_folder': 's3Folder'
     }
 
-    def __init__(self, aws: 'V1AwsDataConnection' =None, check_is_public: 'bool' =None, cluster_ids: 'list[str]' =None, efs: 'V1EfsDataConnection' =None, efs_folder: 'V1EFSFolderDataConnection' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
+    def __init__(self, aws: 'V1AwsDataConnection' =None, check_is_public: 'bool' =None, cluster_ids: 'list[str]' =None, efs: 'V1EfsConfig' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
         """Validate - a model defined in Swagger"""  # noqa: E501
         self._aws = None
         self._check_is_public = None
         self._cluster_ids = None
         self._efs = None
-        self._efs_folder = None
         self._gcp = None
         self._gcs_folder = None
         self._s3_folder = None
@@ -81,8 +78,6 @@ class Validate(object):
             self.cluster_ids = cluster_ids
         if efs is not None:
             self.efs = efs
-        if efs_folder is not None:
-            self.efs_folder = efs_folder
         if gcp is not None:
             self.gcp = gcp
         if gcs_folder is not None:
@@ -154,47 +149,26 @@ class Validate(object):
         self._cluster_ids = cluster_ids
 
     @property
-    def efs(self) -> 'V1EfsDataConnection':
+    def efs(self) -> 'V1EfsConfig':
         """Gets the efs of this Validate.  # noqa: E501
 
 
         :return: The efs of this Validate.  # noqa: E501
-        :rtype: V1EfsDataConnection
+        :rtype: V1EfsConfig
         """
         return self._efs
 
     @efs.setter
-    def efs(self, efs: 'V1EfsDataConnection'):
+    def efs(self, efs: 'V1EfsConfig'):
         """Sets the efs of this Validate.
 
 
         :param efs: The efs of this Validate.  # noqa: E501
-        :type: V1EfsDataConnection
+        :type: V1EfsConfig
         """
 
         self._efs = efs
 
-    @property
-    def efs_folder(self) -> 'V1EFSFolderDataConnection':
-        """Gets the efs_folder of this Validate.  # noqa: E501
-
-
-        :return: The efs_folder of this Validate.  # noqa: E501
-        :rtype: V1EFSFolderDataConnection
-        """
-        return self._efs_folder
-
-    @efs_folder.setter
-    def efs_folder(self, efs_folder: 'V1EFSFolderDataConnection'):
-        """Sets the efs_folder of this Validate.
-
-
-        :param efs_folder: The efs_folder of this Validate.  # noqa: E501
-        :type: V1EFSFolderDataConnection
-        """
-
-        self._efs_folder = efs_folder
-
     @property
     def gcp(self) -> 'V1GcpDataConnection':
         """Gets the gcp of this Validate.  # noqa: E501

lightning_sdk/models.py

@@ -0,0 +1,132 @@
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from lightning_sdk.api import OrgApi, TeamspaceApi, UserApi
+from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _get_authed_user
+
+
+def _get_teamspace_and_path(
+    ts: V1Membership, org_api: OrgApi, user_api: UserApi, authed_user: User
+) -> Tuple[str, Dict[str, Any]]:
+    if ts.owner_type == V1OwnerType.ORGANIZATION:
+        org = org_api._get_org_by_id(ts.owner_id)
+        return f"{org.name}/{ts.name}", {"name": ts.name, "org": org.name}
+
+    if ts.owner_type == V1OwnerType.USER and ts.owner_id != authed_user.id:
+        user = user_api._get_user_by_id(ts.owner_id)  # todo: check also the name
+        return f"{user.username}/{ts.name}", {"name": ts.name, "user": User(name=user.username)}
+
+    if ts.owner_type == V1OwnerType.USER:
+        return f"{authed_user.name}/{ts.name}", {"name": ts.name, "user": authed_user}
+
+    raise RuntimeError(f"Unknown organization type {ts.owner_type}")
+
+
+def _list_teamspaces() -> List[str]:
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    return [
+        _get_teamspace_and_path(ts, org_api, user_api, authed_user)[0]
+        for ts in user_api._get_all_teamspace_memberships("")
+    ]
+
+
+def _get_teamspace(name: str, organization: str) -> Teamspace:
+    """Get a Teamspace object from the SDK."""
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    requested_teamspace = f"{organization}/{name}".lower()
+
+    for ts in user_api._get_all_teamspace_memberships(""):
+        if ts.name != name:
+            continue
+
+        teamspace_path, teamspace = _get_teamspace_and_path(ts, org_api, user_api, authed_user)
+        if requested_teamspace == teamspace_path:
+            return Teamspace(**teamspace)
+
+    options = f"{os.linesep}\t".join(_list_teamspaces())
+    raise RuntimeError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: {os.linesep}\t{options}")
+
+
+def _parse_model_name_and_version(name: str) -> Tuple[str, str, str, str]:
+    """Parse the name argument into its components."""
+    try:
+        org_name, teamspace_name, model_name = name.split("/")
+        parts = model_name.split(":")
+        if len(parts) == 1:
+            return org_name, teamspace_name, parts[0], "latest"
+        if len(parts) == 2:
+            return org_name, teamspace_name, parts[0], parts[1]
+        # The rest of the validation for name and version happens in the backend
+        raise ValueError(
+            "Model version is expected to be in the format `entity/modelname:version` separated by a"
+            f" single colon, but got: {name}"
+        )
+    except ValueError as err:
+        raise ValueError(
+            f"Model name must be in the format 'organization/teamspace/model' but you provided '{name}'."
+        ) from err
+
+
+def download_model(
+    name: str,
+    download_dir: Union[Path, str] = ".",
+    progress_bar: bool = True,
+) -> List[str]:
+    """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>.
+        download_dir: The directory where the Model should be downloaded.
+        progress_bar: Whether to show a progress bar when downloading.
+    """
+    teamspace_owner_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
+
+    download_dir = Path(download_dir)
+
+    api = TeamspaceApi()
+
+    try:
+        return api.download_model_files(
+            name=model_name,
+            version=version,
+            download_dir=download_dir,
+            teamspace_name=teamspace_name,
+            teamspace_owner_name=teamspace_owner_name,
+            progress_bar=progress_bar,
+        )
+    except ApiException as e:
+        # if we get an error, check if the teamspace actually exists (and print the list)
+        # TODO: ideally this would match a specific error about teamspace not being found
+        _ = _get_teamspace(name=teamspace_name, organization=teamspace_owner_name)
+        raise e
+
+
+def upload_model(name: str, path: Union[Path, str] = ".", 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_and_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    teamspace.upload_model(
+        path=path,
+        name=model_name,
+        progress_bar=True,
+        cluster_id=cloud_account,
+    )

lightning_sdk/teamspace.py

@@ -153,7 +153,7 @@ class Teamspace:
 
     def upload_model(
         self,
-        path: str,
+        path: Union[str, Path],
         name: str,
         progress_bar: bool = True,
         cluster_id: Optional[str] = None,
@@ -245,11 +245,17 @@ class Teamspace:
         download_dir = Path(download_dir)
 
         name, version = _parse_model_and_version(name)
+        model_version = self._teamspace_api.get_model_version(name=name, version=version, teamspace_id=self.id)
+        if not model_version.upload_complete:
+            raise RuntimeError(
+                f"Model {name}:{version} is not fully uploaded yet. Please wait until the upload is complete."
+            )
         downloaded_files = self._teamspace_api.download_model_files(
             name=name,
             version=version,
             download_dir=download_dir,
-            teamspace_id=self.id,
+            teamspace_name=self.name,
+            teamspace_owner_name=self.owner.name,
             progress_bar=progress_bar,
         )
 

{lightning_sdk-0.1.37.dist-info → lightning_sdk-0.1.39.dist-info}/METADATA

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