lightning-sdk 0.1.3__py3-none-any.whl → 0.1.46__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/version_uploads_body.py

@@ -0,0 +1,123 @@
+# 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 VersionUploadsBody(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 = {
+        'filepath': 'str'
+    }
+
+    attribute_map = {
+        'filepath': 'filepath'
+    }
+
+    def __init__(self, filepath: 'str' =None):  # noqa: E501
+        """VersionUploadsBody - a model defined in Swagger"""  # noqa: E501
+        self._filepath = None
+        self.discriminator = None
+        if filepath is not None:
+            self.filepath = filepath
+
+    @property
+    def filepath(self) -> 'str':
+        """Gets the filepath of this VersionUploadsBody.  # noqa: E501
+
+
+        :return: The filepath of this VersionUploadsBody.  # noqa: E501
+        :rtype: str
+        """
+        return self._filepath
+
+    @filepath.setter
+    def filepath(self, filepath: 'str'):
+        """Sets the filepath of this VersionUploadsBody.
+
+
+        :param filepath: The filepath of this VersionUploadsBody.  # noqa: E501
+        :type: str
+        """
+
+        self._filepath = filepath
+
+    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(VersionUploadsBody, 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: 'VersionUploadsBody') -> bool:
+        """Returns true if both objects are equal"""
+        if not isinstance(other, VersionUploadsBody):
+            return False
+
+        return self.__dict__ == other.__dict__
+
+    def __ne__(self, other: 'VersionUploadsBody') -> bool:
+        """Returns true if both objects are not equal"""
+        return not self == other

lightning_sdk/lightning_cloud/rest_client.py

@@ -4,7 +4,6 @@ import time
 from functools import wraps
 from typing import Callable, Optional, Any
 
-import click
 import urllib3
 from lightning_sdk.lightning_cloud import env
 from lightning_sdk.lightning_cloud.login import Auth
@@ -26,6 +25,14 @@ from lightning_sdk.lightning_cloud.openapi import (
     UserServiceApi,
     BillingServiceApi,
     EndpointServiceApi,
+    SlurmJobsUserServiceApi,
+    LitLoggerServiceApi,
+    JobsServiceApi,
+    AssistantsServiceApi,
+    StorageServiceApi,
+    DeploymentTemplatesServiceApi,
+    ModelsStoreApi,
+    LitRegistryServiceApi
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.source_code.logs_socket_api import LightningLogsSocketAPI
@@ -78,6 +85,14 @@ class GridRestClient(
         UserServiceApi,
         BillingServiceApi,
         EndpointServiceApi,
+        SlurmJobsUserServiceApi,
+        LitLoggerServiceApi,
+        JobsServiceApi,
+        AssistantsServiceApi,
+        StorageServiceApi,
+        DeploymentTemplatesServiceApi,
+        ModelsStoreApi,
+        LitRegistryServiceApi
 ):
 
     def __init__(self, api_client: Optional[ApiClient] = None):
@@ -112,32 +127,42 @@ def _retry_wrapper(self,
             try:
                 return func(self, *args, **kwargs)
             except (ApiException, urllib3.exceptions.HTTPError) as ex:
-                # retry if the backend fails with all errors except 4xx but not 408 - (Request Timeout)
-                if (isinstance(ex, urllib3.exceptions.HTTPError)
-                        or ex.status in (408, 409)
-                        or not str(ex.status).startswith("4")):
-                    consecutive_errors += 1
-                    backoff_time = _get_next_backoff_time(consecutive_errors)
-
-                    msg = (f"error: {str(ex)}" if isinstance(
-                        ex, urllib3.exceptions.HTTPError) else
-                           f"response: {ex.status}")
-                    logger.debug(
-                        f"The {func.__name__} request failed to reach the server, {msg}."
-                        f" Retrying after {backoff_time} seconds.")
+                if not _should_retry(ex):
+                    raise ex
 
-                    if max_tries is not None and consecutive_errors == max_tries:
-                        raise Exception(
-                            f"The {func.__name__} request failed to reach the server, {msg}."
-                        )
+                consecutive_errors += 1
+                backoff_time = _get_next_backoff_time(consecutive_errors)
 
-                    time.sleep(backoff_time)
-                else:
-                    raise ex
+                msg = (f"error: {str(ex)}" if isinstance(
+                    ex, urllib3.exceptions.HTTPError) else
+                        f"response: {ex.status}")
+                logger.debug(
+                    f"The {func.__name__} request failed to reach the server, {msg}."
+                    f" Retrying after {backoff_time} seconds.")
+
+                if max_tries is not None and consecutive_errors == max_tries:
+                    raise Exception(
+                        f"The {func.__name__} request failed to reach the server, {msg}."
+                    )
+
+                time.sleep(backoff_time)
 
     return wrapped
 
 
+def _should_retry(ex: BaseException) -> bool:
+    if isinstance(ex, urllib3.exceptions.HTTPError):
+        return True
+
+    if str(ex.status).startswith("4") and ex.status not in (400, 401, 404):
+        return True
+
+    if str(ex.status).startswith("5"):
+        return True
+
+    return False
+
+
 class LightningClient(GridRestClient):
     """The LightningClient is a wrapper around the GridRestClient.
 
@@ -176,8 +201,7 @@ def request_auth_warning_wrapper(func):
             return response
         except ApiException as err:
             if err.status == 401:
-                raise click.ClickException(
-                    "Authentication failed. Please run `lightning login`.")
+                raise ConnectionError("Authentication failed. Please run `lightning login`.")
             raise err
 
     return wrap

lightning_sdk/lightning_cloud/source_code/logs_socket_api.py

@@ -17,7 +17,7 @@ class LightningLogsSocketAPI:
 
     def _get_api_token(self):
         token_resp = self._auth_service.auth_service_login(body=V1LoginRequest(
-            username=self._auth.username,
+            username=self._auth.user_id,
             api_key=self._auth.api_key,
         ))
         return token_resp.token

lightning_sdk/lightning_cloud/source_code/tar.py

@@ -5,8 +5,6 @@ import subprocess
 import tarfile
 from typing import Optional, Tuple
 
-import click
-
 MAX_SPLIT_COUNT = 999
 
 
@@ -85,7 +83,7 @@ def get_split_size(total_size: int,
         1 << 31
     )  # max size per part limited by Requests or urllib as shown in ref above
     if total_size > max_size:
-        raise click.ClickException(
+        raise RuntimeError(
             f"The size of the datastore to be uploaded is bigger than our {max_size/(1 << 40):.2f} TBytes limit"
         )
 

lightning_sdk/lightning_cloud/utils/data_connection.py

@@ -1,14 +1,13 @@
 import os
 from time import sleep, time
-from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection
+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
 
 
 def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeout: int = 15) -> None:
     """Utility to add a data connection."""
-    from lightning_sdk.lightning_cloud import rest_client
-
     client = rest_client.LightningClient(retry=False)
 
     project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
@@ -23,6 +22,7 @@ def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeou
         name=bucket_name,
         create_index=True,
         cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
         aws=V1AwsDataConnection(
             source=f"s3://{bucket_name}",
             region=region
@@ -30,11 +30,11 @@ def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeou
     try:
         client.data_connection_service_create_data_connection(body, project_id)
     except (ApiException, urllib3.exceptions.HTTPError) as ex:
-        # Note: This function can be called in a distributed way. 
+        # 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
         # and this request would fail with duplicated key
-        # In this case, it is fine not to raise 
-        if 'duplicate key value violates unique constraint' in str(ex.body):
+        # In this case, it is fine not to raise
+        if isinstance(ex, ApiException) and 'duplicate key value violates unique constraint' in str(ex.body):
             pass
         else:
             raise ex
@@ -43,4 +43,142 @@ def add_s3_connection(bucket_name: str, region: str = "us-east-1", create_timeou
     start = time()
 
     while not os.path.isdir(f"/teamspace/s3_connections/{bucket_name}") and (time() - start) < create_timeout:
-        sleep(1)
+        sleep(1)
+    
+    return
+
+def create_s3_folder(folder_name: str, create_timeout: int = 15) -> None:
+    """
+    Utility function to create a writable S3 folder.
+
+    Parameters:
+        1. folder_name (str): The name of the folder to create.
+        2. create_timeout (int): The timeout for the folder creation.
+    Returns:
+        None
+    """
+    client = rest_client.LightningClient(retry=False)
+
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+    cluster_id = os.getenv("LIGHTNING_CLUSTER_ID")
+
+    # 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
+
+        if existing_folder_name == folder_name and isS3Folder:
+            return
+
+    # If we get here, no matching folder was found, proceed with creation
+    body = Create(
+        name=folder_name,
+        create_resources=True,
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        force=True,
+        writable=True,
+        s3_folder=V1S3FolderDataConnection()
+    )
+    try:
+        client.data_connection_service_create_data_connection(body, project_id)
+    except (ApiException, urllib3.exceptions.HTTPError) as ex:
+        # 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
+        # and this request would fail with duplicated key
+        # In this case, it is fine not to raise
+        if isinstance(ex, ApiException) and 'duplicate key value violates unique constraint' in str(ex.body):
+            pass
+        else:
+            raise ex
+
+    # Wait for the filesystem picks up the new data connection
+    start = time()
+
+    while not os.path.isdir(f"/teamspace/s3_folders/{folder_name}") and (time() - start) < create_timeout:
+        sleep(1)
+
+    return
+
+def create_efs_folder(folder_name: str, region: str) -> 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"
+    """
+    client = rest_client.LightningClient(retry=False)
+
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+    cluster_id = os.getenv("LIGHTNING_CLUSTER_ID")
+
+
+    # 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)
+        isEFS = getattr(connection, 'efs', None) is not None
+
+        if existing_folder_name == folder_name and isEFS:
+            return
+
+    # If we get here, no matching folder was found, proceed with creation
+    body = Create(
+        name=folder_name,
+        create_resources=True,
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        force=True,
+        writable=True,
+        efs=V1EfsConfig(region=region),
+    )
+    try:
+        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
+        # and this request would fail with duplicated key
+        # In this case, it is fine not to raise
+        if'duplicate key value violates unique constraint' in str(e.body):
+            return
+        raise e from None
+
+    except urllib3.exceptions.HTTPError as e:
+        raise e from None
+
+def delete_data_connection(name: str):
+    """Utility to delete a data connection
+
+    Args:
+        name: the name of the data connection to remove
+    """
+    client = rest_client.LightningClient(retry=False)
+
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+
+    # Get existing data connections
+    data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
+
+    data_connection_id = None
+
+    for conn in data_connections:
+        if conn.name == name:
+            data_connection_id = conn.id
+            break
+
+    # data connection wasn't found
+    if data_connection_id is None:
+        return
+
+    try:
+        client.data_connection_service_delete_data_connection(project_id=project_id, id=data_connection_id)
+    except ApiException as e:
+        # TODO: with managed efs folders the contract around when it actually deletes is a bit fragile.
+        #  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

lightning_sdk/machine.py

@@ -1,15 +1,28 @@
-from enum import Enum
+from lightning_sdk.utils.enum import DeprecationEnum
 
 
-class Machine(Enum):
+class Machine(DeprecationEnum):
     """Enum holding all supported machine types for studios."""
 
+    CPU_SMALL = "CPU_SMALL"
     CPU = "CPU"
     DATA_PREP = "DATA_PREP"
+    DATA_PREP_MAX = "DATA_PREP_MAX"
+    DATA_PREP_ULTRA = "DATA_PREP_ULTRA"
     T4 = "T4"
     T4_X_4 = "T4_X_4"
-    V100 = "V100"
-    V100_X_4 = "V100_X_4"
+    L4 = "L4"
+    L4_X_4 = "L4_X_4"
+    L4_X_8 = "L4_X_8"
     A10G = "A10G"
     A10G_X_4 = "A10G_X_4"
+    A10G_X_8 = "A10G_X_8"
+    L40S = "L40S"
+    L40 = "L40", "L40S"  # deprecated
+    L40S_X_4 = "L40S_X_4"
+    L40_X_4 = "L40_X_4", "L40S_X_4"  # deprecated
+    L40S_X_8 = "L40S_X_8"
+    L40_X_8 = "L40_X_8", "L40S_X_8"  # deprecated
     A100_X_8 = "A100_X_8"
+    H100_X_8 = "H100_X_8"
+    H200_X_8 = "H200_X_8"

lightning_sdk/mmt/__init__.py

@@ -0,0 +1,4 @@
+from lightning_sdk.mmt.base import MMTMachine
+from lightning_sdk.mmt.mmt import MMT
+
+__all__ = ["MMT", "MMTMachine"]