lightning-sdk 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. lightning_sdk/__init__.py +1 -1
  2. lightning_sdk/api/deployment_api.py +7 -2
  3. lightning_sdk/api/license_api.py +28 -6
  4. lightning_sdk/api/llm_api.py +118 -26
  5. lightning_sdk/api/studio_api.py +5 -0
  6. lightning_sdk/cli/configure.py +34 -27
  7. lightning_sdk/cli/connect.py +2 -2
  8. lightning_sdk/cli/deploy/_auth.py +11 -19
  9. lightning_sdk/cli/entrypoint.py +20 -2
  10. lightning_sdk/deployment/deployment.py +17 -3
  11. lightning_sdk/lightning_cloud/login.py +2 -2
  12. lightning_sdk/lightning_cloud/openapi/__init__.py +2 -3
  13. lightning_sdk/lightning_cloud/openapi/api/cluster_service_api.py +1 -5
  14. lightning_sdk/lightning_cloud/openapi/api/endpoint_service_api.py +11 -1
  15. lightning_sdk/lightning_cloud/openapi/api/jobs_service_api.py +121 -0
  16. lightning_sdk/lightning_cloud/openapi/api/user_service_api.py +0 -85
  17. lightning_sdk/lightning_cloud/openapi/models/__init__.py +2 -3
  18. lightning_sdk/lightning_cloud/openapi/models/alertingevents_id_body.py +409 -0
  19. lightning_sdk/lightning_cloud/openapi/models/id_codeconfig_body.py +29 -3
  20. lightning_sdk/lightning_cloud/openapi/models/update.py +105 -1
  21. lightning_sdk/lightning_cloud/openapi/models/v1_author.py +201 -0
  22. lightning_sdk/lightning_cloud/openapi/models/v1_blog_post.py +53 -1
  23. lightning_sdk/lightning_cloud/openapi/models/v1_cloud_space.py +27 -1
  24. lightning_sdk/lightning_cloud/openapi/models/v1_cloud_space_environment_template.py +53 -1
  25. lightning_sdk/lightning_cloud/openapi/models/v1_cloud_space_environment_template_config.py +27 -1
  26. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_deletion_options.py +27 -1
  27. lightning_sdk/lightning_cloud/openapi/models/v1_create_cloud_space_environment_template_request.py +105 -1
  28. lightning_sdk/lightning_cloud/openapi/models/v1_create_project_request.py +79 -1
  29. lightning_sdk/lightning_cloud/openapi/models/v1_data_connection.py +79 -1
  30. lightning_sdk/lightning_cloud/openapi/models/v1_deployment_alerting_policy_type.py +1 -0
  31. lightning_sdk/lightning_cloud/openapi/models/v1_get_organization_storage_metadata_response.py +79 -1
  32. lightning_sdk/lightning_cloud/openapi/models/v1_get_project_storage_metadata_response.py +105 -1
  33. lightning_sdk/lightning_cloud/openapi/models/v1_get_user_response.py +79 -1
  34. lightning_sdk/lightning_cloud/openapi/models/v1_get_user_storage_breakdown_response.py +27 -1
  35. lightning_sdk/lightning_cloud/openapi/models/v1_membership.py +27 -1
  36. lightning_sdk/lightning_cloud/openapi/models/v1_message.py +53 -1
  37. lightning_sdk/lightning_cloud/openapi/models/v1_notification_type.py +1 -0
  38. lightning_sdk/lightning_cloud/openapi/models/v1_organization.py +105 -1
  39. lightning_sdk/lightning_cloud/openapi/models/v1_project_storage.py +131 -1
  40. lightning_sdk/lightning_cloud/openapi/models/v1_routing_telemetry.py +27 -1
  41. lightning_sdk/lightning_cloud/openapi/models/v1_storage_asset.py +27 -1
  42. lightning_sdk/lightning_cloud/openapi/models/v1_storage_asset_type.py +2 -0
  43. lightning_sdk/lightning_cloud/openapi/models/v1_transaction.py +27 -1
  44. lightning_sdk/lightning_cloud/openapi/models/v1_update_user_request.py +79 -1
  45. lightning_sdk/lightning_cloud/openapi/models/v1_usage.py +27 -27
  46. lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py +79 -443
  47. lightning_sdk/lightning_cloud/openapi/models/v1_volume.py +499 -31
  48. lightning_sdk/lightning_cloud/rest_client.py +13 -11
  49. lightning_sdk/lightning_cloud/source_code/logs_socket_api.py +8 -3
  50. lightning_sdk/llm/llm.py +52 -1
  51. lightning_sdk/pipeline/pipeline.py +1 -1
  52. lightning_sdk/services/license.py +78 -22
  53. lightning_sdk/services/utilities.py +15 -1
  54. {lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/METADATA +1 -1
  55. {lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/RECORD +59 -60
  56. lightning_sdk/lightning_cloud/openapi/models/v1_ebs.py +0 -279
  57. lightning_sdk/lightning_cloud/openapi/models/v1_get_user_storage_response.py +0 -201
  58. lightning_sdk/lightning_cloud/openapi/models/v1_reservation_billing_session.py +0 -279
  59. {lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/LICENSE +0 -0
  60. {lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/WHEEL +0 -0
  61. {lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/entry_points.txt +0 -0
  62. {lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/top_level.txt +0 -0
lightning_sdk/__init__.py CHANGED
@@ -31,6 +31,6 @@ __all__ = [
31
31
  "User",
32
32
  ]
33
33
 
34
- __version__ = "0.2.19"
34
+ __version__ = "0.2.21"
35
35
  _check_version_and_prompt_upgrade(__version__)
36
36
  _set_tqdm_envvars_noninteractive()
@@ -282,10 +282,11 @@ class DeploymentApi:
282
282
 
283
283
  requires_release = False
284
284
  requires_release |= apply_change(deployment.spec, "image", image)
285
+
285
286
  requires_release |= apply_change(deployment.spec, "entrypoint", entrypoint)
286
287
  requires_release |= apply_change(deployment.spec, "command", command)
287
288
  requires_release |= apply_change(deployment.spec, "env", to_env(env))
288
- requires_release |= apply_change(deployment.spec, "readiness_probe", to_health_check(health_check))
289
+ requires_release |= apply_change(deployment.spec, "readiness_probe", to_health_check(health_check, False))
289
290
  requires_release |= apply_change(deployment.spec, "cluster_id", cloud_account)
290
291
  requires_release |= apply_change(deployment.spec, "spot", spot)
291
292
  requires_release |= apply_change(deployment.spec, "quantity", quantity)
@@ -524,8 +525,12 @@ def to_endpoint(
524
525
 
525
526
 
526
527
  def to_health_check(
527
- health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None
528
+ health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None,
529
+ use_default: bool = True,
528
530
  ) -> Optional[V1JobHealthCheckConfig]:
531
+ if health_check is None and not use_default:
532
+ return None
533
+
529
534
  # Use Default health check if none is provided
530
535
  if not health_check:
531
536
  return V1JobHealthCheckConfig(
@@ -1,11 +1,35 @@
1
+ import os
1
2
  from typing import Optional
3
+ from urllib.parse import urlencode
2
4
 
5
+ from lightning_sdk.lightning_cloud import env
3
6
  from lightning_sdk.lightning_cloud.rest_client import LightningClient
4
7
 
8
+ LICENSE_CODE = os.environ.get("LICENSE_CODE", "d9s79g79ss")
9
+ # https://lightning.ai/home?settings=licenses
10
+ LICENSE_SIGNING_URL = f"{env.LIGHTNING_CLOUD_URL}?settings=licenses"
11
+
12
+
13
+ def generate_url_user_settings(redirect_to: str = LICENSE_SIGNING_URL) -> str:
14
+ params = urlencode({"redirectTo": redirect_to, "okbhrt": LICENSE_CODE})
15
+ return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
16
+
5
17
 
6
18
  class LicenseApi:
7
- def __init__(self) -> None:
8
- self._client = LightningClient(retry=False, max_tries=0)
19
+ _client_authenticated: LightningClient = None
20
+ _client_public: LightningClient = None
21
+
22
+ @property
23
+ def client_public(self) -> LightningClient:
24
+ if not self._client_public:
25
+ self._client_public = LightningClient(retry=False, max_tries=0, with_auth=False)
26
+ return self._client_public
27
+
28
+ @property
29
+ def client_authenticated(self) -> LightningClient:
30
+ if not self._client_authenticated:
31
+ self._client_authenticated = LightningClient(retry=True, max_tries=3, with_auth=True)
32
+ return self._client_authenticated
9
33
 
10
34
  def valid_license(
11
35
  self,
@@ -25,14 +49,12 @@ class LicenseApi:
25
49
  Returns:
26
50
  True if the license key is valid, False otherwise.
27
51
  """
28
- response, code, _ = self._client.product_license_service_validate_product_license_with_http_info(
52
+ response = self.client_public.product_license_service_validate_product_license(
29
53
  license_key=license_key,
30
54
  product_name=product_name,
31
55
  product_version=product_version,
32
56
  product_type=product_type,
33
57
  )
34
- if code != 200:
35
- raise ConnectionError(f"Failed to validate license key: {code} - {response}")
36
58
  return response.valid
37
59
 
38
60
  def list_user_licenses(self, user_id: str) -> list:
@@ -44,5 +66,5 @@ class LicenseApi:
44
66
  Returns:
45
67
  A list of licenses for the user.
46
68
  """
47
- response = self._client.product_license_service_list_user_licenses(user_id=user_id)
69
+ response = self.client_authenticated.product_license_service_list_user_licenses(user_id=user_id)
48
70
  return response.licenses
@@ -1,7 +1,10 @@
1
+ import asyncio
1
2
  import base64
2
3
  import json
3
4
  import os
4
- from typing import Dict, Generator, List, Optional, Union
5
+ import threading
6
+ import warnings
7
+ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
5
8
 
6
9
  from pip._vendor.urllib3 import HTTPResponse
7
10
 
@@ -27,35 +30,40 @@ class LLMApi:
27
30
  result = self._client.assistants_service_list_assistants(user_id=user_id)
28
31
  return result.assistants
29
32
 
33
+ def _parse_stream_line(self, decoded_line: str) -> Optional[V1ConversationResponseChunk]:
34
+ try:
35
+ payload = json.loads(decoded_line)
36
+ result_data = payload.get("result", {})
37
+
38
+ choices = []
39
+ for choice in result_data.get("choices", []):
40
+ delta = choice.get("delta", {})
41
+ choices.append(
42
+ V1ResponseChoice(
43
+ delta=V1ResponseChoiceDelta(**delta),
44
+ finish_reason=choice.get("finishReason"),
45
+ index=choice.get("index"),
46
+ )
47
+ )
48
+
49
+ return V1ConversationResponseChunk(
50
+ choices=choices,
51
+ conversation_id=result_data.get("conversationId"),
52
+ executable=result_data.get("executable"),
53
+ id=result_data.get("id"),
54
+ throughput=result_data.get("throughput"),
55
+ )
56
+ except json.JSONDecodeError:
57
+ warnings.warn("Error decoding JSON:", decoded_line)
58
+ return None
59
+
30
60
  def _stream_chat_response(self, result: HTTPResponse) -> Generator[V1ConversationResponseChunk, None, None]:
31
61
  for line in result.stream():
32
62
  decoded_lines = line.decode("utf-8").strip()
33
63
  for decoded_line in decoded_lines.splitlines():
34
- try:
35
- payload = json.loads(decoded_line)
36
- result_data = payload.get("result", {})
37
-
38
- choices = []
39
- for choice in result_data.get("choices", []):
40
- delta = choice.get("delta", {})
41
- choices.append(
42
- V1ResponseChoice(
43
- delta=V1ResponseChoiceDelta(**delta),
44
- finish_reason=choice.get("finishReason"),
45
- index=choice.get("index"),
46
- )
47
- )
48
-
49
- yield V1ConversationResponseChunk(
50
- choices=choices,
51
- conversation_id=result_data.get("conversationId"),
52
- executable=result_data.get("executable"),
53
- id=result_data.get("id"),
54
- throughput=result_data.get("throughput"),
55
- )
56
-
57
- except json.JSONDecodeError:
58
- print("Error decoding JSON:", decoded_line)
64
+ chunk = self._parse_stream_line(decoded_line)
65
+ if chunk:
66
+ yield chunk
59
67
 
60
68
  def _encode_image_bytes_to_data_url(self, image: str) -> str:
61
69
  with open(image, "rb") as image_file:
@@ -110,6 +118,90 @@ class LLMApi:
110
118
  return result.result
111
119
  return self._stream_chat_response(result)
112
120
 
121
+ async def async_start_conversation(
122
+ self,
123
+ prompt: str,
124
+ system_prompt: Optional[str],
125
+ max_completion_tokens: int,
126
+ assistant_id: str,
127
+ images: Optional[List[str]] = None,
128
+ conversation_id: Optional[str] = None,
129
+ billing_project_id: Optional[str] = None,
130
+ name: Optional[str] = None,
131
+ metadata: Optional[Dict[str, str]] = None,
132
+ stream: bool = False,
133
+ ) -> Union[V1ConversationResponseChunk, AsyncGenerator[V1ConversationResponseChunk, None]]:
134
+ is_internal_conversation = os.getenv("LIGHTNING_INTERNAL_CONVERSATION", "false").lower() == "true"
135
+ body = {
136
+ "message": {
137
+ "author": {"role": "user"},
138
+ "content": [
139
+ {"contentType": "text", "parts": [prompt]},
140
+ ],
141
+ },
142
+ "max_completion_tokens": max_completion_tokens,
143
+ "conversation_id": conversation_id,
144
+ "billing_project_id": billing_project_id,
145
+ "name": name,
146
+ "stream": stream,
147
+ "metadata": metadata or {},
148
+ "internal_conversation": is_internal_conversation,
149
+ }
150
+ if images:
151
+ for image in images:
152
+ url = image
153
+ if not image.startswith("http"):
154
+ url = self._encode_image_bytes_to_data_url(image)
155
+
156
+ body["message"]["content"].append(
157
+ {
158
+ "contentType": "image",
159
+ "parts": [url],
160
+ }
161
+ )
162
+
163
+ if not stream:
164
+ thread = await asyncio.to_thread(
165
+ self._client.assistants_service_start_conversation, body, assistant_id, async_req=True
166
+ )
167
+ result = await asyncio.to_thread(thread.get)
168
+ return result.result
169
+
170
+ conversation_thread = await asyncio.to_thread(
171
+ self._client.assistants_service_start_conversation,
172
+ body,
173
+ assistant_id,
174
+ async_req=True,
175
+ _preload_content=False,
176
+ )
177
+
178
+ return self.stream_response(conversation_thread)
179
+
180
+ async def stream_response(self, thread: Any) -> AsyncGenerator[V1ConversationResponseChunk, None]:
181
+ loop = asyncio.get_event_loop()
182
+ response = await asyncio.to_thread(thread.get)
183
+
184
+ queue = asyncio.Queue()
185
+
186
+ def enqueue() -> None:
187
+ try:
188
+ for line in response:
189
+ decoded_lines = line.decode("utf-8").strip()
190
+ for decoded_line in decoded_lines.splitlines():
191
+ chunk = self._parse_stream_line(decoded_line)
192
+ if chunk:
193
+ asyncio.run_coroutine_threadsafe(queue.put(chunk), loop)
194
+ finally:
195
+ asyncio.run_coroutine_threadsafe(queue.put(None), loop)
196
+
197
+ threading.Thread(target=enqueue, daemon=True).start()
198
+
199
+ while True:
200
+ item = await queue.get()
201
+ if item is None:
202
+ break
203
+ yield item
204
+
113
205
  def list_conversations(self, assistant_id: str) -> List[str]:
114
206
  result = self._client.assistants_service_list_conversations(assistant_id)
115
207
  return result.conversations
@@ -154,6 +154,11 @@ class StudioApi:
154
154
  @backoff.on_exception(backoff.expo, AttributeError, max_tries=10)
155
155
  def _check_code_status_top_up_restore_finished(self, studio_id: str, teamspace_id: str) -> bool:
156
156
  """Retries checking the top_up_restore_finished value of the code status when there's an AttributeError."""
157
+ if (
158
+ self.get_studio_status(studio_id, teamspace_id) is None
159
+ or self.get_studio_status(studio_id, teamspace_id).in_use is None
160
+ ):
161
+ return False
157
162
  startup_status = self.get_studio_status(studio_id, teamspace_id).in_use.startup_status
158
163
  return startup_status and startup_status.top_up_restore_finished
159
164
 
@@ -17,33 +17,10 @@ def configure() -> None:
17
17
  """Configure access to resources on the Lightning AI platform."""
18
18
 
19
19
 
20
- @configure.command(name="ssh")
21
- @click.option(
22
- "--name",
23
- default=None,
24
- help=(
25
- "The name of the studio to obtain SSH config. "
26
- "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
27
- ),
28
- )
29
- @click.option(
30
- "--teamspace",
31
- default=None,
32
- help=(
33
- "The teamspace the studio is part of. "
34
- "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
35
- "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
36
- ),
37
- )
38
- @click.option(
39
- "--overwrite",
40
- is_flag=True,
41
- flag_value=True,
42
- default=False,
43
- help="Whether to overwrite the SSH key and config if they already exist.",
44
- )
45
- def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
46
- """Get SSH config entry for a studio."""
20
+ def _configure_ssh_internal(
21
+ name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False
22
+ ) -> None:
23
+ """Internal function to configure SSH without Click decorators."""
47
24
  auth = Auth()
48
25
  auth.authenticate()
49
26
  console = Console()
@@ -78,6 +55,36 @@ def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite:
78
55
  console.print(f"SSH config updated at {config_path}")
79
56
 
80
57
 
58
+ @configure.command(name="ssh")
59
+ @click.option(
60
+ "--name",
61
+ default=None,
62
+ help=(
63
+ "The name of the studio to obtain SSH config. "
64
+ "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
65
+ ),
66
+ )
67
+ @click.option(
68
+ "--teamspace",
69
+ default=None,
70
+ help=(
71
+ "The teamspace the studio is part of. "
72
+ "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
73
+ "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
74
+ ),
75
+ )
76
+ @click.option(
77
+ "--overwrite",
78
+ is_flag=True,
79
+ flag_value=True,
80
+ default=False,
81
+ help="Whether to overwrite the SSH key and config if they already exist.",
82
+ )
83
+ def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
84
+ """Get SSH config entry for a studio."""
85
+ _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=overwrite)
86
+
87
+
81
88
  def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
82
89
  """Download a file from a URL."""
83
90
  import requests
@@ -4,7 +4,7 @@ from typing import Optional
4
4
 
5
5
  import click
6
6
 
7
- from lightning_sdk.cli.configure import ssh as configure_ssh
7
+ from lightning_sdk.cli.configure import _configure_ssh_internal
8
8
  from lightning_sdk.cli.studios_menu import _StudiosMenu
9
9
 
10
10
 
@@ -22,7 +22,7 @@ def connect() -> None:
22
22
  )
23
23
  def studio(name: Optional[str], teamspace: Optional[str]) -> None:
24
24
  """Connect to a studio via SSH."""
25
- configure_ssh(name=name, teamspace=teamspace, overwrite=False)
25
+ _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=False)
26
26
 
27
27
  menu = _StudiosMenu()
28
28
  studio = menu._get_studio(name=name, teamspace=teamspace)
@@ -2,8 +2,7 @@ import os
2
2
  import time
3
3
  from datetime import datetime
4
4
  from enum import Enum
5
- from typing import Any, List, Optional, TypedDict
6
- from urllib.parse import urlencode
5
+ from typing import List, Optional, TypedDict
7
6
 
8
7
  from rich.console import Console
9
8
  from rich.prompt import Confirm
@@ -11,7 +10,6 @@ from rich.prompt import Confirm
11
10
  from lightning_sdk import Teamspace
12
11
  from lightning_sdk.api import UserApi
13
12
  from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
14
- from lightning_sdk.lightning_cloud import env
15
13
  from lightning_sdk.lightning_cloud.login import Auth, AuthServer
16
14
  from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
17
15
  from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -26,18 +24,7 @@ class _AuthMode(Enum):
26
24
  DEPLOY = "deploy"
27
25
 
28
26
 
29
- class _AuthServer(AuthServer):
30
- def __init__(self, mode: _AuthMode, *args: Any, **kwargs: Any) -> None:
31
- self._mode = mode
32
- super().__init__(*args, **kwargs)
33
-
34
- def get_auth_url(self, port: int) -> str:
35
- redirect_uri = f"http://localhost:{port}/login-complete"
36
- params = urlencode({"redirectTo": redirect_uri, "mode": self._mode.value, "okbhrt": LITSERVE_CODE})
37
- return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
38
-
39
-
40
- class _Auth(Auth):
27
+ class _AuthLitServe(Auth):
41
28
  def __init__(self, mode: _AuthMode, shall_confirm: bool = False) -> None:
42
29
  super().__init__()
43
30
  self._mode = mode
@@ -51,15 +38,20 @@ class _Auth(Auth):
51
38
  if not proceed:
52
39
  raise RuntimeError(
53
40
  "Login cancelled. Please login to Lightning AI to deploy the API. Run `lightning login` to login."
54
- ) from None
41
+ )
55
42
  print("Opening browser for authentication...")
56
43
  print("Please come back to the terminal after logging in.")
57
44
  time.sleep(3)
58
- _AuthServer(self._mode).login_with_browser(self)
45
+ AuthServer({"mode": self._mode, "okbhrt": LITSERVE_CODE}).login_with_browser(self)
59
46
 
60
47
 
61
48
  def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:
62
- auth = _Auth(mode, shall_confirm)
49
+ """Authenticate with Lightning AI.
50
+
51
+ This will open a browser window for authentication.
52
+ If `shall_confirm` is True, it will ask for confirmation before proceeding.
53
+ """
54
+ auth = _AuthLitServe(mode, shall_confirm)
63
55
  auth.authenticate()
64
56
 
65
57
 
@@ -87,7 +79,7 @@ def poll_verified_status(timeout: int = _POLL_TIMEOUT) -> _UserStatus:
87
79
  user_api = UserApi()
88
80
  user = _get_authed_user()
89
81
  start_time = datetime.now()
90
- result = {"onboarded": False, "verified": False}
82
+ result = _UserStatus(onboarded=False, verified=False)
91
83
  while True:
92
84
  user_resp = user_api.get_user(name=user.name)
93
85
  result["onboarded"] = user_resp.status.completed_project_onboarding
@@ -1,10 +1,12 @@
1
1
  import sys
2
+ import traceback
2
3
  from types import TracebackType
3
4
  from typing import Type
4
5
 
5
6
  import click
6
7
  from rich.console import Console
7
8
  from rich.panel import Panel
9
+ from rich.text import Text
8
10
 
9
11
  from lightning_sdk import __version__
10
12
  from lightning_sdk.api.studio_api import _cloud_url
@@ -26,13 +28,29 @@ from lightning_sdk.cli.start import start
26
28
  from lightning_sdk.cli.stop import stop
27
29
  from lightning_sdk.cli.switch import switch
28
30
  from lightning_sdk.cli.upload import upload
31
+ from lightning_sdk.constants import _LIGHTNING_DEBUG
29
32
  from lightning_sdk.lightning_cloud.login import Auth
30
33
 
31
34
 
32
- def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None: # No
35
+ def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:
33
36
  """CLI won't show tracebacks, just print the exception message."""
37
+ # if debug mode, print the traceback using rich
34
38
  console = Console()
35
- console.print(Panel(value))
39
+ if value.args:
40
+ message = str(value.args[0]) if value.args[0] else str(value)
41
+ else:
42
+ message = str(value) or "An unknown error occurred"
43
+
44
+ error_content = Text()
45
+ error_content.append(f"{exception_type.__name__}: ", style="bold red")
46
+ error_content.append(message, style="white")
47
+
48
+ if _LIGHTNING_DEBUG:
49
+ error_content.append("\n\nFull traceback:\n", style="bold yellow")
50
+ tb_lines = traceback.format_exception(exception_type, value, tb)
51
+ error_content.append("".join(tb_lines), style="dim white")
52
+
53
+ console.print(Panel(error_content, title="⚡ Lightning CLI Error", border_style="red"))
36
54
 
37
55
 
38
56
  @click.group(name="lightning", help="Command line interface (CLI) to interact with/manage Lightning AI Studios.")
@@ -276,8 +276,8 @@ class Deployment:
276
276
  cloud_account=cloud_account,
277
277
  machine=machine,
278
278
  image=image,
279
- entrypoint=entrypoint or "",
280
- command=command or "",
279
+ entrypoint=entrypoint,
280
+ command=command,
281
281
  ports=ports,
282
282
  custom_domain=custom_domain,
283
283
  auth=auth,
@@ -340,7 +340,7 @@ class Deployment:
340
340
  return None
341
341
 
342
342
  @property
343
- def readiness_probe(self) -> Optional[Union[HttpHealthCheck, ExecHealthCheck]]:
343
+ def health_check(self) -> Optional[Union[HttpHealthCheck, ExecHealthCheck]]:
344
344
  """The health check to validate the replicas are ready to receive traffic."""
345
345
  if self._deployment:
346
346
  self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
@@ -467,6 +467,20 @@ class Deployment:
467
467
  return self._deployment.spec.image
468
468
  return None
469
469
 
470
+ @property
471
+ def entrypoint(self) -> Optional[str]:
472
+ if self._deployment:
473
+ self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
474
+ return self._deployment.spec.entrypoint
475
+ return None
476
+
477
+ @property
478
+ def command(self) -> Optional[str]:
479
+ if self._deployment:
480
+ self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
481
+ return self._deployment.spec.command
482
+ return None
483
+
470
484
  @property
471
485
  def _session(self) -> Any:
472
486
  if self._request_session is None:
@@ -43,8 +43,8 @@ class Auth:
43
43
  for key in Keys:
44
44
  setattr(self, key.suffix, os.environ.get(key.value, None))
45
45
 
46
- self._with_env_var = bool(
47
- self.user_id and self.api_key) # used by authenticate method
46
+ # used by authenticate method
47
+ self._with_env_var = bool(self.user_id and self.api_key)
48
48
  if self.api_key and not self.user_id:
49
49
  raise ValueError(
50
50
  f"{Keys.USER_ID.value} is missing from env variables. "
@@ -69,6 +69,7 @@ from lightning_sdk.lightning_cloud.openapi.configuration import Configuration
69
69
  from lightning_sdk.lightning_cloud.openapi.models.affiliatelinks_id_body import AffiliatelinksIdBody
70
70
  from lightning_sdk.lightning_cloud.openapi.models.agentmanagedendpoints_id_body import AgentmanagedendpointsIdBody
71
71
  from lightning_sdk.lightning_cloud.openapi.models.agents_id_body import AgentsIdBody
72
+ from lightning_sdk.lightning_cloud.openapi.models.alertingevents_id_body import AlertingeventsIdBody
72
73
  from lightning_sdk.lightning_cloud.openapi.models.alerts_config_billing import AlertsConfigBilling
73
74
  from lightning_sdk.lightning_cloud.openapi.models.alerts_config_studios import AlertsConfigStudios
74
75
  from lightning_sdk.lightning_cloud.openapi.models.app_id_works_body import AppIdWorksBody
@@ -256,6 +257,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_assistant import V1Assistan
256
257
  from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_knowledge_item_status import V1AssistantKnowledgeItemStatus
257
258
  from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_knowledge_status import V1AssistantKnowledgeStatus
258
259
  from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_model_status import V1AssistantModelStatus
260
+ from lightning_sdk.lightning_cloud.openapi.models.v1_author import V1Author
259
261
  from lightning_sdk.lightning_cloud.openapi.models.v1_auto_join_domain_validation import V1AutoJoinDomainValidation
260
262
  from lightning_sdk.lightning_cloud.openapi.models.v1_auto_join_org_response import V1AutoJoinOrgResponse
261
263
  from lightning_sdk.lightning_cloud.openapi.models.v1_autoscaling_spec import V1AutoscalingSpec
@@ -485,7 +487,6 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_drive_status import V1Drive
485
487
  from lightning_sdk.lightning_cloud.openapi.models.v1_drive_type import V1DriveType
486
488
  from lightning_sdk.lightning_cloud.openapi.models.v1_drive_type_spec import V1DriveTypeSpec
487
489
  from lightning_sdk.lightning_cloud.openapi.models.v1_drive_type_status import V1DriveTypeStatus
488
- from lightning_sdk.lightning_cloud.openapi.models.v1_ebs import V1Ebs
489
490
  from lightning_sdk.lightning_cloud.openapi.models.v1_efs_config import V1EfsConfig
490
491
  from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint import V1Endpoint
491
492
  from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint_auth import V1EndpointAuth
@@ -562,7 +563,6 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_get_user_balance_response i
562
563
  from lightning_sdk.lightning_cloud.openapi.models.v1_get_user_notification_preferences_response import V1GetUserNotificationPreferencesResponse
563
564
  from lightning_sdk.lightning_cloud.openapi.models.v1_get_user_response import V1GetUserResponse
564
565
  from lightning_sdk.lightning_cloud.openapi.models.v1_get_user_storage_breakdown_response import V1GetUserStorageBreakdownResponse
565
- from lightning_sdk.lightning_cloud.openapi.models.v1_get_user_storage_response import V1GetUserStorageResponse
566
566
  from lightning_sdk.lightning_cloud.openapi.models.v1_git_credentials import V1GitCredentials
567
567
  from lightning_sdk.lightning_cloud.openapi.models.v1_google_cloud_direct_v1 import V1GoogleCloudDirectV1
568
568
  from lightning_sdk.lightning_cloud.openapi.models.v1_google_cloud_direct_v1_status import V1GoogleCloudDirectV1Status
@@ -831,7 +831,6 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_report_restart_timings_resp
831
831
  from lightning_sdk.lightning_cloud.openapi.models.v1_request_cluster_access_request import V1RequestClusterAccessRequest
832
832
  from lightning_sdk.lightning_cloud.openapi.models.v1_request_cluster_access_response import V1RequestClusterAccessResponse
833
833
  from lightning_sdk.lightning_cloud.openapi.models.v1_request_verification_code_response import V1RequestVerificationCodeResponse
834
- from lightning_sdk.lightning_cloud.openapi.models.v1_reservation_billing_session import V1ReservationBillingSession
835
834
  from lightning_sdk.lightning_cloud.openapi.models.v1_reservation_details import V1ReservationDetails
836
835
  from lightning_sdk.lightning_cloud.openapi.models.v1_resource_tag import V1ResourceTag
837
836
  from lightning_sdk.lightning_cloud.openapi.models.v1_resource_visibility import V1ResourceVisibility
@@ -2726,7 +2726,6 @@ class ClusterServiceApi(object):
2726
2726
  >>> result = thread.get()
2727
2727
 
2728
2728
  :param async_req bool
2729
- :param bool include_pricing:
2730
2729
  :param str cloud_provider:
2731
2730
  :param str project_id:
2732
2731
  :return: V1ListDefaultClusterAcceleratorsResponse
@@ -2749,7 +2748,6 @@ class ClusterServiceApi(object):
2749
2748
  >>> result = thread.get()
2750
2749
 
2751
2750
  :param async_req bool
2752
- :param bool include_pricing:
2753
2751
  :param str cloud_provider:
2754
2752
  :param str project_id:
2755
2753
  :return: V1ListDefaultClusterAcceleratorsResponse
@@ -2757,7 +2755,7 @@ class ClusterServiceApi(object):
2757
2755
  returns the request thread.
2758
2756
  """
2759
2757
 
2760
- all_params = ['include_pricing', 'cloud_provider', 'project_id'] # noqa: E501
2758
+ all_params = ['cloud_provider', 'project_id'] # noqa: E501
2761
2759
  all_params.append('async_req')
2762
2760
  all_params.append('_return_http_data_only')
2763
2761
  all_params.append('_preload_content')
@@ -2778,8 +2776,6 @@ class ClusterServiceApi(object):
2778
2776
  path_params = {}
2779
2777
 
2780
2778
  query_params = []
2781
- if 'include_pricing' in params:
2782
- query_params.append(('includePricing', params['include_pricing'])) # noqa: E501
2783
2779
  if 'cloud_provider' in params:
2784
2780
  query_params.append(('cloudProvider', params['cloud_provider'])) # noqa: E501
2785
2781
  if 'project_id' in params:
@@ -1420,6 +1420,8 @@ class EndpointServiceApi(object):
1420
1420
  :param list[str] ids:
1421
1421
  :param bool active_cloudspaces:
1422
1422
  :param bool active_jobs:
1423
+ :param list[str] cloudspace_ids:
1424
+ :param list[str] job_ids:
1423
1425
  :return: V1ListEndpointsResponse
1424
1426
  If the method is called asynchronously,
1425
1427
  returns the request thread.
@@ -1447,12 +1449,14 @@ class EndpointServiceApi(object):
1447
1449
  :param list[str] ids:
1448
1450
  :param bool active_cloudspaces:
1449
1451
  :param bool active_jobs:
1452
+ :param list[str] cloudspace_ids:
1453
+ :param list[str] job_ids:
1450
1454
  :return: V1ListEndpointsResponse
1451
1455
  If the method is called asynchronously,
1452
1456
  returns the request thread.
1453
1457
  """
1454
1458
 
1455
- all_params = ['project_id', 'cloudspace_id', 'auto_start', 'cluster_id', 'ids', 'active_cloudspaces', 'active_jobs'] # noqa: E501
1459
+ all_params = ['project_id', 'cloudspace_id', 'auto_start', 'cluster_id', 'ids', 'active_cloudspaces', 'active_jobs', 'cloudspace_ids', 'job_ids'] # noqa: E501
1456
1460
  all_params.append('async_req')
1457
1461
  all_params.append('_return_http_data_only')
1458
1462
  all_params.append('_preload_content')
@@ -1492,6 +1496,12 @@ class EndpointServiceApi(object):
1492
1496
  query_params.append(('activeCloudspaces', params['active_cloudspaces'])) # noqa: E501
1493
1497
  if 'active_jobs' in params:
1494
1498
  query_params.append(('activeJobs', params['active_jobs'])) # noqa: E501
1499
+ if 'cloudspace_ids' in params:
1500
+ query_params.append(('cloudspaceIds', params['cloudspace_ids'])) # noqa: E501
1501
+ collection_formats['cloudspaceIds'] = 'multi' # noqa: E501
1502
+ if 'job_ids' in params:
1503
+ query_params.append(('jobIds', params['job_ids'])) # noqa: E501
1504
+ collection_formats['jobIds'] = 'multi' # noqa: E501
1495
1505
 
1496
1506
  header_params = {}
1497
1507