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

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"]

lightning_sdk/mmt/base.py

@@ -0,0 +1,288 @@
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Dict, Optional, Protocol, Tuple, Union
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.status import Status
+    from lightning_sdk.studio import Studio
+    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.user import User
+
+from lightning_sdk.job.base import _BaseJob
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
+
+
+class MMTMachine(Protocol):
+    """A single machine in multi-machine training."""
+
+    @property
+    def name(self) -> str:
+        """The Name of the individual machine. Usually corresponds to the rank."""
+        ...
+
+    @property
+    def machine(self) -> "Machine":
+        """The actual machine type this node is running on."""
+        ...
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        """The path to the artifacts of this job."""
+        ...
+
+    @property
+    def status(self) -> "Status":
+        """The status of this job."""
+        ...
+
+    @property
+    def logs(self) -> str:
+        """The logs of the given machine."""
+        ...
+
+
+class _BaseMMT(_BaseJob):
+    """Base interface to all job types."""
+
+    @classmethod
+    def run(
+        cls,
+        name: str,
+        machine: "Machine",
+        num_machines: int,
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+    ) -> "_BaseMMT":
+        """Run async workloads using a docker image across multiple machines.
+
+        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.
+            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.
+            image: The docker image to run the job with. Mutually exclusive with studio.
+            teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
+            org: The organization owning the teamspace (if any). Defaults to the current organization.
+            user: The user owning the teamspace (if any). Defaults to the current user.
+            cloud_account: The cloud account to run the job on.
+                Defaults to the studio cloud account if running with studio compute env.
+                If not provided will fall back to the teamspaces default cloud account.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+            image_credentials: The credentials used to pull the image. Required if the image is private.
+                This should be the name of the respective credentials secret created on the Lightning AI platform.
+            cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
+                Required if the registry is part of a cloud provider (e.g. ECR).
+            artifacts_local: The path of inside the docker container, you want to persist images from.
+                CAUTION: When setting this to "/", it will effectively erase your container.
+                Only supported for jobs with a docker image compute environment.
+            artifacts_remote: The remote storage to persist your artifacts to.
+                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
+                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
+                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
+                within it.
+                Note that the connection needs to be added to the teamspace already in order for it to be found.
+                Only supported for jobs with a docker image compute environment.
+        """
+        from lightning_sdk.studio import Studio
+
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+
+        if num_machines <= 1:
+            raise ValueError("Multi-Machine training cannot be run with less than 2 Machines")
+
+        if not name:
+            raise ValueError("A job needs to have a name!")
+
+        if image is None:
+            if not isinstance(studio, Studio):
+                studio = Studio(
+                    name=studio, teamspace=teamspace, org=org, user=user, cloud_account=cloud_account, create_ok=False
+                )
+
+            # studio is a Studio instance at this point
+            if teamspace is None:
+                teamspace = studio.teamspace
+            else:
+                teamspace_name = teamspace if isinstance(teamspace, str) else teamspace.name
+
+                if studio.teamspace.name != teamspace_name:
+                    raise ValueError(
+                        "Studio teamspace does not match provided teamspace. "
+                        "Can only run jobs with Studio envs in the teamspace of that Studio."
+                    )
+
+            if cloud_account is None:
+                cloud_account = studio.cloud_account
+
+            if cloud_account != studio.cloud_account:
+                raise ValueError(
+                    "Studio cloud_account does not match provided cloud_account. "
+                    "Can only run jobs with Studio envs in the same cloud_account."
+                )
+
+            if image_credentials is not None:
+                raise ValueError("image_credentials is only supported when using a custom image")
+
+            if cloud_account_auth:
+                raise ValueError("cloud_account_auth is only supported when using a custom image")
+
+            if artifacts_local is not None or artifacts_remote is not None:
+                raise ValueError(
+                    "Specifying artifacts persistence is supported for docker images only. "
+                    "Other jobs will automatically persist artifacts to the teamspace distributed filesystem."
+                )
+
+        else:
+            if studio is not None:
+                raise RuntimeError(
+                    "image and studio are mutually exclusive as both define the environment to run the job in"
+                )
+
+            # they either need to specified both or none of them
+            if bool(artifacts_local) != bool(artifacts_remote):
+                raise ValueError("Artifact persistence requires both artifacts_local and artifacts_remote to be set")
+
+            if artifacts_remote and len(artifacts_remote.split(":")) != 3:
+                raise ValueError(
+                    "Artifact persistence requires exactly three arguments separated by colon of kind "
+                    f"<CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>, got {artifacts_local}"
+                )
+
+        inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
+        inst._submit(
+            num_machines=num_machines,
+            machine=machine,
+            cloud_account=cloud_account,
+            command=command,
+            studio=studio,
+            image=image,
+            env=env,
+            interruptible=interruptible,
+            image_credentials=image_credentials,
+            cloud_account_auth=cloud_account_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        return inst
+
+    @abstractmethod
+    def _submit(
+        self,
+        num_machines: int,
+        machine: "Machine",
+        command: Optional[str] = None,
+        studio: Optional["Studio"] = None,
+        image: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        cloud_account: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> None:
+        """Submit a new multi-machine job to the Lightning AI platform.
+
+        Args:
+            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. 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.
+            image: The docker image to run the job with. Mutually exclusive with studio.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+            cloud_account: The cloud account to run the job on.
+                Defaults to the studio cloud account if running with studio compute env.
+                If not provided will fall back to the teamspaces default cloud account.
+            image_credentials: The credentials used to pull the image. Required if the image is private.
+                This should be the name of the respective credentials secret created on the Lightning AI platform.
+            cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
+                Required if the registry is part of a cloud provider (e.g. ECR).
+            artifacts_local: The path of inside the docker container, you want to persist images from.
+                CAUTION: When setting this to "/", it will effectively erase your container.
+                Only supported for jobs with a docker image compute environment.
+            artifacts_remote: The remote storage to persist your artifacts to.
+                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
+                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
+                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
+                within it.
+                Note that the connection needs to be added to the teamspace already in order for it to be found.
+                Only supported for jobs with a docker image compute environment.
+        """
+
+    @property
+    @abstractmethod
+    def machines(self) -> Tuple[MMTMachine, ...]:
+        """Returns the sub-jobs for each individual instance."""
+
+    @property
+    @abstractmethod
+    def machine(self) -> "Machine":
+        """Returns the machine type this job is running on."""
+
+    @abstractmethod
+    def stop(self) -> None:
+        """Stops the job."""
+
+    @abstractmethod
+    def delete(self) -> None:
+        """Deletes the job.
+
+        Caution: This also deletes all artifacts and snapshots associated with the job.
+        """
+
+    @property
+    @abstractmethod
+    def status(self) -> "Status":
+        """The current status of the job."""
+
+    @property
+    @abstractmethod
+    def artifact_path(self) -> Optional[str]:
+        """Path to the artifacts created by the job within the distributed teamspace filesystem."""
+
+    @property
+    @abstractmethod
+    def snapshot_path(self) -> Optional[str]:
+        """Path to the studio snapshot used to create the job within the distributed teamspace filesystem."""
+
+    @property
+    def share_path(self) -> Optional[str]:
+        """Path to the jobs share path."""
+        return None
+
+    @property
+    def name(self) -> str:
+        """The job's name."""
+        return self._name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        """The teamspace the job is part of."""
+        return self._teamspace
+
+    @property
+    def logs(self) -> str:
+        """Logs of the rank 0 machine."""
+        return self.machines[0].logs
+
+    @abstractmethod
+    def _update_internal_job(self) -> None:
+        pass