lightning-sdk 0.1.38__py3-none-any.whl → 0.1.40__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_query_param.py

@@ -0,0 +1,175 @@
+# coding: utf-8
+
+"""
+    external/v1/auth_service.proto
+
+    No description provided (generated by Swagger Codegen https://github.com/swagger-api/swagger-codegen)  # noqa: E501
+
+    OpenAPI spec version: version not set
+    
+    Generated by: https://github.com/swagger-api/swagger-codegen.git
+
+    NOTE
+    ----
+    standard swagger-codegen-cli for this python client has been modified
+    by custom templates. The purpose of these templates is to include
+    typing information in the API and Model code. Please refer to the
+    main grid repository for more info
+"""
+
+import pprint
+import re  # noqa: F401
+
+from typing import TYPE_CHECKING
+
+import six
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from lightning_sdk.lightning_cloud.openapi.models import *
+
+class V1QueryParam(object):
+    """NOTE: This class is auto generated by the swagger code generator program.
+
+    Do not edit the class manually.
+    """
+    """
+    Attributes:
+      swagger_types (dict): The key is attribute name
+                            and the value is attribute type.
+      attribute_map (dict): The key is attribute name
+                            and the value is json key in definition.
+    """
+    swagger_types = {
+        'description': 'str',
+        'key': 'str',
+        'value': 'str'
+    }
+
+    attribute_map = {
+        'description': 'description',
+        'key': 'key',
+        'value': 'value'
+    }
+
+    def __init__(self, description: 'str' =None, key: 'str' =None, value: 'str' =None):  # noqa: E501
+        """V1QueryParam - a model defined in Swagger"""  # noqa: E501
+        self._description = None
+        self._key = None
+        self._value = None
+        self.discriminator = None
+        if description is not None:
+            self.description = description
+        if key is not None:
+            self.key = key
+        if value is not None:
+            self.value = value
+
+    @property
+    def description(self) -> 'str':
+        """Gets the description of this V1QueryParam.  # noqa: E501
+
+
+        :return: The description of this V1QueryParam.  # noqa: E501
+        :rtype: str
+        """
+        return self._description
+
+    @description.setter
+    def description(self, description: 'str'):
+        """Sets the description of this V1QueryParam.
+
+
+        :param description: The description of this V1QueryParam.  # noqa: E501
+        :type: str
+        """
+
+        self._description = description
+
+    @property
+    def key(self) -> 'str':
+        """Gets the key of this V1QueryParam.  # noqa: E501
+
+
+        :return: The key of this V1QueryParam.  # noqa: E501
+        :rtype: str
+        """
+        return self._key
+
+    @key.setter
+    def key(self, key: 'str'):
+        """Sets the key of this V1QueryParam.
+
+
+        :param key: The key of this V1QueryParam.  # noqa: E501
+        :type: str
+        """
+
+        self._key = key
+
+    @property
+    def value(self) -> 'str':
+        """Gets the value of this V1QueryParam.  # noqa: E501
+
+
+        :return: The value of this V1QueryParam.  # noqa: E501
+        :rtype: str
+        """
+        return self._value
+
+    @value.setter
+    def value(self, value: 'str'):
+        """Sets the value of this V1QueryParam.
+
+
+        :param value: The value of this V1QueryParam.  # noqa: E501
+        :type: str
+        """
+
+        self._value = value
+
+    def to_dict(self) -> dict:
+        """Returns the model properties as a dict"""
+        result = {}
+
+        for attr, _ in six.iteritems(self.swagger_types):
+            value = getattr(self, attr)
+            if isinstance(value, list):
+                result[attr] = list(map(
+                    lambda x: x.to_dict() if hasattr(x, "to_dict") else x,
+                    value
+                ))
+            elif hasattr(value, "to_dict"):
+                result[attr] = value.to_dict()
+            elif isinstance(value, dict):
+                result[attr] = dict(map(
+                    lambda item: (item[0], item[1].to_dict())
+                    if hasattr(item[1], "to_dict") else item,
+                    value.items()
+                ))
+            else:
+                result[attr] = value
+        if issubclass(V1QueryParam, dict):
+            for key, value in self.items():
+                result[key] = value
+
+        return result
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model"""
+        return pprint.pformat(self.to_dict())
+
+    def __repr__(self) -> str:
+        """For `print` and `pprint`"""
+        return self.to_str()
+
+    def __eq__(self, other: 'V1QueryParam') -> bool:
+        """Returns true if both objects are equal"""
+        if not isinstance(other, V1QueryParam):
+            return False
+
+        return self.__dict__ == other.__dict__
+
+    def __ne__(self, other: 'V1QueryParam') -> bool:
+        """Returns true if both objects are not equal"""
+        return not self == other

lightning_sdk/lightning_cloud/openapi/models/{v1_list_managed_models_response.py → v1_resource_visibility.py}

@@ -28,7 +28,7 @@ if TYPE_CHECKING:
     from datetime import datetime
     from lightning_sdk.lightning_cloud.openapi.models import *
 
-class V1ListManagedModelsResponse(object):
+class V1ResourceVisibility(object):
     """NOTE: This class is auto generated by the swagger code generator program.
 
     Do not edit the class manually.
@@ -41,40 +41,40 @@ class V1ListManagedModelsResponse(object):
                             and the value is json key in definition.
     """
     swagger_types = {
-        'models': 'list[V1ManagedModel]'
+        'all_users': 'bool'
     }
 
     attribute_map = {
-        'models': 'models'
+        'all_users': 'allUsers'
     }
 
-    def __init__(self, models: 'list[V1ManagedModel]' =None):  # noqa: E501
-        """V1ListManagedModelsResponse - a model defined in Swagger"""  # noqa: E501
-        self._models = None
+    def __init__(self, all_users: 'bool' =None):  # noqa: E501
+        """V1ResourceVisibility - a model defined in Swagger"""  # noqa: E501
+        self._all_users = None
         self.discriminator = None
-        if models is not None:
-            self.models = models
+        if all_users is not None:
+            self.all_users = all_users
 
     @property
-    def models(self) -> 'list[V1ManagedModel]':
-        """Gets the models of this V1ListManagedModelsResponse.  # noqa: E501
+    def all_users(self) -> 'bool':
+        """Gets the all_users of this V1ResourceVisibility.  # noqa: E501
 
 
-        :return: The models of this V1ListManagedModelsResponse.  # noqa: E501
-        :rtype: list[V1ManagedModel]
+        :return: The all_users of this V1ResourceVisibility.  # noqa: E501
+        :rtype: bool
         """
-        return self._models
+        return self._all_users
 
-    @models.setter
-    def models(self, models: 'list[V1ManagedModel]'):
-        """Sets the models of this V1ListManagedModelsResponse.
+    @all_users.setter
+    def all_users(self, all_users: 'bool'):
+        """Sets the all_users of this V1ResourceVisibility.
 
 
-        :param models: The models of this V1ListManagedModelsResponse.  # noqa: E501
-        :type: list[V1ManagedModel]
+        :param all_users: The all_users of this V1ResourceVisibility.  # noqa: E501
+        :type: bool
         """
 
-        self._models = models
+        self._all_users = all_users
 
     def to_dict(self) -> dict:
         """Returns the model properties as a dict"""
@@ -97,7 +97,7 @@ class V1ListManagedModelsResponse(object):
                 ))
             else:
                 result[attr] = value
-        if issubclass(V1ListManagedModelsResponse, dict):
+        if issubclass(V1ResourceVisibility, dict):
             for key, value in self.items():
                 result[key] = value
 
@@ -111,13 +111,13 @@ class V1ListManagedModelsResponse(object):
         """For `print` and `pprint`"""
         return self.to_str()
 
-    def __eq__(self, other: 'V1ListManagedModelsResponse') -> bool:
+    def __eq__(self, other: 'V1ResourceVisibility') -> bool:
         """Returns true if both objects are equal"""
-        if not isinstance(other, V1ListManagedModelsResponse):
+        if not isinstance(other, V1ResourceVisibility):
             return False
 
         return self.__dict__ == other.__dict__
 
-    def __ne__(self, other: 'V1ListManagedModelsResponse') -> bool:
+    def __ne__(self, other: 'V1ResourceVisibility') -> bool:
         """Returns true if both objects are not equal"""
         return not self == other

lightning_sdk/lightning_cloud/openapi/models/{v1_delete_managed_model_response.py → v1_update_model_visibility_response.py}

@@ -28,7 +28,7 @@ if TYPE_CHECKING:
     from datetime import datetime
     from lightning_sdk.lightning_cloud.openapi.models import *
 
-class V1DeleteManagedModelResponse(object):
+class V1UpdateModelVisibilityResponse(object):
     """NOTE: This class is auto generated by the swagger code generator program.
 
     Do not edit the class manually.
@@ -47,7 +47,7 @@ class V1DeleteManagedModelResponse(object):
     }
 
     def __init__(self):  # noqa: E501
-        """V1DeleteManagedModelResponse - a model defined in Swagger"""  # noqa: E501
+        """V1UpdateModelVisibilityResponse - a model defined in Swagger"""  # noqa: E501
         self.discriminator = None
 
     def to_dict(self) -> dict:
@@ -71,7 +71,7 @@ class V1DeleteManagedModelResponse(object):
                 ))
             else:
                 result[attr] = value
-        if issubclass(V1DeleteManagedModelResponse, dict):
+        if issubclass(V1UpdateModelVisibilityResponse, dict):
             for key, value in self.items():
                 result[key] = value
 
@@ -85,13 +85,13 @@ class V1DeleteManagedModelResponse(object):
         """For `print` and `pprint`"""
         return self.to_str()
 
-    def __eq__(self, other: 'V1DeleteManagedModelResponse') -> bool:
+    def __eq__(self, other: 'V1UpdateModelVisibilityResponse') -> bool:
         """Returns true if both objects are equal"""
-        if not isinstance(other, V1DeleteManagedModelResponse):
+        if not isinstance(other, V1UpdateModelVisibilityResponse):
             return False
 
         return self.__dict__ == other.__dict__
 
-    def __ne__(self, other: 'V1DeleteManagedModelResponse') -> bool:
+    def __ne__(self, other: 'V1UpdateModelVisibilityResponse') -> bool:
         """Returns true if both objects are not equal"""
         return not self == other

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -76,6 +76,7 @@ class V1UserFeatures(object):
         'filesystem_optimisation': 'bool',
         'gcp': 'bool',
         'inference_job_deployment_plugin': 'bool',
+        'instant_capacity_reservation': 'bool',
         'jobs_init': 'bool',
         'jobs_v2': 'bool',
         'landing_studios': 'bool',
@@ -160,6 +161,7 @@ class V1UserFeatures(object):
         'filesystem_optimisation': 'filesystemOptimisation',
         'gcp': 'gcp',
         'inference_job_deployment_plugin': 'inferenceJobDeploymentPlugin',
+        'instant_capacity_reservation': 'instantCapacityReservation',
         'jobs_init': 'jobsInit',
         'jobs_v2': 'jobsV2',
         'landing_studios': 'landingStudios',
@@ -208,7 +210,7 @@ class V1UserFeatures(object):
         'writable_data_connections': 'writableDataConnections'
     }
 
-    def __init__(self, advanced_deployment_autoscaling: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, aws_trainium: 'bool' =None, b2c_experience: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, cluster_proxy: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, crypto_monitoring: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, default_one_cluster: 'bool' =None, deployment_customize_api: 'bool' =None, deployment_data_path: 'bool' =None, deployment_gallery: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_version_visibility: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, dynamic_workload_scheduling: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_efs: 'bool' =None, enable_storage_limits: 'bool' =None, featured_studios_admin: 'bool' =None, filesystem_optimisation: 'bool' =None, gcp: 'bool' =None, inference_job_deployment_plugin: 'bool' =None, jobs_init: 'bool' =None, jobs_v2: 'bool' =None, landing_studios: 'bool' =None, lightning_registry: 'bool' =None, lit_logger: 'bool' =None, lit_logger_storage_v2: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, model_store: 'bool' =None, multiple_deployment_versions: 'bool' =None, multiple_studio_versions: 'bool' =None, org_level_member_permissions: 'bool' =None, plugin_biz_chat: 'bool' =None, plugin_distributed: 'bool' =None, plugin_fiftyone: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_lightning_apps: 'bool' =None, plugin_lightning_apps_distributed: 'bool' =None, plugin_mage_ai: 'bool' =None, plugin_milvus: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_react: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, plugin_weviate: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, restart_ide_on_hang: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, s3_folders: 'bool' =None, show_dev_admin: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, snapshotter_service: 'bool' =None, snowflake_connection: 'bool' =None, spot_v2: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, teamspace_storage_tab: 'bool' =None, use_rclone_mounts_only: 'bool' =None, writable_data_connections: 'bool' =None):  # noqa: E501
+    def __init__(self, advanced_deployment_autoscaling: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, aws_trainium: 'bool' =None, b2c_experience: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, cluster_proxy: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, crypto_monitoring: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, default_one_cluster: 'bool' =None, deployment_customize_api: 'bool' =None, deployment_data_path: 'bool' =None, deployment_gallery: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_version_visibility: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, dynamic_workload_scheduling: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_efs: 'bool' =None, enable_storage_limits: 'bool' =None, featured_studios_admin: 'bool' =None, filesystem_optimisation: 'bool' =None, gcp: 'bool' =None, inference_job_deployment_plugin: 'bool' =None, instant_capacity_reservation: 'bool' =None, jobs_init: 'bool' =None, jobs_v2: 'bool' =None, landing_studios: 'bool' =None, lightning_registry: 'bool' =None, lit_logger: 'bool' =None, lit_logger_storage_v2: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, model_store: 'bool' =None, multiple_deployment_versions: 'bool' =None, multiple_studio_versions: 'bool' =None, org_level_member_permissions: 'bool' =None, plugin_biz_chat: 'bool' =None, plugin_distributed: 'bool' =None, plugin_fiftyone: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_lightning_apps: 'bool' =None, plugin_lightning_apps_distributed: 'bool' =None, plugin_mage_ai: 'bool' =None, plugin_milvus: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_react: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, plugin_weviate: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, restart_ide_on_hang: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, s3_folders: 'bool' =None, show_dev_admin: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, snapshotter_service: 'bool' =None, snowflake_connection: 'bool' =None, spot_v2: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, teamspace_storage_tab: 'bool' =None, use_rclone_mounts_only: 'bool' =None, writable_data_connections: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._advanced_deployment_autoscaling = None
         self._affiliate_links = None
@@ -245,6 +247,7 @@ class V1UserFeatures(object):
         self._filesystem_optimisation = None
         self._gcp = None
         self._inference_job_deployment_plugin = None
+        self._instant_capacity_reservation = None
         self._jobs_init = None
         self._jobs_v2 = None
         self._landing_studios = None
@@ -362,6 +365,8 @@ class V1UserFeatures(object):
             self.gcp = gcp
         if inference_job_deployment_plugin is not None:
             self.inference_job_deployment_plugin = inference_job_deployment_plugin
+        if instant_capacity_reservation is not None:
+            self.instant_capacity_reservation = instant_capacity_reservation
         if jobs_init is not None:
             self.jobs_init = jobs_init
         if jobs_v2 is not None:
@@ -1190,6 +1195,27 @@ class V1UserFeatures(object):
 
         self._inference_job_deployment_plugin = inference_job_deployment_plugin
 
+    @property
+    def instant_capacity_reservation(self) -> 'bool':
+        """Gets the instant_capacity_reservation of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The instant_capacity_reservation of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._instant_capacity_reservation
+
+    @instant_capacity_reservation.setter
+    def instant_capacity_reservation(self, instant_capacity_reservation: 'bool'):
+        """Sets the instant_capacity_reservation of this V1UserFeatures.
+
+
+        :param instant_capacity_reservation: The instant_capacity_reservation of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._instant_capacity_reservation = instant_capacity_reservation
+
     @property
     def jobs_init(self) -> 'bool':
         """Gets the jobs_init of this V1UserFeatures.  # noqa: E501

lightning_sdk/models.py

@@ -0,0 +1,153 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, 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.user import User
+from lightning_sdk.utils.resolve import _get_authed_user
+
+if TYPE_CHECKING:
+    from lightning_sdk.teamspace import Teamspace
+
+
+# TODO: Maybe just have a `Model` object?
+@dataclass
+class UploadedModelInfo:
+    name: str
+    version: str
+    teamspace: str
+    cloud_account: str
+
+
+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."""
+    from lightning_sdk.teamspace import Teamspace
+
+    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,
+    progress_bar: bool = True,
+) -> UploadedModelInfo:
+    """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.
+            If not provided, the default cloud account for the Teamspace will be used.
+        progress_bar: Whether to show a progress bar for the upload.
+    """
+    org_name, teamspace_name, model_name, _ = _parse_model_name_and_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    return teamspace.upload_model(
+        path=path,
+        name=model_name,
+        cloud_account=cloud_account,
+        progress_bar=progress_bar,
+    )

lightning_sdk/teamspace.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, List, Optional, Union
 
 from lightning_sdk.agents import Agent
 from lightning_sdk.api import TeamspaceApi
-from lightning_sdk.api.teamspace_api import UploadedModelInfo
+from lightning_sdk.models import UploadedModelInfo
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
 from lightning_sdk.user import User
@@ -105,6 +105,10 @@ class Teamspace:
 
         return studios
 
+    @property
+    def default_cluster(self) -> str:
+        return self._teamspace.project_settings.preferred_cluster
+
     @property
     def clusters(self) -> List[str]:
         """All clusters associated with that teamspace."""
@@ -153,20 +157,19 @@ class Teamspace:
 
     def upload_model(
         self,
-        path: str,
+        path: Union[str, Path],
         name: str,
+        cloud_account: Optional[str] = None,
         progress_bar: bool = True,
-        cluster_id: Optional[str] = None,
     ) -> UploadedModelInfo:
         """Upload a local checkpoint file to the model store.
 
         Args:
             path: Path to the model file or folder to upload.
             name: Name tag of the model to upload.
+            cloud_account: The name of the cloud account to store the Model in.
+                If not provided, the default cloud account for the Teamspace will be used.
             progress_bar: Whether to show a progress bar for the upload.
-            cluster_id: The name of the cluster to use. Only required if it can't be determined
-                automatically.
-
         """
         if not path:
             raise ValueError("No path provided to upload")
@@ -176,7 +179,7 @@ class Teamspace:
         if not path.exists():
             raise FileNotFoundError(str(path))
 
-        cluster_id = self._teamspace_api._determine_cluster_id(self.id) if cluster_id is None else cluster_id
+        cloud_account = self._teamspace_api._determine_cluster_id(self.id) if cloud_account is None else cloud_account
         filepaths = [path] if path.is_file() else [p for p in path.rglob("*") if p.is_file()]
 
         if not filepaths:
@@ -196,14 +199,14 @@ class Teamspace:
             metadata={"filenames": filenames},
             private=True,
             teamspace_id=self.id,
-            cluster_id=cluster_id,
+            cluster_id=cloud_account,
         )
         self._teamspace_api.upload_model_files(
             model_id=model.model_id,
             version=model.version,
             root_path=root_path,
             filepaths=filepaths,
-            cluster_id=cluster_id,
+            cluster_id=cloud_account,
             teamspace_id=self.id,
             progress_bar=progress_bar,
         )
@@ -216,7 +219,7 @@ class Teamspace:
             name=name,
             version=model.version,
             teamspace=self.name,
-            cluster=cluster_id,
+            cloud_account=cloud_account,
         )
 
     def download_model(
@@ -254,7 +257,8 @@ class Teamspace:
             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.38.dist-info → lightning_sdk-0.1.40.dist-info}/METADATA

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