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

lightning_sdk/services/file_endpoint.py

@@ -0,0 +1,223 @@
+import os
+from pathlib import Path
+from time import sleep
+from typing import Any, Dict, Optional
+
+import requests
+import urllib3
+
+from lightning_sdk.api.utils import _FileUploader
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.lightning_cloud.openapi import CommandArgumentCommandArgumentType
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.services.utilities import _get_cluster, _get_project, _get_service_url
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
+
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+_LIGHTNING_SERVICE_EXECUTION_ID_HEADER = "X-Lightning-Service-Execution-Id"
+_AUTHORIZATION_HEADER = "Authorization"
+
+
+class Client:
+    """This class is used to communicate with the File Endpoint."""
+
+    def __init__(
+        self,
+        name: str,
+        teamspace: Optional[str],
+        cloud_account: Optional[str] = None,
+        cluster_id: Optional[str] = None,  # deprecated in favor of cloud_account
+    ) -> None:
+        """Constructor of the Client.
+
+        Args:
+            name: The name of the Studio File Endpoint Service.
+            teamspace: The name of the teamspace you want to attach the upload data and artifacts to be.
+            cloud_account: The name of the cloud account on which to upload the data.
+
+        """
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster_id)
+
+        self._auth = Auth()
+
+        try:
+            self._auth.authenticate()
+        except ConnectionError as e:
+            raise e
+
+        self._name = name
+        self._teamspace = teamspace
+        self._client = LightningClient()
+        self._project = _get_project(client=self._client, project_name=teamspace)
+        self._cloud_account = _get_cluster(
+            client=self._client, project_id=self._project.project_id, cluster_id=cloud_account
+        )
+        self._file_endpoint = self._client.endpoint_service_get_file_endpoint_by_name(
+            project_id=self._project.project_id, name=self._name
+        )
+
+        self.headers = {_AUTHORIZATION_HEADER: f"Bearer {self._auth.api_key}"}
+
+        self._arguments = []
+        for argument in self._file_endpoint.arguments:
+            self._arguments.append(Argument(**argument.to_dict()))
+
+        self.url = _get_service_url(self._file_endpoint.cloudspace_id, self._file_endpoint.id)
+
+    def run(
+        self,
+        **kwargs: Dict[str, str],
+    ) -> None:
+        """The run method executes the file endpoint.
+
+        Args:
+            kwargs: The keyword arguments associated to the service
+
+        """
+        for argument in self._arguments:
+            if argument.is_file:
+                if argument.name not in kwargs:
+                    raise ValueError(f"This endpoint expects a file for the argument `{argument.name}`.")
+                value = kwargs[argument.name]
+                if not os.path.isfile(value):
+                    raise ValueError(f"This endpoint expects a file for the argument `{argument.name}`.")
+            else:
+                if argument.name not in kwargs:
+                    raise ValueError(f"This endpoint expects a value for the argument `{argument.name}`.")
+                value = kwargs[argument.name]
+                if os.path.isfile(value):
+                    raise ValueError(f"This endpoint doesn't expect a file for `{argument.name}`.")
+
+            argument.value = str(Path(value).resolve())
+
+        missing_names = [v.name for v in self._arguments if v.value is None]
+        if missing_names:
+            raise ValueError(f"You are missing values for the following arguments: {missing_names}")
+
+        # Avoid uploading duplicated files
+        files_to_upload = {}
+        for argument in self._arguments:
+            if not argument.should_upload:
+                continue
+            files_to_upload[argument.value] = argument
+
+        # TODO: Verify if the file exists in Teampace and avoid uploading if already there.
+        for argument in files_to_upload.values():
+            _FileUploader(
+                client=self._client,
+                teamspace_id=self._project.project_id,
+                cloud_account=self._cloud_account.cluster_id,
+                file_path=argument.value,
+                progress_bar=True,
+                remote_path=_sanitize_uploads_remote_path(argument.value),
+            )()
+
+        json = {
+            "teamspace_id": self._project.project_id,
+            "cluster_id": self._cloud_account.cluster_id,
+            "input": {},
+        }
+        for argument in self._arguments:
+            json["input"].update(**argument.to_dict())
+
+        response = requests.post(self.url, json=json, headers=self.headers)
+        if response.status_code != 200:
+            raise Exception(f"The endpoint isn't reachable. Status code: {response.status_code}")
+
+        self.headers = {
+            **self.headers,
+            _LIGHTNING_SERVICE_EXECUTION_ID_HEADER: response.headers[_LIGHTNING_SERVICE_EXECUTION_ID_HEADER],
+        }
+
+        self._check_progress(response.json())
+
+    def _check_progress(self, data: Dict[str, str]) -> Dict[str, str]:
+        """Check the current Studio status."""
+        while True:
+            url = f"{self.url}?run_id={data['run_id']}"
+            response = requests.post(url, headers=self.headers)
+
+            if response.status_code != 200:
+                raise Exception(f"The file endpoint had an error. Status code: {response.status_code}")
+
+            data = response.json()
+
+            # Display the progress status to the user.
+            print(data)
+
+            if data["stage"] == "completed":
+                break
+
+            if data["stage"] == "failed":
+                # TODO: Add more information on why the execution failed.
+                raise RuntimeError("The Studio File Endpoint failed")
+
+            # Wait until making the next request
+            sleep(3)
+
+        return data
+
+
+class Argument:
+    """A holder for the service argument."""
+
+    def __init__(self, name: str, type: CommandArgumentCommandArgumentType, **kwargs: Any) -> None:  # noqa: A002
+        self._name = name
+        self._type = type
+        self._value = None
+        self._kwargs = kwargs
+
+    @property
+    def is_text(self) -> bool:
+        """Whether this argument is of type Text."""
+        return self._type == CommandArgumentCommandArgumentType.TEXT
+
+    @property
+    def is_file(self) -> bool:
+        """Whether this argument is of type File."""
+        return self._type == CommandArgumentCommandArgumentType.FILE
+
+    @property
+    def value(self) -> Any:
+        """Returns the value."""
+        return self._value
+
+    @value.setter
+    def value(self, value: Any) -> None:
+        """Store the value."""
+        if self.is_file and not os.path.exists(value):
+            raise ValueError(f"The argument {self._name} should be a valid file.")
+        self._value = value
+
+    @property
+    def name(self) -> str:
+        """Returns the name."""
+        return self._name
+
+    def to_dict(self) -> Dict[str, str]:
+        """Convert the argument into its OpenAPI dataclass counterpart."""
+        if self.is_text:
+            return {self._name: str(self._value)}
+
+        if self.should_upload:
+            return {self._name: f"/teamspace/Uploads/{os.path.basename(self._value)}"}
+
+        return {self._name: self._value}
+
+    @property
+    def should_upload(self) -> bool:
+        """Whether the file should be uploaded."""
+        if not self.is_file:
+            return False
+
+        value = str(Path(self.value).resolve())
+        if value.startswith("/teamspace/"):
+            return False
+
+        return True
+
+
+def _sanitize_uploads_remote_path(file_path: str) -> str:
+    remote_path = os.path.basename(file_path)
+    return f"/Uploads/{remote_path}"

lightning_sdk/services/finetune/__init__.py

@@ -0,0 +1,35 @@
+from typing import Literal
+
+from lightning_sdk.services.file_endpoint import Client
+
+
+class LLMFinetune(Client):
+    """The LLM Finetune is the client to the LLM Finetune Service Studio.
+
+    Learn more: https://lightning.ai/lightning-ai/studios/llm-finetune-service~01h5rahq6gbhw5m4bzyws0at5h.
+
+    """
+
+    def __init__(self, teamspace: str) -> None:
+        super().__init__(name="lightning-al/llm-finetunes", teamspace=teamspace)
+
+    def run(
+        self,
+        data_path: str,
+        model: Literal["llama-2-7b", "tiny-llama"] = "tiny-llama",
+        mode: Literal["lora", "full"] = "lora",
+        epochs: int = 3,
+        learning_rate: float = 0.0002,
+        micro_batch_size: int = 2,
+        global_batch_size: int = 8,
+    ) -> None:
+        """The run method executes the LLM Finetune Service."""
+        super().run(
+            data_path=data_path,
+            model=model,
+            mode=mode,
+            epochs=epochs,
+            learning_rate=learning_rate,
+            micro_batch_size=micro_batch_size,
+            global_batch_size=global_batch_size,
+        )

lightning_sdk/services/utilities.py

@@ -0,0 +1,111 @@
+import contextlib
+import os
+from typing import Optional
+
+import requests
+import urllib3
+
+from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.lightning_cloud.openapi import V1Membership
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+_CHUNK_SIZE = 1024 * 1024
+
+
+def _get_project(client: LightningClient, project_name: Optional[str] = None) -> V1Membership:
+    """Get a project membership for the user from the backend."""
+    projects = client.projects_service_list_memberships()
+    if len(projects.memberships) == 0:
+        raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
+
+    if project_name is None:
+        return projects.memberships[0]
+
+    matches = []
+    for membership in projects.memberships:
+        if membership.name == project_name or membership.display_name == project_name:
+            matches.append(membership)
+
+    if len(matches) == 1:
+        return matches[0]
+
+    if len(matches) >= 2:
+        raise ValueError(f"We found several teamspaces. Which one do you want to use {[m.name for m in matches]}")
+
+    raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
+
+
+def _get_cluster(client: LightningClient, project_id: str, cluster_id: Optional[str] = None) -> V1Membership:
+    """Get a project membership for the user from the backend."""
+    clusters = client.projects_service_list_project_cluster_bindings(project_id=project_id)
+    if cluster_id:
+        for cluster in clusters.clusters:
+            if cluster.cluster_id == cluster_id:
+                return cluster
+        raise ValueError(
+            f"No valid cluster found with the provided {cluster_id}."
+            f"Found {[c.cluster_id for c in clusters.clusters]}."
+        )
+
+    clusters = sorted(clusters.clusters, key=lambda x: x.created_at)
+    if len(clusters):
+        return clusters[0]
+    return None
+
+
+def _get_service_url(cloud_space_id: str, file_endpoint_id: str) -> str:
+    url = _get_cloud_url()
+    domain = _get_domain(url)
+    protocol = _get_protocol(url)
+    return f"{protocol}//{file_endpoint_id}-{cloud_space_id}.cloudspaces.{domain}"
+
+
+def _get_domain(url: str) -> str:
+    base_url = url.split("//")[1].split("/")[0]
+    if "localhost:9800" in base_url:
+        return "local.litng.ai:8118"
+    if "lightning.ai" in base_url and "localhost:8888" in base_url:
+        return "litng.ai"
+    return base_url
+
+
+def _get_protocol(url: str) -> str:
+    return url.split("//")[0]
+
+
+def download_file(filepath: str, cache_dir: str = "/cache") -> str:
+    """Download the file passed to the service execution by the user."""
+    service_id = os.getenv("LIGHTNING_SERVICE_EXECUTION_ID")
+
+    if service_id is None:
+        raise RuntimeError("The ServiceId is required. Please reach out to lightning.ai team.")
+
+    # Note: Make the function idempotent
+    prefix = os.path.join(cache_dir, service_id)
+    if not filepath.startswith(prefix):
+        saved_filepath = filepath
+        if saved_filepath.startswith("/"):
+            saved_filepath = saved_filepath[1:]
+        saved_filepath = os.path.join(prefix, saved_filepath)
+    else:
+        saved_filepath = filepath
+        filepath = filepath.replace(prefix, "")
+
+    if os.path.exists(saved_filepath):
+        return saved_filepath
+
+    client = LightningClient(retry=False)
+    download_artifacts = client.endpoint_service_download_service_execution_artifact(
+        project_id=os.getenv("LIGHTNING_CLOUD_PROJECT_ID"), id=service_id, page_token=None, filepath=filepath
+    )
+    artifact = download_artifacts.artifacts[0]
+
+    os.makedirs(os.path.dirname(saved_filepath), exist_ok=True)
+
+    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+    with contextlib.suppress(ConnectionError):
+        request = requests.get(artifact.url, stream=True, verify=False)
+        with open(saved_filepath, "wb") as fp:
+            for chunk in request.iter_content(chunk_size=_CHUNK_SIZE):
+                fp.write(chunk)  # type: ignore
+    return saved_filepath

lightning_sdk/status.py

@@ -9,4 +9,5 @@ class Status(Enum):
     Running = 3
     Stopping = 4
     Stopped = 5
-    Failed = 6
+    Completed = 6
+    Failed = 7

lightning_sdk/studio.py

@@ -1,7 +1,9 @@
 import os
+import warnings
 from typing import TYPE_CHECKING, Any, Mapping, Optional, Tuple, Union
 
 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.machine import Machine
 from lightning_sdk.organization import Organization
@@ -9,7 +11,7 @@ from lightning_sdk.owner import Owner
 from lightning_sdk.status import Status
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils import _resolve_teamspace, _setup_logger
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
 
 if TYPE_CHECKING:
     from lightning_sdk.plugin import Plugin
@@ -28,10 +30,9 @@ class Studio:
         teamspace: the name of the teamspace the studio is contained by
         org: the name of the organization owning the :param`teamspace` in case it is owned by an org
         user: the name of the user owning the :param`teamspace` in case it is owned directly by a user instead of an org
-        cluster: the name of the cluster, the studio should be created on.
+        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
-
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
         only one of the arguments can be provided.
@@ -47,13 +48,14 @@ class Studio:
         teamspace: Optional[Union[str, Teamspace]] = None,
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         create_ok: bool = True,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> None:
         self._studio_api = StudioApi()
 
         self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
-        self._cluster = cluster
+        self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
         self._setup_done = False
 
         self._plugins = {}
@@ -68,7 +70,9 @@ class Studio:
                 self._studio = self._studio_api.get_studio(name, self._teamspace.id)
             except ValueError as e:
                 if create_ok:
-                    self._studio = self._studio_api.create_studio(name, self._teamspace.id, cluster=self._cluster)
+                    self._studio = self._studio_api.create_studio(
+                        name, self._teamspace.id, cloud_account=self._cloud_account
+                    )
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e
 
@@ -131,14 +135,32 @@ class Studio:
             return None
         return self._studio_api.get_machine(self._studio.id, self._teamspace.id)
 
-    def start(self, machine: Machine = Machine.CPU) -> None:
+    @property
+    def interruptible(self) -> bool:
+        """Returns whether the Studio is running on a interruptible instance."""
+        if self.status != Status.Running:
+            return None
+
+        return self._studio_api.get_interruptible(self._studio.id, self._teamspace.id)
+
+    @property
+    def cluster(self) -> str:
+        """Returns the cluster the Studio is running on."""
+        warnings.warn("Studio.cluster is deprecated. Use Studio.cloud_account instead", DeprecationWarning)
+        return self.cloud_account
+
+    @property
+    def cloud_account(self) -> str:
+        return self._studio.cluster_id
+
+    def start(self, machine: Union[Machine, str] = Machine.CPU, interruptible: bool = False) -> None:
         """Starts a Studio on the specified machine type (default: CPU-4)."""
         status = self.status
         if status == Status.Running:
-            curr_machine = self.machine
-            if curr_machine != machine:
+            curr_machine = _machine_to_compute_name(self.machine) if self.machine is not None else None
+            if curr_machine != _machine_to_compute_name(machine):
                 raise RuntimeError(
-                    f"Requested to start studio on {machine}, but studio is already running on {curr_machine}."
+                    f"Requested to start studio on {machine}, but studio is already running on {self.machine}."
                     " Consider switching instead!"
                 )
             _logger.info(f"Studio {self.name} is already running")
@@ -146,7 +168,7 @@ class Studio:
 
         if status != Status.Stopped:
             raise RuntimeError(f"Cannot start a studio that is not stopped. Studio {self.name} is {status}.")
-        self._studio_api.start_studio(self._studio.id, self._teamspace.id, machine)
+        self._studio_api.start_studio(self._studio.id, self._teamspace.id, machine, interruptible=interruptible)
 
         self._setup()
 
@@ -166,11 +188,12 @@ class Studio:
         kwargs = self._studio_api.duplicate_studio(self._studio.id, self._teamspace.id, self._teamspace.id)
         return Studio(**kwargs)
 
-    def switch_machine(self, machine: Machine) -> None:
+    def switch_machine(self, machine: Union[Machine, str], interruptible: bool = False) -> None:
         """Switches machine to the provided machine type/.
 
         Args:
             machine: the new machine type to switch to
+            interruptible: determines whether to switch to an interruptible instance
 
         Note:
             this call is blocking until the new machine is provisioned
@@ -181,7 +204,9 @@ class Studio:
             raise RuntimeError(
                 f"Cannot switch machine on a studio that is not running. Studio {self.name} is {status}."
             )
-        self._studio_api.switch_studio_machine(self._studio.id, self._teamspace.id, machine)
+        self._studio_api.switch_studio_machine(
+            self._studio.id, self._teamspace.id, machine, interruptible=interruptible
+        )
 
     def run_with_exit_code(self, *commands: str) -> Tuple[str, int]:
         """Runs given commands on the Studio while returning output and exit code.
@@ -224,9 +249,9 @@ class Studio:
         self._studio_api.upload_file(
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
             file_path=file_path,
-            remote_path=remote_path,
+            remote_path=os.path.normpath(remote_path),
             progress_bar=progress_bar,
         )
 
@@ -240,7 +265,7 @@ class Studio:
             target_path=file_path,
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
         )
 
     def download_folder(self, remote_path: str, target_path: Optional[str] = None) -> None:
@@ -253,8 +278,53 @@ class Studio:
             target_path=target_path,
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
+        )
+
+    @property
+    def auto_sleep(self) -> bool:
+        """Returns if a Studio has auto-sleep enabled."""
+        return not self._studio.code_config.disable_auto_shutdown
+
+    @auto_sleep.setter
+    def auto_sleep(self, value: bool) -> None:
+        if not value and self.machine == Machine.CPU:
+            warnings.warn("Disabling auto-sleep will convert the Studio from free to paid!")
+        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, enabled=value, studio=self._studio)
+        self._update_studio_reference()
+
+    @property
+    def auto_sleep_time(self) -> int:
+        """Returns the time in seconds a Studio has to be idle for auto-sleep to kick in (if enabled)."""
+        return self._studio.code_config.idle_shutdown_seconds
+
+    @auto_sleep_time.setter
+    def auto_sleep_time(self, value: int) -> None:
+        warnings.warn("Setting auto-sleep time will convert the Studio from free to paid!")
+        self._studio_api.update_autoshutdown(
+            self._studio.id, self._teamspace.id, idle_shutdown_seconds=value, studio=self._studio
         )
+        self._update_studio_reference()
+
+    @property
+    def auto_shutdown(self) -> bool:
+        warnings.warn("auto_shutdown is deprecated. Use auto_sleep instead", DeprecationWarning)
+        return self.auto_sleep
+
+    @auto_shutdown.setter
+    def auto_shutdown(self, value: bool) -> None:
+        warnings.warn("auto_shutdown is deprecated. Use auto_sleep instead", DeprecationWarning)
+        self.auto_sleep = value
+
+    @property
+    def auto_shutdown_time(self) -> int:
+        warnings.warn("auto_shutdown_time is deprecated. Use auto_sleep_time instead", DeprecationWarning)
+        return self.auto_sleep_time
+
+    @auto_shutdown_time.setter
+    def auto_shutdown_time(self, value: int) -> None:
+        warnings.warn("auto_shutdown_time is deprecated. Use auto_sleep_time instead", DeprecationWarning)
+        self.auto_sleep_time = value
 
     @property
     def available_plugins(self) -> Mapping[str, str]:
@@ -342,6 +412,9 @@ class Studio:
         """Returns reader friendly representation."""
         return repr(self)
 
+    def _update_studio_reference(self) -> None:
+        self._studio = self._studio_api.get_studio_by_id(studio_id=self._studio.id, teamspace_id=self._teamspace.id)
+
 
 def _internal_status_to_external_status(internal_status: str) -> Status:
     """Converts internal status strings from HTTP requests to external enums."""