lightning-sdk 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -63,8 +63,8 @@ class V1UserFeatures(object):
         'default_one_cluster': 'bool',
         'deployment_alerts': 'bool',
         'deployment_persistent_disk': 'bool',
-        'deployment_reservations': 'bool',
         'dgx_cloud': 'bool',
+        'doc_helper_chat': 'bool',
         'docs_agent': 'bool',
         'drive_v2': 'bool',
         'enable_storage_limits': 'bool',
@@ -85,6 +85,7 @@ class V1UserFeatures(object):
         'multiple_studio_versions': 'bool',
         'org_admin_alerts': 'bool',
         'org_level_member_permissions': 'bool',
+        'org_usage_limits': 'bool',
         'pipelines': 'bool',
         'plugin_distributed': 'bool',
         'plugin_inference': 'bool',
@@ -140,8 +141,8 @@ class V1UserFeatures(object):
         'default_one_cluster': 'defaultOneCluster',
         'deployment_alerts': 'deploymentAlerts',
         'deployment_persistent_disk': 'deploymentPersistentDisk',
-        'deployment_reservations': 'deploymentReservations',
         'dgx_cloud': 'dgxCloud',
+        'doc_helper_chat': 'docHelperChat',
         'docs_agent': 'docsAgent',
         'drive_v2': 'driveV2',
         'enable_storage_limits': 'enableStorageLimits',
@@ -162,6 +163,7 @@ class V1UserFeatures(object):
         'multiple_studio_versions': 'multipleStudioVersions',
         'org_admin_alerts': 'orgAdminAlerts',
         'org_level_member_permissions': 'orgLevelMemberPermissions',
+        'org_usage_limits': 'orgUsageLimits',
         'pipelines': 'pipelines',
         'plugin_distributed': 'pluginDistributed',
         'plugin_inference': 'pluginInference',
@@ -194,7 +196,7 @@ class V1UserFeatures(object):
         'weka': 'weka'
     }
 
-    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byoc_litcr: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloud_space_environment_templates: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_reservations: 'bool' =None, dgx_cloud: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_storage_limits: 'bool' =None, enterprise_compute_admin: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filestore: 'bool' =None, inactive_notify_delete: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, lambda_labs: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, org_admin_alerts: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, stop_ide_container_on_shutdown: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, studios_dashboard: 'bool' =None, studios_dashboard_system_metrics: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None):  # noqa: E501
+    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byoc_litcr: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloud_space_environment_templates: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, dgx_cloud: 'bool' =None, doc_helper_chat: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_storage_limits: 'bool' =None, enterprise_compute_admin: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filestore: 'bool' =None, inactive_notify_delete: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, lambda_labs: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, org_admin_alerts: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, pipelines: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, stop_ide_container_on_shutdown: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, studios_dashboard: 'bool' =None, studios_dashboard_system_metrics: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._affiliate_links = None
         self._agents_v2 = None
@@ -218,8 +220,8 @@ class V1UserFeatures(object):
         self._default_one_cluster = None
         self._deployment_alerts = None
         self._deployment_persistent_disk = None
-        self._deployment_reservations = None
         self._dgx_cloud = None
+        self._doc_helper_chat = None
         self._docs_agent = None
         self._drive_v2 = None
         self._enable_storage_limits = None
@@ -240,6 +242,7 @@ class V1UserFeatures(object):
         self._multiple_studio_versions = None
         self._org_admin_alerts = None
         self._org_level_member_permissions = None
+        self._org_usage_limits = None
         self._pipelines = None
         self._plugin_distributed = None
         self._plugin_inference = None
@@ -315,10 +318,10 @@ class V1UserFeatures(object):
             self.deployment_alerts = deployment_alerts
         if deployment_persistent_disk is not None:
             self.deployment_persistent_disk = deployment_persistent_disk
-        if deployment_reservations is not None:
-            self.deployment_reservations = deployment_reservations
         if dgx_cloud is not None:
             self.dgx_cloud = dgx_cloud
+        if doc_helper_chat is not None:
+            self.doc_helper_chat = doc_helper_chat
         if docs_agent is not None:
             self.docs_agent = docs_agent
         if drive_v2 is not None:
@@ -359,6 +362,8 @@ class V1UserFeatures(object):
             self.org_admin_alerts = org_admin_alerts
         if org_level_member_permissions is not None:
             self.org_level_member_permissions = org_level_member_permissions
+        if org_usage_limits is not None:
+            self.org_usage_limits = org_usage_limits
         if pipelines is not None:
             self.pipelines = pipelines
         if plugin_distributed is not None:
@@ -883,46 +888,46 @@ class V1UserFeatures(object):
         self._deployment_persistent_disk = deployment_persistent_disk
 
     @property
-    def deployment_reservations(self) -> 'bool':
-        """Gets the deployment_reservations of this V1UserFeatures.  # noqa: E501
+    def dgx_cloud(self) -> 'bool':
+        """Gets the dgx_cloud of this V1UserFeatures.  # noqa: E501
 
 
-        :return: The deployment_reservations of this V1UserFeatures.  # noqa: E501
+        :return: The dgx_cloud of this V1UserFeatures.  # noqa: E501
         :rtype: bool
         """
-        return self._deployment_reservations
+        return self._dgx_cloud
 
-    @deployment_reservations.setter
-    def deployment_reservations(self, deployment_reservations: 'bool'):
-        """Sets the deployment_reservations of this V1UserFeatures.
+    @dgx_cloud.setter
+    def dgx_cloud(self, dgx_cloud: 'bool'):
+        """Sets the dgx_cloud of this V1UserFeatures.
 
 
-        :param deployment_reservations: The deployment_reservations of this V1UserFeatures.  # noqa: E501
+        :param dgx_cloud: The dgx_cloud of this V1UserFeatures.  # noqa: E501
         :type: bool
         """
 
-        self._deployment_reservations = deployment_reservations
+        self._dgx_cloud = dgx_cloud
 
     @property
-    def dgx_cloud(self) -> 'bool':
-        """Gets the dgx_cloud of this V1UserFeatures.  # noqa: E501
+    def doc_helper_chat(self) -> 'bool':
+        """Gets the doc_helper_chat of this V1UserFeatures.  # noqa: E501
 
 
-        :return: The dgx_cloud of this V1UserFeatures.  # noqa: E501
+        :return: The doc_helper_chat of this V1UserFeatures.  # noqa: E501
         :rtype: bool
         """
-        return self._dgx_cloud
+        return self._doc_helper_chat
 
-    @dgx_cloud.setter
-    def dgx_cloud(self, dgx_cloud: 'bool'):
-        """Sets the dgx_cloud of this V1UserFeatures.
+    @doc_helper_chat.setter
+    def doc_helper_chat(self, doc_helper_chat: 'bool'):
+        """Sets the doc_helper_chat of this V1UserFeatures.
 
 
-        :param dgx_cloud: The dgx_cloud of this V1UserFeatures.  # noqa: E501
+        :param doc_helper_chat: The doc_helper_chat of this V1UserFeatures.  # noqa: E501
         :type: bool
         """
 
-        self._dgx_cloud = dgx_cloud
+        self._doc_helper_chat = doc_helper_chat
 
     @property
     def docs_agent(self) -> 'bool':
@@ -1344,6 +1349,27 @@ class V1UserFeatures(object):
 
         self._org_level_member_permissions = org_level_member_permissions
 
+    @property
+    def org_usage_limits(self) -> 'bool':
+        """Gets the org_usage_limits of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The org_usage_limits of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._org_usage_limits
+
+    @org_usage_limits.setter
+    def org_usage_limits(self, org_usage_limits: 'bool'):
+        """Sets the org_usage_limits of this V1UserFeatures.
+
+
+        :param org_usage_limits: The org_usage_limits of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._org_usage_limits = org_usage_limits
+
     @property
     def pipelines(self) -> 'bool':
         """Gets the pipelines of this V1UserFeatures.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_validate_data_connection_response.py

@@ -45,6 +45,7 @@ class V1ValidateDataConnectionResponse(object):
         'efs': 'V1EfsConfig',
         'gcp': 'V1GcpDataConnection',
         'gcs_folder': 'V1GCSFolderDataConnection',
+        'r2': 'V1R2DataConnection',
         's3_folder': 'V1S3FolderDataConnection'
     }
 
@@ -53,15 +54,17 @@ class V1ValidateDataConnectionResponse(object):
         'efs': 'efs',
         'gcp': 'gcp',
         'gcs_folder': 'gcsFolder',
+        'r2': 'r2',
         's3_folder': 's3Folder'
     }
 
-    def __init__(self, aws: 'V1AwsDataConnection' =None, efs: 'V1EfsConfig' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
+    def __init__(self, aws: 'V1AwsDataConnection' =None, efs: 'V1EfsConfig' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, r2: 'V1R2DataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
         """V1ValidateDataConnectionResponse - a model defined in Swagger"""  # noqa: E501
         self._aws = None
         self._efs = None
         self._gcp = None
         self._gcs_folder = None
+        self._r2 = None
         self._s3_folder = None
         self.discriminator = None
         if aws is not None:
@@ -72,6 +75,8 @@ class V1ValidateDataConnectionResponse(object):
             self.gcp = gcp
         if gcs_folder is not None:
             self.gcs_folder = gcs_folder
+        if r2 is not None:
+            self.r2 = r2
         if s3_folder is not None:
             self.s3_folder = s3_folder
 
@@ -159,6 +164,27 @@ class V1ValidateDataConnectionResponse(object):
 
         self._gcs_folder = gcs_folder
 
+    @property
+    def r2(self) -> 'V1R2DataConnection':
+        """Gets the r2 of this V1ValidateDataConnectionResponse.  # noqa: E501
+
+
+        :return: The r2 of this V1ValidateDataConnectionResponse.  # noqa: E501
+        :rtype: V1R2DataConnection
+        """
+        return self._r2
+
+    @r2.setter
+    def r2(self, r2: 'V1R2DataConnection'):
+        """Sets the r2 of this V1ValidateDataConnectionResponse.
+
+
+        :param r2: The r2 of this V1ValidateDataConnectionResponse.  # noqa: E501
+        :type: V1R2DataConnection
+        """
+
+        self._r2 = r2
+
     @property
     def s3_folder(self) -> 'V1S3FolderDataConnection':
         """Gets the s3_folder of this V1ValidateDataConnectionResponse.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_weka_data_connection.py

@@ -42,38 +42,33 @@ class V1WekaDataConnection(object):
     """
     swagger_types = {
         'backend_ips': 'list[str]',
-        'cores': 'int',
         'file_system_name': 'str',
-        'management_ip': 'str',
-        'memory_mb': 'int'
+        'memory_mb': 'int',
+        'readonly': 'bool'
     }
 
     attribute_map = {
         'backend_ips': 'backendIps',
-        'cores': 'cores',
         'file_system_name': 'fileSystemName',
-        'management_ip': 'managementIp',
-        'memory_mb': 'memoryMb'
+        'memory_mb': 'memoryMb',
+        'readonly': 'readonly'
     }
 
-    def __init__(self, backend_ips: 'list[str]' =None, cores: 'int' =None, file_system_name: 'str' =None, management_ip: 'str' =None, memory_mb: 'int' =None):  # noqa: E501
+    def __init__(self, backend_ips: 'list[str]' =None, file_system_name: 'str' =None, memory_mb: 'int' =None, readonly: 'bool' =None):  # noqa: E501
         """V1WekaDataConnection - a model defined in Swagger"""  # noqa: E501
         self._backend_ips = None
-        self._cores = None
         self._file_system_name = None
-        self._management_ip = None
         self._memory_mb = None
+        self._readonly = None
         self.discriminator = None
         if backend_ips is not None:
             self.backend_ips = backend_ips
-        if cores is not None:
-            self.cores = cores
         if file_system_name is not None:
             self.file_system_name = file_system_name
-        if management_ip is not None:
-            self.management_ip = management_ip
         if memory_mb is not None:
             self.memory_mb = memory_mb
+        if readonly is not None:
+            self.readonly = readonly
 
     @property
     def backend_ips(self) -> 'list[str]':
@@ -96,27 +91,6 @@ class V1WekaDataConnection(object):
 
         self._backend_ips = backend_ips
 
-    @property
-    def cores(self) -> 'int':
-        """Gets the cores of this V1WekaDataConnection.  # noqa: E501
-
-
-        :return: The cores of this V1WekaDataConnection.  # noqa: E501
-        :rtype: int
-        """
-        return self._cores
-
-    @cores.setter
-    def cores(self, cores: 'int'):
-        """Sets the cores of this V1WekaDataConnection.
-
-
-        :param cores: The cores of this V1WekaDataConnection.  # noqa: E501
-        :type: int
-        """
-
-        self._cores = cores
-
     @property
     def file_system_name(self) -> 'str':
         """Gets the file_system_name of this V1WekaDataConnection.  # noqa: E501
@@ -138,27 +112,6 @@ class V1WekaDataConnection(object):
 
         self._file_system_name = file_system_name
 
-    @property
-    def management_ip(self) -> 'str':
-        """Gets the management_ip of this V1WekaDataConnection.  # noqa: E501
-
-
-        :return: The management_ip of this V1WekaDataConnection.  # noqa: E501
-        :rtype: str
-        """
-        return self._management_ip
-
-    @management_ip.setter
-    def management_ip(self, management_ip: 'str'):
-        """Sets the management_ip of this V1WekaDataConnection.
-
-
-        :param management_ip: The management_ip of this V1WekaDataConnection.  # noqa: E501
-        :type: str
-        """
-
-        self._management_ip = management_ip
-
     @property
     def memory_mb(self) -> 'int':
         """Gets the memory_mb of this V1WekaDataConnection.  # noqa: E501
@@ -180,6 +133,27 @@ class V1WekaDataConnection(object):
 
         self._memory_mb = memory_mb
 
+    @property
+    def readonly(self) -> 'bool':
+        """Gets the readonly of this V1WekaDataConnection.  # noqa: E501
+
+
+        :return: The readonly of this V1WekaDataConnection.  # noqa: E501
+        :rtype: bool
+        """
+        return self._readonly
+
+    @readonly.setter
+    def readonly(self, readonly: 'bool'):
+        """Sets the readonly of this V1WekaDataConnection.
+
+
+        :param readonly: The readonly of this V1WekaDataConnection.  # noqa: E501
+        :type: bool
+        """
+
+        self._readonly = readonly
+
     def to_dict(self) -> dict:
         """Returns the model properties as a dict"""
         result = {}

lightning_sdk/lightning_cloud/openapi/models/validate.py

@@ -48,6 +48,7 @@ class Validate(object):
         'filestore': 'V1FilestoreDataConnection',
         'gcp': 'V1GcpDataConnection',
         'gcs_folder': 'V1GCSFolderDataConnection',
+        'r2': 'V1R2DataConnection',
         's3_folder': 'V1S3FolderDataConnection'
     }
 
@@ -59,10 +60,11 @@ class Validate(object):
         'filestore': 'filestore',
         'gcp': 'gcp',
         'gcs_folder': 'gcsFolder',
+        'r2': 'r2',
         's3_folder': 's3Folder'
     }
 
-    def __init__(self, aws: 'V1AwsDataConnection' =None, check_is_public: 'bool' =None, cluster_ids: 'list[str]' =None, efs: 'V1EfsConfig' =None, filestore: 'V1FilestoreDataConnection' =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, filestore: 'V1FilestoreDataConnection' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, r2: 'V1R2DataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
         """Validate - a model defined in Swagger"""  # noqa: E501
         self._aws = None
         self._check_is_public = None
@@ -71,6 +73,7 @@ class Validate(object):
         self._filestore = None
         self._gcp = None
         self._gcs_folder = None
+        self._r2 = None
         self._s3_folder = None
         self.discriminator = None
         if aws is not None:
@@ -87,6 +90,8 @@ class Validate(object):
             self.gcp = gcp
         if gcs_folder is not None:
             self.gcs_folder = gcs_folder
+        if r2 is not None:
+            self.r2 = r2
         if s3_folder is not None:
             self.s3_folder = s3_folder
 
@@ -237,6 +242,27 @@ class Validate(object):
 
         self._gcs_folder = gcs_folder
 
+    @property
+    def r2(self) -> 'V1R2DataConnection':
+        """Gets the r2 of this Validate.  # noqa: E501
+
+
+        :return: The r2 of this Validate.  # noqa: E501
+        :rtype: V1R2DataConnection
+        """
+        return self._r2
+
+    @r2.setter
+    def r2(self, r2: 'V1R2DataConnection'):
+        """Sets the r2 of this Validate.
+
+
+        :param r2: The r2 of this Validate.  # noqa: E501
+        :type: V1R2DataConnection
+        """
+
+        self._r2 = r2
+
     @property
     def s3_folder(self) -> 'V1S3FolderDataConnection':
         """Gets the s3_folder of this Validate.  # noqa: E501

lightning_sdk/lit_container.py

@@ -80,7 +80,8 @@ class LitContainer:
         tag: str = "latest",
         cloud_account: Optional[str] = None,
         platform: Optional[str] = "linux/amd64",
-    ) -> None:
+        return_final_dict: bool = False,
+    ) -> Optional[Dict]:
         """Upload a container to the docker registry.
 
         Args:
@@ -91,15 +92,24 @@ class LitContainer:
             tag: The tag to use for the container.
             cloud_account: The cloud account where the container is stored.
             platform: The platform the container is meant to run on.
+            return_final_dict: Instructs function to return metadata about container location in platform.
         """
         try:
             teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
         except Exception as e:
             raise ValueError(f"Could not resolve teamspace: {e}") from e
 
-        resp = self._api.upload_container(container, teamspace, tag, cloud_account, platform=platform)
+        resp = self._api.upload_container(
+            container, teamspace, tag, cloud_account, platform=platform, return_final_dict=return_final_dict
+        )
+
+        final_dict = None
         for line in resp:
             print(line)
+            if return_final_dict and isinstance(line, dict) and line.get("finish") is True:
+                final_dict = line
+
+        return final_dict if return_final_dict else None
 
     def download_container(
         self,

lightning_sdk/models.py

@@ -4,6 +4,7 @@ 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 import V1ModelVersionArchive
 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
@@ -85,7 +86,7 @@ def _extend_model_name_with_teamspace(name: str) -> str:
     return f"{teamspace.owner.name}/{teamspace.name}/{name}"
 
 
-def _parse_model_name_and_version(name: str) -> Tuple[str, str, str, str]:
+def _parse_org_teamspace_model_version(name: str) -> Tuple[str, str, str, Optional[str]]:
     """Parse the name argument into its components."""
     try:
         org_name, teamspace_name, model_name = name.split("/")
@@ -119,19 +120,14 @@ def download_model(
         progress_bar: Whether to show a progress bar when downloading.
     """
     name = _extend_model_name_with_teamspace(name)
-    teamspace_owner_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
-    if version is None:
-        version = "default"
-
-    download_dir = Path(download_dir)
-
+    teamspace_owner_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
     api = TeamspaceApi()
 
     try:
         return api.download_model_files(
             name=model_name,
             version=version,
-            download_dir=download_dir,
+            download_dir=Path(download_dir),
             teamspace_name=teamspace_name,
             teamspace_owner_name=teamspace_owner_name,
             progress_bar=progress_bar,
@@ -160,9 +156,11 @@ def upload_model(
         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.
+        metadata: Metadata to attach to the uploaded model.
+            If not provided, an empty dictionary will be used.
     """
     name = _extend_model_name_with_teamspace(name)
-    org_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
+    org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
     teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
     return teamspace.upload_model(
         path=path,
@@ -172,3 +170,34 @@ def upload_model(
         progress_bar=progress_bar,
         metadata=metadata,
     )
+
+
+def delete_model(
+    name: str,
+) -> None:
+    """Delete a model or a version of model.
+
+    Args:
+        name: The name of the model you want to delete or with specified version it deltes only that version.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME> for full model deletion
+            or <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>:<VERSION> for version deletion.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    teamspace.delete_model(name=f"{model_name}:{version}" if version else model_name)
+
+
+def list_model_versions(
+    name: str,
+) -> List[V1ModelVersionArchive]:
+    """List all versions of a model.
+
+    Args:
+        name: The name of the model you want to list versions for.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, _ = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    return teamspace.list_model_versions(name=model_name)

lightning_sdk/serve.py

@@ -248,6 +248,8 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
         cloud_account: Optional[str] = None,
         port: Optional[int] = 8000,
         include_credentials: Optional[bool] = True,
+        cloudspace_id: Optional[str] = None,
+        from_onboarding: bool = False,
     ) -> dict:
         """Run a deployment on the cloud. If the deployment already exists, it will be updated.
 
@@ -264,6 +266,8 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
             cloud_account: The cloud account to run the deployment on. Defaults to None.
             port: The port to run the deployment on. Defaults to 8000.
             include_credentials: Whether to include credentials in the deployment. Defaults to True.
+            cloudspace_id: Connect to a Studio.
+            from_onboarding: Deployment created during onboarding.
 
         Returns:
             dict: The deployment and the URL of the deployment.
@@ -297,6 +301,8 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
             cloud_account=cloud_account,
             ports=[port],
             include_credentials=include_credentials,
+            cloudspace_id=cloudspace_id,
+            from_onboarding=from_onboarding,
         )
 
         return {"deployment": deployment, "url": url}

lightning_sdk/teamspace.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 import lightning_sdk
 from lightning_sdk.agents import Agent
 from lightning_sdk.api import TeamspaceApi
-from lightning_sdk.lightning_cloud.openapi import V1ProjectClusterBinding
+from lightning_sdk.lightning_cloud.openapi import V1Model, V1ModelVersionArchive, V1ProjectClusterBinding
 from lightning_sdk.machine import Machine
 from lightning_sdk.models import UploadedModelInfo
 from lightning_sdk.organization import Organization
@@ -262,6 +262,8 @@ class Teamspace:
         Args:
             path: Path to the model file or folder to upload.
             name: Name tag of the model to upload.
+            version: Version tag of the model to upload.
+                If not provided, the ``vX`` version will be used where X is running index.
             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.
@@ -315,7 +317,7 @@ class Teamspace:
             teamspace_id=self.id,
             progress_bar=progress_bar,
         )
-        self._teamspace_api.complete_model_upload(
+        self._teamspace_api._complete_model_upload(
             model_id=model.model_id,
             version=model.version,
             teamspace_id=self.id,
@@ -387,6 +389,18 @@ class Teamspace:
         name, version = _parse_model_and_version(name)
         self._teamspace_api.delete_model(name=name, version=version, teamspace_id=self.id)
 
+    def list_models(self) -> List[V1Model]:
+        """List all models in the model store."""
+        return self._teamspace_api.list_models(teamspace_id=self.id)
+
+    def list_model_versions(self, name: str) -> List[V1ModelVersionArchive]:
+        """List all versions of a model in the model store."""
+        if ":" in name:
+            raise ValueError(
+                "Model name should not contain a version tag. Please provide the model name without a version."
+            )
+        return self._teamspace_api.list_model_versions(teamspace_id=self.id, model_name=name)
+
 
 def _list_files(path: Union[str, Path]) -> Tuple[List[Path], List[str]]:
     """List all folders in a directory and return them as a list and relative path."""

lightning_sdk/utils/resolve.py

@@ -179,16 +179,23 @@ def skip_studio_init() -> Generator[None, None, None]:
     Studio._skip_init = prev_studio_init_state
 
 
-def _parse_model_and_version(name: str) -> Tuple[str, str]:
+def _parse_model_and_version(name: str) -> Tuple[str, Optional[str]]:
+    """Parse the model name and version from the given string.
+
+    >>> _parse_model_and_version("org/teamspace/modelname")
+    ('org/teamspace/modelname', None)
+    >>> _parse_model_and_version("org/teamspace/modelname:version")
+    ('org/teamspace/modelname', 'version')
+    """
     parts = name.split(":")
     if len(parts) == 1:
-        return parts[0], "default"
+        return parts[0], None
     if len(parts) == 2:
         return 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}"
+        "Model version is expected to be in the format `entity/modelname:version` separated by a single colon,"
+        f" but got: {name}"
     )