lightning-sdk 2025.11.13.post0__py3-none-any.whl

lightning_sdk/api/agents_api.py

@@ -0,0 +1,107 @@
+from typing import List, Optional
+
+from lightning_sdk.lightning_cloud.openapi import (
+    AgentsIdBody,
+    EndpointsIdBody,
+    V1Assistant,
+    V1Endpoint,
+    V1PromptSuggestion,
+    V1UpstreamOpenAI,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+class AgentApi:
+    """Internal API client for handling Agents-related HTTP requests."""
+
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=7)
+
+    def get_agent(self, agent_id: str) -> V1Assistant:
+        """Retrieve the agent by its ID."""
+        try:
+            return self._client.assistants_service_get_assistant(agent_id)
+        except ValueError as e:
+            raise ValueError(f"Agent {agent_id} does not exist") from e
+
+    def delete_agent(self, agent_id: str, teamspace_id: str) -> None:
+        """Delete the agent by its ID."""
+        self._client.assistants_service_delete_assistant(id=agent_id, project_id=teamspace_id)
+
+    def _get_agent_endpoint(self, endpoint_id: str, teampsace_id: str) -> V1Endpoint:
+        return self._client.endpoint_service_get_endpoint(project_id=teampsace_id, ref=endpoint_id)
+
+    def update_agent(
+        self,
+        agent_id: str,
+        teamspace_id: str,
+        name: Optional[str] = None,
+        model: Optional[str] = None,
+        description: Optional[str] = None,
+        prompt_template: Optional[str] = None,
+        prompt_suggestions: Optional[List[str]] = None,
+        knowledge: Optional[str] = None,
+        publish_status: Optional[str] = None,
+        file_uploads_enabled: Optional[bool] = None,
+    ) -> V1Assistant:
+        """Update the agent with provided details."""
+        agent = self.get_agent(agent_id)
+        body = AgentsIdBody(
+            cloudspace_id=agent.cloudspace_id,
+            cluster_id=agent.cluster_id,
+            description=agent.description,
+            endpoint_id=agent.endpoint_id,
+            knowledge=agent.knowledge,
+            model=agent.model,
+            name=agent.name,
+            org_id=agent.org_id,
+            prompt_suggestions=agent.prompt_suggestions,
+            prompt_template=agent.prompt_template,
+            user_id=agent.user_id,
+            publish_status=agent.publish_status,
+            file_uploads_enabled=agent.file_uploads_enabled,
+        )
+
+        if name is not None:
+            body.name = name
+        if model is not None:
+            body.model = model
+        if description is not None:
+            body.description = description
+        if prompt_template is not None:
+            body.prompt_template = prompt_template
+        if prompt_suggestions is not None:
+            formatted_prompt_suggestions = [V1PromptSuggestion(content=suggestion) for suggestion in prompt_suggestions]
+            body.prompt_suggestions = formatted_prompt_suggestions
+        if knowledge is not None:
+            body.knowledge = knowledge
+        if publish_status is not None:
+            body.publish_status = publish_status
+        if file_uploads_enabled is not None:
+            body.file_uploads_enabled = file_uploads_enabled
+
+        return self._client.assistants_service_update_assistant(
+            id=agent_id,
+            project_id=teamspace_id,
+            body=body,
+        )
+
+    def update_agent_endpoint(
+        self, teamspace_id: str, endpoint_id: str, base_url: Optional[str] = None, api_key: Optional[str] = None
+    ) -> V1Endpoint:
+        """Update the agent endpoint with provided details."""
+        endpoint = self._get_agent_endpoint(teampsace_id=teamspace_id, endpoint_id=endpoint_id)
+
+        body = EndpointsIdBody(
+            name=endpoint.name,
+            openai=V1UpstreamOpenAI(api_key=endpoint.openai.api_key, base_url=endpoint.openai.base_url),
+        )
+
+        if base_url is not None:
+            body.openai.base_url = base_url
+        if api_key is not None:
+            body.openai.api_key = api_key
+
+        return self._client.endpoint_service_update_endpoint(
+            project_id=teamspace_id, id=endpoint_id, body=body, async_req=True
+        )

lightning_sdk/api/ai_hub_api.py

@@ -0,0 +1,173 @@
+import traceback
+import warnings
+from typing import Dict, List, Optional, Tuple, Union
+
+import backoff
+
+from lightning_sdk.api.deployment_api import apply_change
+from lightning_sdk.api.utils import _machine_to_compute_name
+from lightning_sdk.lightning_cloud.openapi.models import (
+    CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs,
+    V1Deployment,
+    V1DeploymentTemplate,
+    V1DeploymentTemplateParameter,
+    V1DeploymentTemplateParameterPlacement,
+    V1DeploymentTemplateParameterType,
+    V1JobSpec,
+)
+from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_template_gallery_response import (
+    V1DeploymentTemplateGalleryResponse,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.machine import Machine
+
+
+class AIHubApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=3)
+
+    def api_info(self, api_id: str) -> Tuple[V1DeploymentTemplate, List[Dict[str, str]]]:
+        try:
+            template = self._client.deployment_templates_service_get_deployment_template(api_id)
+        except Exception as e:
+            stack_trace = traceback.format_exc()
+            if "record not found" in stack_trace:
+                raise ValueError(f"api_id={api_id} not found.") from None
+            raise e
+
+        api_arguments = []
+        for param in template.parameter_spec.parameters:
+            default = None
+            if param.type == V1DeploymentTemplateParameterType.INPUT and param.input:
+                default = param.input.default_value
+            if param.type == V1DeploymentTemplateParameterType.SELECT and param.select:
+                default = param.select.options[0]
+            if param.type == V1DeploymentTemplateParameterType.CHECKBOX and param.checkbox:
+                default = (
+                    (param.checkbox.true_value or "True")
+                    if param.checkbox.is_checked
+                    else (param.checkbox.false_value or "False")
+                )
+
+            api_arguments.append(
+                {
+                    "name": param.name,
+                    "short_description": param.short_description,
+                    "required": param.required,
+                    "type": param.type,
+                    "default": default,
+                }
+            )
+        return template, api_arguments
+
+    @backoff.on_predicate(backoff.expo, lambda x: not x, max_tries=5)
+    def list_apis(self, search_query: str) -> List[V1DeploymentTemplateGalleryResponse]:
+        kwargs = {"show_globally_visible": True}
+        return self._client.deployment_templates_service_list_published_deployment_templates(
+            search_query=search_query, **kwargs
+        ).templates
+
+    @staticmethod
+    def _update_parameters(
+        job: V1JobSpec, placements: List[V1DeploymentTemplateParameterPlacement], pattern: str, value: str
+    ) -> None:
+        for placement in placements:
+            if placement == V1DeploymentTemplateParameterPlacement.COMMAND:
+                job.command = job.command.replace(pattern, str(value))
+            if placement == V1DeploymentTemplateParameterPlacement.ENTRYPOINT:
+                job.entrypoint = job.entrypoint.replace(pattern, str(value))
+
+            if placement == V1DeploymentTemplateParameterPlacement.ENV:
+                for e in job.env:
+                    if e.value == pattern:
+                        e.value = str(value)
+
+    @staticmethod
+    def _set_parameters(
+        job: V1JobSpec, parameters: List[V1DeploymentTemplateParameter], api_arguments: Dict[str, str]
+    ) -> V1JobSpec:
+        for p in parameters:
+            if p.name not in api_arguments:
+                if p.type == V1DeploymentTemplateParameterType.INPUT and p.input and p.input.default_value:
+                    api_arguments[p.name] = p.input.default_value
+
+                if p.type == V1DeploymentTemplateParameterType.SELECT and p.select and len(p.select.options) > 0:
+                    api_arguments[p.name] = p.select.options[0]
+
+                if p.type == V1DeploymentTemplateParameterType.CHECKBOX and p.checkbox:
+                    api_arguments[p.name] = (
+                        (p.checkbox.true_value or "") if p.checkbox.is_checked else (p.checkbox.false_value or "")
+                    )
+
+        for p in parameters:
+            name = p.name
+            pattern = f"${{{name}}}"
+            if name in api_arguments:
+                if p.type == V1DeploymentTemplateParameterType.CHECKBOX and p.checkbox:
+                    api_arguments[p.name] = (
+                        (p.checkbox.true_value or "") if api_arguments[name] is True else (p.checkbox.false_value or "")
+                    )
+                AIHubApi._update_parameters(job, p.placements, pattern, api_arguments[name])
+            elif not p.required:
+                AIHubApi._update_parameters(job, p.placements, pattern, "")
+            else:
+                raise ValueError(f"API requires argument '{p.name}' but is not provided with api_arguments.")
+
+        return job
+
+    def run_api(
+        self,
+        template_id: str,
+        project_id: str,
+        cloud_account: str,
+        name: Optional[str],
+        api_arguments: Dict[str, str],
+        machine: Optional[Union[str, Machine]],
+        quantity: Optional[int],
+    ) -> V1Deployment:
+        template = self._client.deployment_templates_service_get_deployment_template(template_id)
+        name = name or template.name
+        template.spec_v2.endpoint.id = None
+
+        # These are needed to ensure templates with a max replicas of 0 will start on creation
+        if template.spec_v2.autoscaling.max_replicas == "0":
+            template.spec_v2.autoscaling.max_replicas = "1"
+        if not template.spec_v2.autoscaling.enabled:
+            template.spec_v2.autoscaling.enabled = True
+
+        AIHubApi._set_parameters(template.spec_v2.job, template.parameter_spec.parameters, api_arguments)
+        if machine and isinstance(machine, Machine):
+            apply_change(template.spec_v2.job, "instance_name", _machine_to_compute_name(machine))
+            apply_change(template.spec_v2.job, "instance_type", _machine_to_compute_name(machine))
+        elif machine and isinstance(machine, str):
+            apply_change(template.spec_v2.job, "instance_name", machine)
+            apply_change(template.spec_v2.job, "instance_type", machine)
+
+        if quantity != template.spec_v2.job.quantity:
+            # If the quantity is different from the published template quantity, override it with warnging
+            warnings.warn(
+                "Overriding the quantity of the template with the provided quantity. "
+                "This may result in unexpected behavior. "
+                f"Please verify the template (https://lightning.ai/lightning-ai/ai-hub/{template_id}) "
+                "and asscoiated parameters before running."
+            )
+            apply_change(template.spec_v2.job, "quantity", quantity)
+
+        # Override the cluster_id with the cloud_account if it is provided
+        if len(cloud_account) > 0:
+            apply_change(template.spec_v2.job, "cluster_id", cloud_account)
+        return self._client.jobs_service_create_deployment(
+            project_id=project_id,
+            body=CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs(
+                autoscaling=template.spec_v2.autoscaling,
+                cluster_id=cloud_account,
+                endpoint=template.spec_v2.endpoint,
+                name=name,
+                replicas=1,
+                spec=template.spec_v2.job,
+                parent_template_id=template_id,
+            ),
+        )
+
+    def delete_api(self, deployment_id: str, teamspace_id: str) -> None:
+        self._client.jobs_service_delete_deployment(project_id=teamspace_id, id=deployment_id)

lightning_sdk/api/base_studio_api.py

@@ -0,0 +1,90 @@
+from typing import Any, List, Optional
+
+from lightning_sdk.lightning_cloud.openapi import (
+    V1CloudSpaceEnvironmentType,
+    V1ListCloudSpaceEnvironmentTemplatesResponse,
+)
+from lightning_sdk.lightning_cloud.openapi.models.update import Update as BaseStudioUpdateBody
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_template import (
+    V1CloudSpaceEnvironmentTemplate,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+class BaseStudioApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(retry=False, max_tries=0)
+
+    def get_base_studio(self, base_studio_id: str, org_id: Optional[str] = None) -> V1CloudSpaceEnvironmentTemplate:
+        """Retrieve the base studio by its ID."""
+        try:
+            return self._client.cloud_space_environment_template_service_get_cloud_space_environment_template(
+                base_studio_id, org_id=org_id or ""
+            )
+        except ValueError as e:
+            raise ValueError(f"Base studio {base_studio_id} does not exist") from e
+
+    def get_all_base_studios(self, org_id: Optional[str]) -> V1ListCloudSpaceEnvironmentTemplatesResponse:
+        """Retrieve all base studios for a given organization."""
+        result = self._client.cloud_space_environment_template_service_list_managed_cloud_space_environment_templates(
+            org_id=org_id or ""
+        )
+        if org_id is not None:
+            org_templates = (
+                self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
+                    org_id=org_id
+                )
+            )
+            result.templates = result.templates + org_templates.templates
+        return result
+
+    def update_base_studio(
+        self,
+        base_studio_id: str,
+        org_id: str,
+        name: Optional[str] = None,
+        allowed_machines: Optional[List[str]] = None,
+        default_machine: Optional[str] = None,
+        disabled: Optional[bool] = None,
+        environment_type: Optional[V1CloudSpaceEnvironmentType] = None,
+        machine_image_version: Optional[str] = None,
+        setup_script_text: Optional[str] = None,
+    ) -> V1CloudSpaceEnvironmentTemplate:
+        base_studio = self.get_base_studio(base_studio_id, org_id)
+
+        # Get the current configuration for the base studio
+        update_body = BaseStudioUpdateBody(
+            org_id=base_studio.org_id,
+            name=base_studio.name,
+            allowed_machines=base_studio.config.allowed_machines,
+            default_machine=base_studio.config.default_machine,
+            environment_type=base_studio.config.environment_type,
+            machine_image_version=base_studio.config.machine_image_version,
+            setup_script_text=base_studio.config.setup_script_text,
+            disabled=base_studio.disabled,
+        )
+
+        # Apply changes only if the new value is not None
+        apply_change(update_body, "name", name)
+        apply_change(update_body, "allowed_machines", allowed_machines)
+        apply_change(update_body, "default_machine", default_machine)
+        apply_change(update_body, "environment_type", environment_type)
+        apply_change(update_body, "machine_image_version", machine_image_version)
+        apply_change(update_body, "setup_script_text", setup_script_text)
+        apply_change(update_body, "disabled", disabled)
+
+        return self._client.cloud_space_environment_template_service_update_cloud_space_environment_template(
+            id=base_studio_id,
+            body=update_body,
+        )
+
+
+def apply_change(spec: Any, key: str, value: Any) -> bool:
+    if value is None:
+        return False
+
+    if getattr(spec, key) != value:
+        setattr(spec, key, value)
+        return True
+
+    return False

lightning_sdk/api/cloud_account_api.py

@@ -0,0 +1,225 @@
+from functools import lru_cache
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1Cluster,
+    V1CloudProvider,
+    V1ClusterType,
+    V1ExternalCluster,
+    V1ListClusterAcceleratorsResponse,
+    V1ListDefaultClusterAcceleratorsResponse,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import CloudProvider
+    from lightning_sdk.teamspace import ConnectionType
+
+
+class CloudAccountApi:
+    """Internal API client for API requests to cluster endpoints."""
+
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=7)
+
+    def get_cloud_account(self, cloud_account_id: str, teamspace_id: str, org_id: str) -> Externalv1Cluster:
+        """Gets the cluster from given params cluster_id, project_id and owner.
+
+        Args:
+            cloud_account_id: the cloud account to get
+            teamspace_id: the teamspace the cloud_account is supposed to be associated with
+            org_id: The owning org of this teamspace
+
+        """
+        res = self._client.cluster_service_get_cluster(id=cloud_account_id, org_id=org_id, project_id=teamspace_id)
+        if not res:
+            raise ValueError(f"CloudAccount {cloud_account_id} does not exist")
+        return res
+
+    @lru_cache(maxsize=None)  # noqa: B019
+    def list_cloud_accounts(self, teamspace_id: str) -> List[V1ExternalCluster]:
+        """Lists the cloud accounts for a given teamspace.
+
+        Args:
+            teamspace_id: The teamspace to list cloud accounts for
+
+        Returns:
+            A list of cloud accounts
+        """
+        res_project = self._client.cluster_service_list_project_clusters(
+            project_id=teamspace_id,
+        )
+        res_global = self._client.cluster_service_list_clusters(
+            project_id=teamspace_id,
+        )
+
+        # can't use set here because the cloud_accounts are not hashable
+        cloud_accounts = []
+        cloud_account_ids = []
+        for cloud_account in res_project.clusters + res_global.clusters:
+            if cloud_account.id not in cloud_account_ids:
+                cloud_accounts.append(cloud_account)
+                cloud_account_ids.append(cloud_account.id)
+
+        return cloud_accounts
+
+    def get_cloud_account_non_org(self, teamspace_id: str, cloud_account_id: str) -> Optional[V1ExternalCluster]:
+        for cluster in self.list_cloud_accounts(teamspace_id=teamspace_id):
+            if cluster.id == cloud_account_id:
+                return cluster
+
+        return None
+
+    @lru_cache(maxsize=None)  # noqa: B019
+    def list_cloud_account_accelerators(
+        self,
+        teamspace_id: str,
+        cloud_account_id: str,
+        org_id: str,
+    ) -> Union[V1ListClusterAcceleratorsResponse, V1ListDefaultClusterAcceleratorsResponse]:
+        """Lists the accelerators for a given cloud account.
+
+        Args:
+            cloud_account_id: cluster ID to list accelerators for
+        """
+        # map cloud_account to provider
+        cloud_provider = None
+        is_default = True
+        for cloud_account in self.list_cloud_accounts(teamspace_id=teamspace_id):
+            if cloud_account.id == cloud_account_id:
+                is_default = cloud_account.spec.cluster_type == V1ClusterType.GLOBAL
+                cloud_provider = self._get_cloud_account_provider(cloud_account)
+                break
+
+        if cloud_provider is None:
+            raise ValueError(
+                f"Cloud Account {cloud_account_id} is not a default cloud account. Are you in the correct teamspace?"
+            )
+
+        if is_default:
+            res = self._list_default_cluster_accelerators(teamspace_id=teamspace_id, cloud_provider=cloud_provider)
+        else:
+            res = self._client.cluster_service_list_cluster_accelerators(
+                id=cloud_account_id,
+                org_id=org_id,
+            )
+
+        if not res:
+            raise ValueError(f"CloudAccount {cloud_account_id} does not exist")
+        return res
+
+    def _list_default_cluster_accelerators(
+        self, teamspace_id: str, cloud_provider: Union[str, "CloudProvider"]
+    ) -> V1ListDefaultClusterAcceleratorsResponse:
+        return self._client.cluster_service_list_default_cluster_accelerators(
+            project_id=teamspace_id, cloud_provider=str(cloud_provider)
+        )
+
+    @lru_cache(maxsize=None)  # noqa: B019
+    def list_global_cloud_accounts(self, teamspace_id: str) -> List[V1ExternalCluster]:
+        """Lists the accelerators for a given teamspace.
+
+        Args:
+            teamspace_id: id of the teamspace to get the associated cloud_accounts for
+        """
+        cloud_accounts = self.list_cloud_accounts(teamspace_id=teamspace_id)
+        if not cloud_accounts:
+            raise ValueError(f"Teamspace {teamspace_id} does not exist")
+        filtered_cloud_accounts = filter(lambda x: x.spec.cluster_type == V1ClusterType.GLOBAL, cloud_accounts)
+        # TODO: remove aggregate filter once finished
+        filtered_cloud_accounts = filter(
+            lambda x: x.spec.driver != V1CloudProvider.LIGHTNING_AGGREGATE, filtered_cloud_accounts
+        )
+        return list(filtered_cloud_accounts)
+
+    def get_cloud_account_provider_mapping(self, teamspace_id: str) -> Dict["CloudProvider", V1ExternalCluster]:
+        """Gets the cloud account <-> provider mapping."""
+        res = self.list_cloud_accounts(teamspace_id=teamspace_id)
+        cloud_accounts = {cloud_account.id: cloud_account for cloud_account in res}
+        providers = {cloud_account.id: self._get_cloud_account_provider(cloud_account) for cloud_account in res}
+
+        mapping = {}
+        for cloud_account_id, provider in providers.items():
+            if provider is not None:
+                mapping[provider] = cloud_accounts[cloud_account_id]
+        return mapping
+
+    @staticmethod
+    def _get_cloud_account_provider(cloud_account: Optional[V1ExternalCluster]) -> Optional["CloudProvider"]:
+        """Determines the cloud provider based on the cloud_account configuration.
+
+        Args:
+            cloud_account: An optional Externalv1Cluster object containing cluster specifications
+
+        Returns:
+            CloudProvider: The determined cloud provider, defaults to AWS if no match is found
+        """
+        from lightning_sdk.machine import CloudProvider
+
+        if not cloud_account:
+            return CloudProvider.AWS
+
+        if cloud_account.spec and cloud_account.spec.driver:
+            if cloud_account.spec.driver == V1CloudProvider.LIGHTNING:
+                return CloudProvider.LIGHTNING
+
+            if cloud_account.spec.driver == V1CloudProvider.DGX:
+                return CloudProvider.DGX
+
+        if cloud_account.spec:
+            if cloud_account.spec.aws_v1:
+                return CloudProvider.AWS
+            if cloud_account.spec.google_cloud_v1:
+                return CloudProvider.GCP
+            if cloud_account.spec.lambda_labs_v1:
+                return CloudProvider.LAMBDA_LABS
+            if cloud_account.spec.voltage_park_v1:
+                return CloudProvider.VOLTAGE_PARK
+            if cloud_account.spec.nebius_v1:
+                return CloudProvider.NEBIUS
+
+        return None
+
+    def resolve_cloud_account(
+        self,
+        teamspace_id: str,
+        cloud_account: Optional[str],
+        cloud_provider: Optional[Union["CloudProvider", str]],
+        default_cloud_account: Optional[str],
+    ) -> Optional[str]:
+        from lightning_sdk.machine import CloudProvider
+
+        if cloud_provider and not isinstance(cloud_provider, CloudProvider):
+            cloud_provider = CloudProvider(cloud_provider)
+
+        if cloud_account:
+            if cloud_provider:
+                cloud_account_resp = self.get_cloud_account_non_org(teamspace_id, cloud_account)
+                cloud_provider_resp = self._get_cloud_account_provider(cloud_account_resp)
+                if cloud_provider_resp != cloud_provider:
+                    raise RuntimeError(
+                        f"Specified both cloud_provider ({cloud_provider}) and cloud_account ({cloud_account} "
+                        f"has cloud provider {cloud_provider_resp}) which don't match!"
+                    )
+
+            return cloud_account
+
+        if cloud_provider:
+            cloud_account_mapping = self.get_cloud_account_provider_mapping(teamspace_id=teamspace_id)
+            if cloud_provider and cloud_provider in cloud_account_mapping:
+                return cloud_account_mapping[cloud_provider].id
+
+        if default_cloud_account:
+            return default_cloud_account
+
+        return None
+
+    @staticmethod
+    def get_cloud_provider_for_connection_type(connection_type: "ConnectionType") -> "CloudProvider":
+        from lightning_sdk.machine import CloudProvider
+        from lightning_sdk.teamspace import ConnectionType
+
+        if connection_type == ConnectionType.EFS:
+            return CloudProvider.AWS
+
+        raise ValueError(f"ConnectionType {ConnectionType} currently not supported!")