lightning-sdk 0.2.15__py3-none-any.whl → 0.2.16__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -48,6 +48,7 @@ class V1UserFeatures(object):
         'auto_join_orgs': 'bool',
         'auto_top_up': 'bool',
         'b2c_experience': 'bool',
+        'big_default_cpus': 'bool',
         'byoc_litcr': 'bool',
         'cap_add': 'list[str]',
         'cap_drop': 'list[str]',
@@ -68,6 +69,7 @@ class V1UserFeatures(object):
         'deployment_requests_tab': 'bool',
         'dgx_cloud': 'bool',
         'docs_agent': 'bool',
+        'down_switch_machine': 'bool',
         'drive_v2': 'bool',
         'enable_storage_limits': 'bool',
         'enterprise_compute_admin': 'bool',
@@ -112,11 +114,11 @@ class V1UserFeatures(object):
         'restartable_jobs': 'bool',
         'runnable_public_studio_page': 'bool',
         'security_docs': 'bool',
+        'seoul_aws_region': 'bool',
         'show_dev_admin': 'bool',
         'single_wallet': 'bool',
         'slurm': 'bool',
         'slurm_machine_selector': 'bool',
-        'stable_vscode_server': 'bool',
         'stop_ide_container_on_shutdown': 'bool',
         'studio_config': 'bool',
         'studio_on_stop': 'bool',
@@ -139,6 +141,7 @@ class V1UserFeatures(object):
         'auto_join_orgs': 'autoJoinOrgs',
         'auto_top_up': 'autoTopUp',
         'b2c_experience': 'b2cExperience',
+        'big_default_cpus': 'bigDefaultCpus',
         'byoc_litcr': 'byocLitcr',
         'cap_add': 'capAdd',
         'cap_drop': 'capDrop',
@@ -159,6 +162,7 @@ class V1UserFeatures(object):
         'deployment_requests_tab': 'deploymentRequestsTab',
         'dgx_cloud': 'dgxCloud',
         'docs_agent': 'docsAgent',
+        'down_switch_machine': 'downSwitchMachine',
         'drive_v2': 'driveV2',
         'enable_storage_limits': 'enableStorageLimits',
         'enterprise_compute_admin': 'enterpriseComputeAdmin',
@@ -203,11 +207,11 @@ class V1UserFeatures(object):
         'restartable_jobs': 'restartableJobs',
         'runnable_public_studio_page': 'runnablePublicStudioPage',
         'security_docs': 'securityDocs',
+        'seoul_aws_region': 'seoulAwsRegion',
         'show_dev_admin': 'showDevAdmin',
         'single_wallet': 'singleWallet',
         'slurm': 'slurm',
         'slurm_machine_selector': 'slurmMachineSelector',
-        'stable_vscode_server': 'stableVscodeServer',
         'stop_ide_container_on_shutdown': 'stopIdeContainerOnShutdown',
         'studio_config': 'studioConfig',
         'studio_on_stop': 'studioOnStop',
@@ -222,7 +226,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, auto_top_up: '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, cloudy_vibe_code: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, concurrent_gpu_limit: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_requests_tab: '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, gcs_folders: '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, lightning_cloud: 'bool' =None, lit_logger: 'bool' =None, manage_storage_costs: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_folders: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, nebius: 'bool' =None, neocloud_studios: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, paygo_free_storage_limit_check: '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, product_license: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: '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, stable_vscode_server: '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, voltage_park: '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, auto_top_up: 'bool' =None, b2c_experience: 'bool' =None, big_default_cpus: '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, cloudy_vibe_code: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, concurrent_gpu_limit: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_requests_tab: 'bool' =None, dgx_cloud: 'bool' =None, docs_agent: 'bool' =None, down_switch_machine: '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, gcs_folders: '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, lightning_cloud: 'bool' =None, lit_logger: 'bool' =None, manage_storage_costs: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_folders: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, nebius: 'bool' =None, neocloud_studios: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, paygo_free_storage_limit_check: '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, product_license: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, seoul_aws_region: '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, voltage_park: '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
@@ -231,6 +235,7 @@ class V1UserFeatures(object):
         self._auto_join_orgs = None
         self._auto_top_up = None
         self._b2c_experience = None
+        self._big_default_cpus = None
         self._byoc_litcr = None
         self._cap_add = None
         self._cap_drop = None
@@ -251,6 +256,7 @@ class V1UserFeatures(object):
         self._deployment_requests_tab = None
         self._dgx_cloud = None
         self._docs_agent = None
+        self._down_switch_machine = None
         self._drive_v2 = None
         self._enable_storage_limits = None
         self._enterprise_compute_admin = None
@@ -295,11 +301,11 @@ class V1UserFeatures(object):
         self._restartable_jobs = None
         self._runnable_public_studio_page = None
         self._security_docs = None
+        self._seoul_aws_region = None
         self._show_dev_admin = None
         self._single_wallet = None
         self._slurm = None
         self._slurm_machine_selector = None
-        self._stable_vscode_server = None
         self._stop_ide_container_on_shutdown = None
         self._studio_config = None
         self._studio_on_stop = None
@@ -327,6 +333,8 @@ class V1UserFeatures(object):
             self.auto_top_up = auto_top_up
         if b2c_experience is not None:
             self.b2c_experience = b2c_experience
+        if big_default_cpus is not None:
+            self.big_default_cpus = big_default_cpus
         if byoc_litcr is not None:
             self.byoc_litcr = byoc_litcr
         if cap_add is not None:
@@ -367,6 +375,8 @@ class V1UserFeatures(object):
             self.dgx_cloud = dgx_cloud
         if docs_agent is not None:
             self.docs_agent = docs_agent
+        if down_switch_machine is not None:
+            self.down_switch_machine = down_switch_machine
         if drive_v2 is not None:
             self.drive_v2 = drive_v2
         if enable_storage_limits is not None:
@@ -455,6 +465,8 @@ class V1UserFeatures(object):
             self.runnable_public_studio_page = runnable_public_studio_page
         if security_docs is not None:
             self.security_docs = security_docs
+        if seoul_aws_region is not None:
+            self.seoul_aws_region = seoul_aws_region
         if show_dev_admin is not None:
             self.show_dev_admin = show_dev_admin
         if single_wallet is not None:
@@ -463,8 +475,6 @@ class V1UserFeatures(object):
             self.slurm = slurm
         if slurm_machine_selector is not None:
             self.slurm_machine_selector = slurm_machine_selector
-        if stable_vscode_server is not None:
-            self.stable_vscode_server = stable_vscode_server
         if stop_ide_container_on_shutdown is not None:
             self.stop_ide_container_on_shutdown = stop_ide_container_on_shutdown
         if studio_config is not None:
@@ -637,6 +647,27 @@ class V1UserFeatures(object):
 
         self._b2c_experience = b2c_experience
 
+    @property
+    def big_default_cpus(self) -> 'bool':
+        """Gets the big_default_cpus of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The big_default_cpus of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._big_default_cpus
+
+    @big_default_cpus.setter
+    def big_default_cpus(self, big_default_cpus: 'bool'):
+        """Sets the big_default_cpus of this V1UserFeatures.
+
+
+        :param big_default_cpus: The big_default_cpus of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._big_default_cpus = big_default_cpus
+
     @property
     def byoc_litcr(self) -> 'bool':
         """Gets the byoc_litcr of this V1UserFeatures.  # noqa: E501
@@ -1057,6 +1088,27 @@ class V1UserFeatures(object):
 
         self._docs_agent = docs_agent
 
+    @property
+    def down_switch_machine(self) -> 'bool':
+        """Gets the down_switch_machine of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The down_switch_machine of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._down_switch_machine
+
+    @down_switch_machine.setter
+    def down_switch_machine(self, down_switch_machine: 'bool'):
+        """Sets the down_switch_machine of this V1UserFeatures.
+
+
+        :param down_switch_machine: The down_switch_machine of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._down_switch_machine = down_switch_machine
+
     @property
     def drive_v2(self) -> 'bool':
         """Gets the drive_v2 of this V1UserFeatures.  # noqa: E501
@@ -1981,6 +2033,27 @@ class V1UserFeatures(object):
 
         self._security_docs = security_docs
 
+    @property
+    def seoul_aws_region(self) -> 'bool':
+        """Gets the seoul_aws_region of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The seoul_aws_region of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._seoul_aws_region
+
+    @seoul_aws_region.setter
+    def seoul_aws_region(self, seoul_aws_region: 'bool'):
+        """Sets the seoul_aws_region of this V1UserFeatures.
+
+
+        :param seoul_aws_region: The seoul_aws_region of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._seoul_aws_region = seoul_aws_region
+
     @property
     def show_dev_admin(self) -> 'bool':
         """Gets the show_dev_admin of this V1UserFeatures.  # noqa: E501
@@ -2065,27 +2138,6 @@ class V1UserFeatures(object):
 
         self._slurm_machine_selector = slurm_machine_selector
 
-    @property
-    def stable_vscode_server(self) -> 'bool':
-        """Gets the stable_vscode_server of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The stable_vscode_server of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._stable_vscode_server
-
-    @stable_vscode_server.setter
-    def stable_vscode_server(self, stable_vscode_server: 'bool'):
-        """Sets the stable_vscode_server of this V1UserFeatures.
-
-
-        :param stable_vscode_server: The stable_vscode_server of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._stable_vscode_server = stable_vscode_server
-
     @property
     def stop_ide_container_on_shutdown(self) -> 'bool':
         """Gets the stop_ide_container_on_shutdown of this V1UserFeatures.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_user_requested_compute_config.py

@@ -42,6 +42,7 @@ class V1UserRequestedComputeConfig(object):
     """
     swagger_types = {
         'affinity_identifier': 'str',
+        'cluster_override': 'str',
         'count': 'int',
         'cpu_image_override': 'str',
         'disk_size': 'int',
@@ -56,6 +57,7 @@ class V1UserRequestedComputeConfig(object):
 
     attribute_map = {
         'affinity_identifier': 'affinityIdentifier',
+        'cluster_override': 'clusterOverride',
         'count': 'count',
         'cpu_image_override': 'cpuImageOverride',
         'disk_size': 'diskSize',
@@ -68,9 +70,10 @@ class V1UserRequestedComputeConfig(object):
         'spot': 'spot'
     }
 
-    def __init__(self, affinity_identifier: 'str' =None, count: 'int' =None, cpu_image_override: 'str' =None, disk_size: 'int' =None, gpu_image_override: 'str' =None, id: 'str' =None, name: 'str' =None, requested_run_duration_seconds: 'str' =None, same_compute_on_resume: 'bool' =None, shm_size: 'int' =None, spot: 'bool' =None):  # noqa: E501
+    def __init__(self, affinity_identifier: 'str' =None, cluster_override: 'str' =None, count: 'int' =None, cpu_image_override: 'str' =None, disk_size: 'int' =None, gpu_image_override: 'str' =None, id: 'str' =None, name: 'str' =None, requested_run_duration_seconds: 'str' =None, same_compute_on_resume: 'bool' =None, shm_size: 'int' =None, spot: 'bool' =None):  # noqa: E501
         """V1UserRequestedComputeConfig - a model defined in Swagger"""  # noqa: E501
         self._affinity_identifier = None
+        self._cluster_override = None
         self._count = None
         self._cpu_image_override = None
         self._disk_size = None
@@ -84,6 +87,8 @@ class V1UserRequestedComputeConfig(object):
         self.discriminator = None
         if affinity_identifier is not None:
             self.affinity_identifier = affinity_identifier
+        if cluster_override is not None:
+            self.cluster_override = cluster_override
         if count is not None:
             self.count = count
         if cpu_image_override is not None:
@@ -128,6 +133,27 @@ class V1UserRequestedComputeConfig(object):
 
         self._affinity_identifier = affinity_identifier
 
+    @property
+    def cluster_override(self) -> 'str':
+        """Gets the cluster_override of this V1UserRequestedComputeConfig.  # noqa: E501
+
+
+        :return: The cluster_override of this V1UserRequestedComputeConfig.  # noqa: E501
+        :rtype: str
+        """
+        return self._cluster_override
+
+    @cluster_override.setter
+    def cluster_override(self, cluster_override: 'str'):
+        """Sets the cluster_override of this V1UserRequestedComputeConfig.
+
+
+        :param cluster_override: The cluster_override of this V1UserRequestedComputeConfig.  # noqa: E501
+        :type: str
+        """
+
+        self._cluster_override = cluster_override
+
     @property
     def count(self) -> 'int':
         """Gets the count of this V1UserRequestedComputeConfig.  # noqa: E501

lightning_sdk/llm/llm.py

@@ -1,3 +1,5 @@
+import os
+import warnings
 from typing import Dict, Generator, List, Optional, Set, Tuple, Union
 
 from lightning_sdk.api.llm_api import LLMApi
@@ -8,7 +10,7 @@ from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
 from lightning_sdk.teamspace import Teamspace
-from lightning_sdk.utils.resolve import _get_authed_user, _resolve_org
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_org, _resolve_teamspace
 
 
 class LLM:
@@ -36,32 +38,60 @@ class LLM:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
         if teamspace is None:
-            if len(possible_teamspaces) == 1:
-                teamspace_name = next(iter(possible_teamspaces.values()))["name"]
-                self._teamspace = Teamspace(name=teamspace_name, org=None, user=user)
-            else:
-                self._teamspace = menu._resolve_teamspace(teamspace)
+            # get current teamspace
+            self._teamspace = _resolve_teamspace(teamspace=None, org=None, user=None)
         else:
             self._teamspace = Teamspace(**menu._get_teamspace_from_name(teamspace, possible_teamspaces))
 
+        if self._teamspace is None:
+            # select the first available teamspace
+            first_teamspace = next(iter(possible_teamspaces.values()), None)
+
+            if first_teamspace:
+                self._teamspace = Teamspace(
+                    name=first_teamspace["name"],
+                    org=first_teamspace["org"],
+                    user=first_teamspace["user"],
+                )
+                warnings.warn(
+                    f"No teamspace given. Using teamspace: {self._teamspace.name}.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
         if self._teamspace is None:
             raise ValueError("Teamspace is required for billing but could not be resolved. ")
 
         self._user = user
 
         self._model_provider, self._model_name = self._parse_model_name(name)
+
+        self._llm_api = LLMApi()
+
         try:
             # check if it is a org model
             self._org = _resolve_org(self._model_provider)
+
+            try:
+                # check if user has access to the org
+                self._org_models = self._build_model_lookup(self._get_org_models())
+            except ApiException:
+                warnings.warn(
+                    f"User is not authenticated to access the model in organization: '{self._model_provider}'.\n"
+                    " Proceeding with appropriate org models, user models, or public models.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                self._model_provider = None
+                raise
         except ApiException:
             if isinstance(self._teamspace.owner, Organization):
                 self._org = self._teamspace.owner
             else:
                 self._org = None
+            self._org_models = self._build_model_lookup(self._get_org_models())
 
-        self._llm_api = LLMApi()
         self._public_models = self._build_model_lookup(self._get_public_models())
-        self._org_models = self._build_model_lookup(self._get_org_models())
         self._user_models = self._build_model_lookup(self._get_user_models())
         self._model = self._get_model()
         self._conversations = {}
@@ -149,20 +179,34 @@ class LLM:
         prompt: str,
         system_prompt: Optional[str] = None,
         max_completion_tokens: Optional[int] = 500,
+        images: Optional[Union[List[str], str]] = None,
         conversation: Optional[str] = None,
+        metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
+        upload_local_images: bool = False,
     ) -> Union[str, Generator[str, None, None]]:
         if conversation and conversation not in self._conversations:
             self._get_conversations()
 
+        if images:
+            if isinstance(images, str):
+                images = [images]
+            for image in images:
+                if not isinstance(image, str):
+                    raise NotImplementedError(f"Image type {type(image)} are not supported yet.")
+                if not image.startswith("http") and upload_local_images:
+                    self._teamspace.upload_file(file_path=image, remote_path=f"images/{os.path.basename(image)}")
+
         conversation_id = self._conversations.get(conversation) if conversation else None
         output = self._llm_api.start_conversation(
             prompt=prompt,
             system_prompt=system_prompt,
             max_completion_tokens=max_completion_tokens,
+            images=images,
             assistant_id=self._model.id,
             conversation_id=conversation_id,
             billing_project_id=self._teamspace.id,
+            metadata=metadata,
             name=conversation,
             stream=stream,
         )

lightning_sdk/studio.py

@@ -1,13 +1,16 @@
 import glob
 import os
 import warnings
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
+from lightning_sdk.api.cluster_api import ClusterApi
 from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.constants import _LIGHTNING_DEBUG
+from lightning_sdk.lightning_cloud.openapi import V1CloudSpaceSourceType
 from lightning_sdk.machine import Machine
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
@@ -24,6 +27,19 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
+class Provider(Enum):
+    # Machine providers based on v1CloudProvider
+    AWS = "AWS"
+    GCP = "GCP"
+    VULTR = "VULTR"
+    LAMBDA_LABS = "LAMBDA_LABS"
+    DGX = "DGX"
+    VOLTAGE_PARK = "VOLTAGE_PARK"
+    NEBIUS = "NEBIUS"
+    CLOUDFLARE = "CLOUDFLARE"
+    LIGHTNING = "LIGHTNING"
+
+
 class Studio:
     """A single Lightning AI Studio.
 
@@ -38,6 +54,8 @@ class Studio:
         cloud_account: the name of the cloud account, the studio should be created on.
             Doesn't matter when the studio already exists.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
+        provider: the provider of the machine, the studio should be created on.
+
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
         only one of the arguments can be provided.
@@ -56,8 +74,11 @@ class Studio:
         cloud_account: Optional[str] = None,
         create_ok: bool = True,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+        provider: Optional[str] = None,
+        source: Optional[V1CloudSpaceSourceType] = None,
     ) -> None:
         self._studio_api = StudioApi()
+        self._cluster_api = ClusterApi()
 
         self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
         self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
@@ -65,6 +86,16 @@ class Studio:
 
         self._plugins = {}
 
+        if provider is not None:
+            if isinstance(provider, str) and provider in Provider.__members__:
+                provider = Provider(provider)
+            else:
+                raise ValueError(f"Invalid provider: {provider}. Must be one of {Provider.__members__.keys()}.")
+            self._cloud_account = self._cluster_api.get_cluster_provider_mapping(
+                self._teamspace.id,
+                self._teamspace.owner.id,
+            )[provider.value]
+
         if name is None:
             studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
             if studio_id is None:
@@ -76,7 +107,7 @@ class Studio:
             except ValueError as e:
                 if create_ok:
                     self._studio = self._studio_api.create_studio(
-                        name, self._teamspace.id, cloud_account=self._cloud_account
+                        name, self._teamspace.id, cloud_account=self._cloud_account, source=source
                     )
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e

lightning_sdk/teamspace.py

@@ -1,8 +1,11 @@
+import glob
 import os
 import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
+from tqdm.auto import tqdm
+
 import lightning_sdk
 from lightning_sdk.agents import Agent
 from lightning_sdk.api import TeamspaceApi
@@ -401,6 +404,71 @@ class Teamspace:
             )
         return self._teamspace_api.list_model_versions(teamspace_id=self.id, model_name=name)
 
+    def upload_file(self, file_path: str, remote_path: Optional[str] = None, progress_bar: bool = True) -> None:
+        """Uploads file to given remote path in the Teamspace drive."""
+        if remote_path is None:
+            remote_path = os.path.split(file_path)[1]
+
+        self._teamspace_api.upload_file(
+            teamspace_id=self._teamspace.id,
+            cloud_account=self.default_cloud_account,
+            file_path=file_path,
+            remote_path=os.path.normpath(remote_path),
+            progress_bar=progress_bar,
+        )
+
+    def upload_folder(self, folder_path: str, remote_path: Optional[str] = None, progress_bar: bool = True) -> None:
+        """Uploads a given folder to a remote path in the Teamspace drive."""
+        if folder_path is None:
+            raise ValueError("Cannot upload a folder that is None.")
+        folder_path = os.path.normpath(folder_path)
+        if os.path.isfile(folder_path):
+            raise NotADirectoryError(f"Cannot upload a file as a folder. '{folder_path}' is a file.")
+        if not os.path.exists(folder_path):
+            raise NotADirectoryError(f"Cannot upload a folder that does not exist. '{folder_path}' is not a directory.")
+        all_files = []
+        for fp in glob.glob(os.path.join(folder_path, "**"), recursive=True):
+            if not os.path.isfile(fp):
+                continue
+            rel_path = os.path.relpath(fp, folder_path)
+            remote_file = os.path.join(remote_path, rel_path) if remote_path else rel_path
+            all_files.append((fp, remote_file))
+
+        if progress_bar:
+            progress_bar = tqdm(total=len(all_files), desc="Uploading files", unit="file")
+        for local_file, remote_path in sorted(all_files, key=lambda p: p[1]):
+            if progress_bar:
+                progress_bar.set_description(f"Uploading {local_file}")
+            self.upload_file(local_file, remote_path=remote_path, progress_bar=False)
+            if progress_bar:
+                progress_bar.update(1)
+        if progress_bar:
+            progress_bar.close()
+
+    def download_file(self, remote_path: str, file_path: Optional[str] = None) -> None:
+        """Downloads a given file in Teamspace drive to a target location."""
+        if file_path is None:
+            file_path = remote_path
+
+        self._teamspace_api.download_file(
+            path=remote_path,
+            target_path=file_path,
+            teamspace_id=self._teamspace.id,
+            cloud_account=self.default_cloud_account,
+        )
+
+    def download_folder(self, remote_path: str, target_path: Optional[str] = None) -> None:
+        """Downloads a folder in the Teamspace drive to a given target path."""
+        if target_path is None:
+            target_path = remote_path
+
+        self._teamspace_api.download_folder(
+            path=remote_path,
+            target_path=target_path,
+            teamspace_id=self._teamspace.id,
+            cloud_account=self.default_cloud_account,
+        )
+
 
 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-0.2.15.dist-info → lightning_sdk-0.2.16.dist-info}/METADATA

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