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

lightning_sdk/studio.py

@@ -1,7 +1,9 @@
 import os
+import warnings
 from typing import TYPE_CHECKING, Any, Mapping, Optional, Tuple, Union
 
 from lightning_sdk.api.studio_api import StudioApi
+from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.machine import Machine
 from lightning_sdk.organization import Organization
@@ -9,7 +11,7 @@ from lightning_sdk.owner import Owner
 from lightning_sdk.status import Status
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils import _resolve_teamspace, _setup_logger
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
 
 if TYPE_CHECKING:
     from lightning_sdk.plugin import Plugin
@@ -28,10 +30,9 @@ class Studio:
         teamspace: the name of the teamspace the studio is contained by
         org: the name of the organization owning the :param`teamspace` in case it is owned by an org
         user: the name of the user owning the :param`teamspace` in case it is owned directly by a user instead of an org
-        cluster: the name of the cluster, the studio should be created on.
+        cloud_account: the name of the cloud account, the studio should be created on.
             Doesn't matter when the studio already exists.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
-
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
         only one of the arguments can be provided.
@@ -47,13 +48,14 @@ class Studio:
         teamspace: Optional[Union[str, Teamspace]] = None,
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         create_ok: bool = True,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> None:
         self._studio_api = StudioApi()
 
         self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
-        self._cluster = cluster
+        self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
         self._setup_done = False
 
         self._plugins = {}
@@ -68,7 +70,9 @@ class Studio:
                 self._studio = self._studio_api.get_studio(name, self._teamspace.id)
             except ValueError as e:
                 if create_ok:
-                    self._studio = self._studio_api.create_studio(name, self._teamspace.id, cluster=self._cluster)
+                    self._studio = self._studio_api.create_studio(
+                        name, self._teamspace.id, cloud_account=self._cloud_account
+                    )
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e
 
@@ -131,14 +135,32 @@ class Studio:
             return None
         return self._studio_api.get_machine(self._studio.id, self._teamspace.id)
 
-    def start(self, machine: Machine = Machine.CPU) -> None:
+    @property
+    def interruptible(self) -> bool:
+        """Returns whether the Studio is running on a interruptible instance."""
+        if self.status != Status.Running:
+            return None
+
+        return self._studio_api.get_interruptible(self._studio.id, self._teamspace.id)
+
+    @property
+    def cluster(self) -> str:
+        """Returns the cluster the Studio is running on."""
+        warnings.warn("Studio.cluster is deprecated. Use Studio.cloud_account instead", DeprecationWarning)
+        return self.cloud_account
+
+    @property
+    def cloud_account(self) -> str:
+        return self._studio.cluster_id
+
+    def start(self, machine: Union[Machine, str] = Machine.CPU, interruptible: bool = False) -> None:
         """Starts a Studio on the specified machine type (default: CPU-4)."""
         status = self.status
         if status == Status.Running:
-            curr_machine = self.machine
-            if curr_machine != machine:
+            curr_machine = _machine_to_compute_name(self.machine) if self.machine is not None else None
+            if curr_machine != _machine_to_compute_name(machine):
                 raise RuntimeError(
-                    f"Requested to start studio on {machine}, but studio is already running on {curr_machine}."
+                    f"Requested to start studio on {machine}, but studio is already running on {self.machine}."
                     " Consider switching instead!"
                 )
             _logger.info(f"Studio {self.name} is already running")
@@ -146,7 +168,7 @@ class Studio:
 
         if status != Status.Stopped:
             raise RuntimeError(f"Cannot start a studio that is not stopped. Studio {self.name} is {status}.")
-        self._studio_api.start_studio(self._studio.id, self._teamspace.id, machine)
+        self._studio_api.start_studio(self._studio.id, self._teamspace.id, machine, interruptible=interruptible)
 
         self._setup()
 
@@ -166,11 +188,12 @@ class Studio:
         kwargs = self._studio_api.duplicate_studio(self._studio.id, self._teamspace.id, self._teamspace.id)
         return Studio(**kwargs)
 
-    def switch_machine(self, machine: Machine) -> None:
+    def switch_machine(self, machine: Union[Machine, str], interruptible: bool = False) -> None:
         """Switches machine to the provided machine type/.
 
         Args:
             machine: the new machine type to switch to
+            interruptible: determines whether to switch to an interruptible instance
 
         Note:
             this call is blocking until the new machine is provisioned
@@ -181,7 +204,9 @@ class Studio:
             raise RuntimeError(
                 f"Cannot switch machine on a studio that is not running. Studio {self.name} is {status}."
             )
-        self._studio_api.switch_studio_machine(self._studio.id, self._teamspace.id, machine)
+        self._studio_api.switch_studio_machine(
+            self._studio.id, self._teamspace.id, machine, interruptible=interruptible
+        )
 
     def run_with_exit_code(self, *commands: str) -> Tuple[str, int]:
         """Runs given commands on the Studio while returning output and exit code.
@@ -224,9 +249,9 @@ class Studio:
         self._studio_api.upload_file(
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
             file_path=file_path,
-            remote_path=remote_path,
+            remote_path=os.path.normpath(remote_path),
             progress_bar=progress_bar,
         )
 
@@ -240,7 +265,7 @@ class Studio:
             target_path=file_path,
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
         )
 
     def download_folder(self, remote_path: str, target_path: Optional[str] = None) -> None:
@@ -253,8 +278,53 @@ class Studio:
             target_path=target_path,
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
+        )
+
+    @property
+    def auto_sleep(self) -> bool:
+        """Returns if a Studio has auto-sleep enabled."""
+        return not self._studio.code_config.disable_auto_shutdown
+
+    @auto_sleep.setter
+    def auto_sleep(self, value: bool) -> None:
+        if not value and self.machine == Machine.CPU:
+            warnings.warn("Disabling auto-sleep will convert the Studio from free to paid!")
+        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, enabled=value, studio=self._studio)
+        self._update_studio_reference()
+
+    @property
+    def auto_sleep_time(self) -> int:
+        """Returns the time in seconds a Studio has to be idle for auto-sleep to kick in (if enabled)."""
+        return self._studio.code_config.idle_shutdown_seconds
+
+    @auto_sleep_time.setter
+    def auto_sleep_time(self, value: int) -> None:
+        warnings.warn("Setting auto-sleep time will convert the Studio from free to paid!")
+        self._studio_api.update_autoshutdown(
+            self._studio.id, self._teamspace.id, idle_shutdown_seconds=value, studio=self._studio
         )
+        self._update_studio_reference()
+
+    @property
+    def auto_shutdown(self) -> bool:
+        warnings.warn("auto_shutdown is deprecated. Use auto_sleep instead", DeprecationWarning)
+        return self.auto_sleep
+
+    @auto_shutdown.setter
+    def auto_shutdown(self, value: bool) -> None:
+        warnings.warn("auto_shutdown is deprecated. Use auto_sleep instead", DeprecationWarning)
+        self.auto_sleep = value
+
+    @property
+    def auto_shutdown_time(self) -> int:
+        warnings.warn("auto_shutdown_time is deprecated. Use auto_sleep_time instead", DeprecationWarning)
+        return self.auto_sleep_time
+
+    @auto_shutdown_time.setter
+    def auto_shutdown_time(self, value: int) -> None:
+        warnings.warn("auto_shutdown_time is deprecated. Use auto_sleep_time instead", DeprecationWarning)
+        self.auto_sleep_time = value
 
     @property
     def available_plugins(self) -> Mapping[str, str]:
@@ -342,6 +412,9 @@ class Studio:
         """Returns reader friendly representation."""
         return repr(self)
 
+    def _update_studio_reference(self) -> None:
+        self._studio = self._studio_api.get_studio_by_id(studio_id=self._studio.id, teamspace_id=self._teamspace.id)
+
 
 def _internal_status_to_external_status(internal_status: str) -> Status:
     """Converts internal status strings from HTTP requests to external enums."""

lightning_sdk/teamspace.py

@@ -1,12 +1,24 @@
-from typing import TYPE_CHECKING, List, Optional, Union
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
+from lightning_sdk.agents import Agent
 from lightning_sdk.api import TeamspaceApi
+from lightning_sdk.models import UploadedModelInfo
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
 from lightning_sdk.user import User
-from lightning_sdk.utils import _get_organizations_for_authed_user, _resolve_org, _resolve_teamspace_name, _resolve_user
+from lightning_sdk.utils.resolve import (
+    _get_organizations_for_authed_user,
+    _parse_model_and_version,
+    _resolve_org,
+    _resolve_teamspace_name,
+    _resolve_user,
+)
 
 if TYPE_CHECKING:
+    from lightning_sdk.job import Job
+    from lightning_sdk.mmt import MMT
     from lightning_sdk.studio import Studio
 
 
@@ -88,19 +100,81 @@ class Teamspace:
         from lightning_sdk.studio import Studio
 
         studios = []
-        clusters = self._teamspace_api.list_clusters(teamspace_id=self.id)
+        clusters = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
         for cl in clusters:
-            _studios = self._teamspace_api.list_studios(teamspace_id=self.id, cluster_id=cl.cluster_id)
+            _studios = self._teamspace_api.list_studios(teamspace_id=self.id, cloud_account=cl.cluster_id)
             for s in _studios:
                 studios.append(Studio(name=s.name, teamspace=self, cluster=cl.cluster_name, create_ok=False))
 
         return studios
 
+    @property
+    def default_cloud_account(self) -> str:
+        return self._teamspace.project_settings.preferred_cluster
+
+    @property
+    def cloud_accounts(self) -> List[str]:
+        """All cloud accounts associated with that teamspace."""
+        clusters = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
+        return [cl.cluster_name for cl in clusters]
+
     @property
     def clusters(self) -> List[str]:
         """All clusters associated with that teamspace."""
-        clusters = self._teamspace_api.list_clusters(teamspace_id=self.id)
-        return [cl.cluster_name for cl in clusters]
+        warnings.warn(
+            "The 'clusters' attribute is deprecated and will be removed in the future. "
+            "Please use the 'cloud_accounts' attribute instead.",
+            DeprecationWarning,
+        )
+        return self.cloud_accounts
+
+    @property
+    def jobs(self) -> Tuple["Job", ...]:
+        from lightning_sdk.job import Job
+        from lightning_sdk.plugin import forced_v1
+
+        jobsv1, jobsv2 = self._teamspace_api.list_jobs(teamspace_id=self.id)
+
+        jobs = []
+
+        for j1 in jobsv1:
+            with forced_v1(Job):
+                # _fetch_job = False to prevent refetching on init since we already got it
+                job = Job(name=j1.name, teamspace=self, _fetch_job=False)
+            job._internal_job._job = j1
+            jobs.append(job)
+
+        for j2 in jobsv2:
+            # _fetch_job = False to prevent refetching on init since we already got it
+            job = Job(name=j2.name, teamspace=self, _fetch_job=False)
+            job._internal_job._job = j2
+            jobs.append(job)
+
+        return tuple(jobs)
+
+    @property
+    def multi_machine_jobs(self) -> Tuple["MMT", ...]:
+        from lightning_sdk.mmt import MMT
+        from lightning_sdk.plugin import forced_v1
+
+        mmtsv1, mmtsv2 = self._teamspace_api.list_mmts(teamspace_id=self.id)
+
+        mmts = []
+
+        for m1 in mmtsv1:
+            with forced_v1(MMT):
+                # _fetch_job = False to prevent refetching on init since we already got it
+                mmt = MMT(name=m1.name, teamspace=self, _fetch_job=False)
+            mmt._internal_mmt._job = m1
+            mmts.append(mmt)
+
+        for m2 in mmtsv2:
+            # _fetch_job = False to prevent refetching on init since we already got it
+            mmt = MMT(name=m2.name, teamspace=self, _fetch_job=False)
+            mmt._internal_mmt._job = m2
+            mmts.append(mmt)
+
+        return tuple(mmts)
 
     def __eq__(self, other: "Teamspace") -> bool:
         """Checks whether the provided other object is equal to this one."""
@@ -116,6 +190,161 @@ class Teamspace:
         """Returns reader friendly representation."""
         return repr(self)
 
+    def create_agent(
+        self,
+        name: str,
+        api_key: str,
+        base_url: str,
+        model: str,
+        org_id: Optional[str] = "",
+        prompt_template: Optional[str] = "",
+        description: Optional[str] = "",
+        prompt_suggestions: Optional[List[str]] = None,
+        file_uploads_enabled: Optional[bool] = None,
+    ) -> "Agent":
+        agent = self._teamspace_api.create_agent(
+            teamspace_id=self.id,
+            name=name,
+            api_key=api_key,
+            base_url=base_url,
+            model=model,
+            org_id=org_id,
+            prompt_template=prompt_template,
+            description=description,
+            prompt_suggestions=prompt_suggestions,
+            file_uploads_enabled=file_uploads_enabled,
+        )
+        return Agent(agent.id)
+
+    def upload_model(
+        self,
+        path: Union[str, Path],
+        name: str,
+        cloud_account: Optional[str] = None,
+        progress_bar: bool = True,
+    ) -> UploadedModelInfo:
+        """Upload a local checkpoint file to the model store.
+
+        Args:
+            path: Path to the model file or folder to upload.
+            name: Name tag of the model to upload.
+            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.
+        """
+        if not path:
+            raise ValueError("No path provided to upload")
+        if not name:
+            raise ValueError("No name provided for the model")
+        path = Path(path).resolve()
+        if not path.exists():
+            raise FileNotFoundError(str(path))
+
+        cloud_account = (
+            self._teamspace_api._determine_cloud_account(self.id) if cloud_account is None else cloud_account
+        )
+        filepaths = [path] if path.is_file() else [p for p in path.rglob("*") if p.is_file()]
+
+        if not filepaths:
+            raise FileNotFoundError(
+                "The path to upload doesn't contain any files. Make sure it points to a file or"
+                f" non-empty folder: {path}"
+            )
+
+        root_path = path
+        if len(filepaths) == 1:
+            root_path = path.parent
+
+        filenames = ",".join(str(f.relative_to(root_path)) for f in filepaths)
+
+        model = self._teamspace_api.create_model(
+            name=name,
+            metadata={"filenames": filenames},
+            private=True,
+            teamspace_id=self.id,
+            cloud_account=cloud_account,
+        )
+        self._teamspace_api.upload_model_files(
+            model_id=model.model_id,
+            version=model.version,
+            root_path=root_path,
+            filepaths=filepaths,
+            cloud_account=cloud_account,
+            teamspace_id=self.id,
+            progress_bar=progress_bar,
+        )
+        self._teamspace_api.complete_model_upload(
+            model_id=model.model_id,
+            version=model.version,
+            teamspace_id=self.id,
+        )
+        return UploadedModelInfo(
+            name=name,
+            version=model.version,
+            teamspace=self.name,
+            cloud_account=cloud_account,
+        )
+
+    def download_model(
+        self,
+        name: str,
+        download_dir: Optional[str] = None,
+        progress_bar: bool = True,
+    ) -> str:
+        """Download a checkpoint from the model store.
+
+        Args:
+            name: Name tag of the model to download. Can optionally also contain a version tag separated by a colon,
+                 e.g. 'modelname:v1'.
+            download_dir: A path to directory where the model should be downloaded. Defaults
+                to the current working directory.
+            progress_bar: Whether to show a progress bar for the download.
+
+        Returns:
+            The absolute path to the downloaded model file or folder.
+
+        """
+        if not name:
+            raise ValueError("No name provided for the model")
+        if download_dir is None:
+            download_dir = Path.cwd()
+        download_dir = Path(download_dir)
+
+        name, version = _parse_model_and_version(name)
+        model_version = self._teamspace_api.get_model_version(name=name, version=version, teamspace_id=self.id)
+        if not model_version.upload_complete:
+            raise RuntimeError(
+                f"Model {name}:{version} is not fully uploaded yet. Please wait until the upload is complete."
+            )
+        downloaded_files = self._teamspace_api.download_model_files(
+            name=name,
+            version=version,
+            download_dir=download_dir,
+            teamspace_name=self.name,
+            teamspace_owner_name=self.owner.name,
+            progress_bar=progress_bar,
+        )
+
+        if not downloaded_files:
+            raise RuntimeError("No files were downloaded. This shouldn't happen, please report a bug.")
+
+        if len(downloaded_files) == 1:
+            downloaded_file = Path(downloaded_files[0])
+            downloaded_path = download_dir / downloaded_file.parts[0]
+            return str(downloaded_path.resolve())
+        return str(Path(download_dir).resolve())
+
+    def delete_model(self, name: str) -> None:
+        """Delete a model from the model store.
+
+        Args:
+            name: Name tag of the model to delete. Can optionally also contain a version tag separated by a colon,
+                 e.g. 'entity/modelname:v1'.
+
+        """
+        name, version = _parse_model_and_version(name)
+        self._teamspace_api.delete_model(name=name, version=version, teamspace_id=self.id)
+
 
 def _resolve_valueerror_message(error: ValueError, owner: Owner, teamspace_name: str) -> ValueError:
     """Resolves the ValueError Message and replaces it with a nicer message."""
@@ -127,18 +356,17 @@ def _resolve_valueerror_message(error: ValueError, owner: Owner, teamspace_name:
             organizations = _get_organizations_for_authed_user()
             message = (
                 f"Teamspace {entire_ts_name} does not exist. "
-                "Is it maybe an organizational Teamspace? You are a member of the following organizations: "
-                f"{[o.name for o in organizations]}. Maybe specify one of these instead "
-                "of your user if the Teamspace belongs to the organization."
+                f"Is {teamspace_name} an organizational Teamspace? You are a member of the following organizations: "
+                f"{[o.name for o in organizations]}. Try specifying the `org` parameter instead "
+                "of `user` if the Teamspace belongs to the organization."
             )
         else:
             # organization teamspace owner
             user = User()
             message = (
                 f"Teamspace {entire_ts_name} does not exist. "
-                f"Is it maybe a user Teamspace. You specified org={owner.name}, "
-                "but maybe the Teamspace is part of your user? "
-                f"Consider specifying user={user.name} instead of your org."
+                f"Is {teamspace_name} a user Teamspace? "
+                f"Consider specifying user={user.name} instead of org={owner.name}."
             )
 
     return ValueError(message, *error.args[1:])

lightning_sdk/user.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from lightning_sdk.api import UserApi
 from lightning_sdk.owner import Owner
-from lightning_sdk.utils import _resolve_user_name
+from lightning_sdk.utils.resolve import _resolve_user_name
 
 
 class User(Owner):

lightning_sdk/utils/dynamic.py

@@ -0,0 +1,61 @@
+from abc import ABCMeta
+from typing import Any, Callable, Tuple, Type
+
+
+class ConditionBaseMeta(ABCMeta):
+    """Metaclass that allows for conditional inheritance.
+
+    This metaclass is used to conditionally inherit from two base classes
+    based on the result of a given condition function.
+
+    Example usage:
+    ```
+    def my_condition():
+        return some_condition  # Define your condition here
+
+    class A:
+        pass
+
+    class B:
+        pass
+
+    class MyNewClass(metaclass=ConditionBaseMeta, condition_func=my_condition, base_true=A, base_false=B):
+        pass
+    ```
+
+    Args:
+        cls: The metaclass itself.
+        name: The name of the class being created.
+        bases: The base classes of the class being created.
+        attrs: The attributes of the class being created.
+        condition_func: The function that determines which base class to inherit from.
+        base_true: The base class to inherit from if the condition function returns True.
+        base_false: The base class to inherit from if the condition function returns False.
+
+    Returns:
+        type: The new class with the appropriate base class.
+    """
+
+    def __new__(
+        cls,
+        name: str,
+        bases: Tuple[type, ...],
+        attrs: dict,
+        condition_func: Callable[[], bool],
+        base_true: Type,
+        base_false: Type,
+    ) -> Type:
+        # Helper function to determine the base class
+        def get_base() -> type:
+            return base_true if condition_func() else base_false
+
+        base = get_base()
+        attrs["_base"] = base
+
+        # Create the final class, inheriting from both the intermediate base and the original bases
+        return super().__new__(cls, name, (base, *bases), attrs)
+
+    def __getattr__(cls, name: str) -> Any:
+        """Get an attribute from the appropriate base class."""
+        # Delegate attribute lookup to the appropriate base class
+        return getattr(cls._base, name)