lightning-sdk 2025.8.1__py3-none-any.whl → 2025.8.6rc1__py3-none-any.whl

lightning_sdk/lightning_cloud/utils/data_connection.py

@@ -1,7 +1,7 @@
 import os
 from time import sleep, time
 from lightning_sdk.lightning_cloud import rest_client
-from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection, V1S3FolderDataConnection, V1EfsConfig, V1GcpDataConnection
+from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection, V1S3FolderDataConnection, V1EfsConfig, V1GcpDataConnection, V1FilestoreDataConnection, V1DataConnectionTier
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 import urllib3
 
@@ -272,6 +272,61 @@ def add_gcs_connection(connection_name: str, bucket_name: str, create_timeout: i
 
     return
 
+def create_filestore_folder(folder_name: str, region: str, capacity_gb: int = 1024, tier: V1DataConnectionTier = V1DataConnectionTier.HDD, create_timeout: int = 600) -> None:
+    """
+    Utility function to create a Filestore folder.
+
+    Args:
+        folder_name: The name of the folder to create.
+        region: the region to create the folder in, e.g. "us-central1".
+        capacity_gb: capacity of the Filestore folder.
+        tier: Filestore folder tier, SSD or HDD.
+        create_timeout (int): The timeout for the folder creation.
+    """
+    client = rest_client.LightningClient(retry=False)
+
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+    cluster_id = os.getenv("LIGHTNING_CLUSTER_ID")
+
+    # Get existing data connections
+    data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
+
+    for connection in data_connections:
+        existing_folder_name = getattr(connection, 'name', None)
+        isFilestore = getattr(connection, 'filestore', None) is not None
+
+        if existing_folder_name == folder_name and isFilestore:
+            return
+
+    # If we get here, no matching folder was found, proceed with creation
+    body = Create(
+        name=folder_name,
+        create_resources=True,
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        filestore=V1FilestoreDataConnection(region=region, capacity_gb=capacity_gb, tier=tier, source=folder_name),
+    )
+    try:
+        connection = client.data_connection_service_create_data_connection(body, project_id)
+    except ApiException as e:
+        # Note: This function can be called in a distributed way.
+        # There is a race condition where one machine might create the entry before another machine
+        # and this request would fail with duplicated key
+        # In this case, it is fine not to raise
+        if'duplicate key value violates unique constraint' in str(e.body):
+            return
+        raise e from None
+
+    except urllib3.exceptions.HTTPError as e:
+        raise e from None
+
+    # Wait for the filesystem picks up the newly added Filestore folder
+    start = time()
+
+    while not os.path.isdir(f"/teamspace/filestore_folders/{folder_name}") and (time() - start) < create_timeout:
+        sleep(1)
+
+
 def delete_data_connection(name: str):
     """Utility to delete a data connection
 

lightning_sdk/llm/llm.py

@@ -1,10 +1,10 @@
-import json
 import os
 from typing import Any, AsyncGenerator, ClassVar, Dict, Generator, List, Optional, Tuple, Union
 
 from lightning_sdk.api import TeamspaceApi, UserApi
 from lightning_sdk.api.llm_api import LLMApi
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
+from lightning_sdk.llm.public_assistants import PUBLIC_MODELS
 
 PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
     "openai": "OpenAI",
@@ -14,17 +14,6 @@ PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
 }
 
 
-def _load_public_assistants() -> Dict[str, Dict[str, Any]]:
-    """Load public assistants from a JSON file."""
-    try:
-        json_path = os.path.join(os.path.dirname(__file__), "public_assistants.json")
-        with open(json_path) as f:
-            return json.load(f)
-    except Exception as e:
-        print(f"[warning] Failed to load public_assistants.json: {e}")
-        return {}
-
-
 class LLM:
     _auth_info_cached: ClassVar[bool] = False
     _cached_auth_info: ClassVar[Dict[str, Optional[str]]] = {}
@@ -138,7 +127,7 @@ class LLM:
             }
             LLM._auth_info_cached = True
             if LLM._public_assistants is None:
-                LLM._public_assistants = _load_public_assistants()
+                LLM._public_assistants = PUBLIC_MODELS
         # Always assign to the current instance
         self._teamspace_name = LLM._cached_auth_info["teamspace_name"]
         self._teamspace_id = LLM._cached_auth_info["teamspace_id"]

lightning_sdk/llm/public_assistants.py

@@ -0,0 +1,17 @@
+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},
+    "lightning-ai/gpt-oss-20b": {
+        "id": "ast_01k1y2ywfh9zzexjzrva6y96ms",
+        "context_length": 128000,
+    },
+    "lightning-ai/gpt-oss-120b": {
+        "id": "ast_01k1z7ajxbz7wq62xj5bc0mrcs",
+        "context_length": 128000,
+    },
+}

lightning_sdk/pipeline/steps.py

@@ -59,6 +59,7 @@ class DeploymentStep:
         include_credentials: Optional[bool] = None,
         max_runtime: Optional[int] = None,
         wait_for: Optional[Union[str, List[str]]] = DEFAULT,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
     ) -> None:
         self.name = name
         self.studio = _get_studio(studio)
@@ -99,11 +100,28 @@ class DeploymentStep:
         self.include_credentials = include_credentials or True
         self.max_runtime = max_runtime
         self.wait_for = wait_for
+        self.cloud_provider = cloud_provider
 
     def to_proto(
         self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
     ) -> V1PipelineStep:
-        _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+        machine_image_version = None
+
+        studio = _get_studio(self.studio)
+        if isinstance(studio, Studio):
+            machine_image_version = studio._studio.machine_image_version
+
+            if self.cloud_account is None:
+                self.cloud_account = studio.cloud_account
+            elif studio.cloud_account != self.cloud_account:
+                raise ValueError("The provided cloud account doesn't match the studio")
+
+        resolved_cloud_account = CloudAccountApi().resolve_cloud_account(
+            teamspace.id, self.cloud_account, self.cloud_provider, teamspace.default_cloud_account
+        )
+
+        _validate_cloud_account(cloud_account, resolved_cloud_account, shared_filesystem)
+
         return V1PipelineStep(
             name=self.name,
             type=V1PipelineStepType.DEPLOYMENT,
@@ -127,6 +145,7 @@ class DeploymentStep:
                     cloudspace_id=self.studio._studio.id if self.studio else None,
                     include_credentials=self.include_credentials,
                     max_runtime=self.max_runtime,
+                    machine_image_version=machine_image_version,
                 ),
                 strategy=to_strategy(self.release_strategy),
             ),
@@ -186,8 +205,12 @@ class JobStep:
     def to_proto(
         self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
     ) -> V1PipelineStep:
+        machine_image_version = None
+
         studio = _get_studio(self.studio)
         if isinstance(studio, Studio):
+            machine_image_version = studio._studio.machine_image_version
+
             if self.cloud_account is None:
                 self.cloud_account = studio.cloud_account
             elif studio.cloud_account != self.cloud_account:
@@ -215,6 +238,7 @@ class JobStep:
             artifacts_local=None,
             artifacts_remote=None,
             max_runtime=self.max_runtime,
+            machine_image_version=machine_image_version,
         )
 
         return V1PipelineStep(
@@ -277,8 +301,12 @@ class MMTStep:
     def to_proto(
         self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
     ) -> V1PipelineStep:
+        machine_image_version = None
+
         studio = _get_studio(self.studio)
         if isinstance(studio, Studio):
+            machine_image_version = studio._studio.machine_image_version
+
             if self.cloud_account is None:
                 self.cloud_account = studio.cloud_account
             elif studio.cloud_account != self.cloud_account:
@@ -303,6 +331,7 @@ class MMTStep:
             artifacts_local=None,  # deprecated in favor of path_mappings
             artifacts_remote=None,  # deprecated in favor of path_mappings
             max_runtime=self.max_runtime,
+            machine_image_version=machine_image_version,
         )
 
         return V1PipelineStep(

lightning_sdk/services/file_endpoint.py

@@ -8,7 +8,6 @@ import urllib3
 
 from lightning_sdk.api.utils import _FileUploader
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.lightning_cloud.openapi import CommandArgumentCommandArgumentType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.services.utilities import _get_cluster, _get_project, _get_service_url
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
@@ -162,7 +161,7 @@ class Client:
 class Argument:
     """A holder for the service argument."""
 
-    def __init__(self, name: str, type: CommandArgumentCommandArgumentType, **kwargs: Any) -> None:  # noqa: A002
+    def __init__(self, name: str, type: str, **kwargs: Any) -> None:  # noqa: A002
         self._name = name
         self._type = type
         self._value = None
@@ -171,12 +170,12 @@ class Argument:
     @property
     def is_text(self) -> bool:
         """Whether this argument is of type Text."""
-        return self._type == CommandArgumentCommandArgumentType.TEXT
+        return self._type == "Text"
 
     @property
     def is_file(self) -> bool:
         """Whether this argument is of type File."""
-        return self._type == CommandArgumentCommandArgumentType.FILE
+        return self._type == "File"
 
     @property
     def value(self) -> Any:

{lightning_sdk-2025.8.1.dist-info → lightning_sdk-2025.8.6rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.8.1
+Version: 2025.8.6rc1
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License