lightning-sdk 0.1.2__py3-none-any.whl → 0.1.46__py3-none-any.whl

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,
+        )

lightning_sdk/{utils.py → utils/resolve.py}

@@ -2,7 +2,7 @@ import logging
 import os
 import warnings
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Generator, List, Optional, Union
+from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
 
 from lightning_sdk.api import TeamspaceApi, UserApi
 from lightning_sdk.machine import Machine
@@ -13,6 +13,9 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 
+_LIGHTNING_SERVICE_EXECUTION_ID_KEY = "LIGHTNING_SERVICE_EXECUTION_ID"
+
+
 def _setup_logger(name: str, level: int = logging.INFO) -> logging.Logger:
     _logger = logging.getLogger(name)
     _handler = logging.StreamHandler()
@@ -43,6 +46,24 @@ def _resolve_deprecated_cloud_compute(machine: Machine, cloud_compute: Optional[
     return machine
 
 
+def _resolve_deprecated_cluster(cloud_account: Optional[str], cluster: Optional[str]) -> Optional[str]:
+    if cluster is not None:
+        if cloud_account is not None:
+            raise ValueError(
+                "Cannot use both 'cluster' and 'cloud_account' at the same time."
+                "Please don't set the 'cluster' as it will be deprecated!"
+            )
+
+        warnings.warn(
+            "The 'cluster' argument will be deprecated in the future! "
+            "Please consider using the 'cloud_account' argument instead!",
+            DeprecationWarning,
+        )
+        return cluster
+
+    return cloud_account
+
+
 def _resolve_org_name(name: Optional[str]) -> Optional[str]:
     if name is None:
         name = os.environ.get("LIGHTNING_ORG", "") or None
@@ -54,7 +75,6 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 
     if isinstance(org, Organization):
         return org
-
     org = _resolve_org_name(org)
 
     if org is None:
@@ -68,7 +88,6 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 def _resolve_user_name(name: Optional[str]) -> Optional[str]:
     if name is None:
         name = os.environ.get("LIGHTNING_USERNAME", "") or None
-
     return name
 
 
@@ -79,7 +98,6 @@ def _resolve_user(user: Optional[Union[str, "User"]]) -> Optional["User"]:
         return user
 
     user = _resolve_user_name(user)
-
     if user is None:
         return None
 
@@ -131,6 +149,12 @@ def _get_organizations_for_authed_user() -> List["Organization"]:
     return [Organization(_org.name) for _org in _orgs]
 
 
+def _get_teamspace_names_for_authed_user() -> List[str]:
+    """Returns Teamspace's names the current Authed user is a member of."""
+    teamspaces = UserApi()._get_all_teamspace_memberships("")
+    return sorted([ts.name for ts in teamspaces])
+
+
 def _get_authed_user() -> "User":
     from lightning_sdk.user import User
 
@@ -150,3 +174,16 @@ def skip_studio_init() -> Generator[None, None, None]:
     yield
 
     Studio._skip_init = prev_studio_init_state
+
+
+def _parse_model_and_version(name: str) -> Tuple[str, str]:
+    parts = name.split(":")
+    if len(parts) == 1:
+        return parts[0], "latest"
+    if len(parts) == 2:
+        return parts[0], parts[1]
+    # The rest of the validation for name and version happens in the backend
+    raise ValueError(
+        "Model version is expected to be in the format `entity/modelname:version` separated by a"
+        f" single colon, but got: {name}"
+    )

lightning_sdk-0.1.46.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Grid.ai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{lightning_sdk-0.1.2.dist-info → lightning_sdk-0.1.46.dist-info}/METADATA

@@ -1,14 +1,37 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.1.2
+Version: 0.1.46
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
+License: MIT License
+        
+        Copyright (c) 2024 Grid.ai
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
 Keywords: deep learning,machine learning,pytorch,AI
 Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: backoff
-Requires-Dist: click
 Requires-Dist: fastapi
 Requires-Dist: pyjwt
 Requires-Dist: python-multipart
@@ -22,6 +45,9 @@ Requires-Dist: tqdm
 Requires-Dist: fire
 Requires-Dist: simple-term-menu
 Requires-Dist: lightning-utilities
+Provides-Extra: serve
+Requires-Dist: litserve>=0.2.5; extra == "serve"
+Requires-Dist: docker; extra == "serve"
 
 # Lightning SDK
 
@@ -70,8 +96,8 @@ s.install_plugin("jobs")
 s.install_plugin("multi-machine-training")
 
 # run the resulting plugins to start 1 job and 1 multi-machine training
-s.installed_plugins["jobs"].run("python my_dummy_file", name="my_first_job", cloud_compute=Machine.A10G)
-s.installed_plugins["multi-machine-training"].run("python my_dummy_file", name="my_first_mmt", cloud_compute=Machine.T4, num_instances=42)
+s.installed_plugins["jobs"].run("python my_dummy_file", name="my_first_job", machine=Machine.A10G)
+s.installed_plugins["multi-machine-training"].run("python my_dummy_file", name="my_first_mmt", machine=Machine.T4, num_instances=42)
 
 print("Stopping Studio")
 s.stop()