lightning-sdk 0.2.22__py3-none-any.whl → 0.2.24rc0__py3-none-any.whl

lightning_sdk/pipeline/{types.py → steps.py}

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from lightning_sdk.api.deployment_api import (
     AutoScaleConfig,
@@ -20,52 +20,64 @@ 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 Studio
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.teamspace import Teamspace
     from lightning_sdk.user import User
 
 
-from lightning_sdk.pipeline.utils import DEFAULT
-
-
-class Deployment:
+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[List[float]] = None,
-        release_strategy: Optional["ReleaseStrategy"] = 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,
-        env: Union[List[Union["Secret", "Env"]], Dict[str, str], None] = 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,
+        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,
         wait_for: Optional[Union[str, List[str]]] = DEFAULT,
     ) -> None:
         self.name = name
-        self.machine = machine
+        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="CPU" if machine.is_cpu() else "GPU",
+                    name=autoscaling_metric_name,
                     target=80,
                 )
             ],
@@ -74,22 +86,26 @@ class Deployment:
         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 ""
+        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.wait_for = wait_for
 
-    def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+    def to_proto(
+        self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+    ) -> V1PipelineStep:
         _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
         return V1PipelineStep(
             name=self.name,
             type=V1PipelineStepType.DEPLOYMENT,
-            wait_for=to_wait_for(self.wait_for),
+            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),
@@ -106,18 +122,20 @@ class Deployment:
                     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,
                 ),
                 strategy=to_strategy(self.release_strategy),
             ),
         )
 
 
-class Job:
+class JobStep:
     # Note: This class is only temporary while pipeline is wip
 
     def __init__(
         self,
-        machine: Union["Machine", str],
+        machine: Optional[Union["Machine", str]] = None,
         name: Optional[str] = None,
         command: Optional[str] = None,
         studio: Union["Studio", str, None] = None,
@@ -135,14 +153,21 @@ class Job:
         wait_for: Union[str, List[str], None] = DEFAULT,
     ) -> None:
         self.name = name
-        self.machine = machine
+        self.machine = machine or Machine.CPU
         self.command = command
-        self.studio = studio
+        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
+        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
@@ -151,7 +176,9 @@ class Job:
         self.path_mappings = path_mappings
         self.wait_for = wait_for
 
-    def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+    def to_proto(
+        self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+    ) -> V1PipelineStep:
         studio = _get_studio(self.studio)
         if isinstance(studio, Studio):
             if self.cloud_account is None:
@@ -181,12 +208,12 @@ class Job:
         return V1PipelineStep(
             name=self.name,
             type=V1PipelineStepType.JOB,
-            wait_for=to_wait_for(self.wait_for),
+            wait_for=_to_wait_for(self.wait_for),
             job=body,
         )
 
 
-class MMT:
+class MMTStep:
     # Note: This class is only temporary while pipeline is wip
 
     def __init__(
@@ -209,16 +236,22 @@ class MMT:
         path_mappings: Optional[Dict[str, str]] = None,
         wait_for: Optional[Union[str, List[str]]] = DEFAULT,
     ) -> None:
-        self.machine = machine
+        self.machine = machine or Machine.CPU
         self.num_machines = num_machines
         self.name = name
         self.command = command
-        self.studio = studio
+        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
+        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
@@ -227,7 +260,9 @@ class MMT:
         self.path_mappings = path_mappings
         self.wait_for = wait_for
 
-    def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+    def to_proto(
+        self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
+    ) -> V1PipelineStep:
         studio = _get_studio(self.studio)
         if isinstance(studio, Studio):
             if self.cloud_account is None:
@@ -258,37 +293,23 @@ class MMT:
         return V1PipelineStep(
             name=self.name,
             type=V1PipelineStepType.MMT,
-            wait_for=to_wait_for(self.wait_for),
+            wait_for=_to_wait_for(self.wait_for),
             mmt=body,
         )
 
 
-def to_wait_for(wait_for: Optional[Union[str, List[str]]]) -> Optional[List[str]]:
-    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 _validate_cloud_account(pipeline_cloud_account: str, step_cloud_account: str, shared_filesystem: bool) -> None:
-    if not shared_filesystem:
-        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 wait_for to be on the same cluster."
-            f" Found {pipeline_cloud_account} and {step_cloud_account}"
-        )
-
+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 _get_studio(studio: Union["Studio", str, None]) -> Union[Studio, None]:
-    if studio is None:
-        return None
+    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
 
-    if isinstance(studio, Studio):
-        return studio
 
-    return Studio(studio)
+__all__ = ["JobStep", "MMTStep", "DeploymentStep", "DeploymentReleaseStep"]

lightning_sdk/pipeline/utils.py

@@ -1,6 +1,12 @@
-from typing import List
+from typing import Any, List, Literal, Optional, Union
 
-from lightning_sdk.lightning_cloud.openapi.models import V1PipelineStep
+from lightning_sdk.lightning_cloud.openapi.models import (
+    V1JobSpec,
+    V1PipelineStep,
+    V1PipelineStepType,
+    V1SharedFilesystem,
+)
+from lightning_sdk.studio import Studio
 
 DEFAULT = "DEFAULT"
 
@@ -20,7 +26,7 @@ def prepare_steps(steps: List["V1PipelineStep"]) -> List["V1PipelineStep"]:
         else:
             raise ValueError(f"A step with the name {current_step.name} already exists.")
 
-    if steps[0].wait_for != DEFAULT:
+    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 = []
@@ -52,3 +58,59 @@ def prepare_steps(steps: List["V1PipelineStep"]) -> List["V1PipelineStep"]:
                     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/sandbox.py

@@ -0,0 +1,157 @@
+import ast
+import types
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional, Type, Union
+
+from lightning_sdk.machine import Machine
+from lightning_sdk.organization import Organization
+from lightning_sdk.status import Status
+from lightning_sdk.studio import Studio
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+
+
+@dataclass
+class Output:
+    text: str
+    exit_code: int
+
+
+class _Sandbox:
+    """Sandbox runs AI generated code safely and discards the machine after use.
+
+    Users can run any arbitrary code in a sandbox with sudo permissions.
+
+    Args:
+        name: The name of the sandbox.
+        machine: The machine to use for the sandbox.
+        interruptible: Whether the sandbox is interruptible.
+        teamspace: The teamspace to use for the sandbox.
+        org: The organization to use for the sandbox.
+        user: The user to use for the sandbox.
+        cloud_account: The cloud account to use for the sandbox.
+        disable_secrets: If true, user secrets such as LIGHTNING_API_KEY are not stored in the sandbox.
+
+    Example:
+        with Sandbox() as sandbox:
+            output = sandbox.run("python --version")
+            print(output.text)
+    """
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        machine: Optional[str] = None,
+        interruptible: Optional[bool] = None,
+        teamspace: Optional[Union[str, Teamspace]] = None,
+        org: Optional[Union[str, Organization]] = None,
+        user: Optional[Union[str, User]] = None,
+        cloud_account: Optional[str] = None,
+        disable_secrets: bool = True,
+    ) -> None:
+        if name is None:
+            timestr = datetime.now().strftime("%b-%d-%H_%M")
+            name = f"sandbox-{timestr}"
+
+        self._machine = machine or Machine.CPU
+        self._interruptible = interruptible
+        self._studio = Studio(
+            name=name,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            cloud_account=cloud_account,
+            disable_secrets=disable_secrets,
+        )
+
+    @property
+    def is_running(self) -> bool:
+        return self._studio.status == Status.Running
+
+    def start(self) -> None:
+        """Starts the sandbox if it is not already running."""
+        if self._studio.status == Status.Running:
+            raise RuntimeError(
+                "Cannot start sandbox: it is already running.\n\n"
+                "To ensure proper lifecycle management, either:\n"
+                "  • Avoid calling `start()` multiple times manually, or\n"
+                "  • Use the sandbox as a context manager:\n"
+                "      with Sandbox() as sandbox:\n"
+                "          # your code here\n"
+            )
+
+        if self._studio.status == Status.Pending:
+            raise RuntimeError("Cannot start sandbox: it is already starting. Wait for it to finish starting.")
+        self._studio.start(machine=self._machine, interruptible=self._interruptible)
+
+    def delete(self) -> None:
+        """Deletes the sandbox if it is not already deleted."""
+        if self._studio.status == Status.NotCreated:
+            raise RuntimeError("Cannot delete sandbox: it is not created.")
+        self._studio.delete()
+
+    def run(self, command: str) -> Output:
+        """Runs the command and returns the output."""
+        output, exit_code = self._studio.run_with_exit_code(command)
+        if exit_code != 0:
+            raise Exception(f"Command failed with exit code {exit_code}: {output}")
+        return Output(text=output, exit_code=exit_code)
+
+    @staticmethod
+    def _validate_python_code(code: str) -> None:
+        """Validates Python code for syntax errors.
+
+        Args:
+            code: The Python code string to validate
+
+        Raises:
+            SyntaxError: If the code has syntax errors
+            ValueError: If the code is empty or only whitespace
+            IndentationError: If the code has improper indentation
+
+        Note:
+            This validation only catches syntax-level errors. Runtime errors like
+            NameError or TypeError can only be caught when the code is actually executed.
+        """
+        if not code or not code.strip():
+            raise ValueError("Code cannot be empty or only whitespace")
+
+        try:
+            ast.parse(code)
+        except (SyntaxError, IndentationError) as e:
+            error_type = type(e).__name__
+            raise type(e)(f"Invalid Python {error_type.lower()}: {e}") from e
+
+    def run_python_code(self, code: str) -> Output:
+        """Runs the python code and returns the output.
+
+        Args:
+            code: The Python code string to execute
+
+        Returns:
+            Output: The result of executing the code
+
+        Raises:
+            SyntaxError: If the code has syntax errors
+            ValueError: If the code is empty or only whitespace
+        """
+        # Validate the code before running
+        self._validate_python_code(code)
+
+        command = f"python - <<EOF\n{code}\nEOF"
+        return self.run(command)
+
+    def __enter__(self) -> "_Sandbox":
+        """Starts the sandbox if it is not running and returns the sandbox."""
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[types.TracebackType] = None,
+    ) -> None:
+        """Deletes the sandbox after use."""
+        self._studio.delete()

lightning_sdk/services/license.py

@@ -39,8 +39,8 @@ MESSAGE_GUIDE_SIGN_LICENSE = (
     "│ Details: license key ({license_strats}...{license_ends}) for package {package_name:<56} │\n"
     "│ Please make sure you have signed the license agreement and set the license key.                           │\n"
     "│                                                                                                           │\n"
-    f"│ Sign the license agreement here (if you dont have Lightning account, you will asked to create one):       │\n"
-    f"│   {generate_url_user_settings():<102}\n"
+    "│ Sign the license agreement here (if you dont have Lightning account, you will asked to create one):       │\n"
+    "│   {link}\n"
     "│                                                                                                           │\n"
     "│ Once you have the license key, you may need to reinstall this package to activate it. Use the commands:   │\n"
     "│                                                                                                           │\n"
@@ -57,6 +57,7 @@ def generate_message_guide_sign_license(package_name: str, reason: str, license_
     if not license_key:
         license_key = "." * 64
     return MESSAGE_GUIDE_SIGN_LICENSE.format(
+        link=generate_url_user_settings(name=package_name),
         package_name=package_name,
         reason=reason.ljust(102, " "),
         license_strats=license_key[:5],
@@ -273,12 +274,17 @@ class LightningLicense:
     @property
     def license_key(self) -> Optional[str]:
         """Get the license key."""
+        name = self.product_name.replace("-", "_")
+        if not self._license_key:
+            # If the license key is not set, first try to find it env variables
+            self._license_key = os.environ.get(f"LIGHTNING_{name.upper()}_LICENSE_KEY", None)
+        if not self._license_key:
+            # If the license key is not set, second try to find it in the package root
+            self._license_key = self._find_package_license_key(name)
+        # If not found, try to find it in the user home
         if not self._license_key:
-            # If the license key is not set, fist try to find it in the package root
-            self._license_key = self._find_package_license_key(self.product_name.replace("-", "_"))
             # If not found, try to find it in the user home
-            if not self._license_key:
-                self._license_key = self._find_user_license_key(self.product_name)
+            self._license_key = self._find_user_license_key(self.product_name)
         return self._license_key
 
     @property

lightning_sdk/studio.py

@@ -76,6 +76,7 @@ class Studio:
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
         provider: Optional[str] = None,
         source: Optional[V1CloudSpaceSourceType] = None,
+        disable_secrets: bool = False,
     ) -> None:
         self._studio_api = StudioApi()
         self._cluster_api = ClusterApi()
@@ -83,9 +84,13 @@ class Studio:
         self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
         self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
         self._setup_done = False
+        self._disable_secrets = disable_secrets
 
         self._plugins = {}
 
+        if self._teamspace is None:
+            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
+
         if provider is not None:
             if isinstance(provider, str) and provider in Provider.__members__:
                 provider = Provider(provider)
@@ -107,7 +112,11 @@ class Studio:
             except ValueError as e:
                 if create_ok:
                     self._studio = self._studio_api.create_studio(
-                        name, self._teamspace.id, cloud_account=self._cloud_account, source=source
+                        name,
+                        self._teamspace.id,
+                        cloud_account=self._cloud_account,
+                        source=source,
+                        disable_secrets=self._disable_secrets,
                     )
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e

{lightning_sdk-0.2.22.dist-info → lightning_sdk-0.2.24rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.2.22
+Version: 0.2.24rc0
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License