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

lightning_sdk/services/file_endpoint.py

@@ -1,100 +1,142 @@
-import json
 import os
+from pathlib import Path
 from time import sleep
 from typing import Any, Dict, Optional
-from uuid import uuid4
 
 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,
-    ProjectIdServiceexecutionBody,
-    V1CommandArgument,
-)
+from lightning_sdk.lightning_cloud.openapi import CommandArgumentCommandArgumentType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
-from lightning_sdk.services.uploader import _ServiceFileUploader
-from lightning_sdk.services.utilities import _get_project
-from lightning_sdk.teamspace import Teamspace
-from lightning_sdk.utils import _resolve_teamspace
+from lightning_sdk.services.utilities import _get_cluster, _get_project, _get_service_url
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
 
-_FILE_TO_UPLOADS_KEY = "files_to_upload"
-_DOWNLOAD_IDS_KEY = "download_ids"
-_LIGHTNING_SERVICES_PIPELINE_ID = uuid4().hex
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
+_LIGHTNING_SERVICE_EXECUTION_ID_HEADER = "X-Lightning-Service-Execution-Id"
+_AUTHORIZATION_HEADER = "Authorization"
 
-class FileEndpoint:
-    # TODO: To be merged with the new Client
+
+class Client:
     """This class is used to communicate with the File Endpoint."""
 
     def __init__(
         self,
-        url: str,
+        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 FileEndpoint.
+        """Constructor of the Client.
 
         Args:
-            url: The url of the Studio File Endpoint Service
+            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.
 
         """
-        self.url = url
+        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, args: Optional[Dict[str, str]] = None, files: Optional[Dict[str, str]] = None, output_dir: str = "results"
+        self,
+        **kwargs: Dict[str, str],
     ) -> None:
         """The run method executes the file endpoint.
 
         Args:
-            args: The arguments sent as json data to the file endpoint
-            files: The files to be uploaded to the file endpoint
-            output_dir: The directory output where the artifacts files will be downloaded.
+            kwargs: The keyword arguments associated to the service
 
         """
-        if files is None:
-            files = {}
-        if args is None:
-            args = {}
-
-        response = requests.post(self.url, json=args)
-
-        if response.status_code != 200:
-            raise Exception(f"The endpoint isn't reachable. Status code: {response.status_code}")
+        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}`.")
 
-        data = response.json()
+            argument.value = str(Path(value).resolve())
 
-        if _FILE_TO_UPLOADS_KEY in data:
-            self._upload_files(data, files)
+        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}")
 
-        data = self._check_progress(data)
+        # Avoid uploading duplicated files
+        files_to_upload = {}
+        for argument in self._arguments:
+            if not argument.should_upload:
+                continue
+            files_to_upload[argument.value] = argument
 
-        if _DOWNLOAD_IDS_KEY in data:
-            self._download_files(data, output_dir)
+        # 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),
+            )()
 
-    def _upload_files(self, data: Dict[str, Any], files: Dict[str, str]) -> None:
-        """Upload the files to the Studio."""
-        # TODO: Move this to pre-signed URLs
-        files_to_upload = data[_FILE_TO_UPLOADS_KEY]
+        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())
 
-        if len(files) != len(files_to_upload):
-            raise ValueError(
-                f"This endpoint is expecting {len(files_to_upload)} files to be uploaded. Found only {files}."
-            )
+        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}")
 
-        for file_to_upload in files_to_upload:
-            upload_id = file_to_upload["upload_id"]
-            name = file_to_upload["name"]
-            url = f"{self.url}?upload_id={upload_id}"
-            with open(files[name], "rb") as f:
-                response = requests.post(url, files={upload_id: f})
+        self.headers = {
+            **self.headers,
+            _LIGHTNING_SERVICE_EXECUTION_ID_HEADER: response.headers[_LIGHTNING_SERVICE_EXECUTION_ID_HEADER],
+        }
 
-            if response.status_code != 200:
-                raise Exception(f"Failed to upload the file {name}. Status code: {response.status_code}")
+        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)
+            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}")
@@ -112,25 +154,10 @@ class FileEndpoint:
                 raise RuntimeError("The Studio File Endpoint failed")
 
             # Wait until making the next request
-            sleep(1)
+            sleep(3)
 
         return data
 
-    def _download_files(self, data: Dict[str, str], output_dir: str) -> None:
-        """Download the artifact files."""
-        os.makedirs(output_dir, exist_ok=True)
-
-        for download_id in data[_DOWNLOAD_IDS_KEY]:
-            url = f"{self.url}?download_id={download_id}"
-
-            with requests.post(url, stream=True) as r:
-                r.raise_for_status()
-                filename = r.headers["Content-Disposition"].split("filename=")[1]
-                filename = os.path.basename(filename)
-                with open(os.path.join(output_dir, filename), "wb") as f:
-                    for chunk in r.iter_content(chunk_size=8192):
-                        f.write(chunk)
-
 
 class Argument:
     """A holder for the service argument."""
@@ -168,153 +195,29 @@ class Argument:
         """Returns the name."""
         return self._name
 
-    def to_openapi(self) -> V1CommandArgument:
+    def to_dict(self) -> Dict[str, str]:
         """Convert the argument into its OpenAPI dataclass counterpart."""
         if self.is_text:
-            return V1CommandArgument(name=self._name, type=self._type, value=str(self._value))
-
-        value = self._value.replace("/teamspace/studios/this_studio/", "")
-        return V1CommandArgument(name=self._name, type=self._type, value=value)
-
+            return {self._name: str(self._value)}
 
-class Client:
-    """This class is used to communicate with the Service Endpoint."""
-
-    def __init__(
-        self,
-        id: str,  # noqa: A002
-        teamspace: Optional[str] = None,
-    ) -> None:
-        """Constructor of the FileEndpoint.
-
-        Args:
-            id: The id of the Studio File Endpoint Service.
-            teamspace: The name of the Teamspace in which the Service
+        if self.should_upload:
+            return {self._name: f"/teamspace/Uploads/{os.path.basename(self._value)}"}
 
-        """
-        auth = Auth()
+        return {self._name: self._value}
 
-        try:
-            auth.authenticate()
-        except ConnectionError as e:
-            raise e
-
-        self.id = id
-
-        self._client = LightningClient()
-        if teamspace is not None:
-            teamspace = _resolve_teamspace(teamspace, org=None, user=None)
-
-        self._teamspace_id = (
-            teamspace.id if isinstance(teamspace, Teamspace) else _get_project(client=self._client).project_id
-        )
-
-        self._file_endpoint = self._client.endpoint_service_get_file_endpoint(project_id=self._teamspace_id, id=self.id)
-        self._arguments = []
-
-        for argument in self._file_endpoint.arguments:
-            self._arguments.append(Argument(**argument.to_dict()))
-
-    def run(
-        self,
-        pipeline_id: Optional[str] = _LIGHTNING_SERVICES_PIPELINE_ID,
-        download_artifacts: bool = True,
-        **kwargs: Dict[str, str],
-    ) -> None:
-        """The run method executes the file endpoint.
-
-        Args:
-            pipeline_id: The ID of the current pipeline
-            download_artifacts: Whether to download the artifacts associated to the service execution.
-            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 = value
-
-        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}")
-
-        service_execution = self._client.endpoint_service_create_service_execution(
-            project_id=self._teamspace_id,
-            body=ProjectIdServiceexecutionBody(
-                file_endpoint_id=self.id,
-                pipeline_id=pipeline_id,
-                arguments=[argument.to_openapi() for argument in self._arguments],
-            ),
-        )
-
-        for argument in self._arguments:
-            if argument.is_text:
-                continue
-
-            upload_id = None
-            for service_argument in service_execution.arguments:
-                if argument.name == service_argument.name:
-                    upload_id = service_argument.id
-
-            if upload_id is None:
-                raise ValueError("This shouldn't have happened. Please, contact lightning.ai.")
-
-            # TODO: Move the logic to use pre-defined pre-signed URLs
-            _ServiceFileUploader(
-                client=self._client,
-                teamspace_id=self._teamspace_id,
-                service_execution_id=service_execution.id,
-                upload_id=upload_id,
-                file_path=argument.value,
-                progress_bar=True,
-            )()
-
-        self._client.endpoint_service_run_service_execution(
-            project_id=self._teamspace_id,
-            id=service_execution.id,
-            body={},
-        )
-
-        sleep(1)
-
-        result = {}
-
-        while True:
-            service_execution_status = self._client.endpoint_service_get_service_execution_status(
-                project_id=self._teamspace_id,
-                id=service_execution.id,
-                body={},
-            )
-
-            if "phase" in service_execution_status.data:
-                phase = service_execution_status.data["phase"]
-
-                if "FAILED" in phase:
-                    print("Failed executing. Please, contact lightning.ai")
-                    raise ValueError(f"Failed executing {service_execution.id}. Please, contact lightning.ai")
+    @property
+    def should_upload(self) -> bool:
+        """Whether the file should be uploaded."""
+        if not self.is_file:
+            return False
 
-                if "COMPLETED" in phase:
-                    if "result" in service_execution_status.data:
-                        service_execution_status.data["result"] = json.loads(service_execution_status.data["result"])
-                    print(service_execution_status)
-                    break
+        value = str(Path(self.value).resolve())
+        if value.startswith("/teamspace/"):
+            return False
 
-            print(service_execution_status)
+        return True
 
-            sleep(3)
 
-        if download_artifacts and "result" in service_execution_status.data:
-            result = service_execution_status.data["result"]
-            for k, v in result.items():
-                print(k, v)
+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

@@ -1,17 +1,17 @@
 from typing import Literal
 
-from lightning_sdk.services.file_endpoint import FileEndpoint
+from lightning_sdk.services.file_endpoint import Client
 
 
-class LLMFinetune(FileEndpoint):
+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) -> None:
-        super().__init__(url="https://finetune-01hra53x9nzbhc774s2ecp7bcp.cloudspaces.litng.ai")
+    def __init__(self, teamspace: str) -> None:
+        super().__init__(name="lightning-al/llm-finetunes", teamspace=teamspace)
 
     def run(
         self,
@@ -22,16 +22,14 @@ class LLMFinetune(FileEndpoint):
         learning_rate: float = 0.0002,
         micro_batch_size: int = 2,
         global_batch_size: int = 8,
-        output_dir: str = "results",
     ) -> None:
         """The run method executes the LLM Finetune Service."""
-        args = {
-            "model": str(model),
-            "mode": str(mode),
-            "epochs": str(epochs),
-            "learning_rate": str(learning_rate),
-            "micro_batch_size": str(micro_batch_size),
-            "global_batch_size": str(global_batch_size),
-        }
-        files = {"data_path": data_path}
-        super().run(args, files, output_dir)
+        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

@@ -1,38 +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
 
-LIGHTNING_CLOUD_PROJECT_ID = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+_CHUNK_SIZE = 1024 * 1024
 
 
-def _get_project(client: LightningClient, project_id: Optional[str] = None, verbose: bool = True) -> V1Membership:
+def _get_project(client: LightningClient, project_name: Optional[str] = None) -> V1Membership:
     """Get a project membership for the user from the backend."""
-    if project_id is None:
-        project_id = LIGHTNING_CLOUD_PROJECT_ID
-
-    if project_id is not None:
-        project = client.projects_service_get_project(project_id)
-        if not project:
-            raise ValueError(
-                "Environment variable `LIGHTNING_CLOUD_PROJECT_ID` is set but could not find an associated project."
-            )
-        return V1Membership(
-            name=project.name,
-            display_name=project.display_name,
-            description=project.description,
-            created_at=project.created_at,
-            project_id=project.id,
-            owner_id=project.owner_id,
-            owner_type=project.owner_type,
-            quotas=project.quotas,
-            updated_at=project.updated_at,
-        )
-
     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 len(projects.memberships) > 1 and verbose:
-        print(f"Defaulting to the project: {projects.memberships[0].name}")
-    return projects.memberships[0]
+
+    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

@@ -4,9 +4,14 @@ from enum import Enum
 class Status(Enum):
     """Enum holding all possible studio status types."""
 
-    NotCreated = 1
-    Pending = 2
-    Running = 3
-    Stopping = 4
-    Stopped = 5
-    Failed = 6
+    NotCreated = "NotCreated"
+    Pending = "Pending"
+    Running = "Running"
+    Stopping = "Stopping"
+    Stopped = "Stopped"
+    Completed = "Completed"
+    Failed = "Failed"
+
+    def __str__(self) -> str:
+        """String representation of the enum."""
+        return self.value