lightning-sdk 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -85,6 +85,7 @@ class V1UserFeatures(object):
         'mmt_strategy_selector': 'bool',
         'multicloud_saas': 'bool',
         'multiple_studio_versions': 'bool',
+        'nerf_fs_nonpaying': 'bool',
         'org_admin_alerts': 'bool',
         'org_level_member_permissions': 'bool',
         'org_usage_limits': 'bool',
@@ -110,6 +111,7 @@ class V1UserFeatures(object):
         'slurm_machine_selector': 'bool',
         'stop_ide_container_on_shutdown': 'bool',
         'studio_config': 'bool',
+        'studio_deployment': 'bool',
         'studio_on_stop': 'bool',
         'studio_version_visibility': 'bool',
         'studios_dashboard': 'bool',
@@ -167,6 +169,7 @@ class V1UserFeatures(object):
         'mmt_strategy_selector': 'mmtStrategySelector',
         'multicloud_saas': 'multicloudSaas',
         'multiple_studio_versions': 'multipleStudioVersions',
+        'nerf_fs_nonpaying': 'nerfFsNonpaying',
         'org_admin_alerts': 'orgAdminAlerts',
         'org_level_member_permissions': 'orgLevelMemberPermissions',
         'org_usage_limits': 'orgUsageLimits',
@@ -192,6 +195,7 @@ class V1UserFeatures(object):
         'slurm_machine_selector': 'slurmMachineSelector',
         'stop_ide_container_on_shutdown': 'stopIdeContainerOnShutdown',
         'studio_config': 'studioConfig',
+        'studio_deployment': 'studioDeployment',
         'studio_on_stop': 'studioOnStop',
         'studio_version_visibility': 'studioVersionVisibility',
         'studios_dashboard': 'studiosDashboard',
@@ -204,7 +208,7 @@ class V1UserFeatures(object):
         'weka': 'weka'
     }
 
-    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byoc_litcr: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloud_space_environment_templates: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, concurrent_gpu_limit: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, dgx_cloud: 'bool' =None, doc_helper_chat: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_storage_limits: 'bool' =None, enterprise_compute_admin: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filestore: 'bool' =None, gcp_local_disk_binding: 'bool' =None, inactive_notify_delete: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, lambda_labs: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, org_admin_alerts: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, pipelines: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, stop_ide_container_on_shutdown: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, studios_dashboard: 'bool' =None, studios_dashboard_system_metrics: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, voltage_park: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None):  # noqa: E501
+    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byoc_litcr: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloud_space_environment_templates: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, concurrent_gpu_limit: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, dgx_cloud: 'bool' =None, doc_helper_chat: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_storage_limits: 'bool' =None, enterprise_compute_admin: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filestore: 'bool' =None, gcp_local_disk_binding: 'bool' =None, inactive_notify_delete: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, lambda_labs: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, org_admin_alerts: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, pipelines: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, stop_ide_container_on_shutdown: 'bool' =None, studio_config: 'bool' =None, studio_deployment: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, studios_dashboard: 'bool' =None, studios_dashboard_system_metrics: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, voltage_park: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._affiliate_links = None
         self._agents_v2 = None
@@ -250,6 +254,7 @@ class V1UserFeatures(object):
         self._mmt_strategy_selector = None
         self._multicloud_saas = None
         self._multiple_studio_versions = None
+        self._nerf_fs_nonpaying = None
         self._org_admin_alerts = None
         self._org_level_member_permissions = None
         self._org_usage_limits = None
@@ -275,6 +280,7 @@ class V1UserFeatures(object):
         self._slurm_machine_selector = None
         self._stop_ide_container_on_shutdown = None
         self._studio_config = None
+        self._studio_deployment = None
         self._studio_on_stop = None
         self._studio_version_visibility = None
         self._studios_dashboard = None
@@ -374,6 +380,8 @@ class V1UserFeatures(object):
             self.multicloud_saas = multicloud_saas
         if multiple_studio_versions is not None:
             self.multiple_studio_versions = multiple_studio_versions
+        if nerf_fs_nonpaying is not None:
+            self.nerf_fs_nonpaying = nerf_fs_nonpaying
         if org_admin_alerts is not None:
             self.org_admin_alerts = org_admin_alerts
         if org_level_member_permissions is not None:
@@ -424,6 +432,8 @@ class V1UserFeatures(object):
             self.stop_ide_container_on_shutdown = stop_ide_container_on_shutdown
         if studio_config is not None:
             self.studio_config = studio_config
+        if studio_deployment is not None:
+            self.studio_deployment = studio_deployment
         if studio_on_stop is not None:
             self.studio_on_stop = studio_on_stop
         if studio_version_visibility is not None:
@@ -1369,6 +1379,27 @@ class V1UserFeatures(object):
 
         self._multiple_studio_versions = multiple_studio_versions
 
+    @property
+    def nerf_fs_nonpaying(self) -> 'bool':
+        """Gets the nerf_fs_nonpaying of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The nerf_fs_nonpaying of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._nerf_fs_nonpaying
+
+    @nerf_fs_nonpaying.setter
+    def nerf_fs_nonpaying(self, nerf_fs_nonpaying: 'bool'):
+        """Sets the nerf_fs_nonpaying of this V1UserFeatures.
+
+
+        :param nerf_fs_nonpaying: The nerf_fs_nonpaying of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._nerf_fs_nonpaying = nerf_fs_nonpaying
+
     @property
     def org_admin_alerts(self) -> 'bool':
         """Gets the org_admin_alerts of this V1UserFeatures.  # noqa: E501
@@ -1894,6 +1925,27 @@ class V1UserFeatures(object):
 
         self._studio_config = studio_config
 
+    @property
+    def studio_deployment(self) -> 'bool':
+        """Gets the studio_deployment of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The studio_deployment of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._studio_deployment
+
+    @studio_deployment.setter
+    def studio_deployment(self, studio_deployment: 'bool'):
+        """Sets the studio_deployment of this V1UserFeatures.
+
+
+        :param studio_deployment: The studio_deployment of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._studio_deployment = studio_deployment
+
     @property
     def studio_on_stop(self) -> 'bool':
         """Gets the studio_on_stop of this V1UserFeatures.  # noqa: E501

lightning_sdk/llm/llm.py

@@ -1,41 +1,64 @@
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple, Union
 
+from lightning_sdk.api import UserApi
 from lightning_sdk.api.llm_api import LLMApi
+from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import V1Assistant
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.organization import Organization
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
 
 
 class LLM:
-    def __init__(self, name: str) -> None:
+    def __init__(
+        self,
+        name: str,
+        user: Union[str, "User", None] = None,
+        org: Union[str, "Organization", None] = None,
+    ) -> None:
+        self._auth = Auth()
+        self._user = None
+
+        try:
+            self._auth.authenticate()
+            self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
+        except ConnectionError as e:
+            raise e
+
+        self._name = name
+        try:
+            self._user = _resolve_user(self._user or user)
+        except ValueError:
+            self._user = None
+
         self._name = name
         self._org, self._model_name = self._parse_model_name(name)
+        try:
+            # check if it is a org model
+            self._org = _resolve_org(self._org or org)
+        except ApiException:
+            self._org = None
+
         self._llm_api = LLMApi()
-        self._models = self._build_model_lookup(self._llm_api.list_models())
-        self._model_exists()
-        self._public_models = self._build_public_model_lookup(self._get_public_models())
+        self._public_models = self._build_model_lookup(self._get_public_models())
+        self._org_models = self._build_model_lookup(self._get_org_models())
+        self._user_models = self._build_model_lookup(self._get_user_models())
         self._model = self._get_model()
+        self._conversations = {}
 
     def _parse_model_name(self, name: str) -> Tuple[str, str]:
         parts = name.split("/")
-        if len(parts) != 2:
-            raise ValueError(f"Model name must be in the format `organization/model_name`, but got '{name}'.")
-        return parts[0], parts[1]
+        if len(parts) == 1:
+            # a user model or a org model
+            return None, parts[0]
+        if len(parts) == 2:
+            return parts[0], parts[1]
+        raise ValueError(
+            f"Model name must be in the format `organization/model_name` or `model_name`, but got '{name}'."
+        )
 
     def _build_model_lookup(self, endpoints: List[str]) -> Dict[str, Set[str]]:
-        return {endpoint.id: {model.name for model in endpoint.models_metadata} for endpoint in endpoints}
-
-    def _model_exists(self) -> bool:
-        if self._org not in self._models:
-            raise ValueError(
-                f"Model provider {self._org} not found. Available models providers: {list(self._models.keys())}"
-            )
-
-        if self._model_name not in self._models[self._org]:
-            raise ValueError(
-                f"Model {self._model_name} not found. Available models by {self._org}: {self._models[self._org]}"
-            )
-        return True
-
-    def _build_public_model_lookup(self, endpoints: List[str]) -> Dict[str, Set[str]]:
         result = {}
         for endpoint in endpoints:
             result.setdefault(endpoint.model, []).append(endpoint)
@@ -44,13 +67,94 @@ class LLM:
     def _get_public_models(self) -> List[str]:
         return self._llm_api.get_public_models()
 
-    def _get_model(self, public_model: bool = True) -> V1Assistant:
-        # TODO figure out how to identify if model is public or not
-        if not public_model:
-            raise NotImplementedError("Non-public models are not supported yet.")
+    def _get_org_models(self) -> List[str]:
+        return self._llm_api.get_org_models(self._org.id) if self._org else []
+
+    def _get_user_models(self) -> List[str]:
+        return self._llm_api.get_user_models(self._user.id) if self._user else []
+
+    def _get_model(self) -> V1Assistant:
         # TODO how to handle multiple models with same model type? For now, just use the first one
-        return self._public_models.get(self._model_name)[0]
+        if self._model_name in self._public_models:
+            return self._public_models.get(self._model_name)[0]
+        if self._model_name in self._org_models:
+            return self._org_models.get(self._model_name)[0]
+        if self._model_name in self._user_models:
+            return self._user_models.get(self._model_name)[0]
+
+        available_models = []
+        if self._public_models:
+            available_models.append(f"Public Models: {', '.join(self._public_models.keys())}")
 
-    def chat(self, prompt: str, system_prompt: Optional[str] = None) -> str:
-        output = self._llm_api.start_conversation(prompt, system_prompt, self._model.id)
+        if self._org and self._org_models:
+            available_models.append(f"Org ({self._org.name}) Models: {', '.join(self._org_models.keys())}")
+
+        if self._user and self._user_models:
+            available_models.append(f"User ({self._user.name}) Models: {', '.join(self._user_models.keys())}")
+
+        available_models_str = "\n".join(available_models)
+        raise ValueError(f"Model '{self._model_name}' not found. \nAvailable models: \n{available_models_str}")
+
+    def _get_conversations(self) -> Dict[str, str]:
+        # TODO: after updating backend, this will fetch conversations from backend
+        # conversations = self._llm_api.list_conversations(assistant_id=self._model.id)
+        return self._conversations
+
+    def _fetch_conversations(self) -> None:
+        self._conversations = self._get_conversations()
+
+    def chat(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        max_completion_tokens: Optional[int] = 500,
+        conversation: Optional[str] = None,
+    ) -> str:
+        if conversation and conversation not in self._conversations:
+            self._fetch_conversations()
+
+        conversation_id = self._conversations.get(conversation) if conversation else None
+        output = self._llm_api.start_conversation(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            max_completion_tokens=max_completion_tokens,
+            assistant_id=self._model.id,
+            conversation_id=conversation_id,
+        )
+        if conversation and not conversation_id:
+            self._conversations[conversation] = output.conversation_id
         return output.choices[0].delta.content
+
+    def list_conversations(self) -> List[Dict]:
+        self._fetch_conversations()
+        return list(self._conversations.keys())
+
+    def _get_conversation_messages(self, conversation_id: str) -> Optional[str]:
+        return self._llm_api.get_conversation(assistant_id=self._model.id, conversation_id=conversation_id)
+
+    def get_history(self, conversation: str) -> Optional[List[Dict]]:
+        # TODO: after updating backend, this will fetch conversation from backend
+        if conversation not in self._conversations:
+            self._fetch_conversations()
+
+        if conversation not in self._conversations:
+            raise ValueError(
+                f"Conversation '{conversation}' not found. \nAvailable conversations: {self._conversations.keys()}"
+            )
+
+        messages = self._get_conversation_messages(self._conversations[conversation])
+        history = []
+        for message in messages:
+            if message.author.role == "user":
+                history.append({"role": "user", "content": message.content[0].parts[0]})
+            elif message.author.role == "assistant":
+                history.append({"role": "assistant", "content": message.content[0].parts[0]})
+        return history
+
+    def reset_conversation(self, conversation: str) -> None:
+        if conversation in self._conversations:
+            self._llm_api.reset_conversation(
+                assistant_id=self._model.id,
+                conversation_id=self._conversations[conversation],
+            )
+            del self._conversations[conversation]

lightning_sdk/plugin.py

@@ -395,6 +395,25 @@ class SlurmJobsPlugin(_Plugin):
         return resp
 
 
+class CustomPortPlugin(_Plugin):
+    """Plugin handling the port of a given service."""
+
+    _plugin_run_name = "Custom Port"
+    _slug_name = "custom-port"
+
+    def run(self, name: Optional[str] = None, port: int = 8000) -> str:
+        """Starts a new port to the given Studio."""
+        if name is None:
+            name = _run_name("port")
+
+        return self._studio._studio_api.start_new_port(
+            teamspace_id=self._studio._teamspace.id,
+            studio_id=self._studio._studio.id,
+            name=name,
+            port=port,
+        )
+
+
 @runtime_checkable
 class _RunnablePlugin(Protocol):
     _plugin_run_name: str

lightning_sdk/studio.py

@@ -1,7 +1,10 @@
+import glob
 import os
 import warnings
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
+from tqdm.auto import tqdm
+
 from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.constants import _LIGHTNING_DEBUG
@@ -265,6 +268,34 @@ class Studio:
             progress_bar=progress_bar,
         )
 
+    def upload_folder(self, folder_path: str, remote_path: Optional[str] = None, progress_bar: bool = True) -> None:
+        """Uploads a given folder to a remote path on the Studio."""
+        if folder_path is None:
+            raise ValueError("Cannot upload a folder that is None.")
+        folder_path = os.path.normpath(folder_path)
+        if os.path.isfile(folder_path):
+            raise NotADirectoryError(f"Cannot upload a file as a folder. '{folder_path}' is a file.")
+        if not os.path.exists(folder_path):
+            raise NotADirectoryError(f"Cannot upload a folder that does not exist. '{folder_path}' is not a directory.")
+        all_files = []
+        for fp in glob.glob(os.path.join(folder_path, "**"), recursive=True):
+            if not os.path.isfile(fp):
+                continue
+            rel_path = os.path.relpath(fp, folder_path)
+            remote_file = os.path.join(remote_path, rel_path) if remote_path else rel_path
+            all_files.append((fp, remote_file))
+
+        if progress_bar:
+            progress_bar = tqdm(total=len(all_files), desc="Uploading files", unit="file")
+        for local_file, remote_path in sorted(all_files, key=lambda p: p[1]):
+            if progress_bar:
+                progress_bar.set_description(f"Uploading {local_file}")
+            self.upload_file(local_file, remote_path=remote_path, progress_bar=False)
+            if progress_bar:
+                progress_bar.update(1)
+        if progress_bar:
+            progress_bar.close()
+
     def download_file(self, remote_path: str, file_path: Optional[str] = None) -> None:
         """Downloads a file from the Studio to a given target path."""
         if file_path is None:
@@ -445,6 +476,7 @@ class Studio:
     def _add_plugin(self, plugin_name: str) -> None:
         """Adds the just installed plugin to the internal list of plugins."""
         from lightning_sdk.plugin import (
+            CustomPortPlugin,
             InferenceServerPlugin,
             JobsPlugin,
             MultiMachineTrainingPlugin,
@@ -458,6 +490,7 @@ class Studio:
             "jobs": JobsPlugin,
             "multi-machine-training": MultiMachineTrainingPlugin,
             "inference-server": InferenceServerPlugin,
+            "custom-port": CustomPortPlugin,
         }.get(plugin_name, Plugin)
 
         description = self._list_installed_plugins()[plugin_name]

{lightning_sdk-0.2.12.dist-info → lightning_sdk-0.2.14.dist-info}/METADATA

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