lightning-sdk 0.1.3__py3-none-any.whl → 0.1.46__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,10 +1,20 @@
+import warnings
+from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional, 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.studio import Studio
@@ -88,19 +98,33 @@ 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
 
     def __eq__(self, other: "Teamspace") -> bool:
         """Checks whether the provided other object is equal to this one."""
@@ -116,6 +140,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 +306,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)

lightning_sdk/utils/enum.py

@@ -0,0 +1,116 @@
+from enum import Enum, EnumMeta, unique
+from types import MappingProxyType
+from typing import Any, Dict, Iterator, Union
+
+
+# We silence annotation errors as the function signatures are exactly the same in python core enum.py
+# modified from https://stackoverflow.com/a/62309159
+class OnAccess(EnumMeta):
+    """Runs a user-specified function whenever member is accessed."""
+
+    def __getattribute__(cls, name: str) -> Any:  # noqa: N805
+        """Called when an attribute is accessed."""
+        obj = super().__getattribute__(name)
+        if isinstance(obj, Enum) and obj._on_access:
+            obj._on_access()
+        return _resolve_deprecated_enum_fields(obj)
+
+    def __getitem__(cls, name: str) -> Any:  # noqa: N805
+        """Called when an item is accessed."""
+        member = super().__getitem__(name)
+        if isinstance(member, Enum) and member._on_access:
+            member._on_access()
+
+        return _resolve_deprecated_enum_fields(member)
+
+    def __call__(
+        cls,  # noqa: N805
+        value,  # noqa: ANN001
+        names=None,  # noqa: ANN001
+        *,
+        module=None,  # noqa: ANN001
+        qualname=None,  # noqa: ANN001
+        type=None,  # noqa: A002, ANN001
+        start=1,  # noqa: ANN001
+    ) -> Any:
+        obj = super().__call__(value, names, module=module, qualname=qualname, type=type, start=start)
+        if isinstance(obj, Enum) and obj._on_access:
+            obj._on_access()
+        return _resolve_deprecated_enum_fields(obj)
+
+    @staticmethod
+    def _filter_func(member: "DeprecationEnum") -> bool:
+        return not member.deprecated
+
+    def __iter__(cls) -> Iterator["DeprecationEnum"]:  # noqa: N805
+        """Returns members in definition order.
+
+        Excludes deprecated members
+
+        """
+        return filter(cls._filter_func, (cls._member_map_[name] for name in cls._member_names_))
+
+    def __len__(cls) -> int:  # noqa: N805
+        """Length of the enum excluding deprecated members."""
+        return len(cls.__members__)
+
+    @property
+    def __members__(cls) -> Dict[str, "DeprecationEnum"]:  # noqa: N805
+        """Returns a mapping of member name->value.
+
+        This mapping lists all enum members, including aliases. Note that this is a read-only view of the internal
+        mapping.
+
+        Excludes deprecated members.
+
+        """
+        return MappingProxyType(
+            {name: member for name, member in cls._member_map_.items() if not cls._filter_func(member)}
+        )
+
+    def __reversed__(cls) -> Iterator["DeprecationEnum"]:  # noqa: N805
+        """Returns members in reverse definition order.
+
+        Excludes deprecated members
+
+        """
+        return filter(cls._filter_func, (cls._member_map_[name] for name in reversed(cls._member_names_)))
+
+
+def _resolve_deprecated_enum_fields(value: Union[Any, "DeprecationEnum"]) -> Union[Any, "DeprecationEnum"]:
+    """Resolves deprecated enum fields to their preferred values."""
+    if getattr(value, "deprecated", False):
+        return value.__class__[value.preferred_way]
+    return value
+
+
+# NOTE: even deprecated enums need to have unique values.
+# otherwise the deprecation warning will not be triggered
+# as duplicated values are handled as aliases towards the original value
+# and the original value is not accessed. However, _resolve_deprecated_enum_fields
+# will handle this case and return the preferred value
+@unique
+class DeprecationEnum(Enum, metaclass=OnAccess):
+    def __new__(cls, value, *args) -> Any:  # noqa: ANN001, ANN002
+        member = object.__new__(cls)
+        member._value_ = value
+        member.deprecated = False
+        member._on_access = None
+
+        if args:
+            member.preferred_way = args[0]
+            member._on_access = member.deprecate
+            member.deprecated = True
+
+        return member
+
+    def deprecate(self) -> None:
+        import warnings
+
+        clsname = self.__class__.__name__
+
+        warnings.warn(
+            f"{clsname}.{self.name} is deprecated. Use {clsname}.{self.preferred_way} instead!",
+            DeprecationWarning,
+            stacklevel=3,
+        )