lightning-sdk 0.1.56__py3-none-any.whl → 0.1.58__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 (41) hide show
  1. lightning_sdk/__init__.py +3 -2
  2. lightning_sdk/api/lit_container_api.py +82 -20
  3. lightning_sdk/cli/ai_hub.py +61 -10
  4. lightning_sdk/cli/configure.py +110 -65
  5. lightning_sdk/cli/connect.py +32 -16
  6. lightning_sdk/cli/delete.py +81 -32
  7. lightning_sdk/cli/download.py +177 -90
  8. lightning_sdk/cli/entrypoint.py +44 -16
  9. lightning_sdk/cli/generate.py +48 -16
  10. lightning_sdk/cli/inspect.py +43 -3
  11. lightning_sdk/cli/list.py +130 -41
  12. lightning_sdk/cli/run.py +0 -6
  13. lightning_sdk/cli/teamspace_menu.py +1 -1
  14. lightning_sdk/cli/upload.py +13 -3
  15. lightning_sdk/helpers.py +20 -0
  16. lightning_sdk/job/job.py +1 -1
  17. lightning_sdk/lightning_cloud/openapi/__init__.py +5 -0
  18. lightning_sdk/lightning_cloud/openapi/api/data_connection_service_api.py +105 -0
  19. lightning_sdk/lightning_cloud/openapi/api/jobs_service_api.py +113 -0
  20. lightning_sdk/lightning_cloud/openapi/api/lit_logger_service_api.py +4 -4
  21. lightning_sdk/lightning_cloud/openapi/models/__init__.py +5 -0
  22. lightning_sdk/lightning_cloud/openapi/models/agents_id_body.py +105 -1
  23. lightning_sdk/lightning_cloud/openapi/models/deployments_id_body.py +29 -3
  24. lightning_sdk/lightning_cloud/openapi/models/id_visibility_body1.py +149 -0
  25. lightning_sdk/lightning_cloud/openapi/models/model_id_visibility_body.py +27 -1
  26. lightning_sdk/lightning_cloud/openapi/models/setup.py +149 -0
  27. lightning_sdk/lightning_cloud/openapi/models/v1_assistant.py +105 -1
  28. lightning_sdk/lightning_cloud/openapi/models/v1_deployment.py +29 -3
  29. lightning_sdk/lightning_cloud/openapi/models/v1_gcp_data_connection_setup.py +123 -0
  30. lightning_sdk/lightning_cloud/openapi/models/v1_job_spec.py +27 -1
  31. lightning_sdk/lightning_cloud/openapi/models/v1_setup_data_connection_response.py +123 -0
  32. lightning_sdk/lightning_cloud/openapi/models/v1_update_deployment_visibility_response.py +97 -0
  33. lightning_sdk/lightning_cloud/openapi/models/v1_update_metrics_stream_visibility_response.py +27 -1
  34. lightning_sdk/lightning_cloud/openapi/models/v1_update_model_visibility_response.py +27 -1
  35. lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py +79 -1
  36. {lightning_sdk-0.1.56.dist-info → lightning_sdk-0.1.58.dist-info}/METADATA +2 -1
  37. {lightning_sdk-0.1.56.dist-info → lightning_sdk-0.1.58.dist-info}/RECORD +41 -36
  38. {lightning_sdk-0.1.56.dist-info → lightning_sdk-0.1.58.dist-info}/LICENSE +0 -0
  39. {lightning_sdk-0.1.56.dist-info → lightning_sdk-0.1.58.dist-info}/WHEEL +0 -0
  40. {lightning_sdk-0.1.56.dist-info → lightning_sdk-0.1.58.dist-info}/entry_points.txt +0 -0
  41. {lightning_sdk-0.1.56.dist-info → lightning_sdk-0.1.58.dist-info}/top_level.txt +0 -0
lightning_sdk/__init__.py CHANGED
@@ -1,7 +1,7 @@
1
1
  from lightning_sdk.agents import Agent
2
2
  from lightning_sdk.ai_hub import AIHub
3
3
  from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__ # noqa: F401
4
- from lightning_sdk.helpers import _check_version_and_prompt_upgrade
4
+ from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
5
5
  from lightning_sdk.job import Job
6
6
  from lightning_sdk.machine import Machine
7
7
  from lightning_sdk.mmt import MMT
@@ -29,5 +29,6 @@ __all__ = [
29
29
  "AIHub",
30
30
  ]
31
31
 
32
- __version__ = "0.1.56"
32
+ __version__ = "0.1.58"
33
33
  _check_version_and_prompt_upgrade(__version__)
34
+ _set_tqdm_envvars_noninteractive()
@@ -1,4 +1,5 @@
1
- from typing import Any, Callable, Generator, List
1
+ import time
2
+ from typing import Any, Callable, Dict, Generator, Iterator, List
2
3
 
3
4
  import docker
4
5
  import requests
@@ -12,11 +13,11 @@ from lightning_sdk.teamspace import Teamspace
12
13
 
13
14
  class LCRAuthFailedError(Exception):
14
15
  def __init__(self) -> None:
15
- super().__init__(
16
- "Failed to authenticate with Lightning Container Registry. Please login manually "
17
- "using the following command:\n "
18
- "echo $LIGHTNING_API_KEY | docker login litcr.io --username=LIGHTNING_USERNAME --password-stdin"
19
- )
16
+ super().__init__("Failed to authenticate with Lightning Container Registry. Please try again.")
17
+
18
+
19
+ class DockerPushError(Exception):
20
+ pass
20
21
 
21
22
 
22
23
  def retry_on_lcr_auth_failure(func: Callable) -> Callable:
@@ -24,7 +25,7 @@ def retry_on_lcr_auth_failure(func: Callable) -> Callable:
24
25
  try:
25
26
  return func(self, *args, **kwargs)
26
27
  except LCRAuthFailedError:
27
- self.authenticate()
28
+ self.authenticate(reauth=True)
28
29
  return func(self, *args, **kwargs)
29
30
 
30
31
  return wrapper
@@ -37,15 +38,35 @@ class LitContainerApi:
37
38
  try:
38
39
  self._docker_client = docker.from_env()
39
40
  self._docker_client.ping()
41
+ self._docker_auth_config = {}
40
42
  except docker.errors.DockerException as e:
41
43
  raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
42
44
 
43
- def authenticate(self) -> bool:
44
- authed_user = self._client.auth_service_get_user()
45
- username = authed_user.username
46
- api_key = authed_user.api_key
47
- resp = self._docker_client.login(username, password=api_key, registry=_get_registry_url())
48
- return resp["Status"] == "Login Succeeded"
45
+ def authenticate(self, reauth: bool = False) -> bool:
46
+ try:
47
+ authed_user = self._client.auth_service_get_user()
48
+ username = authed_user.username
49
+ api_key = authed_user.api_key
50
+ registry = _get_registry_url()
51
+ resp = self._docker_client.login(username, password=api_key, registry=registry, reauth=reauth)
52
+
53
+ if (
54
+ resp.get("username", None) == username
55
+ and resp.get("password", None) == api_key
56
+ and resp.get("serveraddress", None) == registry
57
+ ):
58
+ self._docker_auth_config = {"username": username, "password": api_key}
59
+ return True
60
+
61
+ # This is a new 200 response auth attempt from the client.
62
+ if "Status" in resp and resp["Status"] == "Login Succeeded":
63
+ self._docker_auth_config = {"username": username, "password": api_key}
64
+ return True
65
+
66
+ return False
67
+ except Exception as e:
68
+ print(f"Authentication error: {e} resp: {resp}")
69
+ return False
49
70
 
50
71
  def list_containers(self, project_id: str) -> List:
51
72
  project = self._client.lit_registry_service_get_lit_project_registry(project_id)
@@ -61,7 +82,15 @@ class LitContainerApi:
61
82
  try:
62
83
  self._docker_client.images.get(container)
63
84
  except docker.errors.ImageNotFound:
64
- raise ValueError(f"Container {container} does not exist") from None
85
+ try:
86
+ self._docker_client.images.pull(container, tag)
87
+ self._docker_client.images.get(container)
88
+ except docker.errors.APIError as e:
89
+ raise ValueError(f"Could not pull container {container}") from e
90
+ except docker.errors.ImageNotFound as e:
91
+ raise ValueError(f"Container {container} does not exist") from e
92
+ except Exception as e:
93
+ raise ValueError(f"Unable to upload {container}") from e
65
94
 
66
95
  registry_url = _get_registry_url()
67
96
  container_basename = container.split("/")[-1]
@@ -69,22 +98,55 @@ class LitContainerApi:
69
98
  tagged = self._docker_client.api.tag(container, repository, tag)
70
99
  if not tagged:
71
100
  raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
72
- lines = self._docker_client.api.push(repository, stream=True, decode=True)
73
- for line in lines:
74
- if "errorDetail" in line and ("authorization failed" in line["error"] or "unauth" in line["error"]):
75
- raise LCRAuthFailedError()
76
- yield line
101
+ yield from self._push_with_retry(repository)
77
102
  yield {
78
103
  "finish": True,
79
104
  "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container_basename}",
80
105
  }
81
106
 
107
+ def _push_with_retry(self, repository: str, max_retries: int = 3) -> Iterator[Dict[str, Any]]:
108
+ def is_auth_error(error_msg: str) -> bool:
109
+ auth_errors = ["unauthorized", "authentication required", "unauth"]
110
+ return any(err in error_msg.lower() for err in auth_errors)
111
+
112
+ def is_timeout_error(error_msg: str) -> bool:
113
+ timeout_errors = ["proxyconnect tcp", "i/o timeout"]
114
+ return any(err in error_msg.lower() for err in timeout_errors)
115
+
116
+ for attempt in range(max_retries):
117
+ try:
118
+ if attempt > 0:
119
+ # This is important, if we don't set reauth here then we just keep using the
120
+ # same authentication context that we know just failed.
121
+ self.authenticate(reauth=True)
122
+ time.sleep(2)
123
+
124
+ lines = self._docker_client.api.push(repository, stream=True, decode=True, auth_config=self._docker_auth_config)
125
+ for line in lines:
126
+ if isinstance(line, dict) and "error" in line:
127
+ error = line["error"]
128
+ if is_auth_error(error) or is_timeout_error(error):
129
+ if attempt < max_retries - 1:
130
+ break
131
+ raise DockerPushError(f"Max retries reached: {error}")
132
+ raise DockerPushError(f"Push error: {error}")
133
+ yield line
134
+ else:
135
+ return
136
+
137
+ except docker.errors.APIError as e:
138
+ if (is_auth_error(str(e)) or is_timeout_error(str(e))) and attempt < max_retries - 1:
139
+ continue
140
+ raise DockerPushError(f"Push failed: {e}") from e
141
+
142
+ raise DockerPushError("Max push retries reached")
143
+
82
144
  @retry_on_lcr_auth_failure
83
145
  def download_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
84
146
  registry_url = _get_registry_url()
85
147
  repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
86
148
  try:
87
- self._docker_client.images.pull(repository, tag=tag)
149
+ self._docker_client.images.pull(repository, tag=tag, auth_config=self._docker_auth_config)
88
150
  except requests.exceptions.HTTPError as e:
89
151
  if "unauthorized" in e.response.text:
90
152
  raise LCRAuthFailedError() from e
@@ -1,4 +1,6 @@
1
- from typing import List, Optional
1
+ from typing import Optional
2
+
3
+ import click
2
4
 
3
5
  from lightning_sdk.ai_hub import AIHub
4
6
  from lightning_sdk.cli.studios_menu import _StudiosMenu
@@ -7,10 +9,7 @@ from lightning_sdk.cli.studios_menu import _StudiosMenu
7
9
  class _AIHub(_StudiosMenu):
8
10
  """Interact with Lightning Studio - AI Hub."""
9
11
 
10
- def __init__(self) -> None:
11
- self._hub = AIHub()
12
-
13
- def api_info(self, api_id: str) -> dict:
12
+ def api_info(self, api_id: str) -> None:
14
13
  """Get full API template info such as input details.
15
14
 
16
15
  Example:
@@ -19,15 +18,15 @@ class _AIHub(_StudiosMenu):
19
18
  Args:
20
19
  api_id: The ID of the API for which information is requested.
21
20
  """
22
- return self._hub.api_info(api_id)
21
+ return api_info(api_id=api_id)
23
22
 
24
- def list_apis(self, search: Optional[str] = None) -> List[dict]:
23
+ def list_apis(self, search: Optional[str] = None) -> None:
25
24
  """List API templates available in the AI Hub.
26
25
 
27
26
  Args:
28
27
  search: Search for API templates by name.
29
28
  """
30
- return self._hub.list_apis(search=search)
29
+ return list_apis(search=search)
31
30
 
32
31
  def deploy(
33
32
  self,
@@ -36,7 +35,7 @@ class _AIHub(_StudiosMenu):
36
35
  name: Optional[str] = None,
37
36
  teamspace: Optional[str] = None,
38
37
  org: Optional[str] = None,
39
- ) -> dict:
38
+ ) -> None:
40
39
  """Deploy an API template from the AI Hub.
41
40
 
42
41
  Args:
@@ -46,4 +45,56 @@ class _AIHub(_StudiosMenu):
46
45
  teamspace: Teamspace to deploy the API to. Defaults to user's default teamspace.
47
46
  org: Organization to deploy the API to. Defaults to user's default organization.
48
47
  """
49
- return self._hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
48
+ return deploy(api_id=api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
49
+
50
+
51
+ @click.group(name="aihub")
52
+ def aihub() -> None:
53
+ """Interact with Lightning Studio - AI Hub."""
54
+
55
+
56
+ # @aihub.command(name="api-info")
57
+ # @click.argument("api-id")
58
+ def api_info(api_id: str) -> None:
59
+ """Get full API template info such as input details.
60
+
61
+ Example:
62
+ lightning aihub api_info API-ID
63
+
64
+ API-ID: The ID of the API for which information is requested.
65
+ """
66
+ ai_hub = AIHub()
67
+ ai_hub.api_info(api_id)
68
+
69
+
70
+ # @aihub.command(name="list-apis")
71
+ # @click.option("--search", default=None, help="Search for API templates by name.")
72
+ def list_apis(search: Optional[str]) -> None:
73
+ """List API templates available in the AI Hub."""
74
+ ai_hub = AIHub()
75
+ ai_hub.list_apis(search=search)
76
+
77
+
78
+ # @aihub.command(name="deploy")
79
+ # @click.argument("api-id")
80
+ # @click.option(
81
+ # "--cloud-account",
82
+ # default=None,
83
+ # help="Cloud Account to deploy the API to. Defaults to user's default cloud account.",
84
+ # )
85
+ # @click.option("--name", default=None, help="Name of the deployed API. Defaults to the name of the API template.")
86
+ # @click.option(
87
+ # "--teamspace",
88
+ # default=None,
89
+ # help="Teamspace to deploy the API to. Defaults to user's default teamspace.",
90
+ # )
91
+ # @click.option(
92
+ # "--org",
93
+ # default=None,
94
+ # help="Organization to deploy the API to. Defaults to user's default organization.",
95
+ # )
96
+ def deploy(
97
+ api_id: str, cloud_account: Optional[str], name: Optional[str], teamspace: Optional[str], org: Optional[str]
98
+ ) -> None:
99
+ ai_hub = AIHub()
100
+ ai_hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
@@ -1,14 +1,98 @@
1
+ import os
1
2
  import platform
2
3
  import uuid
3
4
  from pathlib import Path
4
5
  from typing import Optional, Union
5
6
 
7
+ import click
6
8
  from rich.console import Console
7
9
 
8
- from lightning_sdk.cli.generate import _Generate
10
+ from lightning_sdk.cli.generate import _generate_ssh_config
11
+ from lightning_sdk.cli.studios_menu import _StudiosMenu
9
12
  from lightning_sdk.lightning_cloud.login import Auth
10
13
 
11
14
 
15
+ class _Configure(_StudiosMenu):
16
+ """Configure lightning products."""
17
+
18
+ def ssh(self, name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
19
+ """Get SSH config entry for a studio.
20
+
21
+ Args:
22
+ name: The name of the studio to obtain SSH config.
23
+ If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
24
+ teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
25
+ If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
26
+ overwrite: Whether to overwrite the SSH key and config if they already exist.
27
+ """
28
+ ssh(name=name, teamspace=teamspace, overwrite=overwrite)
29
+
30
+
31
+ @click.group(name="configure")
32
+ def configure() -> None:
33
+ """Configure access to resources on the Lightning AI platform."""
34
+
35
+
36
+ # @configure.command(name="ssh")
37
+ # @click.option(
38
+ # "--name",
39
+ # default=None,
40
+ # help=(
41
+ # "The name of the studio to obtain SSH config. "
42
+ # "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
43
+ # ),
44
+ # )
45
+ # @click.option(
46
+ # "--teamspace",
47
+ # default=None,
48
+ # help=(
49
+ # "The teamspace the studio is part of. "
50
+ # "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
51
+ # "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
52
+ # ),
53
+ # )
54
+ # @click.option(
55
+ # "--overwrite",
56
+ # is_flag=True,
57
+ # flag_value=True,
58
+ # help="Whether to overwrite the SSH key and config if they already exist.",
59
+ # )
60
+ def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
61
+ """Get SSH config entry for a studio."""
62
+ auth = Auth()
63
+ auth.authenticate()
64
+ console = Console()
65
+ ssh_dir = Path.home() / ".ssh"
66
+ ssh_dir.mkdir(parents=True, exist_ok=True)
67
+
68
+ key_path = ssh_dir / "lightning_rsa"
69
+ config_path = ssh_dir / "config"
70
+
71
+ # Check if the SSH key already exists
72
+ if key_path.exists() and (key_path.with_suffix(".pub")).exists() and not overwrite:
73
+ console.print(f"SSH key already exists at {key_path}")
74
+ else:
75
+ _download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name="lightning_rsa", overwrite=overwrite)
76
+ console.print(f"SSH key generated and saved to {key_path}")
77
+
78
+ # Check if the SSH config already contains the required configuration
79
+ menu = _StudiosMenu()
80
+ studio = menu._get_studio(name=name, teamspace=teamspace)
81
+ config_content = _generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name)
82
+ if config_path.exists():
83
+ with config_path.open("r") as config_file:
84
+ # check if the host already exists in the config
85
+ if f"Host {studio.name}" in config_file.read():
86
+ console.print("SSH config already contains the required configuration.")
87
+ return
88
+
89
+ with config_path.open("a") as config_file:
90
+ config_file.write(os.linesep)
91
+ config_file.write(config_content)
92
+ config_file.write(os.linesep)
93
+ console.print(f"SSH config updated at {config_path}")
94
+
95
+
12
96
  def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
13
97
  """Download a file from a URL."""
14
98
  import requests
@@ -26,67 +110,28 @@ def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Op
26
110
  local_path.chmod(0o600)
27
111
 
28
112
 
29
- class _Configure(_Generate):
30
- """Configure lightning products."""
31
-
32
- @staticmethod
33
- def _download_ssh_keys(
34
- api_key: str,
35
- key_id: str = "",
36
- ssh_home: Union[str, Path] = "",
37
- ssh_key_name: str = "lightning_rsa",
38
- overwrite: bool = False,
39
- ) -> None:
40
- if not ssh_home:
41
- ssh_home = Path.home() / ".ssh"
42
- elif isinstance(ssh_home, str):
43
- ssh_home = Path(ssh_home)
44
- if not key_id:
45
- key_id = str(uuid.uuid4())
46
-
47
- path_key = ssh_home / ssh_key_name
48
- path_pub = ssh_home / f"{ssh_key_name}.pub"
49
-
50
- # todo: consider hitting the API to get the key pair directly instead of using wget
51
- _download_file(
52
- f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
53
- path_key,
54
- overwrite=overwrite,
55
- chmod=0o600,
56
- )
57
- _download_file(f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}", path_pub, overwrite=overwrite)
58
-
59
- def ssh(self, overwrite: bool = False, ssh_key_name: str = "lightning_rsa") -> None:
60
- """Get SSH config entry for a studio.
61
-
62
- Args:
63
- overwrite: Whether to overwrite the SSH key and config if they already exist.
64
- ssh_key_name: The name of the SSH key to generate
65
- """
66
- auth = Auth()
67
- auth.authenticate()
68
- console = Console()
69
- ssh_dir = Path.home() / ".ssh"
70
- ssh_dir.mkdir(parents=True, exist_ok=True)
71
-
72
- key_path = ssh_dir / ssh_key_name
73
- config_path = ssh_dir / "config"
74
-
75
- # Check if the SSH key already exists
76
- if key_path.exists() and (key_path.with_suffix(".pub")).exists() and not overwrite:
77
- console.print(f"SSH key already exists at {key_path}")
78
- else:
79
- self._download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name=ssh_key_name, overwrite=overwrite)
80
- console.print(f"SSH key generated and saved to {key_path}")
81
-
82
- # Check if the SSH config already contains the required configuration
83
- config_content = self._generate_ssh_config(str(key_path))
84
- if config_path.exists():
85
- with config_path.open("r") as config_file:
86
- if config_content.strip() in config_file.read():
87
- console.print("SSH config already contains the required configuration.")
88
- return
89
-
90
- with config_path.open("a") as config_file:
91
- config_file.write(config_content)
92
- console.print(f"SSH config updated at {config_path}")
113
+ def _download_ssh_keys(
114
+ api_key: str,
115
+ key_id: str = "",
116
+ ssh_home: Union[str, Path] = "",
117
+ ssh_key_name: str = "lightning_rsa",
118
+ overwrite: bool = False,
119
+ ) -> None:
120
+ if not ssh_home:
121
+ ssh_home = Path.home() / ".ssh"
122
+ elif isinstance(ssh_home, str):
123
+ ssh_home = Path(ssh_home)
124
+ if not key_id:
125
+ key_id = str(uuid.uuid4())
126
+
127
+ path_key = ssh_home / ssh_key_name
128
+ path_pub = ssh_home / f"{ssh_key_name}.pub"
129
+
130
+ # todo: consider hitting the API to get the key pair directly instead of using wget
131
+ _download_file(
132
+ f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
133
+ path_key,
134
+ overwrite=overwrite,
135
+ chmod=0o600,
136
+ )
137
+ _download_file(f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}", path_pub, overwrite=overwrite)
@@ -2,11 +2,13 @@ import subprocess
2
2
  import sys
3
3
  from typing import Optional
4
4
 
5
- from lightning_sdk.cli.configure import _Configure
6
- from lightning_sdk.lightning_cloud.login import Auth
5
+ import click
7
6
 
7
+ from lightning_sdk.cli.configure import ssh as configure_ssh
8
+ from lightning_sdk.cli.studios_menu import _StudiosMenu
8
9
 
9
- class _Connect(_Configure):
10
+
11
+ class _Connect:
10
12
  """Connect to lightning products."""
11
13
 
12
14
  def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
@@ -16,16 +18,30 @@ class _Connect(_Configure):
16
18
  name: The name of the studio to connect to.
17
19
  teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
18
20
  """
19
- auth = Auth()
20
- auth.authenticate() # this is maybe not needed
21
- studio = self._get_studio(name=name, teamspace=teamspace)
22
- host = "ssh.lightning.ai"
23
- username = f"s_{studio._studio.id}"
24
-
25
- self.ssh(overwrite=False)
26
-
27
- try:
28
- subprocess.run(["ssh", f"{username}@{host}"])
29
- except Exception as ex:
30
- print(f"Failed to establish SSH connection: {ex}")
31
- sys.exit(1)
21
+ return studio(name=name, teamspace=teamspace)
22
+
23
+
24
+ @click.group(name="connect")
25
+ def connect() -> None:
26
+ """Connect to lightning products."""
27
+
28
+
29
+ # @connect.command(name="studio")
30
+ # @click.option("--name", default=None, help="The name of the studio to connect to.")
31
+ # @click.option(
32
+ # "--teamspace",
33
+ # default=None,
34
+ # help="The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
35
+ # )
36
+ def studio(name: Optional[str], teamspace: Optional[str]) -> None:
37
+ """Connect to a studio via SSH."""
38
+ configure_ssh(name=name, teamspace=teamspace, overwrite=False)
39
+
40
+ menu = _StudiosMenu()
41
+ studio = menu._get_studio(name=name, teamspace=teamspace)
42
+
43
+ try:
44
+ subprocess.run(["ssh", studio.name])
45
+ except Exception as ex:
46
+ print(f"Failed to establish SSH connection: {ex}")
47
+ sys.exit(1)
@@ -1,5 +1,6 @@
1
1
  from typing import Optional
2
2
 
3
+ import click
3
4
  from rich.console import Console
4
5
 
5
6
  from lightning_sdk.cli.exceptions import StudioCliError
@@ -20,15 +21,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
20
21
  teamspace: The teamspace to delete the container from. Should be specified as {owner}/{name}
21
22
  If not provided, can be selected in an interactive menu.
22
23
  """
23
- api = LitContainer()
24
- resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
25
- try:
26
- api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
27
- Console().print(f"Container {container} deleted successfully.")
28
- except Exception as e:
29
- raise StudioCliError(
30
- f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
31
- ) from None
24
+ delete_container(container=container, teamspace=teamspace)
32
25
 
33
26
  def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
34
27
  """Delete a job.
@@ -40,10 +33,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
40
33
  If not specified can be selected interactively.
41
34
 
42
35
  """
43
- job = super().job(name=name, teamspace=teamspace)
44
-
45
- job.delete()
46
- Console().print(f"Successfully deleted {job.name}!")
36
+ job(name=name, teamspace=teamspace)
47
37
 
48
38
  def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
49
39
  """Delete a multi-machine job.
@@ -55,10 +45,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
55
45
  If not specified can be selected interactively.
56
46
 
57
47
  """
58
- mmt = super().mmt(name=name, teamspace=teamspace)
59
-
60
- mmt.delete()
61
- Console().print(f"Successfully deleted {mmt.name}!")
48
+ mmt(name=name, teamspace=teamspace)
62
49
 
63
50
  def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
64
51
  """Delete an existing studio.
@@ -70,18 +57,80 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
70
57
  teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
71
58
  If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
72
59
  """
73
- if teamspace is not None:
74
- ts_splits = teamspace.split("/")
75
- if len(ts_splits) != 2:
76
- raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
77
- owner, teamspace = ts_splits
78
- else:
79
- owner, teamspace = None, None
80
-
81
- try:
82
- studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
83
- except (RuntimeError, ValueError):
84
- studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
85
-
86
- studio.delete()
87
- Console().print("Studio successfully deleted")
60
+ studio(name=name, teamspace=teamspace)
61
+
62
+
63
+ @click.group()
64
+ def delete() -> None:
65
+ """Delete resources on the Lightning AI platform."""
66
+
67
+
68
+ # @delete.command(name="container")
69
+ # @click.option("--container", help="The name of the container to delete.")
70
+ # @click.option("--teamspace", default=None, help=("The teamspace to delete the container from. "
71
+ # "Should be specified as {owner}/{name} "
72
+ # "If not provided, can be selected in an interactive menu."),)
73
+ def delete_container(container: str, teamspace: Optional[str] = None) -> None:
74
+ """Delete the docker container CONTAINER."""
75
+ api = LitContainer()
76
+ menu = _TeamspacesMenu()
77
+ resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
78
+ try:
79
+ api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
80
+ Console().print(f"Container {container} deleted successfully.")
81
+ except Exception as e:
82
+ raise StudioCliError(
83
+ f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
84
+ ) from None
85
+
86
+
87
+ # @delete.command(name="job")
88
+ # @click.option("--name", help="The name of the job to delete.")
89
+ # @click.option("--teamspace", default=None, help=("The teamspace to delete the job from. "
90
+ # "Should be specified as {owner}/{name} "
91
+ # "If not provided, can be selected in an interactive menu."),)
92
+ def job(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
93
+ """Delete a job."""
94
+ menu = _JobAndMMTAction()
95
+ job = menu.job(name=name, teamspace=teamspace)
96
+
97
+ job.delete()
98
+ Console().print(f"Successfully deleted {job.name}!")
99
+
100
+
101
+ # @delete.command(name="mmt")
102
+ # @click.option("--name", help="The name of the multi-machine job to delete.")
103
+ # @click.option("--teamspace", default=None, help=("The teamspace to delete the job from. "
104
+ # "Should be specified as {owner}/{name} "
105
+ # "If not provided, can be selected in an interactive menu."),)
106
+ def mmt(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
107
+ """Delete a multi-machine job."""
108
+ menu = _JobAndMMTAction()
109
+ mmt = menu.mmt(name=name, teamspace=teamspace)
110
+
111
+ mmt.delete()
112
+ Console().print(f"Successfully deleted {mmt.name}!")
113
+
114
+
115
+ # @delete.command(name="studio")
116
+ # @click.option("--name", help="The name of the studio to delete.")
117
+ # @click.option("--teamspace", default=None, help=("The teamspace to delete the studio from. "
118
+ # "Should be specified as {owner}/{name} "
119
+ # "If not provided, can be selected in an interactive menu."),)
120
+ def studio(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
121
+ """Delete an existing studio."""
122
+ if teamspace is not None:
123
+ ts_splits = teamspace.split("/")
124
+ if len(ts_splits) != 2:
125
+ raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
126
+ owner, teamspace = ts_splits
127
+ else:
128
+ owner, teamspace = None, None
129
+
130
+ try:
131
+ studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
132
+ except (RuntimeError, ValueError):
133
+ studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
134
+
135
+ studio.delete()
136
+ Console().print("Studio successfully deleted")