Skip to content

Commit 9497994

Browse files
jvstmejvstme
committed
Exports CLI
1 parent 39d1d6b commit 9497994

5 files changed

Lines changed: 229 additions & 0 deletions

File tree

‎src/dstack/_internal/cli/commands/export.py‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
import argparse
2+
from typing import Any, Union
3+
4+
from rich.table import Table
5+
6+
from dstack._internal.cli.commands import APIBaseCommand
7+
from dstack._internal.cli.services.completion import ExportNameCompleter
8+
from dstack._internal.cli.utils.common import add_row_from_dict, confirm_ask, console
9+
from dstack._internal.core.models.exports import Export
10+
11+
12+
class ExportCommand(APIBaseCommand):
13+
NAME = "export"
14+
DESCRIPTION = "Manage exports"
15+
16+
def _register(self):
17+
super()._register()
18+
self._parser.set_defaults(subfunc=self._list)
19+
subparsers = self._parser.add_subparsers(dest="action")
20+
21+
list_parser = subparsers.add_parser(
22+
"list", help="List exports", formatter_class=self._parser.formatter_class
23+
)
24+
list_parser.set_defaults(subfunc=self._list)
25+
26+
create_parser = subparsers.add_parser(
27+
"create", help="Create an export", formatter_class=self._parser.formatter_class
28+
)
29+
create_parser.add_argument(
30+
"name",
31+
help="The name of the export",
32+
)
33+
create_parser.add_argument(
34+
"--importer",
35+
action="append",
36+
dest="importers",
37+
help="Importer project name (can be specified multiple times)",
38+
default=[],
39+
)
40+
create_parser.add_argument(
41+
"--fleet",
42+
action="append",
43+
dest="fleets",
44+
help="Fleet name to export (can be specified multiple times)",
45+
default=[],
46+
)
47+
create_parser.set_defaults(subfunc=self._create)
48+
49+
update_parser = subparsers.add_parser(
50+
"update", help="Update an export", formatter_class=self._parser.formatter_class
51+
)
52+
update_parser.add_argument(
53+
"name",
54+
help="The name of the export",
55+
).completer = ExportNameCompleter() # type: ignore[attr-defined]
56+
update_parser.add_argument(
57+
"--add-importer",
58+
action="append",
59+
dest="add_importers",
60+
help="Importer project name to add (can be specified multiple times)",
61+
default=[],
62+
)
63+
update_parser.add_argument(
64+
"--remove-importer",
65+
action="append",
66+
dest="remove_importers",
67+
help="Importer project name to remove (can be specified multiple times)",
68+
default=[],
69+
)
70+
update_parser.add_argument(
71+
"--add-fleet",
72+
action="append",
73+
dest="add_fleets",
74+
help="Fleet name to add (can be specified multiple times)",
75+
default=[],
76+
)
77+
update_parser.add_argument(
78+
"--remove-fleet",
79+
action="append",
80+
dest="remove_fleets",
81+
help="Fleet name to remove (can be specified multiple times)",
82+
default=[],
83+
)
84+
update_parser.set_defaults(subfunc=self._update)
85+
86+
delete_parser = subparsers.add_parser(
87+
"delete", help="Delete an export", formatter_class=self._parser.formatter_class
88+
)
89+
delete_parser.add_argument(
90+
"name",
91+
help="The name of the export",
92+
).completer = ExportNameCompleter() # type: ignore[attr-defined]
93+
delete_parser.add_argument(
94+
"-y", "--yes", help="Don't ask for confirmation", action="store_true"
95+
)
96+
delete_parser.set_defaults(subfunc=self._delete)
97+
98+
def _command(self, args: argparse.Namespace):
99+
super()._command(args)
100+
args.subfunc(args)
101+
102+
def _list(self, args: argparse.Namespace):
103+
exports = self.api.client.exports.list(self.api.project)
104+
print_exports_table(exports)
105+
106+
def _create(self, args: argparse.Namespace):
107+
with console.status("Creating export..."):
108+
export = self.api.client.exports.create(
109+
project_name=self.api.project,
110+
name=args.name,
111+
importer_projects=args.importers,
112+
exported_fleets=args.fleets,
113+
)
114+
print_exports_table([export])
115+
116+
def _update(self, args: argparse.Namespace):
117+
with console.status("Updating export..."):
118+
export = self.api.client.exports.update(
119+
project_name=self.api.project,
120+
name=args.name,
121+
add_importer_projects=args.add_importers,
122+
remove_importer_projects=args.remove_importers,
123+
add_exported_fleets=args.add_fleets,
124+
remove_exported_fleets=args.remove_fleets,
125+
)
126+
print_exports_table([export])
127+
128+
def _delete(self, args: argparse.Namespace):
129+
if not args.yes and not confirm_ask(f"Delete the export [code]{args.name}[/]?"):
130+
console.print("\nExiting...")
131+
return
132+
133+
with console.status("Deleting export..."):
134+
self.api.client.exports.delete(project_name=self.api.project, name=args.name)
135+
136+
console.print(f"Export [code]{args.name}[/] deleted")
137+
138+
139+
def print_exports_table(exports: list[Export]):
140+
table = Table(box=None)
141+
table.add_column("NAME", no_wrap=True)
142+
table.add_column("FLEETS")
143+
table.add_column("IMPORTERS")
144+
145+
for export in exports:
146+
fleets = (
147+
", ".join([f.name for f in export.exported_fleets]) if export.exported_fleets else "-"
148+
)
149+
importers = ", ".join([i.project_name for i in export.imports]) if export.imports else "-"
150+
151+
row: dict[Union[str, int], Any] = {
152+
"NAME": export.name,
153+
"FLEETS": fleets,
154+
"IMPORTERS": importers,
155+
}
156+
add_row_from_dict(table, row)
157+
158+
console.print(table)
159+
console.print()

‎src/dstack/_internal/cli/main.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from dstack._internal.cli.commands.completion import CompletionCommand
1010
from dstack._internal.cli.commands.delete import DeleteCommand
1111
from dstack._internal.cli.commands.event import EventCommand
12+
from dstack._internal.cli.commands.export import ExportCommand
1213
from dstack._internal.cli.commands.fleet import FleetCommand
1314
from dstack._internal.cli.commands.gateway import GatewayCommand
1415
from dstack._internal.cli.commands.init import InitCommand
@@ -66,6 +67,7 @@ def main():
6667
AttachCommand.register(subparsers)
6768
DeleteCommand.register(subparsers)
6869
EventCommand.register(subparsers)
70+
ExportCommand.register(subparsers)
6971
FleetCommand.register(subparsers)
7072
GatewayCommand.register(subparsers)
7173
InitCommand.register(subparsers)

‎src/dstack/_internal/cli/services/completion.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,11 @@ def fetch_resource_names(self, api: Client) -> Iterable[str]:
8080
return [r.name for r in api.client.secrets.list(api.project)]
8181

8282

83+
class ExportNameCompleter(BaseAPINameCompleter):
84+
def fetch_resource_names(self, api: Client) -> Iterable[str]:
85+
return [r.name for r in api.client.exports.list(api.project)]
86+
87+
8388
class ProjectNameCompleter(BaseCompleter):
8489
"""
8590
Completer for local project names.

‎src/dstack/api/server/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from dstack.api.server._auth import AuthAPIClient
1818
from dstack.api.server._backends import BackendsAPIClient
1919
from dstack.api.server._events import EventsAPIClient
20+
from dstack.api.server._exports import ExportsAPIClient
2021
from dstack.api.server._files import FilesAPIClient
2122
from dstack.api.server._fleets import FleetsAPIClient
2223
from dstack.api.server._gateways import GatewaysAPIClient
@@ -50,6 +51,7 @@ class APIClient:
5051
logs: operations with logs
5152
gateways: operations with gateways
5253
volumes: operations with volumes
54+
exports: operations with exports
5355
files: operations with files
5456
"""
5557

@@ -126,6 +128,10 @@ def gateways(self) -> GatewaysAPIClient:
126128
def volumes(self) -> VolumesAPIClient:
127129
return VolumesAPIClient(self._request, self._logger)
128130

131+
@property
132+
def exports(self) -> ExportsAPIClient:
133+
return ExportsAPIClient(self._request, self._logger)
134+
129135
@property
130136
def files(self) -> FilesAPIClient:
131137
return FilesAPIClient(self._request, self._logger)

‎src/dstack/api/server/_exports.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
from typing import List
2+
3+
from pydantic import parse_obj_as
4+
5+
from dstack._internal.core.models.exports import Export
6+
from dstack._internal.server.schemas.exports import (
7+
CreateExportRequest,
8+
DeleteExportRequest,
9+
UpdateExportRequest,
10+
)
11+
from dstack.api.server._group import APIClientGroup
12+
13+
14+
class ExportsAPIClient(APIClientGroup):
15+
def list(self, project_name: str) -> List[Export]:
16+
resp = self._request(f"/api/project/{project_name}/exports/list")
17+
return parse_obj_as(List[Export.__response__], resp.json())
18+
19+
def create(
20+
self,
21+
project_name: str,
22+
name: str,
23+
*,
24+
importer_projects: List[str] = [],
25+
exported_fleets: List[str] = [],
26+
) -> Export:
27+
body = CreateExportRequest(
28+
name=name,
29+
importer_projects=importer_projects,
30+
exported_fleets=exported_fleets,
31+
)
32+
resp = self._request(f"/api/project/{project_name}/exports/create", body=body.json())
33+
return parse_obj_as(Export.__response__, resp.json())
34+
35+
def update(
36+
self,
37+
project_name: str,
38+
name: str,
39+
*,
40+
add_importer_projects: List[str] = [],
41+
remove_importer_projects: List[str] = [],
42+
add_exported_fleets: List[str] = [],
43+
remove_exported_fleets: List[str] = [],
44+
) -> Export:
45+
body = UpdateExportRequest(
46+
name=name,
47+
add_importer_projects=add_importer_projects,
48+
remove_importer_projects=remove_importer_projects,
49+
add_exported_fleets=add_exported_fleets,
50+
remove_exported_fleets=remove_exported_fleets,
51+
)
52+
resp = self._request(f"/api/project/{project_name}/exports/update", body=body.json())
53+
return parse_obj_as(Export.__response__, resp.json())
54+
55+
def delete(self, project_name: str, name: str) -> None:
56+
body = DeleteExportRequest(name=name)
57+
self._request(f"/api/project/{project_name}/exports/delete", body=body.json())

0 commit comments

Comments
 (0)