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

lightning_sdk/__init__.py

@@ -1,21 +1,31 @@
+from lightning_sdk.agents import Agent
+from lightning_sdk.ai_hub import AIHub
+from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noqa: F401
+from lightning_sdk.helpers import _check_version_and_prompt_upgrade
+from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
 from lightning_sdk.organization import Organization
-from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin
+from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
 from lightning_sdk.status import Status
 from lightning_sdk.studio import Studio
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
 
 __all__ = [
+    "Job",
     "JobsPlugin",
     "Machine",
     "MultiMachineTrainingPlugin",
     "Organization",
     "Plugin",
+    "SlurmJobsPlugin",
     "Status",
     "Studio",
     "Teamspace",
     "User",
+    "Agent",
+    "AIHub",
 ]
 
-__version__ = "0.1.3"
+__version__ = "0.1.46"
+_check_version_and_prompt_upgrade(__version__)

lightning_sdk/agents.py

@@ -0,0 +1,46 @@
+from typing import List, Optional
+
+from lightning_sdk.api.agents_api import AgentApi
+
+
+class Agent:
+    def __init__(self, agent_id: str) -> None:
+        self.id = agent_id
+        self._agent_api = AgentApi()
+
+        try:
+            self._agent = self._agent_api.get_agent(agent_id)
+        except ValueError as e:
+            raise ValueError(f"Agent {agent_id}") from e
+
+    def update(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        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,
+    ) -> None:
+        """Update the agent and its endpoint."""
+        self._agent_api.update_agent_endpoint(
+            teamspace_id=self._agent.project_id, endpoint_id=self._agent.endpoint_id, base_url=base_url, api_key=api_key
+        )
+        agent = self._agent_api.update_agent(
+            agent_id=self.id,
+            teamspace_id=self._agent.project_id,
+            name=name,
+            model=model,
+            description=description,
+            prompt_template=prompt_template,
+            prompt_suggestions=prompt_suggestions,
+            knowledge=knowledge,
+            publish_status=publish_status,
+        )
+        self._agent = agent
+
+    def delete(self) -> None:
+        self._agent_api.delete_agent(agent_id=self.id, teamspace_id=self._agent.project_id)

lightning_sdk/ai_hub.py

@@ -0,0 +1,185 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from urllib.parse import quote
+
+from lightning_sdk.api import AIHubApi, UserApi
+from lightning_sdk.lightning_cloud import login
+from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace
+
+if TYPE_CHECKING:
+    from lightning_sdk import Organization, Teamspace
+
+
+class AIHub:
+    """An interface to interact with the AI Hub.
+
+    Example:
+        from lightning_sdk import AIHub
+        hub = AIHub()
+
+        # List public API templates
+        api_list = hub.list_apis()
+
+        # Get detailed information about an API template
+        api_info =  hub.api_info("temp_xxxx")
+
+        # Deploy an API template
+        deployment = hub.deploy("temp_xxxx")
+    """
+
+    def __init__(self) -> None:
+        self._api = AIHubApi()
+        self._auth = None
+
+    def api_info(self, api_id: str) -> dict:
+        """Get full API template info such as input details.
+
+        Example:
+            ai_hub = AIHub()
+            api_info = ai_hub.api_info("temp_xxxx")
+
+        Args:
+            api_id: The ID of the API for which information is requested.
+
+        Returns:
+            A dictionary containing detailed information about the API,
+            including its name, description, creation and update timestamps,
+            parameters, tags, job specifications, and autoscaling settings.
+        """
+        template, api_arguments = self._api.api_info(api_id)
+
+        return {
+            "name": template.name,
+            "description": template.description,
+            "created_at": template.created_at,
+            "updated_at": template.updated_at,
+            "api_arguments": api_arguments,
+            "tags": [tag.name for tag in template.tags],
+            "job": {
+                "image": template.spec_v2.job.image,
+                "interruptible": template.spec_v2.job.spot,
+                "instance_type": template.spec_v2.job.instance_type,
+                "resources": template.spec_v2.job.resources,
+            },
+            "autoscaling": {
+                "enabled": template.spec_v2.autoscaling.enabled,
+                "min_replicas": template.spec_v2.autoscaling.min_replicas,
+                "max_replicas": template.spec_v2.autoscaling.max_replicas,
+            },
+        }
+
+    def list_apis(
+        self,
+        search: Optional[str] = None,
+    ) -> List[Dict[str, str]]:
+        """Get a list of public AI Hub API templates.
+
+        Example:
+            ai_hub = AIHub()
+            api_list = ai_hub.list_apis(search="Llama")
+
+        Args:
+            search: A search query to filter the list of APIs. Defaults to None.
+
+        Returns:
+            A list of dictionaries, each containing information about an API,
+            such as its ID, name, description, creator's username, and creation timestamp.
+        """
+        search_query = search or ""
+        api_templates = self._api.list_apis(search_query=search_query)
+        results = []
+        for template in api_templates:
+            result = {
+                "id": template.id,
+                "name": template.name,
+                "description": template.description,
+                "creator_username": template.creator_username,
+            }
+            results.append(result)
+        return results
+
+    def _authenticate(
+        self,
+        teamspace: Optional[Union[str, "Teamspace"]] = None,
+        org: Optional[Union[str, "Organization"]] = None,
+        user: Optional[Union[str, "User"]] = None,
+    ) -> "Teamspace":
+        if self._auth is None:
+            self._auth = login.Auth()
+        try:
+            self._auth.authenticate()
+            user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
+        except ConnectionError as e:
+            raise e
+
+        org = _resolve_org(org)
+        teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user if org is None else None)
+        if teamspace is None:
+            raise ValueError("You need to pass a teamspace or an org for your deployment.")
+        return teamspace
+
+    def run(
+        self,
+        api_id: str,
+        api_arguments: Optional[Dict[str, Any]] = None,
+        name: Optional[str] = None,
+        cloud_account: Optional[str] = None,
+        teamspace: Optional[Union[str, "Teamspace"]] = None,
+        org: Optional[Union[str, "Organization"]] = None,
+    ) -> Dict[str, Union[str, bool]]:
+        """Deploy an API from the AI Hub.
+
+        Example:
+            from lightning_sdk import AIHub
+            hub = AIHub()
+            deployment = hub.run("temp_xxxx")
+
+            # Using API arguments
+            api_arugments = {"model": "unitary/toxic-bert", "batch_size": 10, "token": "lit_xxxx"}
+            deployment = hub.run("temp_xxxx", api_arugments=api_arugments)
+
+        Args:
+            api_id: The ID of the AIHub template you want to run.
+            api_arguments: Additional API argument, such as model name, or batch size.
+            name: Name for the deployed API. Defaults to None.
+            cloud_account: The cloud account where you want to run the template, such as "lightning-public-prod".
+                Defaults to None.
+            teamspace: The team or group for deployment. Defaults to None.
+            org: The organization for deployment. Defaults to None.
+
+        Returns:
+            A dictionary containing the name of the deployed API,
+            the URL to access it, and whether it is interruptible.
+
+        Raises:
+            ValueError: If a teamspace or organization is not provided.
+            ConnectionError: If there is an issue with logging in.
+        """
+        teamspace = self._authenticate(teamspace, org)
+        teamspace_id = teamspace.id
+
+        api_arguments = api_arguments or {}
+        deployment = self._api.run_api(
+            template_id=api_id,
+            cloud_account=cloud_account or "",
+            project_id=teamspace_id,
+            name=name,
+            api_arguments=api_arguments,
+        )
+        url = (
+            quote(
+                f"{LIGHTNING_CLOUD_URL}/{teamspace._org.name}/{teamspace.name}/jobs/{deployment.name}",
+                safe=":/()",
+            )
+            + "?app_id=deployment"
+        )
+        print("Deployment available at:", url)
+
+        return {
+            "id": deployment.id,
+            "name": deployment.name,
+            "deployment_url": url,
+            "api_endpoint": deployment.status.urls[0],
+            "interruptible": deployment.spec.spot,
+        }

lightning_sdk/api/__init__.py

@@ -1,3 +1,5 @@
+from lightning_sdk.api.agents_api import AgentApi
+from lightning_sdk.api.ai_hub_api import AIHubApi
 from lightning_sdk.api.org_api import OrgApi
 from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.teamspace_api import TeamspaceApi
@@ -8,4 +10,6 @@ __all__ = [
     "StudioApi",
     "TeamspaceApi",
     "UserApi",
+    "AgentApi",
+    "AIHubApi",
 ]

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,130 @@
+import traceback
+from typing import Dict, List, Optional, Tuple
+
+import backoff
+
+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
+
+
+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 "True")
+                        if p.checkbox.is_checked
+                        else (p.checkbox.false_value or "False")
+                    )
+
+        for p in parameters:
+            name = p.name
+            pattern = f"${{{name}}}"
+            if name in api_arguments:
+                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]
+    ) -> V1Deployment:
+        template = self._client.deployment_templates_service_get_deployment_template(template_id)
+        name = name or template.name
+        template.spec_v2.endpoint.id = None
+
+        AIHubApi._set_parameters(template.spec_v2.job, template.parameter_spec.parameters, api_arguments)
+        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=0,
+                spec=template.spec_v2.job,
+            ),
+        )