lightning-sdk 2025.11.13.post0__py3-none-any.whl

lightning_sdk/lightning_cloud/source_code/tar.py

@@ -0,0 +1,194 @@
+from dataclasses import dataclass
+import math
+import os
+import subprocess
+import tarfile
+from typing import Optional, Tuple
+
+MAX_SPLIT_COUNT = 999
+
+
+def get_dir_size_and_count(source_dir: str,
+                           prefix: Optional[str] = None) -> Tuple[int, int]:
+    """
+    Get size and file count of a directory
+
+    Parameters
+    ----------
+    source_dir: str
+        Directory path
+
+    Returns
+    -------
+    Tuple[int, int]
+        Size in megabytes and file count
+    """
+    size = 0
+    count = 0
+    for root, _, files in os.walk(source_dir, topdown=True):
+        for f in files:
+            if prefix and not f.startswith(prefix):
+                continue
+
+            full_path = os.path.join(root, f)
+            size += os.path.getsize(full_path)
+            count += 1
+
+    return (size, count)
+
+
+@dataclass
+class TarResults:
+    """
+    This class holds the results of running tar_path.
+
+    Attributes
+    ----------
+    before_size: int
+        The total size of the original directory files in bytes
+    after_size: int
+        The total size of the compressed and tarred split files in bytes
+    """
+    before_size: int
+    after_size: int
+
+
+def get_split_size(total_size: int,
+                   minimum_split_size: int = 1024 * 1000 * 20,
+                   max_split_count: int = MAX_SPLIT_COUNT) -> int:
+    """
+    Calculate the split size we should use to split the multipart upload of an object to a bucket.  We are limited
+    to 1000 max parts as the way we are using ListMultipartUploads.
+    More info
+        https://github.com/gridai/grid/pull/5267
+        https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html#mpu-process
+        https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListMultipartUploads.html
+        https://github.com/psf/requests/issues/2717#issuecomment-724725392 Python or requests has a limit of 2**31 bytes for a single file upload.
+
+    Parameters
+    ----------
+    minimum_split_size: int
+        The minimum split size to use
+    max_split_count: int
+        The maximum split count
+    total_size: int
+        Total size of the file to split
+
+    Returns
+    -------
+    int
+        Split size
+    """
+    max_size = max_split_count * (
+        1 << 31
+    )  # max size per part limited by Requests or urllib as shown in ref above
+    if total_size > max_size:
+        raise RuntimeError(
+            f"The size of the datastore to be uploaded is bigger than our {max_size/(1 << 40):.2f} TBytes limit"
+        )
+
+    split_size = minimum_split_size
+    split_count = math.ceil(total_size / split_size)
+    if split_count > max_split_count:
+        # Adjust the split size based on max split count
+        split_size = math.ceil(total_size / max_split_count)
+
+    return split_size
+
+
+def tar_path(source_path: str,
+             target_file: str,
+             compression: bool = False) -> TarResults:
+    """
+    Create tar from directory using `tar`
+
+    Parameters
+    ----------
+    source_path: str
+        Source directory or file
+    target_file
+        Target tar file
+    compression: bool, default False
+        Enable compression, which is disabled by default.
+
+    Returns
+    -------
+    TarResults
+        Results that holds file counts and sizes
+    """
+    if os.path.isdir(source_path):
+        before_size, _ = get_dir_size_and_count(source_path)
+    else:
+        before_size = os.path.getsize(source_path)
+
+    try:
+        _tar_path_subprocess(source_path, target_file, compression)
+    except subprocess.CalledProcessError:
+        _tar_path_python(source_path, target_file, compression)
+
+    after_size = os.stat(target_file).st_size
+    return TarResults(before_size=before_size, after_size=after_size)
+
+
+def _tar_path_python(source_path: str,
+                     target_file: str,
+                     compression: bool = False) -> None:
+    """
+    Create tar from directory using `python`
+
+    Parameters
+    ----------
+    source_path: str
+        Source directory or file
+    target_file
+        Target tar file
+    compression: bool, default False
+        Enable compression, which is disabled by default.
+    """
+    file_mode = "w:gz" if compression else "w:"
+
+    with tarfile.open(target_file, file_mode) as tar:
+        if os.path.isdir(source_path):
+            tar.add(str(source_path), arcname=".")
+        elif os.path.isfile(source_path):
+            file_info = tarfile.TarInfo(os.path.basename(str(source_path)))
+            tar.addfile(file_info, open(source_path))
+
+
+def _tar_path_subprocess(source_path: str,
+                         target_file: str,
+                         compression: bool = False) -> None:
+    """
+    Create tar from directory using `tar`
+
+    Parameters
+    ----------
+    source_path: str
+        Source directory or file
+    target_file
+        Target tar file
+    compression: bool, default False
+        Enable compression, which is disabled by default.
+    """
+    # Only add compression when users explicitly request it.
+    # We do this because it takes too long to compress
+    # large datastores.
+    tar_flags = '-cvf'
+    if compression:
+        tar_flags = '-zcvf'
+    if os.path.isdir(source_path):
+        command = f"tar -C {source_path} {tar_flags} {target_file} ./"
+    else:
+        abs_path = os.path.abspath(source_path)
+        parent_dir = os.path.dirname(abs_path)
+        base_name = os.path.basename(abs_path)
+        command = f"tar -C {parent_dir} {tar_flags} {target_file} {base_name}"
+
+    subprocess.check_call(command,
+                          stdout=subprocess.DEVNULL,
+                          stderr=subprocess.DEVNULL,
+                          shell=True,
+                          env={
+                              "GZIP": "-9",
+                              "COPYFILE_DISABLE": "1"
+                          })

lightning_sdk/lightning_cloud/source_code/uploader.py

@@ -0,0 +1,97 @@
+import time
+
+import requests
+from requests.adapters import HTTPAdapter
+from rich.progress import BarColumn, Progress, TextColumn
+from urllib3.util.retry import Retry
+
+
+class FileUploader:
+    """
+    This class uploads a source file with presigned url to S3
+
+    Attributes
+    ----------
+    source_file: str
+        Source file to upload
+    presigned_url: str
+        Presigned urls dictionary, with key as part number and values as urls
+    retries: int
+        Amount of retries when requests encounter an error
+    total_size: int
+        Size of all files to upload
+    name: str
+        Name of this upload to display progress
+    """
+    workers: int = 8
+    retries: int = 10000
+    disconnect_retry_wait_seconds: int = 5
+
+    progress = Progress(
+        TextColumn("[bold blue]{task.description}", justify="left"),
+        BarColumn(bar_width=None),
+        "[self.progress.percentage]{task.percentage:>3.1f}%")
+
+    def __init__(self, presigned_url: str, source_file: str, total_size: int,
+                 name: str):
+        self.presigned_url = presigned_url
+        self.source_file = source_file
+        self.total_size = total_size
+        self.name = name
+
+    @staticmethod
+    def upload_s3_data(url: str, data: bytes, retries: int,
+                       disconnect_retry_wait_seconds: int) -> str:
+        """
+        Send data to s3 url
+
+        Parameters
+        ----------
+        url: str
+            S3 url string to send data to
+        data: bytes
+             Bytes of data to send to S3
+        retries: int
+            Amount of retries
+        disconnect_retry_wait_seconds: int
+            Amount of seconds between disconnect retry
+
+        Returns
+        -------
+        str
+            ETag from response
+        """
+        disconnect_retries = retries
+        while disconnect_retries > 0:
+            try:
+                retries = Retry(total=10)
+                with requests.Session() as s:
+                    s.mount('https://', HTTPAdapter(max_retries=retries))
+                    response = s.put(url, data=data)
+                    if 'ETag' not in response.headers:
+                        raise ValueError(
+                            f"Unexpected response from S3, response: {response.content}"
+                        )
+                    return response.headers['ETag']
+            except BrokenPipeError:
+                time.sleep(disconnect_retry_wait_seconds)
+                disconnect_retries -= 1
+
+        raise ValueError("Unable to upload file after multiple attempts")
+
+    def upload(self) -> None:
+        """
+        Upload files from source dir into target path in S3
+        """
+        task_id = self.progress.add_task("upload",
+                                         filename=self.name,
+                                         total=self.total_size)
+        self.progress.start()
+        try:
+            with open(self.source_file, 'rb') as f:
+                data = f.read()
+            self.upload_s3_data(self.presigned_url, data, self.retries,
+                                self.disconnect_retry_wait_seconds)
+            self.progress.update(task_id, advance=len(data))
+        finally:
+            self.progress.stop()

lightning_sdk/lightning_cloud/utils/__init__.py

@@ -0,0 +1,3 @@
+from lightning_sdk.lightning_cloud.utils.data_connection import add_s3_connection
+
+__all__ = ["add_s3_connection"] 

lightning_sdk/lightning_cloud/utils/data_connection.py

@@ -0,0 +1,411 @@
+import os
+from time import sleep, time
+from lightning_sdk.lightning_cloud import rest_client
+from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection, V1S3FolderDataConnection, V1EfsConfig, V1GcpDataConnection, V1FilestoreDataConnection, V1DataConnectionTier, V1R2DataConnection
+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."""
+    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 == bucket_name):
+        return
+
+    body = Create(
+        name=bucket_name,
+        create_index=True,
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        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:
+        # 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_connections/{bucket_name}") and (time() - start) < create_timeout:
+        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, 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)
+
+    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:
+        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
+        # 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
+
+    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 add_gcs_connection(connection_name: str, bucket_name: str, create_timeout: int = 15) -> None:
+    """
+    Utility function to add a GCS data connection.
+
+    Parameters:
+        1. connection_name (str): The name of the data connection.
+        2. bucket_name (str): The name of the bucket to attach.
+        3. create_timeout (int): The timeout for the data connectio 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 and ensure there is no an existing one with the same name and type
+    data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
+
+    for connection in data_connections:
+        existing_connection_name = getattr(connection, 'name', None)
+        isGCSConnection = getattr(connection, 'gcp', None) is not None
+
+        # Same name and type already exists
+        if existing_connection_name == connection_name and isGCSConnection:
+            return
+
+    body = Create(
+        name=connection_name,
+        create_index=True,
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        gcp=V1GcpDataConnection(
+            source=f"gs://{bucket_name}",
+        ))
+    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 newly added GCS data connection
+    start = time()
+
+    while not os.path.isdir(f"/teamspace/gcs_connections/{bucket_name}") and (time() - start) < create_timeout:
+        sleep(1)
+
+    return
+
+def create_filestore_folder(folder_name: str, region: str, capacity_gb: int = 1024, tier: V1DataConnectionTier = V1DataConnectionTier.HDD, create_timeout: int = 600) -> None:
+    """
+    Utility function to create a Filestore folder.
+
+    Args:
+        folder_name: The name of the folder to create.
+        region: the region to create the folder in, e.g. "us-central1".
+        capacity_gb: capacity of the Filestore folder.
+        tier: Filestore folder tier, SSD or HDD.
+        create_timeout (int): The timeout for the folder creation.
+    """
+    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)
+        isFilestore = getattr(connection, 'filestore', None) is not None
+
+        if existing_folder_name == folder_name and isFilestore:
+            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],
+        filestore=V1FilestoreDataConnection(region=region, capacity_gb=capacity_gb, tier=tier, source=folder_name),
+    )
+    try:
+        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
+        # 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
+
+    # Wait for the filesystem picks up the newly added Filestore folder
+    start = time()
+
+    while not os.path.isdir(f"/teamspace/filestore_folders/{folder_name}") and (time() - start) < create_timeout:
+        sleep(1)
+
+
+def create_cloud_agnostic_folder(folder_name: str, create_timeout: int = 30) -> None:
+    """
+    Utility function to create a Cloud-Agnostic (R2) folder.
+
+    Args:
+        folder_name: The name of the folder to create.
+        create_timeout (int): The timeout for the folder creation.
+    """
+    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)
+        isCloudAgnostic = getattr(connection, 'r2_folder', None) is not None
+
+        if existing_folder_name == folder_name and isCloudAgnostic:
+            return
+
+    # If we get here, no matching folder was found, proceed with creation
+    body = Create(
+        name=folder_name,
+        create_resources=True,
+        r2=V1R2DataConnection(name=folder_name),
+    )
+    try:
+        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
+        # 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
+
+    # Wait for the filesystem picks up the newly added Cloud-Agnostic folder
+    start = time()
+
+    while not os.path.isdir(f"/teamspace/lightning_storage/{folder_name}") and (time() - start) < create_timeout:
+        sleep(1)
+
+
+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