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

lightning_sdk/models.py

@@ -0,0 +1,153 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+from lightning_sdk.api import OrgApi, TeamspaceApi, UserApi
+from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _get_authed_user
+
+if TYPE_CHECKING:
+    from lightning_sdk.teamspace import Teamspace
+
+
+# TODO: Maybe just have a `Model` object?
+@dataclass
+class UploadedModelInfo:
+    name: str
+    version: str
+    teamspace: str
+    cloud_account: str
+
+
+def _get_teamspace_and_path(
+    ts: V1Membership, org_api: OrgApi, user_api: UserApi, authed_user: User
+) -> Tuple[str, Dict[str, Any]]:
+    if ts.owner_type == V1OwnerType.ORGANIZATION:
+        org = org_api._get_org_by_id(ts.owner_id)
+        return f"{org.name}/{ts.name}", {"name": ts.name, "org": org.name}
+
+    if ts.owner_type == V1OwnerType.USER and ts.owner_id != authed_user.id:
+        user = user_api._get_user_by_id(ts.owner_id)  # todo: check also the name
+        return f"{user.username}/{ts.name}", {"name": ts.name, "user": User(name=user.username)}
+
+    if ts.owner_type == V1OwnerType.USER:
+        return f"{authed_user.name}/{ts.name}", {"name": ts.name, "user": authed_user}
+
+    raise RuntimeError(f"Unknown organization type {ts.owner_type}")
+
+
+def _list_teamspaces() -> List[str]:
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    return [
+        _get_teamspace_and_path(ts, org_api, user_api, authed_user)[0]
+        for ts in user_api._get_all_teamspace_memberships("")
+    ]
+
+
+def _get_teamspace(name: str, organization: str) -> "Teamspace":
+    """Get a Teamspace object from the SDK."""
+    from lightning_sdk.teamspace import Teamspace
+
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    requested_teamspace = f"{organization}/{name}".lower()
+
+    for ts in user_api._get_all_teamspace_memberships(""):
+        if ts.name != name:
+            continue
+
+        teamspace_path, teamspace = _get_teamspace_and_path(ts, org_api, user_api, authed_user)
+        if requested_teamspace == teamspace_path:
+            return Teamspace(**teamspace)
+
+    options = f"{os.linesep}\t".join(_list_teamspaces())
+    raise RuntimeError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: {os.linesep}\t{options}")
+
+
+def _parse_model_name_and_version(name: str) -> Tuple[str, str, str, str]:
+    """Parse the name argument into its components."""
+    try:
+        org_name, teamspace_name, model_name = name.split("/")
+        parts = model_name.split(":")
+        if len(parts) == 1:
+            return org_name, teamspace_name, parts[0], "latest"
+        if len(parts) == 2:
+            return org_name, teamspace_name, parts[0], parts[1]
+        # The rest of the validation for name and version happens in the backend
+        raise ValueError(
+            "Model version is expected to be in the format `entity/modelname:version` separated by a"
+            f" single colon, but got: {name}"
+        )
+    except ValueError as err:
+        raise ValueError(
+            f"Model name must be in the format 'organization/teamspace/model' but you provided '{name}'."
+        ) from err
+
+
+def download_model(
+    name: str,
+    download_dir: Union[Path, str] = ".",
+    progress_bar: bool = True,
+) -> List[str]:
+    """Download a Model.
+
+    Args:
+        name: The name of the Model you want to download.
+        This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+        download_dir: The directory where the Model should be downloaded.
+        progress_bar: Whether to show a progress bar when downloading.
+    """
+    teamspace_owner_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
+
+    download_dir = Path(download_dir)
+
+    api = TeamspaceApi()
+
+    try:
+        return api.download_model_files(
+            name=model_name,
+            version=version,
+            download_dir=download_dir,
+            teamspace_name=teamspace_name,
+            teamspace_owner_name=teamspace_owner_name,
+            progress_bar=progress_bar,
+        )
+    except ApiException as e:
+        # if we get an error, check if the teamspace actually exists (and print the list)
+        # TODO: ideally this would match a specific error about teamspace not being found
+        _ = _get_teamspace(name=teamspace_name, organization=teamspace_owner_name)
+        raise e
+
+
+def upload_model(
+    name: str,
+    path: Union[Path, str] = ".",
+    cloud_account: Optional[str] = None,
+    progress_bar: bool = True,
+) -> UploadedModelInfo:
+    """Upload a Model.
+
+    Args:
+        name: The name of the Model you want to upload.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+        path: The path to the file or directory you want to upload. Defaults to the current directory.
+        cloud_account: The name of the cloud account to store the Model in.
+            If not provided, the default cloud account for the Teamspace will be used.
+        progress_bar: Whether to show a progress bar for the upload.
+    """
+    org_name, teamspace_name, model_name, _ = _parse_model_name_and_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    return teamspace.upload_model(
+        path=path,
+        name=model_name,
+        cloud_account=cloud_account,
+        progress_bar=progress_bar,
+    )

lightning_sdk/organization.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from lightning_sdk.api import OrgApi
 from lightning_sdk.owner import Owner
-from lightning_sdk.utils import _resolve_org_name
+from lightning_sdk.utils.resolve import _resolve_org_name
 
 
 class Organization(Owner):

lightning_sdk/plugin.py

@@ -1,14 +1,23 @@
 import datetime
 import logging
+import os
+import warnings
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generator, Optional, Protocol, Union, runtime_checkable
 
+from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
 from lightning_sdk.studio import Studio
-from lightning_sdk.utils import _resolve_deprecated_cloud_compute, _setup_logger
+from lightning_sdk.utils.resolve import (
+    _LIGHTNING_SERVICE_EXECUTION_ID_KEY,
+    _resolve_deprecated_cloud_compute,
+    _setup_logger,
+)
 
 if TYPE_CHECKING:
     from lightning_sdk.lightning_cloud.openapi import Externalv1LightningappInstance
+    from lightning_sdk.mmt import MMT
 
 _logger = _setup_logger(__name__)
 
@@ -80,10 +89,20 @@ class _Plugin(ABC):
 
 
 class Plugin(_Plugin):
-    """Plugin class to handle arbitrary plugins on the Studio."""
+    """Plugin class to handle arbitrary plugins on the Studio.
+
+    While plugins can be installed and uninstalled via the SDK, most of them are designed to be run from the Studio's UI
+    only. Calling run just won't do anything in most cases.
+
+    """
 
     def run(self) -> str:
         """Executes the command of the plugin on the Studio."""
+        warnings.warn(
+            "While plugins can be installed and uninstalled via the SDK, most of them are designed to "
+            "be run from the Studio's UI only. "
+            "Calling run just won't do anything in most cases."
+        )
         if self._has_been_executed:
             logging.info("This plugin has already been executed and can be run only once per Studio.")
             return None
@@ -108,25 +127,32 @@ class JobsPlugin(_Plugin):
         name: Optional[str] = None,
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
-    ) -> "Externalv1LightningappInstance":
-        """Launches an asynchronous job."""
-        if name is None:
+        interruptible: bool = False,
+    ) -> Job:
+        """Launches an asynchronous job.
+
+        Args:
+            command: The command to be executed.
+            name: The name of the job.
+            machine: The machine to run the job on.
+            interruptible: Whether to run the job on an interruptible machine.
+                These are cheaper but can be preempted at any time.
+        """
+        if not name:
             name = _run_name("job")
 
         machine = _resolve_deprecated_cloud_compute(machine, cloud_compute)
 
-        resp = self._studio._studio_api.create_job(
-            entrypoint=command,
+        return Job.run(
             name=name,
             machine=machine,
-            studio_id=self._studio._studio.id,
-            teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
+            command=command,
+            studio=self._studio,
+            teamspace=self._studio.teamspace,
+            cloud_account=self._studio.cloud_account,
+            interruptible=interruptible,
         )
 
-        _logger.info(_success_message(resp, self))
-        return resp
-
 
 class MultiMachineTrainingPlugin(_Plugin):
     """Plugin handling multi-machine-training jobs."""
@@ -141,28 +167,35 @@ class MultiMachineTrainingPlugin(_Plugin):
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
         num_instances: int = 2,
-        strategy: str = "parallel",
-    ) -> "Externalv1LightningappInstance":
-        """Launches an asynchronous multi-machine-training."""
-        if name is None:
+        interruptible: bool = False,
+    ) -> "MMT":
+        """Launches an asynchronous multi-machine-training.
+
+        Args:
+            command: The command to be executed.
+            name: The name of the job.
+            machine: The machine to run the job on.
+            num_instances: The number of instances to run the job on.
+            interruptible: Whether to run the job on an interruptible machine.
+                These are cheaper but can be preempted at any time.
+        """
+        from lightning_sdk.mmt import MMT
+
+        if not name:
             name = _run_name("dist-run")
 
         machine = _resolve_deprecated_cloud_compute(machine, cloud_compute)
 
-        # TODO: assert num_instances >=2
-        resp = self._studio._studio_api.create_multi_machine_job(
-            entrypoint=command,
-            name=name,
-            num_instances=num_instances,
-            machine=machine,
-            strategy=strategy,
-            studio_id=self._studio._studio.id,
-            teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
-        )
-
-        _logger.info(_success_message(resp, self))
-        return resp
+        with forced_v1(MMT) as v1mmt:
+            return v1mmt.run(
+                name=name,
+                num_machines=num_instances,
+                machine=machine,
+                command=command,
+                studio=self._studio,
+                teamspace=self._studio.teamspace,
+                interruptible=interruptible,
+            )
 
 
 class MultiMachineDataPrepPlugin(_Plugin):
@@ -178,9 +211,19 @@ class MultiMachineDataPrepPlugin(_Plugin):
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
         num_instances: int = 2,
-    ) -> "Externalv1LightningappInstance":
-        """Launches an asynchronous multi-machine-processing-job."""
-        if name is None:
+        interruptible: bool = False,
+    ) -> Job:
+        """Launches an asynchronous multi-machine-data-processing job.
+
+        Args:
+            command: The command to be executed.
+            name: The name of the job.
+            machine: The machine to run the job on.
+            num_instances: The number of instances to run the job on.
+            interruptible: Whether to run the job on an interruptible machine.
+                These are cheaper but can be preempted at any time.
+        """
+        if not name:
             name = _run_name("data-prep")
 
         machine = _resolve_deprecated_cloud_compute(machine, cloud_compute)
@@ -192,11 +235,12 @@ class MultiMachineDataPrepPlugin(_Plugin):
             machine=machine,
             studio_id=self._studio._studio.id,
             teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
+            cloud_account=self._studio.cloud_account,
+            interruptible=interruptible,
         )
 
-        _logger.info(_success_message(resp, self))
-        return resp
+        with forced_v1(Job) as v1_job:
+            return v1_job(resp.name, self._studio.teamspace)
 
 
 class InferenceServerPlugin(_Plugin):
@@ -218,6 +262,7 @@ class InferenceServerPlugin(_Plugin):
         max_batch_size: int = 4,
         timeout_batching: float = 0.3,
         endpoint: str = "/predict",
+        interruptible: bool = False,
     ) -> "Externalv1LightningappInstance":
         """Launches an asynchronous inference server."""
         if name is None:
@@ -238,7 +283,119 @@ class InferenceServerPlugin(_Plugin):
             endpoint=endpoint,
             studio_id=self._studio._studio.id,
             teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
+            cloud_account=self._studio.cloud_account,
+            interruptible=interruptible,
+        )
+
+        _logger.info(_success_message(resp, self))
+        with forced_v1(Job) as v1_job:
+            return v1_job(resp.name, self._studio.teamspace)
+
+
+class SlurmJobsPlugin(_Plugin):
+    """Plugin handling asynchronous SLURM jobs."""
+
+    _plugin_run_name = "slurm"
+    _slug_name = "slurm"
+
+    def run(
+        self,
+        command: str,
+        name: Optional[str] = None,
+        cluster_id: Optional[str] = None,
+        work_dir: str = "/home/lightning_manager",
+        num_gpus: int = 1,
+        sync_env: bool = True,
+        cache_id: Optional[str] = None,
+    ) -> "Externalv1LightningappInstance":
+        """Launches an asynchronous SLURM job.
+
+        Args:
+            command: The command to be passed to the SLURM Job.
+            name: The name of the SLURM Job.
+            cluster_id: The name of the SLURM Cluster to submit the job on.
+                If the cluster_id isn't provided, the oldest running SLURM cluster will be selected.
+            work_dir: The position where the the files will be created on the SLURM cluster.
+            num_gpus: The number of GPUs requested.
+            sync_env: Whether to force an environement sync.
+            cache_id: A string to avoid re-downloading the Studio files to the SLURM cluster.
+                If you update your files and don't change the cache_id, they won't be used.
+
+        """
+        from lightning_sdk.lightning_cloud.openapi import SlurmJobsBody
+
+        if work_dir == "":
+            raise ValueError("The argument `work_dir` needs to be a proper path on the SLURM Cluster.")
+
+        if num_gpus <= 0:
+            raise ValueError("The argument `num_gpus` needs to be strictly positive.")
+
+        if name is None:
+            name = _run_name("slurm")
+
+        client = self._studio._studio_api._client
+
+        clusters = client.cluster_service_list_project_clusters(project_id=self._studio._teamspace.id).clusters
+        slurm_clusters = [cluster for cluster in clusters if cluster.spec.slurm_v1 is not None]
+        running_slurm_clusters = [
+            cluster for cluster in slurm_clusters if cluster.status.phase == "CLUSTER_STATE_RUNNING"
+        ]
+        running_slurm_clusters = sorted(running_slurm_clusters, key=lambda x: x.created_at)
+
+        if not running_slurm_clusters:
+            raise RuntimeError(
+                "You don't have any running SLURM clusters associated to this project. "
+                "Please, check your Teamspace Cloud Account."
+            )
+
+        selected_cluster = None
+
+        if cluster_id:
+            for cluster in running_slurm_clusters:
+                if cluster.cluster_id == cluster_id:
+                    selected_cluster = cluster
+                    break
+
+            if not selected_cluster:
+                raise ValueError(f"The provided cluster {cluster_id} wasn't found.")
+        else:
+            selected_cluster = running_slurm_clusters[0]
+
+        service_id = os.getenv(_LIGHTNING_SERVICE_EXECUTION_ID_KEY)
+
+        # TODO: Move this to the BE
+        envs = [
+            f"LIGHTNING_CLOUD_PROJECT_ID={os.getenv('LIGHTNING_CLOUD_PROJECT_ID')}",
+            f"LIGHTNING_USERNAME={os.getenv('LIGHTNING_USERNAME')}",
+            f"LIGHTNING_USER_ID={os.getenv('LIGHTNING_USER_ID')}",
+            f"LIGHTNING_API_KEY={os.getenv('LIGHTNING_API_KEY')}",
+            f"LIGHTNING_CLOUD_URL={os.getenv('LIGHTNING_CLOUD_URL')}",
+        ]
+
+        if service_id:
+            envs.append(f"{_LIGHTNING_SERVICE_EXECUTION_ID_KEY}={service_id}")
+
+        if "&&" in command:
+            # We are adding the env varaibles to the latest command
+            splits = command.split("&&")
+            splits[-1] = " ".join(envs) + " " + splits[-1]
+            command = " && ".join(splits)
+        else:
+            command = " ".join(envs) + " " + command
+
+        resp = client.slurm_jobs_user_service_create_user_slurm_job(
+            project_id=self._studio._teamspace.id,
+            body=SlurmJobsBody(
+                cloudspace_id=self._studio._studio.id,
+                cluster_id=selected_cluster.id,
+                command=command,
+                name=name,
+                sync_env=sync_env,
+                work_dir=work_dir,
+                service_id=service_id,
+                num_gpus=num_gpus,
+                cache_id=cache_id,
+            ),
         )
 
         _logger.info(_success_message(resp, self))
@@ -257,7 +414,7 @@ class _RunnablePlugin(Protocol):
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
         **kwargs: Any,
-    ) -> "Externalv1LightningappInstance":
+    ) -> Union["Externalv1LightningappInstance", Job]:
         ...
 
 
@@ -266,6 +423,15 @@ def _run_name(plugin_type: str) -> str:
     return f"{plugin_type}-{datetime.datetime.now().strftime('%b-%d-%H_%M')}"
 
 
-def _success_message(resp: "Externalv1LightningappInstance", plugin_instance: _RunnablePlugin) -> str:
+def _success_message(resp: Union["Externalv1LightningappInstance", Job], plugin_instance: _RunnablePlugin) -> str:
     """Compiles the success message for a given runnable plugin."""
     return f"{plugin_instance._plugin_run_name} {resp.name} was successfully launched. View it at https://lightning.ai/{plugin_instance._studio.owner.name}/{plugin_instance._studio.teamspace.name}/studios/{plugin_instance.studio}/app?app_id={plugin_instance._slug_name}&job_name={resp.name}"
+
+
+@contextmanager
+def forced_v1(cls: Any) -> Generator[Any, None, None]:
+    """Forces to use the v1 version of a class when using a class with multiple backends."""
+    orig_val = getattr(cls, "_force_v1", False)
+    cls._force_v1 = True
+    yield cls
+    cls._force_v1 = orig_val

lightning_sdk/services/__init__.py

@@ -0,0 +1,5 @@
+from lightning_sdk.services.file_endpoint import Client
+from lightning_sdk.services.finetune import LLMFinetune
+from lightning_sdk.services.utilities import download_file
+
+__all__ = ["LLMFinetune", "Client", "download_file"]