lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/pipeline/steps.py

@@ -0,0 +1,365 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from lightning_sdk.api.deployment_api import (
+    AutoScaleConfig,
+    AutoScalingMetric,
+    BasicAuth,
+    Env,
+    ExecHealthCheck,
+    HttpHealthCheck,
+    ReleaseStrategy,
+    Secret,
+    TokenAuth,
+    to_autoscaling,
+    to_endpoint,
+    to_spec,
+    to_strategy,
+)
+from lightning_sdk.job.v2 import JobApiV2
+from lightning_sdk.lightning_cloud.openapi.models import (
+    V1CreateDeploymentRequest,
+    V1PipelineStep,
+    V1PipelineStepType,
+    V1SharedFilesystem,
+)
+from lightning_sdk.machine import Machine
+from lightning_sdk.mmt.v2 import MMTApiV2
+from lightning_sdk.pipeline.utils import DEFAULT, _get_studio, _to_wait_for, _validate_cloud_account
+from lightning_sdk.studio import CloudAccountApi, Studio
+
+if TYPE_CHECKING:
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.teamspace import CloudProvider, Teamspace
+    from lightning_sdk.user import User
+
+
+class DeploymentStep:
+    # Note: This class is only temporary while pipeline is wip
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        studio: Optional[Union[str, Studio]] = None,
+        machine: Optional["Machine"] = None,
+        image: Optional[str] = None,
+        autoscale: Optional[AutoScaleConfig] = None,
+        ports: Optional[Union[float, List[float]]] = None,
+        release_strategy: Optional[ReleaseStrategy] = None,
+        entrypoint: Optional[str] = None,
+        command: Optional[str] = None,
+        commands: Optional[List[str]] = None,
+        env: Union[List[Union[Secret, Env]], Dict[str, str], None] = None,
+        spot: Optional[bool] = None,
+        replicas: Optional[int] = None,
+        health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None,
+        auth: Optional[Union[BasicAuth, TokenAuth]] = None,
+        cloud_account: Optional[str] = None,
+        custom_domain: Optional[str] = None,
+        quantity: Optional[int] = None,
+        include_credentials: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
+        wait_for: Optional[Union[str, List[str]]] = DEFAULT,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
+    ) -> None:
+        self.name = name
+        self.studio = _get_studio(studio)
+        if cloud_account and studio and cloud_account != studio.cloud_account != cloud_account:
+            raise ValueError(
+                f"The provided cloud account `{cloud_account}` doesn't match"
+                f" the Studio cloud account {self.studio.cloud_account}"
+            )
+
+        self.machine = machine or Machine.CPU
+        self.image = image
+        autoscaling_metric_name = (
+            ("CPU" if self.machine.is_cpu() else "GPU") if isinstance(self.machine, Machine) else "CPU"
+        )
+        self.autoscale = autoscale or AutoScaleConfig(
+            min_replicas=0,
+            max_replicas=1,
+            target_metrics=[
+                AutoScalingMetric(
+                    name=autoscaling_metric_name,
+                    target=80,
+                )
+            ],
+        )
+        self.ports = ports
+        self.release_strategy = release_strategy
+        self.entrypoint = entrypoint
+        self.command = command
+        self.commands = commands
+        self.env = env
+        self.spot = spot
+        self.replicas = replicas or 1
+        self.health_check = health_check
+        self.auth = auth
+        self.cloud_account = cloud_account or "" if self.studio is None else self.studio.cloud_account
+        self.custom_domain = custom_domain
+        self.quantity = quantity
+        self.include_credentials = include_credentials or True
+        self.max_runtime = max_runtime
+        self.wait_for = wait_for
+        self.cloud_provider = cloud_provider
+
+    def to_proto(
+        self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+    ) -> V1PipelineStep:
+        machine_image_version = None
+
+        studio = _get_studio(self.studio)
+        if isinstance(studio, Studio):
+            machine_image_version = studio._studio.machine_image_version
+
+            if self.cloud_account is None:
+                self.cloud_account = studio.cloud_account
+            elif studio.cloud_account != self.cloud_account:
+                raise ValueError("The provided cloud account doesn't match the studio")
+
+        resolved_cloud_account = CloudAccountApi().resolve_cloud_account(
+            teamspace.id, self.cloud_account, self.cloud_provider, teamspace.default_cloud_account
+        )
+
+        _validate_cloud_account(cloud_account, resolved_cloud_account, shared_filesystem)
+
+        return V1PipelineStep(
+            name=self.name,
+            type=V1PipelineStepType.DEPLOYMENT,
+            wait_for=_to_wait_for(self.wait_for),
+            deployment=V1CreateDeploymentRequest(
+                autoscaling=to_autoscaling(self.autoscale, self.replicas),
+                endpoint=to_endpoint(self.ports, self.auth, self.custom_domain),
+                name=self.name,
+                project_id=teamspace.id,
+                replicas=self.replicas,
+                spec=to_spec(
+                    cloud_account=self.cloud_account or cloud_account,
+                    command=self.command,
+                    entrypoint=self.entrypoint,
+                    env=self.env,
+                    image=self.image,
+                    spot=self.spot,
+                    machine=self.machine,
+                    health_check=self.health_check,
+                    quantity=self.quantity,
+                    cloudspace_id=self.studio._studio.id if self.studio else None,
+                    include_credentials=self.include_credentials,
+                    max_runtime=self.max_runtime,
+                    machine_image_version=machine_image_version,
+                ),
+                strategy=to_strategy(self.release_strategy),
+            ),
+        )
+
+
+class JobStep:
+    # Note: This class is only temporary while pipeline is wip
+
+    def __init__(
+        self,
+        machine: Optional[Union["Machine", str]] = None,
+        name: Optional[str] = None,
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Union[str, None] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
+        wait_for: Union[str, List[str], None] = DEFAULT,
+        reuse_snapshot: bool = True,
+    ) -> None:
+        self.name = name
+        self.machine = machine or Machine.CPU
+        self.command = command
+        self.studio = _get_studio(studio)
+
+        if cloud_account and self.studio and cloud_account != self.studio.cloud_account != cloud_account:
+            raise ValueError(
+                f"The provided cloud account `{cloud_account}` doesn't match"
+                f" the Studio cloud account {self.studio.cloud_account}"
+            )
+
+        self.image = image
+        self.teamspace = teamspace
+        self.org = org
+        self.user = user
+        self.cloud_account = cloud_account or "" if self.studio is None else self.studio.cloud_account
+        self.cloud_provider = cloud_provider
+        self.env = env
+        self.interruptible = interruptible
+        self.image_credentials = image_credentials
+        self.cloud_account_auth = cloud_account_auth
+        self.entrypoint = entrypoint
+        self.path_mappings = path_mappings
+        self.max_runtime = max_runtime
+        self.wait_for = wait_for
+        self.reuse_snapshot = reuse_snapshot
+
+    def to_proto(
+        self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+    ) -> V1PipelineStep:
+        machine_image_version = None
+
+        studio = _get_studio(self.studio)
+        if isinstance(studio, Studio):
+            machine_image_version = studio._studio.machine_image_version
+
+            if self.cloud_account is None:
+                self.cloud_account = studio.cloud_account
+            elif studio.cloud_account != self.cloud_account:
+                raise ValueError("The provided cloud account doesn't match the studio")
+
+        resolved_cloud_account = CloudAccountApi().resolve_cloud_account(
+            teamspace.id, self.cloud_account, self.cloud_provider, teamspace.default_cloud_account
+        )
+
+        _validate_cloud_account(cloud_account, resolved_cloud_account, shared_filesystem)
+
+        body = JobApiV2._create_job_body(
+            name=self.name,
+            command=self.command,
+            cloud_account=resolved_cloud_account or cloud_account,
+            studio_id=studio._studio.id if isinstance(studio, Studio) else None,
+            image=self.image,
+            machine=self.machine,
+            interruptible=self.interruptible,
+            env=self.env,
+            image_credentials=self.image_credentials,
+            cloud_account_auth=self.cloud_account_auth,
+            entrypoint=self.entrypoint,
+            path_mappings=self.path_mappings,
+            artifacts_local=None,
+            artifacts_remote=None,
+            max_runtime=self.max_runtime,
+            machine_image_version=machine_image_version,
+            reuse_snapshot=self.reuse_snapshot,
+        )
+
+        return V1PipelineStep(
+            name=self.name,
+            type=V1PipelineStepType.JOB,
+            wait_for=_to_wait_for(self.wait_for),
+            job=body,
+        )
+
+
+class MMTStep:
+    # Note: This class is only temporary while pipeline is wip
+
+    def __init__(
+        self,
+        name: str,
+        machine: Union["Machine", str],
+        num_machines: Optional[int] = 2,
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
+        wait_for: Optional[Union[str, List[str]]] = DEFAULT,
+        reuse_snapshot: bool = True,
+    ) -> None:
+        self.machine = machine or Machine.CPU
+        self.num_machines = num_machines
+        self.name = name
+        self.command = command
+        self.studio = _get_studio(studio)
+
+        if cloud_account and self.studio and cloud_account != self.studio.cloud_account != cloud_account:
+            raise ValueError(
+                f"The provided cloud account `{cloud_account}` doesn't match"
+                f" the Studio cloud account {self.studio.cloud_account}"
+            )
+        self.image = image
+        self.teamspace = teamspace
+        self.org = org
+        self.user = user
+        self.cloud_account = cloud_account or "" if self.studio is None else self.studio.cloud_account
+        self.env = env
+        self.interruptible = interruptible
+        self.image_credentials = image_credentials
+        self.cloud_account_auth = cloud_account_auth
+        self.entrypoint = entrypoint
+        self.path_mappings = path_mappings
+        self.max_runtime = max_runtime
+        self.wait_for = wait_for
+        self.reuse_snapshot = reuse_snapshot
+
+    def to_proto(
+        self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+    ) -> V1PipelineStep:
+        machine_image_version = None
+
+        studio = _get_studio(self.studio)
+        if isinstance(studio, Studio):
+            machine_image_version = studio._studio.machine_image_version
+
+            if self.cloud_account is None:
+                self.cloud_account = studio.cloud_account
+            elif studio.cloud_account != self.cloud_account:
+                raise ValueError("The provided cloud account doesn't match the studio")
+
+        _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+
+        body = MMTApiV2._create_mmt_body(
+            name=self.name,
+            num_machines=self.num_machines,
+            command=self.command,
+            cloud_account=self.cloud_account or cloud_account,
+            studio_id=studio._studio.id if isinstance(studio, Studio) else None,
+            image=self.image,
+            machine=self.machine,
+            interruptible=self.interruptible,
+            env=self.env,
+            image_credentials=self.image_credentials,
+            cloud_account_auth=self.cloud_account_auth,
+            entrypoint=self.entrypoint,
+            path_mappings=self.path_mappings,
+            artifacts_local=None,  # deprecated in favor of path_mappings
+            artifacts_remote=None,  # deprecated in favor of path_mappings
+            max_runtime=self.max_runtime,
+            machine_image_version=machine_image_version,
+            reuse_snapshot=self.reuse_snapshot,
+        )
+
+        return V1PipelineStep(
+            name=self.name,
+            type=V1PipelineStepType.MMT,
+            wait_for=_to_wait_for(self.wait_for),
+            mmt=body,
+        )
+
+
+class DeploymentReleaseStep(DeploymentStep):
+    def __init__(self, *args: Any, deployment_name: Optional[str] = None, **kwargs: Any) -> None:
+        if not deployment_name:
+            raise ValueError("The deployment name is required")
+        self._deployment_name = deployment_name
+        super().__init__(*args, **kwargs)
+
+    def to_proto(self, *args: Any, **kwargs: Any) -> V1PipelineStep:
+        proto: V1PipelineStep = super().to_proto(*args, **kwargs)
+        proto.deployment.name = self._deployment_name
+        proto.deployment.pipeline_reuse_deployment_between_runs = True
+        return proto
+
+
+__all__ = ["JobStep", "MMTStep", "DeploymentStep", "DeploymentReleaseStep"]

lightning_sdk/pipeline/utils.py

@@ -0,0 +1,116 @@
+from typing import Any, List, Literal, Optional, Union
+
+from lightning_sdk.lightning_cloud.openapi.models import (
+    V1JobSpec,
+    V1PipelineStep,
+    V1PipelineStepType,
+    V1SharedFilesystem,
+)
+from lightning_sdk.studio import Studio
+
+DEFAULT = "DEFAULT"
+
+
+def prepare_steps(steps: List["V1PipelineStep"]) -> List["V1PipelineStep"]:
+    """The prepare_steps function is responsible for creating dependencies between steps.
+
+    The dependencies are based on whether a step wait_for to be executed before another.
+    """
+    name_to_step = {}
+    name_to_idx = {}
+
+    for current_step_idx, current_step in enumerate(steps):
+        if current_step.name not in name_to_step:
+            name_to_step[current_step.name] = current_step
+            name_to_idx[current_step.name] = current_step_idx
+        else:
+            raise ValueError(f"A step with the name {current_step.name} already exists.")
+
+    if steps[0].wait_for not in [None, DEFAULT, []]:
+        raise ValueError("The first step isn't allowed to receive `wait_for=...`.")
+
+    steps[0].wait_for = []
+
+    # This implements a linear dependency between the steps as the default behaviour
+    for current_step_idx, current_step in reversed(list(enumerate(steps))):
+        if current_step_idx == 0:
+            continue
+
+        if current_step.wait_for == DEFAULT:
+            prev_step_idx = current_step_idx - 1
+            wait_for = []
+            while prev_step_idx > -1:
+                prev_step = steps[prev_step_idx]
+                wait_for.insert(0, steps[prev_step_idx].name)
+                if prev_step.wait_for != []:
+                    break
+                prev_step_idx -= 1
+            current_step.wait_for = wait_for
+        else:
+            for name in current_step.wait_for:
+                if current_step.name == name:
+                    raise ValueError(f"You can only reference prior steps. Found {current_step.name}")
+
+                if name not in name_to_step:
+                    raise ValueError(f"The step {current_step_idx} doesn't have a valid wait_for. Found {name}")
+
+                if name_to_idx[name] >= name_to_idx[current_step.name]:
+                    raise ValueError("You can only reference prior steps")
+
+    return steps
+
+
+def _get_studio(studio: Union["Studio", str, None]) -> Union[Studio, None]:
+    if studio is None:
+        return None
+
+    if isinstance(studio, Studio):
+        return studio
+
+    return Studio(studio)
+
+
+def _validate_cloud_account(
+    pipeline_cloud_account: str, step_cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+) -> None:
+    shared_filesystem_enable = (
+        shared_filesystem.enabled if isinstance(shared_filesystem, V1SharedFilesystem) else shared_filesystem
+    )
+    if not shared_filesystem_enable:
+        return
+
+    if pipeline_cloud_account != "" and step_cloud_account != "" and pipeline_cloud_account != step_cloud_account:
+        raise ValueError(
+            "With shared filesystem enabled, all the pipeline steps requires to be on the same cluster."
+            f" Found {pipeline_cloud_account} and {step_cloud_account}"
+        )
+
+
+def _to_wait_for(wait_for: Optional[Union[str, List[str]]]) -> Optional[Union[List[str], Literal["DEFAULT"]]]:
+    if wait_for == DEFAULT:
+        return wait_for
+
+    if wait_for is None:
+        return []
+
+    return wait_for if isinstance(wait_for, list) else [wait_for]
+
+
+def _get_cloud_account(steps: List[V1PipelineStep]) -> Optional[str]:
+    if len(steps) == 0:
+        return None
+
+    cluster_ids: set[str] = set()
+    for step in steps:
+        job_spec = _get_spec(step)
+        cluster_ids.add(job_spec.cluster_id)
+
+    return sorted(cluster_ids)[0]
+
+
+def _get_spec(step: Any) -> V1JobSpec:
+    if step.type == V1PipelineStepType.DEPLOYMENT:
+        return step.deployment.spec
+    if step.type == V1PipelineStepType.MMT:
+        return step.mmt.spec
+    return step.job.spec

lightning_sdk/plugin.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Generator, Optional, Protocol, Union, run
 from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
 from lightning_sdk.studio import Studio
+from lightning_sdk.utils.logging import TrackCallsABCMeta
 from lightning_sdk.utils.resolve import (
     _LIGHTNING_SERVICE_EXECUTION_ID_KEY,
     _resolve_deprecated_cloud_compute,
@@ -22,7 +23,7 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
-class _Plugin(ABC):
+class _Plugin(ABC, metaclass=TrackCallsABCMeta):
     """Abstract Plugin class defining the API.
 
     Args:
@@ -128,6 +129,7 @@ class JobsPlugin(_Plugin):
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
         interruptible: bool = False,
+        reuse_snapshot: bool = True,
     ) -> Job:
         """Launches an asynchronous job.
 
@@ -137,6 +139,8 @@ class JobsPlugin(_Plugin):
             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.
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         if not name:
             name = _run_name("job")
@@ -151,6 +155,7 @@ class JobsPlugin(_Plugin):
             teamspace=self._studio.teamspace,
             cloud_account=self._studio.cloud_account,
             interruptible=interruptible,
+            reuse_snapshot=reuse_snapshot,
         )
 
 
@@ -186,16 +191,15 @@ class MultiMachineTrainingPlugin(_Plugin):
 
         machine = _resolve_deprecated_cloud_compute(machine, cloud_compute)
 
-        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,
-            )
+        return MMT.run(
+            name=name,
+            num_machines=num_instances,
+            machine=machine,
+            command=command,
+            studio=self._studio,
+            teamspace=self._studio.teamspace,
+            interruptible=interruptible,
+        )
 
 
 class MultiMachineDataPrepPlugin(_Plugin):
@@ -304,7 +308,6 @@ class SlurmJobsPlugin(_Plugin):
         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":
@@ -316,7 +319,6 @@ class SlurmJobsPlugin(_Plugin):
             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.
@@ -327,9 +329,6 @@ class SlurmJobsPlugin(_Plugin):
         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")
 
@@ -393,7 +392,6 @@ class SlurmJobsPlugin(_Plugin):
                 sync_env=sync_env,
                 work_dir=work_dir,
                 service_id=service_id,
-                num_gpus=num_gpus,
                 cache_id=cache_id,
             ),
         )
@@ -402,6 +400,25 @@ class SlurmJobsPlugin(_Plugin):
         return resp
 
 
+class CustomPortPlugin(_Plugin):
+    """Plugin handling the port of a given service."""
+
+    _plugin_run_name = "Custom Port"
+    _slug_name = "custom-port"
+
+    def run(self, name: Optional[str] = None, port: int = 8000) -> str:
+        """Starts a new port to the given Studio."""
+        if name is None:
+            name = _run_name("port")
+
+        return self._studio._studio_api.start_new_port(
+            teamspace_id=self._studio._teamspace.id,
+            studio_id=self._studio._studio.id,
+            name=name,
+            port=port,
+        )
+
+
 @runtime_checkable
 class _RunnablePlugin(Protocol):
     _plugin_run_name: str
@@ -432,6 +449,8 @@ def _success_message(resp: Union["Externalv1LightningappInstance", Job], plugin_
 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
+    try:
+        cls._force_v1 = True
+        yield cls
+    finally:
+        cls._force_v1 = orig_val