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

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

@@ -1,4 +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"]
+__all__ = ["LLMFinetune", "Client", "download_file"]