lightning-sdk 0.2.15__py3-none-any.whl → 0.2.16__py3-none-any.whl

lightning_sdk/__init__.py

@@ -31,6 +31,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "0.2.15"
+__version__ = "0.2.16"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/base_studio_api.py

@@ -16,11 +16,17 @@ class BaseStudioApi:
         """Retrieve the base studio by its ID."""
         try:
             return self._client.cloud_space_environment_template_service_get_cloud_space_environment_template(
-                base_studio_id, org_id
+                base_studio_id, org_id=org_id
             )
         except ValueError as e:
             raise ValueError(f"Base studio {base_studio_id} does not exist") from e
 
+    def get_all_base_studios(self, org_id: str) -> List[V1CloudSpaceEnvironmentTemplate]:
+        """Retrieve all base studios for a given organization."""
+        return self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
+            org_id=org_id
+        )
+
     def update_base_studio(
         self,
         base_studio_id: str,

lightning_sdk/api/cluster_api.py

@@ -1,4 +1,11 @@
-from lightning_sdk.lightning_cloud.openapi import Externalv1Cluster
+from typing import Dict, List, Optional
+
+from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1Cluster,
+    V1CloudProvider,
+    V1ClusterType,
+    V1ListClusterAcceleratorsResponse,
+)
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -20,3 +27,78 @@ class ClusterApi:
         if not res:
             raise ValueError(f"Cluster {cluster_id} does not exist")
         return res
+
+    def list_cluster_accelerators(self, cluster_id: str, org_id: str) -> V1ListClusterAcceleratorsResponse:
+        """Lists the accelerators for a given cluster.
+
+        :param cluster_id: cluster ID test
+        :param project_id: the project the cluster is supposed to be associated with
+        :param org_id: The owning org of this cluster
+        """
+        res = self._client.cluster_service_list_cluster_accelerators(
+            id=cluster_id,
+            org_id=org_id,
+        )
+        if not res:
+            raise ValueError(f"Cluster {cluster_id} does not exist")
+        return res
+
+    def list_global_clusters(self, project_id: str, org_id: str) -> List[Externalv1Cluster]:
+        """Lists the accelerators for a given project.
+
+        :param project_id: project ID test
+        :param org_id: The owning org of this project
+        """
+        res = self._client.cluster_service_list_clusters(
+            project_id=project_id,
+            org_id=org_id,
+        )
+        if not res:
+            raise ValueError(f"Project {project_id} does not exist")
+        filtered_clusters = filter(lambda x: x.spec.cluster_type == V1ClusterType.GLOBAL, res.clusters)
+        return list(filtered_clusters)
+
+    def get_cluster_provider_mapping(self, project_id: str, org_id: str) -> Dict[V1CloudProvider, str]:
+        """Gets the cluster provider mapping."""
+        res = self.list_global_clusters(
+            project_id=project_id,
+            org_id=org_id,
+        )
+        return {self._get_cluster_provider(cluster): cluster.id for cluster in res}
+
+    def _get_cluster_provider(self, cluster: Optional[Externalv1Cluster]) -> V1CloudProvider:
+        """Determines the cloud provider based on the cluster configuration.
+
+        Args:
+            cluster: An optional Externalv1Cluster object containing cluster specifications
+
+        Returns:
+            V1CloudProvider: The determined cloud provider, defaults to AWS if no match is found
+        """
+        if not cluster:
+            return V1CloudProvider.AWS
+
+        if (
+            cluster.spec
+            and cluster.spec.driver
+            and cluster.spec.driver in [V1CloudProvider.LIGHTNING, V1CloudProvider.DGX]
+        ):
+            return cluster.spec.driver
+
+        if cluster.spec:
+            if cluster.spec.aws_v1:
+                return V1CloudProvider.AWS
+            if cluster.spec.google_cloud_v1:
+                return V1CloudProvider.GCP
+            if cluster.spec.lambda_labs_v1:
+                return V1CloudProvider.LAMBDA_LABS
+            if cluster.spec.vultr_v1:
+                return V1CloudProvider.VULTR
+            if cluster.spec.slurm_v1:
+                return V1CloudProvider.SLURM
+            if cluster.spec.voltage_park_v1:
+                return V1CloudProvider.VOLTAGE_PARK
+            if cluster.spec.nebius_v1:
+                return V1CloudProvider.NEBIUS
+
+        return V1CloudProvider.AWS

lightning_sdk/api/llm_api.py

@@ -1,5 +1,6 @@
+import base64
 import json
-from typing import Generator, List, Optional, Union
+from typing import Dict, Generator, List, Optional, Union
 
 from pip._vendor.urllib3 import HTTPResponse
 
@@ -55,25 +56,29 @@ class LLMApi:
                 except json.JSONDecodeError:
                     print("Error decoding JSON:", decoded_line)
 
+    def _encode_image_bytes_to_data_url(self, image: str, mime_type: str = "image/jpeg") -> str:
+        with open(image, "rb") as image_file:
+            b64 = base64.b64encode(image_file.read()).decode("utf-8")
+            return f"data:{mime_type};base64,{b64}"
+
     def start_conversation(
         self,
         prompt: str,
         system_prompt: Optional[str],
         max_completion_tokens: int,
         assistant_id: str,
+        images: Optional[List[str]] = None,
         conversation_id: Optional[str] = None,
         billing_project_id: Optional[str] = None,
         name: Optional[str] = None,
+        metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
     ) -> Union[V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None]]:
         body = {
             "message": {
                 "author": {"role": "user"},
                 "content": [
-                    {
-                        "contentType": "text",
-                        "parts": [prompt],
-                    }
+                    {"contentType": "text", "parts": [prompt]},
                 ],
             },
             "max_completion_tokens": max_completion_tokens,
@@ -81,7 +86,21 @@ class LLMApi:
             "billing_project_id": billing_project_id,
             "name": name,
             "stream": stream,
+            "metadata": metadata or {},
         }
+        if images:
+            for image in images:
+                url = image
+                if not image.startswith("http"):
+                    url = self._encode_image_bytes_to_data_url(image)
+
+                body["message"]["content"].append(
+                    {
+                        "contentType": "image",
+                        "parts": [url],
+                    }
+                )
+
         result = self._client.assistants_service_start_conversation(body, assistant_id, _preload_content=not stream)
         if not stream:
             return result.result

lightning_sdk/api/studio_api.py

@@ -35,6 +35,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpace,
     V1CloudSpaceInstanceConfig,
     V1CloudSpaceSeedFile,
+    V1CloudSpaceSourceType,
     V1CloudSpaceState,
     V1EndpointType,
     V1GetCloudSpaceInstanceStatusResponse,
@@ -110,6 +111,7 @@ class StudioApi:
         name: str,
         teamspace_id: str,
         cloud_account: Optional[str] = None,
+        source: Optional[V1CloudSpaceSourceType] = None,
     ) -> V1CloudSpace:
         """Create a Studio with a given name in a given Teamspace on a possibly given cloud_account."""
         body = ProjectIdCloudspacesBody(
@@ -117,6 +119,7 @@ class StudioApi:
             name=name,
             display_name=name,
             seed_files=[V1CloudSpaceSeedFile(path="main.py", contents="print('Hello, Lightning World!')\n")],
+            source=source,
         )
         studio = self._client.cloud_space_service_create_cloud_space(body, teamspace_id)
 

lightning_sdk/api/teamspace_api.py

@@ -1,10 +1,20 @@
 import os
+import tempfile
+import zipfile
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
+import requests
 from tqdm.auto import tqdm
 
-from lightning_sdk.api.utils import _download_model_files, _DummyBody, _get_model_version, _ModelFileUploader
+from lightning_sdk.api.utils import (
+    _download_model_files,
+    _DummyBody,
+    _FileUploader,
+    _get_model_version,
+    _ModelFileUploader,
+    _resolve_teamspace_remote_path,
+)
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     Externalv1LightningappInstance,
@@ -17,6 +27,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1ClusterAccelerator,
     V1Endpoint,
     V1Job,
+    V1LoginRequest,
     V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
@@ -331,3 +342,118 @@ class TeamspaceApi:
             model_id = self.get_model(teamspace_id=teamspace_id, model_name=model_name).id
         response = self.models_api.models_store_list_model_versions(project_id=teamspace_id, model_id=model_id)
         return response.versions
+
+    def upload_file(
+        self,
+        teamspace_id: str,
+        cloud_account: str,
+        file_path: str,
+        remote_path: str,
+        progress_bar: bool,
+    ) -> None:
+        """Uploads file to given remote path in the Teamspace drive."""
+        _FileUploader(
+            client=self._client,
+            teamspace_id=teamspace_id,
+            cloud_account=cloud_account,
+            file_path=file_path,
+            remote_path=_resolve_teamspace_remote_path(remote_path),
+            progress_bar=progress_bar,
+        )()
+
+    def download_file(
+        self,
+        path: str,
+        target_path: str,
+        teamspace_id: str,
+        cloud_account: str,
+        progress_bar: bool = True,
+    ) -> None:
+        """Downloads a given file in Teamspace drive to a target location."""
+        # TODO: Update this endpoint to permit basic auth
+        auth = Auth()
+        auth.authenticate()
+        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
+
+        query_params = {
+            "clusterId": cloud_account,
+            "key": _resolve_teamspace_remote_path(path),
+            "token": token,
+        }
+
+        r = requests.get(
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
+            params=query_params,
+            stream=True,
+        )
+        total_length = int(r.headers.get("content-length"))
+
+        if progress_bar:
+            pbar = tqdm(
+                desc=f"Downloading {os.path.split(path)[1]}",
+                total=total_length,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1000,
+            )
+
+            pbar_update = pbar.update
+        else:
+            pbar_update = lambda x: None
+
+        target_dir = os.path.split(target_path)[0]
+        if target_dir:
+            os.makedirs(target_dir, exist_ok=True)
+        with open(target_path, "wb") as f:
+            for chunk in r.iter_content(chunk_size=4096 * 8):
+                f.write(chunk)
+                pbar_update(len(chunk))
+
+    def download_folder(
+        self,
+        path: str,
+        target_path: str,
+        teamspace_id: str,
+        cloud_account: str,
+        progress_bar: bool = True,
+    ) -> None:
+        """Downloads a given folder from Teamspace drive to a target location."""
+        # TODO: Update this endpoint to permit basic auth
+        auth = Auth()
+        auth.authenticate()
+        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
+
+        query_params = {
+            "clusterId": cloud_account,
+            "prefix": _resolve_teamspace_remote_path(path),
+            "token": token,
+        }
+
+        r = requests.get(
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
+            params=query_params,
+            stream=True,
+        )
+
+        if progress_bar:
+            pbar = tqdm(
+                desc=f"Downloading {os.path.split(path)[1]}",
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1000,
+            )
+
+            pbar_update = pbar.update
+        else:
+            pbar_update = lambda x: None
+
+        if target_path:
+            os.makedirs(target_path, exist_ok=True)
+
+        with tempfile.TemporaryFile() as f:
+            for chunk in r.iter_content(chunk_size=4096 * 8):
+                f.write(chunk)
+                pbar_update(len(chunk))
+
+            with zipfile.ZipFile(f) as z:
+                z.extractall(target_path)

lightning_sdk/api/utils.py

@@ -355,6 +355,10 @@ def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
     return f"/cloudspaces/{studio_id}/code/content/{path.replace('/teamspace/studios/this_studio/', '')}"
 
 
+def _resolve_teamspace_remote_path(path: str) -> str:
+    return f"/Uploads/{path.replace('/teamspace/studios/this_studio/', '')}"
+
+
 _DOWNLOAD_REQUEST_CHUNK_SIZE = 10 * _BYTES_PER_MB
 _DOWNLOAD_MIN_CHUNK_SIZE = 100 * _BYTES_PER_KB
 

lightning_sdk/base_studio.py

@@ -45,7 +45,12 @@ class BaseStudio:
 
         self._base_studio_api = BaseStudioApi()
 
-        self._base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+        if name is not None:
+            base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+
+            if base_studio is None:
+                raise ValueError(f"Base studio with name {name} does not exist in organization {self._org.name}")
+            self._base_studio = base_studio
 
     def update(
         self,
@@ -68,3 +73,11 @@ class BaseStudio:
             setup_script_text=setup_script_text,
             disabled=disabled,
         )
+
+    def list(self) -> List[V1CloudSpaceEnvironmentTemplate]:
+        """List all base studios in the organization.
+
+        Returns:
+            List[V1CloudSpaceEnvironmentTemplate]: A list of base studio templates.
+        """
+        return self._base_studio_api.get_all_base_studios(self._org.id)

lightning_sdk/cli/create.py

@@ -7,9 +7,12 @@ import click
 from rich.console import Console
 
 from lightning_sdk import Machine, Studio
+from lightning_sdk.api.cluster_api import ClusterApi
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.studio import Provider
 
 _MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+_PROVIDER_VALUES = tuple([provider.value for provider in Provider])
 
 
 @click.group("create")
@@ -44,8 +47,18 @@ def create() -> None:
         "or fall back to the teamspace default."
     ),
 )
+@click.option(
+    "--provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help="The provider to create the studio on. If --cloud-account is specified, this option is prioritized.",
+)
 def studio(
-    name: str, teamspace: Optional[str] = None, start: Optional[str] = None, cloud_account: Optional[str] = None
+    name: str,
+    teamspace: Optional[str] = None,
+    start: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    provider: Optional[str] = None,
 ) -> None:
     """Create a new studio on the Lightning AI platform.
 
@@ -57,6 +70,13 @@ def studio(
     menu = _TeamspacesMenu()
     teamspace_resolved = menu._resolve_teamspace(teamspace)
 
+    if provider is not None:
+        cluster_api = ClusterApi()
+        cloud_account = cluster_api.get_cluster_provider_mapping(
+            teamspace_resolved.id,
+            teamspace_resolved.owner.id,
+        )[provider]
+
     # default cloud account to current studios cloud account if run from studio
     # else it will fall back to teamspace default in the backend
     if cloud_account is None:

lightning_sdk/cli/deploy/_auth.py

@@ -0,0 +1,189 @@
+import os
+import time
+from datetime import datetime
+from enum import Enum
+from typing import Any, List, Optional, TypedDict
+from urllib.parse import urlencode
+
+from rich.console import Console
+from rich.prompt import Confirm
+
+from lightning_sdk import Teamspace
+from lightning_sdk.api import UserApi
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.lightning_cloud import env
+from lightning_sdk.lightning_cloud.login import Auth, AuthServer
+from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_teamspace
+
+LITSERVE_CODE = os.environ.get("LITSERVE_CODE", "j39bzk903h")
+_POLL_TIMEOUT = 120
+
+
+class _AuthMode(Enum):
+    DEVBOX = "dev"
+    DEPLOY = "deploy"
+
+
+class _AuthServer(AuthServer):
+    def __init__(self, mode: _AuthMode, *args: Any, **kwargs: Any) -> None:
+        self._mode = mode
+        super().__init__(*args, **kwargs)
+
+    def get_auth_url(self, port: int) -> str:
+        redirect_uri = f"http://localhost:{port}/login-complete"
+        params = urlencode({"redirectTo": redirect_uri, "mode": self._mode.value, "okbhrt": LITSERVE_CODE})
+        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+
+
+class _Auth(Auth):
+    def __init__(self, mode: _AuthMode, shall_confirm: bool = False) -> None:
+        super().__init__()
+        self._mode = mode
+        self._shall_confirm = shall_confirm
+
+    def _run_server(self) -> None:
+        if self._shall_confirm:
+            proceed = Confirm.ask(
+                "Authenticating with Lightning AI. This will open a browser window. Continue?", default=True
+            )
+            if not proceed:
+                raise RuntimeError(
+                    "Login cancelled. Please login to Lightning AI to deploy the API. Run `lightning login` to login."
+                ) from None
+        print("Opening browser for authentication...")
+        print("Please come back to the terminal after logging in.")
+        time.sleep(3)
+        _AuthServer(self._mode).login_with_browser(self)
+
+
+def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:
+    auth = _Auth(mode, shall_confirm)
+    auth.authenticate()
+
+
+def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
+    if teamspace is None:
+        user = _get_authed_user()
+        menu = _TeamspacesMenu()
+        possible_teamspaces = menu._get_possible_teamspaces(user)
+        if len(possible_teamspaces) == 1:
+            name = next(iter(possible_teamspaces.values()))["name"]
+            return Teamspace(name=name, org=org, user=user)
+
+        return menu._resolve_teamspace(teamspace)
+
+    return _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+
+
+class _UserStatus(TypedDict):
+    verified: bool
+    onboarded: bool
+
+
+def poll_verified_status(timeout: int = _POLL_TIMEOUT) -> _UserStatus:
+    """Polls the verified status of the user until it is True or a timeout occurs."""
+    user_api = UserApi()
+    user = _get_authed_user()
+    start_time = datetime.now()
+    result = {"onboarded": False, "verified": False}
+    while True:
+        user_resp = user_api.get_user(name=user.name)
+        result["onboarded"] = user_resp.status.completed_project_onboarding
+        result["verified"] = user_resp.status.verified
+        if user_resp.status.verified:
+            return result
+        if (datetime.now() - start_time).total_seconds() > timeout:
+            break
+        time.sleep(5)
+    return result
+
+
+class _OnboardingStatus(Enum):
+    NOT_VERIFIED = "not_verified"
+    ONBOARDING = "onboarding"
+    ONBOARDED = "onboarded"
+
+
+class _Onboarding:
+    def __init__(self, console: Console) -> None:
+        self.console = console
+        self.user = _get_authed_user()
+        self.user_api = UserApi()
+        self.client = LightningClient(max_tries=7)
+
+    @property
+    def verified(self) -> bool:
+        return self.user_api.get_user(name=self.user.name).status.verified
+
+    @property
+    def is_onboarded(self) -> bool:
+        return self.user_api.get_user(name=self.user.name).status.completed_project_onboarding
+
+    @property
+    def can_join_org(self) -> bool:
+        return len(self.client.organizations_service_list_joinable_organizations().joinable_organizations) > 0
+
+    @property
+    def status(self) -> _OnboardingStatus:
+        if not self.verified:
+            return _OnboardingStatus.NOT_VERIFIED
+        if self.is_onboarded:
+            return _OnboardingStatus.ONBOARDED
+        return _OnboardingStatus.ONBOARDING
+
+    def _wait_user_onboarding(self, timeout: int = _POLL_TIMEOUT) -> None:
+        """Wait for user onboarding if they can join the teamspace otherwise move to select a teamspace."""
+        status = self.status
+        if status == _OnboardingStatus.ONBOARDED:
+            return
+
+        self.console.print("Waiting for account setup. Visit lightning.ai")
+        start_time = datetime.now()
+        while self.status != _OnboardingStatus.ONBOARDED:
+            time.sleep(5)
+            if self.is_onboarded:
+                return
+            if (datetime.now() - start_time).total_seconds() > timeout:
+                break
+
+        raise RuntimeError("Timed out waiting for onboarding status")
+
+    def get_cloudspace_id(self, teamspace: Teamspace) -> Optional[str]:
+        cloudspaces: List[V1CloudSpace] = self.client.cloud_space_service_list_cloud_spaces(teamspace.id).cloudspaces
+        cloudspaces = sorted(cloudspaces, key=lambda cloudspace: cloudspace.created_at, reverse=True)
+        if len(cloudspaces) == 0:
+            raise RuntimeError("Error creating deployment! Finish account setup at lightning.ai first.")
+        # get the first cloudspace
+        cloudspace = cloudspaces[0]
+        if "scratch-studio" in cloudspace.name or "scratch-studio" in cloudspace.display_name:
+            return cloudspace.id
+        return None
+
+    def select_teamspace(self, teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
+        """Select a teamspace while onboarding.
+
+        If user is being onboarded and can't join any org, the teamspace it will be resolved to the default
+         personal teamspace.
+        If user is being onboarded and can join an org then it will select default teamspace from the org.
+        """
+        if self.is_onboarded:
+            return select_teamspace(teamspace, org, user)
+
+        # Run only when user hasn't completed onboarding yet.
+        menu = _TeamspacesMenu()
+        self._wait_user_onboarding()
+        # Onboarding has been completed - user already selected organization if they could
+        possible_teamspaces = menu._get_possible_teamspaces(self.user)
+        if len(possible_teamspaces) == 1:
+            # User didn't select any org
+            value = next(iter(possible_teamspaces.values()))
+            return Teamspace(name=value["name"], org=value["org"], user=value["user"])
+
+        for _, value in possible_teamspaces.items():
+            # User select an org
+            # Onboarding teamspace will be the default teamspace in the selected org
+            if value["org"]:
+                return Teamspace(name=value["name"], org=value["org"], user=value["user"])
+        raise RuntimeError("Unable to select teamspace. Visit lightning.ai")