lightning-sdk 0.1.39__py3-none-any.whl → 0.1.41__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_rule_resource.py

@@ -72,6 +72,7 @@ class V1RuleResource(object):
     DEPLOYMENT = "deployment"
     MODEL = "model"
     SLURMJOB = "slurmJob"
+    MULTIMACHINEJOB = "multiMachineJob"
     """
     Attributes:
       swagger_types (dict): The key is attribute name

lightning_sdk/lightning_cloud/openapi/models/v1_secret_type.py

@@ -40,6 +40,7 @@ class V1SecretType(object):
     UNSPECIFIED = "SECRET_TYPE_UNSPECIFIED"
     DOCKER_REGISTRY = "SECRET_TYPE_DOCKER_REGISTRY"
     GCP_CREDENTIALS = "SECRET_TYPE_GCP_CREDENTIALS"
+    SNOWFLAKE_CREDENTIALS = "SECRET_TYPE_SNOWFLAKE_CREDENTIALS"
     """
     Attributes:
       swagger_types (dict): The key is attribute name

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

@@ -51,7 +51,7 @@ class V1UserFeatures(object):
         'b2c_experience': 'bool',
         'cap_add': 'list[str]',
         'cap_drop': 'list[str]',
-        'capacity_reservation': 'bool',
+        'capacity_reservation_byoc': 'bool',
         'capacity_reservation_dry_run': 'bool',
         'cluster_proxy': 'bool',
         'code_tab': 'bool',
@@ -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',
@@ -135,7 +136,7 @@ class V1UserFeatures(object):
         'b2c_experience': 'b2cExperience',
         'cap_add': 'capAdd',
         'cap_drop': 'capDrop',
-        'capacity_reservation': 'capacityReservation',
+        'capacity_reservation_byoc': 'capacityReservationByoc',
         'capacity_reservation_dry_run': 'capacityReservationDryRun',
         'cluster_proxy': 'clusterProxy',
         'code_tab': 'codeTab',
@@ -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_byoc: '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
@@ -220,7 +222,7 @@ class V1UserFeatures(object):
         self._b2c_experience = None
         self._cap_add = None
         self._cap_drop = None
-        self._capacity_reservation = None
+        self._capacity_reservation_byoc = None
         self._capacity_reservation_dry_run = None
         self._cluster_proxy = None
         self._code_tab = 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
@@ -312,8 +315,8 @@ class V1UserFeatures(object):
             self.cap_add = cap_add
         if cap_drop is not None:
             self.cap_drop = cap_drop
-        if capacity_reservation is not None:
-            self.capacity_reservation = capacity_reservation
+        if capacity_reservation_byoc is not None:
+            self.capacity_reservation_byoc = capacity_reservation_byoc
         if capacity_reservation_dry_run is not None:
             self.capacity_reservation_dry_run = capacity_reservation_dry_run
         if cluster_proxy is not 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:
@@ -666,25 +671,25 @@ class V1UserFeatures(object):
         self._cap_drop = cap_drop
 
     @property
-    def capacity_reservation(self) -> 'bool':
-        """Gets the capacity_reservation of this V1UserFeatures.  # noqa: E501
+    def capacity_reservation_byoc(self) -> 'bool':
+        """Gets the capacity_reservation_byoc of this V1UserFeatures.  # noqa: E501
 
 
-        :return: The capacity_reservation of this V1UserFeatures.  # noqa: E501
+        :return: The capacity_reservation_byoc of this V1UserFeatures.  # noqa: E501
         :rtype: bool
         """
-        return self._capacity_reservation
+        return self._capacity_reservation_byoc
 
-    @capacity_reservation.setter
-    def capacity_reservation(self, capacity_reservation: 'bool'):
-        """Sets the capacity_reservation of this V1UserFeatures.
+    @capacity_reservation_byoc.setter
+    def capacity_reservation_byoc(self, capacity_reservation_byoc: 'bool'):
+        """Sets the capacity_reservation_byoc of this V1UserFeatures.
 
 
-        :param capacity_reservation: The capacity_reservation of this V1UserFeatures.  # noqa: E501
+        :param capacity_reservation_byoc: The capacity_reservation_byoc of this V1UserFeatures.  # noqa: E501
         :type: bool
         """
 
-        self._capacity_reservation = capacity_reservation
+        self._capacity_reservation_byoc = capacity_reservation_byoc
 
     @property
     def capacity_reservation_dry_run(self) -> 'bool':
@@ -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

@@ -1,14 +1,26 @@
 import os
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+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.teamspace import Teamspace
 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
@@ -38,8 +50,10 @@ def _list_teamspaces() -> List[str]:
     ]
 
 
-def _get_teamspace(name: str, organization: str) -> Teamspace:
+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()
@@ -113,20 +127,27 @@ def download_model(
         raise e
 
 
-def upload_model(name: str, path: Union[Path, str] = ".", cloud_account: Optional[str] = None) -> None:
+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>.
+            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)
-    teamspace.upload_model(
+    return teamspace.upload_model(
         path=path,
         name=model_name,
-        progress_bar=True,
-        cluster_id=cloud_account,
+        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."""
@@ -155,18 +159,17 @@ class Teamspace:
         self,
         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(
@@ -290,18 +293,17 @@ def _resolve_valueerror_message(error: ValueError, owner: Owner, teamspace_name:
             organizations = _get_organizations_for_authed_user()
             message = (
                 f"Teamspace {entire_ts_name} does not exist. "
-                "Is it maybe an organizational Teamspace? You are a member of the following organizations: "
-                f"{[o.name for o in organizations]}. Maybe specify one of these instead "
-                "of your user if the Teamspace belongs to the organization."
+                f"Is {teamspace_name} an organizational Teamspace? You are a member of the following organizations: "
+                f"{[o.name for o in organizations]}. Try specifying the `org` parameter instead "
+                "of `user` if the Teamspace belongs to the organization."
             )
         else:
             # organization teamspace owner
             user = User()
             message = (
                 f"Teamspace {entire_ts_name} does not exist. "
-                f"Is it maybe a user Teamspace. You specified org={owner.name}, "
-                "but maybe the Teamspace is part of your user? "
-                f"Consider specifying user={user.name} instead of your org."
+                f"Is {teamspace_name} a user Teamspace? "
+                f"Consider specifying user={user.name} instead of org={owner.name}."
             )
 
     return ValueError(message, *error.args[1:])

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

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