lightning-sdk 2025.8.6rc2__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/teamspace.py

@@ -1,35 +1,59 @@
 import glob
 import os
 import warnings
+from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
 import lightning_sdk
 from lightning_sdk.agents import Agent
-from lightning_sdk.api import TeamspaceApi
-from lightning_sdk.lightning_cloud.openapi import V1Model, V1ModelVersionArchive, V1ProjectClusterBinding
-from lightning_sdk.machine import Machine
+from lightning_sdk.api import CloudAccountApi, TeamspaceApi
+from lightning_sdk.lightning_cloud.openapi import V1ClusterType, V1Model, V1ModelVersionArchive, V1ProjectClusterBinding
+from lightning_sdk.machine import CloudProvider, Machine
 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.logging import TrackCallsMeta
 from lightning_sdk.utils.resolve import (
     _get_organizations_for_authed_user,
     _parse_model_and_version,
     _resolve_org,
     _resolve_teamspace_name,
     _resolve_user,
+    skip_studio_init,
 )
 
 if TYPE_CHECKING:
     from lightning_sdk.job import Job
     from lightning_sdk.mmt import MMT
-    from lightning_sdk.studio import Studio
+    from lightning_sdk.studio import VM, Studio
 
 
-class Teamspace:
+class FolderLocation(Enum):
+    AWS = "AWS"
+    GCP = "GCP"
+    CLOUD_AGNOSTIC = "CLOUD_AGNOSTIC"
+
+    def __str__(self) -> str:
+        """Converts the FolderLocation to a str."""
+        return self.value
+
+
+class ConnectionType(Enum):
+    EFS = "EFS"
+    S3 = "S3"
+    GCS = "GCS"
+    FILESTORE = "FILESTORE"
+
+    def __str__(self) -> str:
+        """Converts the FolderLocation to a str."""
+        return self.value
+
+
+class Teamspace(metaclass=TrackCallsMeta):
     """A teamspace is a collection of Studios, Clusters, Members and an associated Budget.
 
     Args:
@@ -53,6 +77,7 @@ class Teamspace:
         user: Optional[Union[str, User]] = None,
     ) -> None:
         self._teamspace_api = TeamspaceApi()
+        self._cloud_account_api = CloudAccountApi()
 
         name = _resolve_teamspace_name(name)
 
@@ -70,6 +95,20 @@ class Teamspace:
             self._user = _resolve_user(user)
             self._org = _resolve_org(org)
 
+        # If still no user or org resolved, try config defaults
+        if self._user is None and self._org is None:
+            from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+            config = Config()
+            owner_type = config.get_value(DefaultConfigKeys.teamspace_owner_type)
+            owner_name = config.get_value(DefaultConfigKeys.teamspace_owner)
+
+            if owner_type and owner_name:
+                if owner_type.lower() == "organization":
+                    self._org = _resolve_org(owner_name)
+                elif owner_type.lower() == "user":
+                    self._user = _resolve_user(owner_name)
+
         self._owner: Owner
         if self._user is None and self._org is None:
             raise RuntimeError(
@@ -106,12 +145,25 @@ class Teamspace:
         """All studios within that teamspace."""
         from lightning_sdk.studio import Studio
 
+        return self._get_studios(Studio)
+
+    @property
+    def vms(self) -> List["VM"]:
+        from lightning_sdk.studio import VM
+
+        return [x for x in self._get_studios(VM) if isinstance(x, VM)]
+
+    def _get_studios(self, target_cls: type) -> List[Union["Studio", "VM"]]:
         studios = []
         cloud_accounts = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
         for cl in cloud_accounts:
             _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))
+                with skip_studio_init():
+                    studio = target_cls(name=s.name, teamspace=self, cluster=cl.cluster_name, create_ok=False)
+                    studio._studio = s
+                    studio._teamspace = self
+                    studios.append(studio)
 
         return studios
 
@@ -195,6 +247,24 @@ class Teamspace:
 
         return tuple(mmts)
 
+    @property
+    def secrets(self) -> Dict[str, str]:
+        """All (encrypted) secrets for the teamspace.
+
+        Note:
+            Once created, the secret values are encrypted and cannot be viewed here anymore.
+        """
+        return self._teamspace_api.get_secrets(self.id)
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set the (encrypted) secrets for the teamspace."""
+        if not self._teamspace_api.verify_secret_name(key):
+            raise ValueError(
+                "Secret keys must only contain alphanumeric characters and underscores and not begin with a number."
+            )
+
+        self._teamspace_api.set_secret(self.id, key, value)
+
     def list_machines(self, cloud_account: Optional[str] = None) -> List[Machine]:
         if cloud_account is None:
             cloud_account = os.getenv("LIGHTNING_CLUSTER_ID") or self.default_cloud_account
@@ -478,6 +548,96 @@ class Teamspace:
             cloud_account=self.default_cloud_account,
         )
 
+    def new_folder(
+        self, name: str, location: Optional[FolderLocation] = None, cloud_account: Optional[str] = None
+    ) -> None:
+        """Create a new folder in this Teamspace.
+
+        Args:
+            name: The name of the folder. Folders will be accesible under `/teamspace/folders/<name>`
+            location: The location of the folder. Defaults to cloud agnostic.
+            cloud_account: The cloud account to create the folder in. Not used for cloud agnostic folders.
+        """
+        if cloud_account is None:
+            cloud_account = self.default_cloud_account
+
+        cloud_accounts = self._cloud_account_api.list_cloud_accounts(self.id)
+        resolved_cloud_accounts = [
+            external_cloud for external_cloud in cloud_accounts if external_cloud.id == cloud_account
+        ]
+
+        if len(resolved_cloud_accounts) == 0:
+            raise ValueError(f"Cloud account not found: {cloud_account}")
+
+        resolved_cloud_account = resolved_cloud_accounts[0]
+
+        # if the cloud account is global, default to agnostic
+        if location is None and resolved_cloud_account.spec.cluster_type == V1ClusterType.GLOBAL:
+            location = FolderLocation.CLOUD_AGNOSTIC
+
+        # if it's global, then default to agnostic, and aws / gcp otherwise if set
+        if (
+            location is not None
+            and location != FolderLocation.CLOUD_AGNOSTIC
+            and resolved_cloud_account.spec.cluster_type == V1ClusterType.GLOBAL
+        ):
+            providers = self._cloud_account_api.get_cloud_account_provider_mapping(self.id)
+
+            if location == FolderLocation.AWS:
+                resolved_cloud_account = providers[CloudProvider.AWS]
+            elif location == FolderLocation.GCP:
+                resolved_cloud_account = providers[CloudProvider.GCP]
+
+        if location == FolderLocation.CLOUD_AGNOSTIC:
+            self._teamspace_api.new_folder(self.id, name, None)
+        else:
+            self._teamspace_api.new_folder(self.id, name, resolved_cloud_account)
+
+        return
+
+    def new_connection(
+        self,
+        name: str,
+        source: str,
+        connection_type: ConnectionType,
+        writable: bool = True,
+        cloud_account: Optional[str] = None,
+        region: Optional[str] = None,
+    ) -> None:
+        """Add an existing data source to this Teamspace.
+
+        Args:
+          name: the name under which your data will be available in this Teamspace
+          source: the source spec of your data. Format depends on the type of data to connect.
+            For EFS, this should be the filsystem id
+          connection_type: the kind of data to connect to this Teamspace
+          writable: whether to support write-back to this data source. If False, the data is connected as read-only
+          cloud_account: which cloud-account to connect to the data source to.
+            If not specified, will retrieve the cloud-account that matches the required provider type
+            starting with private cloud accounts and falling back to public if necessary.
+          region: which provider region this data is in. Required for some connection types only.
+        """
+        provider_for_connection = self._cloud_account_api.get_cloud_provider_for_connection_type(connection_type)
+
+        if connection_type == ConnectionType.EFS and region is None:
+            raise ValueError("Region must be specified")
+
+        cloud_account = self._cloud_account_api.resolve_cloud_account(
+            self.id, cloud_account, provider_for_connection, None
+        )
+
+        cloud_accounts = self._cloud_account_api.list_cloud_accounts(self.id)
+        resolved_cloud_accounts = [
+            external_cloud for external_cloud in cloud_accounts if external_cloud.id == cloud_account
+        ]
+
+        if len(resolved_cloud_accounts) == 0:
+            raise ValueError(f"Cloud account not found: {cloud_account}")
+
+        resolved_cloud_account = resolved_cloud_accounts[0]
+
+        self._teamspace_api.new_connection(self.id, name, source, resolved_cloud_account, writable, region or "")
+
 
 def _list_files(path: Union[str, Path]) -> Tuple[List[Path], List[str]]:
     """List all folders in a directory and return them as a list and relative path."""

lightning_sdk/user.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 
 from lightning_sdk.api import UserApi
 from lightning_sdk.owner import Owner
@@ -37,6 +37,24 @@ class User(Owner):
         """The user's ID."""
         return self._user.id
 
+    @property
+    def secrets(self) -> Dict[str, str]:
+        """All (encrypted) secrets for the user.
+
+        Note:
+            Once created, the secret values are encrypted and cannot be viewed here anymore.
+        """
+        return self._user_api.get_secrets()
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set a (encrypted) secret for the user."""
+        if not self._user_api.verify_secret_name(key):
+            raise ValueError(
+                "Secret keys must only contain alphanumeric characters and underscores and not begin with a number."
+            )
+
+        self._user_api.set_secret(key, value)
+
     def __repr__(self) -> str:
         """Returns reader friendly representation."""
         return f"User(name={self.name})"

lightning_sdk/utils/config.py

@@ -0,0 +1,179 @@
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, Mapping, Optional, Sequence
+
+import yaml
+
+_DEFAULT_CONFIG_FILE_PATH = "~/.lightning/config.yaml"
+
+
+@dataclass(frozen=True)
+class DefaultConfigKeys:
+    """Default configuration keys for the Lightning SDK."""
+
+    organization: str = "organization.name"
+    user: str = "user.name"
+
+    teamspace_name: str = "teamspace.name"
+    teamspace_owner: str = "teamspace.owner"
+    teamspace_owner_type: str = "teamspace.owner_type"
+
+    machine: str = "machine.name"
+
+    studio: str = "studio.name"
+
+    cloud_account: str = "cloud_account.name"
+    cloud_provider: str = "cloud_provider.name"
+
+    license: str = "license"
+
+
+class ConfigProxy:
+    def __init__(self, root: "Config", *path: str) -> None:
+        self._root = root
+        self._path = path  # list of keys from root
+
+    def __getattr__(self, name: str) -> "ConfigProxy":
+        """Returns a reference to a nested ConfigProxy object a level deeper in the config hierarchy.
+
+        Args:
+            name: the name of the attribute to access, which corresponds to a key in the config.
+
+        Returns:
+            ConfigProxy: the next ConfigProxy object for the attribute.
+        """
+        # Build a deeper path and return a new proxy
+        return ConfigProxy(self._root, *self._path, name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Sets the name attribute to value at the current hierarchy level.
+
+        Args:
+            name: the attribute name to set, which corresponds to a key in the config.
+            value: the value to set for the given attribute name in the config.
+        """
+        if name in ("_root", "_path"):  # internal attributes
+            super().__setattr__(name, value)
+        else:
+            # Assign a nested value in the root config
+            self._root._set_nested([*self._path, name], value)
+
+
+class Config:
+    def __init__(self, config_file: Optional[str] = None) -> None:
+        """Config class to manage configuration settings for the lightning SDK and CLI.
+
+        Args:
+            config_file: the file path where the configuration is stored.
+            If None, defaults to "~/.lightning/config.yaml".
+        """
+        if config_file is None:
+            config_file = _DEFAULT_CONFIG_FILE_PATH
+        self._config_file = os.path.expanduser(config_file)
+
+    def _load_config(self) -> Dict[str, Any]:
+        if not os.path.exists(self._config_file):
+            return {}  # Return empty dict if config doesn't exist
+        with open(self._config_file) as f:
+            return yaml.safe_load(f) or {}
+
+    def _save_config(self, config: Dict[str, Any]) -> None:
+        os.makedirs(os.path.dirname(self._config_file), exist_ok=True)
+        config = _unflatten_dict(config)
+        with open(self._config_file, "w") as f:
+            yaml.safe_dump(config, f, default_flow_style=False, sort_keys=True)
+
+    def _set_nested(self, keys: Sequence[str], value: str) -> None:
+        config = self._load_config()
+        curr = config
+        for k in keys[:-1]:
+            if k not in curr or not isinstance(curr[k], Dict):
+                curr[k] = {}
+            curr = curr[k]
+        curr[keys[-1]] = value
+        self._save_config(config)
+
+    def get_value(self, key_path: str) -> Optional[str]:
+        """Gets a value from the config using dot notation.
+
+        Args:
+            key_path: the dot-separated path to the config value (e.g. "teamspace.name")
+
+        Returns:
+            The config value if found, None otherwise
+        """
+        return self._get_value_type(key_path, str)
+
+    def get_sub_config(self, key_path: str) -> Optional[Mapping[str, Any]]:
+        """Gets a subconfig from the config using dot notation.
+
+        Args:
+            key_path: the dot-separated path to the subconfig (e.g. "license")
+
+        Returns:
+            The subconfig if found, None otherwise
+        """
+        return self._get_value_type(key_path, Mapping)
+
+    def _get_value_type(self, key_path: str, subtype: type) -> Optional[Any]:
+        config = self._load_config()
+        if not isinstance(config, Mapping):
+            return None
+
+        keys = key_path.split(".")
+        curr = config
+        for k in keys:
+            if not isinstance(curr, dict) or k not in curr:
+                return None
+            curr = curr[k]
+        return curr if isinstance(curr, subtype) else None
+
+    def __getattr__(self, name: str) -> ConfigProxy:
+        """Returns a proxy to the actual values to allow for nested access.
+
+        Args:
+            name: the name of the value to retrieve.
+
+        Returns:
+            ConfigProxy: a proxy object that allows nested access to the configuration.
+        """
+        return ConfigProxy(self, name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Sets the name attribute to value at the root level."""
+        if name in ("_config_file",):  # internal attributes
+            super().__setattr__(name, value)
+        else:
+            # Assign a value at the root level
+            self._set_nested([name], value)
+
+    def __repr__(self) -> str:
+        """Returns a string representation of the config."""
+        return str(self)
+
+    def __str__(self) -> str:
+        """Returns a string representation of the config."""
+        return yaml.dump(
+            {"Config": {"config_file": self._config_file, **self._load_config()}},
+            indent=4,
+            sort_keys=True,
+        )
+
+    def get(self, key: str) -> Optional[str]:
+        return self.get_value(key)
+
+    def set(self, key: str, value: str) -> None:
+        self._set_nested([key], value)
+
+
+def _unflatten_dict(flat_dict: Dict[str, Any]) -> Dict[str, Any]:
+    unflattened_dict = {}
+    for key, value in flat_dict.items():
+        keys = key.split(".")
+        curr = unflattened_dict
+        for k in keys[:-1]:
+            if k not in curr:
+                curr[k] = {}
+            curr = curr[k]
+        curr[keys[-1]] = value
+    return unflattened_dict

lightning_sdk/utils/license.py

@@ -0,0 +1,13 @@
+from functools import lru_cache
+
+from lightning_sdk.api.license_api import LicenseApi
+
+
+class License:
+    def __init__(self, license_key: str, product_name: str) -> None:
+        self.license_key = license_key
+        self.product_name = product_name
+
+    @lru_cache(maxsize=1)  # noqa: B019
+    def validate(self) -> bool:
+        return LicenseApi(self.license_key).validate_license(self.license_key, self.product_name)

lightning_sdk/utils/logging.py

@@ -0,0 +1,79 @@
+import functools
+import inspect
+import time
+import traceback
+from abc import ABCMeta
+from contextlib import suppress
+from typing import Callable, Dict, Tuple
+
+from lightning_sdk.__version__ import __version__
+from lightning_sdk.lightning_cloud.openapi import V1CreateSDKCommandHistoryRequest, V1SDKCommandHistorySeverity
+from lightning_sdk.lightning_cloud.openapi.models.v1_sdk_command_history_type import V1SDKCommandHistoryType
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+def track_calls() -> Callable[..., any]:
+    def decorator(func: Callable[..., any]) -> Callable[..., any]:
+        @functools.wraps(func)
+        def wrapper(*args: Tuple[any, ...], **kwargs: Dict[str, any]) -> any:
+            start_time = time.time()
+            sig = inspect.signature(func)
+            bound_args = sig.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+
+            args_str = ", ".join(f"{k}: {v}" for k, v in bound_args.arguments.items() if k != "self")
+            message = f"VERSION: {__version__} | ARGS: {args_str} "
+
+            body = V1CreateSDKCommandHistoryRequest(
+                command=func.__qualname__,
+                message=message,
+                project_id=None,
+                severity=V1SDKCommandHistorySeverity.INFO,
+                type=V1SDKCommandHistoryType.SDK,
+            )
+
+            try:
+                return func(*args, **kwargs)
+
+            except Exception as e:
+                body.severity = V1SDKCommandHistorySeverity.ERROR
+                body.message += f" | Error: {type(e).__name__}: {e!s} | Traceback: {traceback.format_exc(limit=3)}"
+                body.duration = int(time.time() - start_time)
+                with suppress(Exception):
+                    client = LightningClient(retry=False, max_tries=0)
+                    client.s_dk_command_history_service_create_sdk_command_history(body=body)
+                raise
+
+        return wrapper
+
+    return decorator
+
+
+class TrackCallsMeta(type):
+    def __new__(cls, name: any, bases: any, attrs: any) -> type.__new__:
+        for attr_name, attr_value in attrs.items():
+            if attr_name.startswith("__") and attr_name not in ("__init__", "__call__"):
+                attrs[attr_name] = attr_value
+                continue
+
+            # Noisy
+            if attr_name in ("_cls_name"):
+                attrs[attr_name] = attr_value
+                continue
+
+            if callable(attr_value):
+                attrs[attr_name] = track_calls()(attr_value)
+
+            elif isinstance(attr_value, property):
+                fget = track_calls()(attr_value.fget) if attr_value.fget else None
+                fset = track_calls()(attr_value.fset) if attr_value.fset else None
+                fdel = track_calls()(attr_value.fdel) if attr_value.fdel else None
+                attrs[attr_name] = property(fget, fset, fdel)
+
+            else:
+                attrs[attr_name] = attr_value
+        return super().__new__(cls, name, bases, attrs)
+
+
+class TrackCallsABCMeta(ABCMeta, TrackCallsMeta):
+    """Combined metaclass for classes that need both ABC and TrackCalls functionality."""