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

lightning_sdk/llm/llm.py

@@ -1,10 +1,12 @@
 import os
-from typing import Any, AsyncGenerator, ClassVar, Dict, Generator, List, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Any, AsyncGenerator, ClassVar, Dict, Generator, List, Literal, Optional, Tuple, Union
 
 from lightning_sdk.api import TeamspaceApi, UserApi
 from lightning_sdk.api.llm_api import LLMApi, authenticate
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
 from lightning_sdk.llm.public_assistants import PUBLIC_MODELS
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace
 
 PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
     "openai": "OpenAI",
@@ -14,6 +16,38 @@ PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
 }
 
 
+@dataclass
+class ModelMetadata:
+    name: str
+    provider: str
+    status: str
+    context_length: int
+    max_completion_tokens: Optional[int]
+    prompt_price: float
+    completion_price: float
+    capabilities: Dict[str, bool]
+    throughput: float
+    time_to_first_token: float
+
+    def __str__(self) -> str:
+        """Return a user-friendly string representation of the model metadata.
+
+        Returns:
+            str: A formatted multi-line string containing model information including
+                 name, provider, status, context length, pricing, performance metrics,
+                 and key capabilities (images and files support).
+        """
+        return f"""
+Model: {self.name}
+Provider: {self.provider}
+Status: {self.status}
+Context Length: {self.context_length:,} tokens
+Pricing: ${self.prompt_price:.2e}/prompt token, ${self.completion_price:.2e}/completion token
+Performance: {self.throughput:.1f} tokens/sec, {self.time_to_first_token:.1f}ms TTFT
+Capabilities: Images={self.capabilities.get('images', False)}, Files={self.capabilities.get('files', False)}
+        """.strip()
+
+
 class LLM:
     _auth_info_cached: ClassVar[bool] = False
     _cached_auth_info: ClassVar[Dict[str, Optional[str]]] = {}
@@ -31,33 +65,33 @@ class LLM:
     ) -> None:
         """Initializes the LLM instance with teamspace information, which is required for billing purposes.
 
-        Teamspace information is resolved through the following methods:
-        1. `.lightning/credentials.json` - Attempts to retrieve the teamspace from the local credentials file.
-        2. Environment Variables - Checks for `LIGHTNING_*` environment variables.
-        3. User Authentication - Redirects the user to the login page if teamspace information is not found.
-
         Args:
             name (str): The name of the model or resource.
-            teamspace (Optional[str]): The specified teamspace for billing. If not provided, it will be resolved
-                                       through the above methods.
+            teamspace (Optional[str]): The specified teamspace for billing. If not provided, it will
+                                        use the available default teamspace.
             enable_async (Optional[bool]): Enable async requests
 
         Raises:
             ValueError: If teamspace information cannot be resolved.
         """
         teamspace_name = None
+        teamspace_owner = None
         if teamspace:
-            try:
-                owner, teamspace_name = teamspace.split("/", maxsplit=1)
-            except ValueError as e:
-                raise ValueError(
-                    f"Invalid teamspace format: '{teamspace}'. "
-                    "Teamspace should be specified as '{teamspace_owner}/{teamspace_name}' "
-                    "(e.g., 'my-org/my-teamspace')."
-                ) from e
+            if "/" in teamspace:
+                try:
+                    teamspace_owner, teamspace_name = teamspace.split("/", maxsplit=1)
+                except ValueError as e:
+                    raise ValueError(
+                        f"Invalid teamspace format: '{teamspace}'. "
+                        "Teamspace should be specified as '{org}/{teamspace_name}' or '{org}'"
+                        "(e.g., 'my-org/my-teamspace', 'my-org')."
+                    ) from e
+            else:
+                # org is given
+                teamspace_name = teamspace
 
         self._model_provider, self._model_name = self._parse_model_name(name)
-        self._get_auth_info(teamspace_name)
+        self._get_auth_info(teamspace_owner, teamspace_name)
 
         self._enable_async = enable_async
 
@@ -69,6 +103,7 @@ class LLM:
         self._context_length = None
         self._model_id = self._get_model_id()
         self._conversations = {}
+        self._metadata = None
 
     @property
     def name(self) -> str:
@@ -78,44 +113,116 @@ class LLM:
     def provider(self) -> str:
         return self._model_provider
 
-    def context_length(self, model: Optional[str] = None) -> Optional[int]:
-        if model is None:
-            return self._context_length
+    @property
+    def metadata(self) -> ModelMetadata:
+        if self._metadata is None:
+            model = self._llm_api.get_model_metadata(self._teamspace_id, self._model_name)
+            abilities = (
+                model.abilities
+                or type(
+                    "obj",
+                    (object,),
+                    {"can_receive_images": False, "can_receive_files": False, "can_call_hub_deployment": False},
+                )()
+            )
+
+            self._metadata = ModelMetadata(
+                name=self._model_name,
+                provider=self._model_provider,
+                status=model.status,
+                context_length=int(model.context_length),
+                max_completion_tokens=int(model.max_completion_tokens) if model.max_completion_tokens != "0" else None,
+                prompt_price=model.prompt_token_price,
+                completion_price=model.completion_token_price,
+                capabilities={
+                    "images": abilities.can_receive_images,
+                    "files": abilities.can_receive_files,
+                    "hub_deployment": abilities.can_call_hub_deployment,
+                },
+                throughput=model.throughput,
+                time_to_first_token=model.time_to_first_token,
+            )
+        return self._metadata
+
+    @property
+    def context_length(self) -> Optional[int]:
+        """Context length for the current model."""
+        if self._context_length is None:
+            try:
+                self._context_length = self.metadata.context_length
+            except Exception as e:
+                raise ValueError(f"Cannot access context length: {e}") from e
+        return self._context_length
 
+    def get_context_length(self, model: Optional[str] = None) -> Optional[int]:
+        """Get context length for the given model."""
         context_info = self._public_assistants.get(model)
-        if context_info is None or "context_length" not in context_info:
-            raise ValueError(f"Cannot access context length of model '{model}'.")
+        if context_info and "context_length" in context_info:
+            return int(context_info["context_length"])
 
-        return int(context_info["context_length"])
+        try:
+            temp_metadata = self._llm_api.get_model_metadata(self._teamspace_id, model)
+            return int(temp_metadata.context_length)
+        except Exception as e:
+            raise ValueError(f"Cannot access context length of model '{model}': {e}") from e
 
-    def _get_auth_info(self, teamspace_name: Optional[str] = None) -> None:
-        # TODO: Validate user input teamspace name
+    def _get_auth_info(self, teamspace_owner: Optional[str] = None, teamspace_name: Optional[str] = None) -> None:
         if not LLM._auth_info_cached:
-            if teamspace_name is None:
-                # studio users
-                teamspace_name = os.environ.get("LIGHTNING_TEAMSPACE", None)
-
-            if teamspace_name is None:
-                # local users with no given teamspace
+            if teamspace_owner and teamspace_name:
+                # org with specific teamspace
+                try:
+                    t = _resolve_teamspace(teamspace=teamspace_name, org=None, user=teamspace_owner)
+                except Exception:
+                    try:
+                        t = _resolve_teamspace(teamspace=teamspace_name, org=teamspace_owner, user=None)
+                    except Exception as err:
+                        raise ValueError(
+                            f"Teamspace {teamspace_owner}/{teamspace_name} not found."
+                            "Please verify owner name (username or organization) and the teamspace name are correct."
+                        ) from err
+
+                    os.environ["LIGHTNING_TEAMSPACE"] = t.name
+                    os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = t.id
+
+            elif teamspace_name and teamspace_owner is None:
+                # if only org name is given, use the default teamspace
                 try:
-                    authenticate(model=f"{self.provider}/{self.name}")
+                    org = _resolve_org(teamspace_name)
                     teamspace_api = TeamspaceApi()
-                    user_api = UserApi()
-                    authed_user = user_api._client.auth_service_get_user()
-                    default_teamspace = teamspace_api.list_teamspaces(owner_id=authed_user.id)[0]
-                    teamspace_name = default_teamspace.name
-                    teamspace_id = default_teamspace.id
-                    os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = teamspace_id
-                    os.environ["LIGHTNING_TEAMSPACE"] = teamspace_name
+                    teamspace = teamspace_api.list_teamspaces(org.id)[0]
+
                 except Exception as err:
-                    # throw an appropriate error that guides users to login through the platform
                     raise ValueError(
-                        "Teamspace information is missing. "
-                        "If this is your first time using LitAI, please log in at https://lightning.ai/sign-up "
-                        "and re-run your script, or set the environment variable LIGHTNING_TEAMSPACE=<your-teamspace>."
+                        f"Organization {teamspace_name} not found. Please verify the organization name."
                     ) from err
+                os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = teamspace.id
+                os.environ["LIGHTNING_TEAMSPACE"] = teamspace.name
+
+            else:
+                if teamspace_name is None:
+                    # studio users
+                    teamspace_name = os.environ.get("LIGHTNING_TEAMSPACE", None)
+
+                if teamspace_name is None:
+                    # local users with no given teamspace
+                    try:
+                        authenticate(model=f"{self.provider}/{self.name}")
+                        teamspace_api = TeamspaceApi()
+                        user_api = UserApi()
+                        authed_user = user_api._client.auth_service_get_user()
+                        default_teamspace = teamspace_api.list_teamspaces(owner_id=authed_user.id)[0]
+                        teamspace_name = default_teamspace.name
+                        teamspace_id = default_teamspace.id
+                        os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = teamspace_id
+                        os.environ["LIGHTNING_TEAMSPACE"] = teamspace_name
+                    except Exception as err:
+                        # throw an appropriate error that guides users to login through the platform
+                        raise ValueError(
+                            "Teamspace information is missing. "
+                            "If this is your first time using LitAI, please log in at https://lightning.ai/sign-up "
+                            "and re-run your script, or set environment variable LIGHTNING_TEAMSPACE=<your-teamspace>."
+                        ) from err
 
-            # TODO: when teamspace_name is given, we don't know the teamspace_id yet
             # TODO: if LIGHTNING_CLOUD_PROJECT_ID does not exist, we have to get the id from the teamspace name
 
             LLM._cached_auth_info = {
@@ -138,16 +245,21 @@ class LLM:
         self._cloud_url = LLM._cached_auth_info["cloud_url"]
         self._org = None
 
-    def _parse_model_name(self, name: str) -> Tuple[str, str]:
-        parts = name.split("/")
-        if len(parts) == 1:
-            # a user model or a org model
-            return None, parts[0]
-        if len(parts) == 2:
-            return parts[0].lower(), parts[1]
-        raise ValueError(
-            f"Model name must be in the format `organization/model_name` or `model_name`, but got '{name}'."
-        )
+    @staticmethod
+    def _parse_model_name(name: str) -> Tuple[str, str]:
+        """Parses the model name into provider and model name.
+
+        >>> LLM._parse_model_name("openai/v1/gpt-3.5-turbo")
+        ('openai', 'v1/gpt-3.5-turbo')
+        """
+        if "/" not in name:
+            raise ValueError(
+                f"Invalid model name format: '{name}'. "
+                "Model name must be in the format `provider/model_name`."
+                "(e.g., 'lightning-ai/gpt-oss-20b')"
+            )
+        provider, model_name = name.split("/", maxsplit=1)
+        return provider.lower(), model_name
 
     # returns the assistant ID
     def _get_model_id(self) -> str:
@@ -250,12 +362,13 @@ class LLM:
         self,
         prompt: str,
         system_prompt: Optional[str] = None,
-        max_completion_tokens: Optional[int] = 500,
+        max_completion_tokens: Optional[int] = None,
         images: Optional[Union[List[str], str]] = None,
         conversation: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
         full_response: bool = False,
+        reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None,
         **kwargs: Any,
     ) -> Union[str, AsyncGenerator[str, None]]:
         conversation_id = self._conversations.get(conversation) if conversation else None
@@ -284,17 +397,21 @@ class LLM:
         self,
         prompt: str,
         system_prompt: Optional[str] = None,
-        max_completion_tokens: Optional[int] = 500,
+        max_completion_tokens: Optional[int] = None,
         images: Optional[Union[List[str], str]] = None,
         conversation: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
         full_response: bool = False,
         tools: Optional[List[Dict[str, Any]]] = None,
+        reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None,
         **kwargs: Any,
     ) -> Union[
         V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None], str, Generator[str, None, None]
     ]:
+        if reasoning_effort is not None and reasoning_effort not in ["none", "low", "medium", "high"]:
+            raise ValueError("reasoning_effort must be 'none', 'low', 'medium', 'high', or None")
+
         if conversation and conversation not in self._conversations:
             self._get_conversations()
 
@@ -317,6 +434,7 @@ class LLM:
                 metadata,
                 stream,
                 full_response,
+                reasoning_effort,
                 **kwargs,
             )
 
@@ -332,6 +450,7 @@ class LLM:
             name=conversation,
             stream=stream,
             tools=tools,
+            reasoning_effort=reasoning_effort,
             **kwargs,
         )
         if not stream:

lightning_sdk/llm/public_assistants.py

@@ -1,11 +1,36 @@
 PUBLIC_MODELS = {
-    "openai/gpt-4o": {"id": "ast_01jdjds71fs8gt47jexzed4czs", "context_length": 128000},
-    "openai/gpt-4": {"id": "ast_01jd38ze6tjbrcd4942nhz41zn", "context_length": 8192},
-    "openai/o3-mini": {"id": "ast_01jz3t13fhnjhh11t1k8b5gyp1", "context_length": 128000},
-    "anthropic/claude-3-5-sonnet-20240620": {"id": "ast_01jd3923a6p98rqwh3dpj686pq", "context_length": 200000},
-    "google/gemini-2.5-pro": {"id": "ast_01jz3tdb1fhey798k95pv61v57", "context_length": 1048576},
-    "google/gemini-2.5-flash": {"id": "ast_01jz3thxskg4fcdk4xhkjkym5a", "context_length": 8000},
-    "google/gemini-2.5-flash-lite-preview-06-17": {"id": "ast_01jz3thxskg4fcdk4xhkjkym5b", "context_length": 8000},
+    "openai/gpt-4o": {
+        "id": "ast_01jdjds71fs8gt47jexzed4czs",
+        "context_length": 128000
+    },
+    "openai/gpt-4": {
+        "id": "ast_01jd38ze6tjbrcd4942nhz41zn",
+        "context_length": 8192
+    },
+    "openai/o3-mini": {
+        "id": "ast_01jz3t13fhnjhh11t1k8b5gyp1",
+        "context_length": 128000
+    },
+    "anthropic/claude-3-5-sonnet-20240620": {
+        "id": "ast_01jd3923a6p98rqwh3dpj686pq",
+        "context_length": 200000
+    },
+    "google/gemini-2.5-pro": {
+        "id": "ast_01jz3tdb1fhey798k95pv61v57",
+        "context_length": 1048576
+    },
+    "google/gemini-2.5-flash": {
+        "id": "ast_01jz3thxskg4fcdk4xhkjkym5a",
+        "context_length": 1048576
+    },
+    "google/gemini-2.5-flash-lite-preview-06-17": {
+        "id": "ast_01jz3thxskg4fcdk4xhkjkym5b",
+        "context_length": 1048576
+    },
+    "lightning-ai/DeepSeek-V3.1": {
+        "id": "ast_01k378z15k0msktaxjpfrfwqqh",
+        "context_length": 163840
+    },
     "lightning-ai/gpt-oss-20b": {
         "id": "ast_01k1y2ywfh9zzexjzrva6y96ms",
         "context_length": 128000,
@@ -14,4 +39,16 @@ PUBLIC_MODELS = {
         "id": "ast_01k1z7ajxbz7wq62xj5bc0mrcs",
         "context_length": 128000,
     },
+    "openai/gpt-5": {
+        "id": "ast_01jdjds71fs8gt47jexzed4cza",
+        "context_length": 400000
+    },
+    "openai/gpt-5-mini": {
+        "id": "ast_01jdjds71fs8gt47jexzed4czb",
+        "context_length": 400000
+    },
+    "openai/gpt-5-nano": {
+        "id": "ast_01jdjds71fs8gt47jexzed4czc",
+        "context_length": 400000
+    },
 }

lightning_sdk/machine.py

@@ -6,7 +6,6 @@ from typing import Any, ClassVar, Optional, Tuple
 class CloudProvider(Enum):
     AWS = "AWS"
     GCP = "GCP"
-    VULTR = "VULTR"
     LAMBDA_LABS = "LAMBDA_LABS"
     DGX = "DGX"
     VOLTAGE_PARK = "VOLTAGE_PARK"
@@ -35,6 +34,7 @@ class Machine:
 
     # supported GPU types
     # supported T4 variations
+    T4_SMALL: ClassVar["Machine"]
     T4: ClassVar["Machine"]
     T4_X_2: ClassVar["Machine"]
     T4_X_4: ClassVar["Machine"]
@@ -116,7 +116,7 @@ class Machine:
         )
         for m in possible_values:
             for machine_id in [machine, *additional_machine_ids]:
-                if machine_id in (
+                if machine_id and machine_id in (
                     getattr(m, "name", None),
                     getattr(m, "instance_type", None),
                     getattr(m, "slug", None),
@@ -127,6 +127,21 @@ class Machine:
             return cls(machine, *additional_machine_ids)
         return cls(machine, machine, machine)
 
+    @classmethod
+    def _from_accelerator(cls, accelerator: Any) -> "Machine":
+        if accelerator.accelerator_type == "GPU":
+            accelerator_resources_count = accelerator.resources.gpu
+        else:
+            accelerator_resources_count = accelerator.resources.cpu
+
+        return Machine.from_str(
+            accelerator.slug_multi_cloud,
+            accelerator.slug,
+            accelerator.instance_id,
+            accelerator.secondary_instance_id,
+            f"lit-{accelerator.family.lower()}-{accelerator_resources_count}",
+        )
+
 
 # CPU machines
 # default CPU machines
@@ -148,6 +163,7 @@ Machine.DATA_PREP_ULTRA = Machine(
 
 # GPU machines
 # available T4 machines
+Machine.T4_SMALL = Machine(name="T4_SMALL", slug="lit-t4-1-small", family="T4", accelerator_count=1)
 Machine.T4 = Machine(name="T4", slug="lit-t4-1", family="T4", accelerator_count=1)
 Machine.T4_X_2 = Machine(name="T4_X_2", slug="lit-t4-2", family="T4", accelerator_count=2)
 Machine.T4_X_4 = Machine(name="T4_X_4", slug="lit-t4-4", family="T4", accelerator_count=4)
@@ -202,3 +218,6 @@ Machine.H200 = Machine(name="H200", slug="lit-h200x-1", family="H200", accelerat
 Machine.H200_X_8 = Machine(name="H200_X_8", slug="lit-h200x-8", family="H200", accelerator_count=8)
 # available B200 machines
 Machine.B200_X_8 = Machine(name="B200_X_8", slug="lit-b200x-8", family="B200", accelerator_count=8)
+
+
+DEFAULT_MACHINE = Machine.CPU.name

lightning_sdk/mmt/base.py

@@ -75,6 +75,7 @@ class _BaseMMT(_BaseJob):
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+        reuse_snapshot: bool = True,
     ) -> "_BaseMMT":
         """Run async workloads using a docker image across multiple machines.
 
@@ -119,6 +120,8 @@ class _BaseMMT(_BaseJob):
                 Irrelevant for most machines, required for some of the top-end machines on GCP.
                 If in doubt, set it. Won't have an effect on machines not requiring it.
                 Defaults to 3h
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         from lightning_sdk.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
@@ -214,6 +217,7 @@ class _BaseMMT(_BaseJob):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         return inst
 
@@ -236,6 +240,7 @@ class _BaseMMT(_BaseJob):
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
         max_runtime: Optional[int] = None,
+        reuse_snapshot: bool = True,
     ) -> None:
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -271,6 +276,8 @@ class _BaseMMT(_BaseJob):
                 Irrelevant for most machines, required for some of the top-end machines on GCP.
                 If in doubt, set it. Won't have an effect on machines not requiring it.
                 Defaults to 3h
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
 
     @property

lightning_sdk/mmt/mmt.py

@@ -20,9 +20,9 @@ _logger = _setup_logger(__name__)
 class MMT(_BaseMMT):
     """Class to submit and manage multi-machine jobs on the Lightning AI Platform."""
 
-    _force_v1: (
-        bool
-    ) = False  # required for studio plugin still working correctly as v2 currently does not support the studio env
+    _force_v1: bool = (
+        False  # required for studio plugin still working correctly as v2 currently does not support the studio env
+    )
 
     def __init__(
         self,
@@ -102,6 +102,7 @@ class MMT(_BaseMMT):
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+        reuse_snapshot: bool = True,
     ) -> "MMT":
         """Run async workloads using a docker image across multiple machines.
 
@@ -141,6 +142,8 @@ class MMT(_BaseMMT):
                     }
                 If the path inside the connection is omitted it's assumed to be the root path of that connection.
                 Only applicable when submitting docker jobs.
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         ret_val = super().run(
             name=name,
@@ -164,6 +167,7 @@ class MMT(_BaseMMT):
             artifacts_remote=artifacts_remote,
             cluster=cluster,  # deprecated in favor of cloud_account
             max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         # required for typing with "MMT"
         assert isinstance(ret_val, cls)
@@ -191,6 +195,7 @@ class MMT(_BaseMMT):
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
+        reuse_snapshot: bool = True,
     ) -> "MMT":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -231,6 +236,8 @@ class MMT(_BaseMMT):
                 Irrelevant for most machines, required for some of the top-end machines on GCP.
                 If in doubt, set it. Won't have an effect on machines not requiring it.
                 Defaults to 3h
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
@@ -249,6 +256,7 @@ class MMT(_BaseMMT):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         return self
 

lightning_sdk/mmt/v1.py

@@ -58,6 +58,7 @@ class _MMTV1(_BaseMMT):
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        reuse_snapshot: bool = True,
     ) -> "_MMTV1":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -96,7 +97,8 @@ class _MMTV1(_BaseMMT):
                 Irrelevant for most machines, required for some of the top-end machines on GCP.
                 If in doubt, set it. Won't have an effect on machines not requiring it.
                 Defaults to 3h
-
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         raise NotImplementedError("Cannot submit new mmts with MMTV1!")
 

lightning_sdk/mmt/v2.py

@@ -58,6 +58,7 @@ class _MMTV2(_BaseMMT):
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
+        reuse_snapshot: bool = True,
     ) -> "_MMTV2":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -98,6 +99,8 @@ class _MMTV2(_BaseMMT):
                 Irrelevant for most machines, required for some of the top-end machines on GCP.
                 If in doubt, set it. Won't have an effect on machines not requiring it.
                 Defaults to 3h
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -141,6 +144,7 @@ class _MMTV2(_BaseMMT):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         self._job = submitted
         self._name = submitted.name

lightning_sdk/owner.py

@@ -2,12 +2,13 @@ from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, List
 
 from lightning_sdk.api import TeamspaceApi
+from lightning_sdk.utils.logging import TrackCallsABCMeta
 
 if TYPE_CHECKING:
     from lightning_sdk.teamspace import Teamspace
 
 
-class Owner(ABC):
+class Owner(ABC, metaclass=TrackCallsABCMeta):
     """Represents an owner of teamspaces and studios."""
 
     def __init__(self) -> None:

lightning_sdk/pipeline/steps.py

@@ -175,6 +175,7 @@ class JobStep:
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         wait_for: Union[str, List[str], None] = DEFAULT,
+        reuse_snapshot: bool = True,
     ) -> None:
         self.name = name
         self.machine = machine or Machine.CPU
@@ -201,6 +202,7 @@ class JobStep:
         self.path_mappings = path_mappings
         self.max_runtime = max_runtime
         self.wait_for = wait_for
+        self.reuse_snapshot = reuse_snapshot
 
     def to_proto(
         self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
@@ -239,6 +241,7 @@ class JobStep:
             artifacts_remote=None,
             max_runtime=self.max_runtime,
             machine_image_version=machine_image_version,
+            reuse_snapshot=self.reuse_snapshot,
         )
 
         return V1PipelineStep(
@@ -272,6 +275,7 @@ class MMTStep:
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         wait_for: Optional[Union[str, List[str]]] = DEFAULT,
+        reuse_snapshot: bool = True,
     ) -> None:
         self.machine = machine or Machine.CPU
         self.num_machines = num_machines
@@ -297,6 +301,7 @@ class MMTStep:
         self.path_mappings = path_mappings
         self.max_runtime = max_runtime
         self.wait_for = wait_for
+        self.reuse_snapshot = reuse_snapshot
 
     def to_proto(
         self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
@@ -332,6 +337,7 @@ class MMTStep:
             artifacts_remote=None,  # deprecated in favor of path_mappings
             max_runtime=self.max_runtime,
             machine_image_version=machine_image_version,
+            reuse_snapshot=self.reuse_snapshot,
         )
 
         return V1PipelineStep(