lightning-sdk 0.1.51__py3-none-any.whl → 0.1.54__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_resource_visibility.py

@@ -41,25 +41,51 @@ class V1ResourceVisibility(object):
                             and the value is json key in definition.
     """
     swagger_types = {
+        'all_org_users': 'bool',
         'all_users': 'bool',
         'user_ids': 'list[str]'
     }
 
     attribute_map = {
+        'all_org_users': 'allOrgUsers',
         'all_users': 'allUsers',
         'user_ids': 'userIds'
     }
 
-    def __init__(self, all_users: 'bool' =None, user_ids: 'list[str]' =None):  # noqa: E501
+    def __init__(self, all_org_users: 'bool' =None, all_users: 'bool' =None, user_ids: 'list[str]' =None):  # noqa: E501
         """V1ResourceVisibility - a model defined in Swagger"""  # noqa: E501
+        self._all_org_users = None
         self._all_users = None
         self._user_ids = None
         self.discriminator = None
+        if all_org_users is not None:
+            self.all_org_users = all_org_users
         if all_users is not None:
             self.all_users = all_users
         if user_ids is not None:
             self.user_ids = user_ids
 
+    @property
+    def all_org_users(self) -> 'bool':
+        """Gets the all_org_users of this V1ResourceVisibility.  # noqa: E501
+
+
+        :return: The all_org_users of this V1ResourceVisibility.  # noqa: E501
+        :rtype: bool
+        """
+        return self._all_org_users
+
+    @all_org_users.setter
+    def all_org_users(self, all_org_users: 'bool'):
+        """Sets the all_org_users of this V1ResourceVisibility.
+
+
+        :param all_org_users: The all_org_users of this V1ResourceVisibility.  # noqa: E501
+        :type: bool
+        """
+
+        self._all_org_users = all_org_users
+
     @property
     def all_users(self) -> 'bool':
         """Gets the all_users of this V1ResourceVisibility.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -62,12 +62,14 @@ class V1UserFeatures(object):
         'deployment_customize_api': 'bool',
         'deployment_data_path': 'bool',
         'deployment_persistent_disk': 'bool',
+        'deployment_version_upgrade': 'bool',
         'deployment_version_visibility': 'bool',
         'docs_agent': 'bool',
         'drive_v2': 'bool',
         'enable_crypto_crackdown': 'bool',
         'enable_efs': 'bool',
         'enable_storage_limits': 'bool',
+        'fair_share': 'bool',
         'featured_studios_admin': 'bool',
         'filesystem_optimisation': 'bool',
         'inference_job_deployment_plugin': 'bool',
@@ -146,12 +148,14 @@ class V1UserFeatures(object):
         'deployment_customize_api': 'deploymentCustomizeApi',
         'deployment_data_path': 'deploymentDataPath',
         'deployment_persistent_disk': 'deploymentPersistentDisk',
+        'deployment_version_upgrade': 'deploymentVersionUpgrade',
         'deployment_version_visibility': 'deploymentVersionVisibility',
         'docs_agent': 'docsAgent',
         'drive_v2': 'driveV2',
         'enable_crypto_crackdown': 'enableCryptoCrackdown',
         'enable_efs': 'enableEfs',
         'enable_storage_limits': 'enableStorageLimits',
+        'fair_share': 'fairShare',
         'featured_studios_admin': 'featuredStudiosAdmin',
         'filesystem_optimisation': 'filesystemOptimisation',
         'inference_job_deployment_plugin': 'inferenceJobDeploymentPlugin',
@@ -208,7 +212,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, 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, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_customize_api: 'bool' =None, deployment_data_path: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_version_visibility: 'bool' =None, docs_agent: 'bool' =None, drive_v2: '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, inference_job_deployment_plugin: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: '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, mmt_v2: 'bool' =None, multiple_deployment_versions: 'bool' =None, multiple_studio_versions: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: '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, security_docs: '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, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: '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, 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, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_customize_api: 'bool' =None, deployment_data_path: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_version_upgrade: 'bool' =None, deployment_version_visibility: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_efs: 'bool' =None, enable_storage_limits: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filesystem_optimisation: 'bool' =None, inference_job_deployment_plugin: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: '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, mmt_v2: 'bool' =None, multiple_deployment_versions: 'bool' =None, multiple_studio_versions: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: '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, security_docs: '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, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: '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
@@ -231,12 +235,14 @@ class V1UserFeatures(object):
         self._deployment_customize_api = None
         self._deployment_data_path = None
         self._deployment_persistent_disk = None
+        self._deployment_version_upgrade = None
         self._deployment_version_visibility = None
         self._docs_agent = None
         self._drive_v2 = None
         self._enable_crypto_crackdown = None
         self._enable_efs = None
         self._enable_storage_limits = None
+        self._fair_share = None
         self._featured_studios_admin = None
         self._filesystem_optimisation = None
         self._inference_job_deployment_plugin = None
@@ -334,6 +340,8 @@ class V1UserFeatures(object):
             self.deployment_data_path = deployment_data_path
         if deployment_persistent_disk is not None:
             self.deployment_persistent_disk = deployment_persistent_disk
+        if deployment_version_upgrade is not None:
+            self.deployment_version_upgrade = deployment_version_upgrade
         if deployment_version_visibility is not None:
             self.deployment_version_visibility = deployment_version_visibility
         if docs_agent is not None:
@@ -346,6 +354,8 @@ class V1UserFeatures(object):
             self.enable_efs = enable_efs
         if enable_storage_limits is not None:
             self.enable_storage_limits = enable_storage_limits
+        if fair_share is not None:
+            self.fair_share = fair_share
         if featured_studios_admin is not None:
             self.featured_studios_admin = featured_studios_admin
         if filesystem_optimisation is not None:
@@ -896,6 +906,27 @@ class V1UserFeatures(object):
 
         self._deployment_persistent_disk = deployment_persistent_disk
 
+    @property
+    def deployment_version_upgrade(self) -> 'bool':
+        """Gets the deployment_version_upgrade of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The deployment_version_upgrade of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._deployment_version_upgrade
+
+    @deployment_version_upgrade.setter
+    def deployment_version_upgrade(self, deployment_version_upgrade: 'bool'):
+        """Sets the deployment_version_upgrade of this V1UserFeatures.
+
+
+        :param deployment_version_upgrade: The deployment_version_upgrade of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._deployment_version_upgrade = deployment_version_upgrade
+
     @property
     def deployment_version_visibility(self) -> 'bool':
         """Gets the deployment_version_visibility of this V1UserFeatures.  # noqa: E501
@@ -1022,6 +1053,27 @@ class V1UserFeatures(object):
 
         self._enable_storage_limits = enable_storage_limits
 
+    @property
+    def fair_share(self) -> 'bool':
+        """Gets the fair_share of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The fair_share of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._fair_share
+
+    @fair_share.setter
+    def fair_share(self, fair_share: 'bool'):
+        """Sets the fair_share of this V1UserFeatures.
+
+
+        :param fair_share: The fair_share of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._fair_share = fair_share
+
     @property
     def featured_studios_admin(self) -> 'bool':
         """Gets the featured_studios_admin of this V1UserFeatures.  # noqa: E501

lightning_sdk/lightning_cloud/utils/data_connection.py

@@ -4,7 +4,7 @@ from lightning_sdk.lightning_cloud import rest_client
 from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection, V1S3FolderDataConnection, V1EfsConfig
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 import urllib3
-
+import os
 
 def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeout: int = 15) -> None:
     """Utility to add a data connection."""
@@ -26,7 +26,7 @@ def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeou
         aws=V1AwsDataConnection(
             source=f"s3://{bucket_name}",
             region=region
-    ))
+        ))
     try:
         client.data_connection_service_create_data_connection(body, project_id)
     except (ApiException, urllib3.exceptions.HTTPError) as ex:
@@ -44,7 +44,7 @@ def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeou
 
     while not os.path.isdir(f"/teamspace/s3_connections/{bucket_name}") and (time() - start) < create_timeout:
         sleep(1)
-    
+
     return
 
 def create_s3_folder(folder_name: str, create_timeout: int = 15) -> None:
@@ -64,7 +64,7 @@ def create_s3_folder(folder_name: str, create_timeout: int = 15) -> None:
 
     # Get existing data connections
     data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
-    
+
     for connection in data_connections:
         existing_folder_name = getattr(connection, 'name', None)
         isS3Folder = getattr(connection, 's3_folder', None) is not None
@@ -102,13 +102,14 @@ def create_s3_folder(folder_name: str, create_timeout: int = 15) -> None:
 
     return
 
-def create_efs_folder(folder_name: str, region: str) -> None:
+def create_efs_folder(folder_name: str, region: str, create_timeout: int = 60) -> None:
     """
     Utility function to create a EFS folder.
 
     Args:
         folder_name: The name of the folder to create.
         region: the region to create the efs in. Could be something like "us-east-1"
+        create_timeout: Timeout for creating the efs folder. Defaults to 60 seconds
     """
     client = rest_client.LightningClient(retry=False)
 
@@ -137,7 +138,7 @@ def create_efs_folder(folder_name: str, region: str) -> None:
         efs=V1EfsConfig(region=region),
     )
     try:
-        client.data_connection_service_create_data_connection(body, project_id)
+        connection = client.data_connection_service_create_data_connection(body, project_id)
     except ApiException as e:
         # Note: This function can be called in a distributed way.
         # There is a race condition where one machine might create the entry before another machine
@@ -150,6 +151,72 @@ def create_efs_folder(folder_name: str, region: str) -> None:
     except urllib3.exceptions.HTTPError as e:
         raise e from None
 
+    start = time()
+    while True:
+        if time() - start > create_timeout:
+            print(f"Dataconnection {connection.name} didn't become accessible withing {create_timeout} seconds!")
+
+
+        data_connection = client.data_connection_service_get_data_connection(project_id=project_id, id=connection.id)
+        if data_connection.accessible:
+            break
+
+        sleep(1)
+
+def add_efs_connection(name: str, filesystem_id: str, region: str = "us-east-1", create_timeout: int = 60) -> None:
+    """Utility to add an existing EFS filesystem as a data connection.
+
+    Args:
+        name: The name to give to the data connection
+        filesystem_id: The ID of the existing EFS filesystem (e.g., 'fs-1234567')
+        region: AWS region where the EFS filesystem exists (default: 'us-east-1')
+        create_timeout: Timeout in seconds to wait for the connection to be accessible (default: 60)
+    """
+    client = rest_client.LightningClient(retry=False)
+
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+    cluster_id = os.getenv("LIGHTNING_CLUSTER_ID")
+
+    data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
+
+    if any(d for d in data_connections if d.name == name):
+        return
+
+    body = Create(
+        name=name,
+        create_resources=False,  # Don't create new EFS since we're connecting to existing one
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        force=True,
+        writable=True,
+        efs=V1EfsConfig(
+            region=region,
+            file_system_id=filesystem_id,
+        )
+    )
+
+    try:
+        connection = client.data_connection_service_create_data_connection(body, project_id)
+    except (ApiException, urllib3.exceptions.HTTPError) as ex:
+        if isinstance(ex, ApiException) and 'duplicate key value violates unique constraint' in str(ex.body):
+            pass
+        else:
+            raise ex
+
+    start = time()
+    while True:
+        if time() - start > create_timeout:
+            print(f"Dataconnection {connection.name} didn't become accessible withing {create_timeout} seconds!")
+            break
+
+        data_connection = client.data_connection_service_get_data_connection(project_id=project_id, id=connection.id)
+        if data_connection.accessible:
+            break
+
+        sleep(1)
+
+    return
+
 def delete_data_connection(name: str):
     """Utility to delete a data connection
 
@@ -181,4 +248,5 @@ def delete_data_connection(name: str):
         #  It may exhaust the attempts before the connection is actually unmounted from the studio.
         #  for now it's best to actually stop the studio and all other things where the connection
         #  is mounted before trying to delete it
-        raise e from None
+        raise e from None
+

lightning_sdk/mmt/mmt.py

@@ -2,7 +2,6 @@ from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 from lightning_sdk.api.user_api import UserApi
-from lightning_sdk.job.job import _has_jobs_v2
 from lightning_sdk.mmt.base import MMTMachine, _BaseMMT
 from lightning_sdk.mmt.v1 import _MMTV1
 from lightning_sdk.mmt.v2 import _MMTV2
@@ -21,16 +20,18 @@ _logger = _setup_logger(__name__)
 
 @lru_cache(maxsize=None)
 def _has_mmt_v2() -> bool:
-    # users need both mmtv2 and jobsv2 flags in order for mmtv2 to work correctly
-    if not _has_jobs_v2():
-        return False
-
     api = UserApi()
     try:
-        return api._get_feature_flags().mmt_v2
+        feature_flags = api._get_feature_flags()
     except Exception:
         return False
 
+    try:
+        return feature_flags.mmt_v2
+    except AttributeError:
+        # Feature flag doesn't exist anymore, so return True
+        return True
+
 
 class MMT(_BaseMMT):
     """Class to submit and manage multi-machine jobs on the Lightning AI Platform."""
@@ -120,7 +121,7 @@ class MMT(_BaseMMT):
         Args:
             name: The name of the job. Needs to be unique within the teamspace.
             machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
-            num_machine: The number of machines to run on.
+            num_machines: The number of machines to run on.
             command: The command to run inside your job. Required if using a studio. Optional if using an image.
                 If not provided for images, will run the container entrypoint and default command.
             studio: The studio env to run the job with. Mutually exclusive with image.

lightning_sdk/plugin.py

@@ -432,6 +432,8 @@ def _success_message(resp: Union["Externalv1LightningappInstance", Job], plugin_
 def forced_v1(cls: Any) -> Generator[Any, None, None]:
     """Forces to use the v1 version of a class when using a class with multiple backends."""
     orig_val = getattr(cls, "_force_v1", False)
-    cls._force_v1 = True
-    yield cls
-    cls._force_v1 = orig_val
+    try:
+        cls._force_v1 = True
+        yield cls
+    finally:
+        cls._force_v1 = orig_val

lightning_sdk/studio.py

@@ -1,6 +1,6 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
 from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.utils import _machine_to_compute_name
@@ -14,6 +14,8 @@ from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
 
 if TYPE_CHECKING:
+    from lightning_sdk.job import Job
+    from lightning_sdk.mmt import MMT
     from lightning_sdk.plugin import Plugin
 
 _logger = _setup_logger(__name__)
@@ -281,6 +283,71 @@ class Studio:
             cloud_account=self._studio.cluster_id,
         )
 
+    def run_job(
+        self,
+        name: str,
+        machine: Union["Machine", str],
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+    ) -> "Job":
+        """Run async workloads using the compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        """
+        from lightning_sdk.job import Job
+
+        return Job.run(
+            name=name,
+            machine=machine,
+            command=command,
+            studio=self,
+            image=None,
+            teamspace=self.teamspace,
+            cloud_account=self.cloud_account,
+            env=env,
+            interruptible=interruptible,
+        )
+
+    def run_mmt(
+        self,
+        name: str,
+        num_machines: int,
+        machine: Union["Machine", str],
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+    ) -> "MMT":
+        """Run async workloads using the compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            num_machines: The number of machines to run on.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        """
+        from lightning_sdk.mmt import MMT
+
+        return MMT.run(
+            name=name,
+            num_machines=num_machines,
+            machine=machine,
+            command=command,
+            studio=self,
+            image=None,
+            teamspace=self.teamspace,
+            cloud_account=self.cloud_account,
+            env=env,
+            interruptible=interruptible,
+        )
+
     @property
     def auto_sleep(self) -> bool:
         """Returns if a Studio has auto-sleep enabled."""

{lightning_sdk-0.1.51.dist-info → lightning_sdk-0.1.54.dist-info}/METADATA

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