diff --git a/src/dstack/_internal/cli/commands/gateway.py b/src/dstack/_internal/cli/commands/gateway.py index 1a61c48856..384feaba63 100644 --- a/src/dstack/_internal/cli/commands/gateway.py +++ b/src/dstack/_internal/cli/commands/gateway.py @@ -11,7 +11,12 @@ confirm_ask, console, ) -from dstack._internal.cli.utils.gateway import get_gateways_table, print_gateways_table +from dstack._internal.cli.utils.gateway import ( + get_gateways_table, + print_gateways_json, + print_gateways_table, +) +from dstack._internal.core.errors import CLIError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.gateways import GatewayConfiguration from dstack._internal.utils.logging import get_logger @@ -43,6 +48,19 @@ def _register(self): parser.add_argument( "-v", "--verbose", action="store_true", help="Show more information" ) + parser.add_argument( + "--format", + choices=["plain", "json"], + default="plain", + help="Output format (default: plain)", + ) + parser.add_argument( + "--json", + action="store_const", + const="json", + dest="format", + help="Output in JSON format (equivalent to --format json)", + ) create_parser = subparsers.add_parser( "create", @@ -91,9 +109,15 @@ def _command(self, args: argparse.Namespace): args.subfunc(args) def _list(self, args: argparse.Namespace): + if args.watch and args.format == "json": + raise CLIError("JSON output is not supported together with --watch") + gateways = self.api.client.gateways.list(self.api.project) if not args.watch: - print_gateways_table(gateways, verbose=args.verbose) + if args.format == "json": + print_gateways_json(gateways, project=self.api.project) + else: + print_gateways_table(gateways, verbose=args.verbose) return try: diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index d173982dab..33f3656d02 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/src/dstack/_internal/cli/commands/offer.py @@ -1,6 +1,6 @@ import argparse from pathlib import Path -from typing import List +from typing import List, Literal, cast from dstack._internal.cli.commands import APIBaseCommand from dstack._internal.cli.services.args import cpu_spec, disk_spec, gpu_spec @@ -13,8 +13,8 @@ from dstack._internal.cli.utils.run import print_offers_json, print_run_plan from dstack._internal.core.errors import CLIError from dstack._internal.core.models.configurations import ApplyConfigurationType, TaskConfiguration +from dstack._internal.core.models.gpus import GpuGroup from dstack._internal.core.models.runs import RunSpec -from dstack._internal.server.schemas.gpus import GpuGroup from dstack.api.utils import load_profile @@ -130,7 +130,12 @@ def _command(self, args: argparse.Namespace): else: if args.group_by: gpus = self._list_gpus(args, run_spec) - print_gpu_json(gpus, run_spec, args.group_by, self.api.project) + print_gpu_json( + gpus, + run_spec, + cast(List[Literal["gpu", "backend", "region", "count"]], args.group_by), + self.api.project, + ) else: run_plan = self.api.client.runs.get_plan( self.api.project, diff --git a/src/dstack/_internal/cli/commands/ps.py b/src/dstack/_internal/cli/commands/ps.py index 9e5e5484af..c94b5410f3 100644 --- a/src/dstack/_internal/cli/commands/ps.py +++ b/src/dstack/_internal/cli/commands/ps.py @@ -10,6 +10,7 @@ LIVE_TABLE_REFRESH_RATE_PER_SEC, console, ) +from dstack._internal.core.errors import CLIError class PsCommand(APIBaseCommand): @@ -43,12 +44,31 @@ def _register(self): type=int, default=None, ) + self._parser.add_argument( + "--format", + choices=["plain", "json"], + default="plain", + help="Output format (default: plain)", + ) + self._parser.add_argument( + "--json", + action="store_const", + const="json", + dest="format", + help="Output in JSON format (equivalent to --format json)", + ) def _command(self, args: argparse.Namespace): super()._command(args) + if args.watch and args.format == "json": + raise CLIError("JSON output is not supported together with --watch") + runs = self.api.runs.list(all=args.all, limit=args.last) if not args.watch: - console.print(run_utils.get_runs_table(runs, verbose=args.verbose)) + if args.format == "json": + run_utils.print_runs_json(self.api.project, runs) + else: + console.print(run_utils.get_runs_table(runs, verbose=args.verbose)) return try: diff --git a/src/dstack/_internal/cli/models/__init__.py b/src/dstack/_internal/cli/models/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/dstack/_internal/cli/models/gateways.py b/src/dstack/_internal/cli/models/gateways.py new file mode 100644 index 0000000000..94dfa88982 --- /dev/null +++ b/src/dstack/_internal/cli/models/gateways.py @@ -0,0 +1,16 @@ +from typing import List + +from dstack._internal.core.models.common import CoreConfig, generate_dual_core_model +from dstack._internal.core.models.gateways import Gateway +from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent + + +class GatewayCommandOutputConfig(CoreConfig): + json_dumps = pydantic_orjson_dumps_with_indent + + +class GatewayCommandOutput(generate_dual_core_model(GatewayCommandOutputConfig)): + """JSON output model for `dstack gateway` command.""" + + project: str + gateways: List[Gateway] diff --git a/src/dstack/_internal/cli/models/offers.py b/src/dstack/_internal/cli/models/offers.py new file mode 100644 index 0000000000..56d5e21ea7 --- /dev/null +++ b/src/dstack/_internal/cli/models/offers.py @@ -0,0 +1,47 @@ +from typing import List, Literal, Optional + +from dstack._internal.core.models.common import CoreConfig, generate_dual_core_model +from dstack._internal.core.models.gpus import GpuGroup +from dstack._internal.core.models.instances import InstanceOfferWithAvailability +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent + + +class OfferRequirementsConfig(CoreConfig): + json_dumps = pydantic_orjson_dumps_with_indent + + +class OfferRequirements(generate_dual_core_model(OfferRequirementsConfig)): + """Profile/requirements output model for CLI commands.""" + + resources: ResourcesSpec + max_price: Optional[float] = None + spot: Optional[bool] = None + reservation: Optional[str] = None + + +class OfferCommandOutputConfig(CoreConfig): + json_dumps = pydantic_orjson_dumps_with_indent + + +class OfferCommandOutput(generate_dual_core_model(OfferCommandOutputConfig)): + """JSON output model for `dstack offer` command.""" + + project: str + user: str + requirements: OfferRequirements + offers: List[InstanceOfferWithAvailability] + total_offers: int + + +class OfferCommandGroupByGpuOutputConfig(CoreConfig): + json_dumps = pydantic_orjson_dumps_with_indent + + +class OfferCommandGroupByGpuOutput(generate_dual_core_model(OfferCommandGroupByGpuOutputConfig)): + """JSON output model for `dstack offer` command with GPU grouping.""" + + project: str + requirements: OfferRequirements + group_by: List[Literal["gpu", "backend", "region", "count"]] + gpus: List[GpuGroup] diff --git a/src/dstack/_internal/cli/models/runs.py b/src/dstack/_internal/cli/models/runs.py new file mode 100644 index 0000000000..db951b752d --- /dev/null +++ b/src/dstack/_internal/cli/models/runs.py @@ -0,0 +1,16 @@ +from typing import List + +from dstack._internal.core.models.common import CoreConfig, generate_dual_core_model +from dstack._internal.core.models.runs import Run +from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent + + +class PsCommandOutputConfig(CoreConfig): + json_dumps = pydantic_orjson_dumps_with_indent + + +class PsCommandOutput(generate_dual_core_model(PsCommandOutputConfig)): + """JSON output model for `dstack ps` command.""" + + project: str + runs: List[Run] diff --git a/src/dstack/_internal/cli/utils/gateway.py b/src/dstack/_internal/cli/utils/gateway.py index bad09579e4..9d9aae0f9a 100644 --- a/src/dstack/_internal/cli/utils/gateway.py +++ b/src/dstack/_internal/cli/utils/gateway.py @@ -2,6 +2,7 @@ from rich.table import Table +from dstack._internal.cli.models.gateways import GatewayCommandOutput from dstack._internal.cli.utils.common import add_row_from_dict, console from dstack._internal.core.models.gateways import Gateway from dstack._internal.utils.common import DateFormatter, pretty_date @@ -13,6 +14,15 @@ def print_gateways_table(gateways: List[Gateway], verbose: bool = False): console.print() +def print_gateways_json(gateways: List[Gateway], project: str) -> None: + """Print gateways information in JSON format.""" + output = GatewayCommandOutput( + project=project, + gateways=gateways, + ) + print(output.json()) + + def get_gateways_table( gateways: List[Gateway], verbose: bool = False, diff --git a/src/dstack/_internal/cli/utils/gpu.py b/src/dstack/_internal/cli/utils/gpu.py index 0ba30d623d..89638cb62f 100644 --- a/src/dstack/_internal/cli/utils/gpu.py +++ b/src/dstack/_internal/cli/utils/gpu.py @@ -1,66 +1,37 @@ import shutil -from typing import List +from typing import List, Literal from rich.table import Table +from dstack._internal.cli.models.offers import OfferCommandGroupByGpuOutput, OfferRequirements from dstack._internal.cli.utils.common import console +from dstack._internal.core.models.gpus import GpuGroup from dstack._internal.core.models.profiles import SpotPolicy from dstack._internal.core.models.runs import Requirements, RunSpec, get_policy_map -from dstack._internal.server.schemas.gpus import GpuGroup -def print_gpu_json(gpus, run_spec, group_by_cli, api_project): +def print_gpu_json( + gpus: List[GpuGroup], + run_spec: RunSpec, + group_by: List[Literal["gpu", "backend", "region", "count"]], + project: str, +): """Print GPU information in JSON format.""" - req = Requirements( + req = OfferRequirements( resources=run_spec.configuration.resources, max_price=run_spec.merged_profile.max_price, spot=get_policy_map(run_spec.merged_profile.spot_policy, default=SpotPolicy.AUTO), reservation=run_spec.configuration.reservation, ) - if req.spot is None: - spot_policy = "auto" - elif req.spot: - spot_policy = "spot" - else: - spot_policy = "on-demand" - - output = { - "project": api_project, - "user": "admin", # TODO: Get actual user name - "resources": req.resources.dict(), - "spot_policy": spot_policy, - "max_price": req.max_price, - "reservation": run_spec.configuration.reservation, - "group_by": group_by_cli, - "gpus": [], - } - - for gpu_group in gpus: - gpu_data = { - "name": gpu_group.name, - "memory_mib": gpu_group.memory_mib, - "vendor": gpu_group.vendor.value, - "availability": [av.value for av in gpu_group.availability], - "spot": gpu_group.spot, - "count": {"min": gpu_group.count.min, "max": gpu_group.count.max}, - "price": {"min": gpu_group.price.min, "max": gpu_group.price.max}, - } - - if gpu_group.backend: - gpu_data["backend"] = gpu_group.backend.value - if gpu_group.backends: - gpu_data["backends"] = [b.value for b in gpu_group.backends] - if gpu_group.region: - gpu_data["region"] = gpu_group.region - if gpu_group.regions: - gpu_data["regions"] = gpu_group.regions - - output["gpus"].append(gpu_data) - - import json + output = OfferCommandGroupByGpuOutput( + project=project, + requirements=req, + group_by=group_by, + gpus=gpus, + ) - print(json.dumps(output, indent=2)) + print(output.json()) def print_gpu_table(gpus: List[GpuGroup], run_spec: RunSpec, group_by: List[str], project: str): diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py index bd77185bdf..68dc828f79 100644 --- a/src/dstack/_internal/cli/utils/run.py +++ b/src/dstack/_internal/cli/utils/run.py @@ -4,6 +4,8 @@ from rich.markup import escape from rich.table import Table +from dstack._internal.cli.models.offers import OfferCommandOutput, OfferRequirements +from dstack._internal.cli.models.runs import PsCommandOutput from dstack._internal.cli.utils.common import NO_OFFERS_WARNING, add_row_from_dict, console from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.configurations import DevEnvironmentConfiguration @@ -14,6 +16,7 @@ ) from dstack._internal.core.models.profiles import ( DEFAULT_RUN_TERMINATION_IDLE_TIME, + SpotPolicy, TerminationPolicy, ) from dstack._internal.core.models.runs import ( @@ -24,6 +27,7 @@ ProbeSpec, RunPlan, RunStatus, + get_policy_map, ) from dstack._internal.core.models.runs import ( Run as CoreRun, @@ -43,33 +47,31 @@ def print_offers_json(run_plan: RunPlan, run_spec): """Print offers information in JSON format.""" job_plan = run_plan.job_plans[0] - output = { - "project": run_plan.project_name, - "user": run_plan.user, - "resources": job_plan.job_spec.requirements.resources.dict(), - "max_price": (job_plan.job_spec.requirements.max_price), - "spot": run_spec.configuration.spot_policy, - "reservation": run_plan.run_spec.configuration.reservation, - "offers": [], - "total_offers": job_plan.total_offers, - } + requirements = OfferRequirements( + resources=job_plan.job_spec.requirements.resources, + max_price=job_plan.job_spec.requirements.max_price, + spot=get_policy_map(run_spec.configuration.spot_policy, default=SpotPolicy.AUTO), + reservation=run_plan.run_spec.configuration.reservation, + ) - for offer in job_plan.offers: - output["offers"].append( - { - "backend": ("ssh" if offer.backend.value == "remote" else offer.backend.value), - "region": offer.region, - "instance_type": offer.instance.name, - "resources": offer.instance.resources.dict(), - "spot": offer.instance.resources.spot, - "price": float(offer.price), - "availability": offer.availability.value, - } - ) + output = OfferCommandOutput( + project=run_plan.project_name, + user=run_plan.user, + requirements=requirements, + offers=job_plan.offers, + total_offers=job_plan.total_offers, + ) + + print(output.json()) - import json - print(json.dumps(output, indent=2)) +def print_runs_json(project: str, runs: List[Run]) -> None: + """Print runs information in JSON format.""" + output = PsCommandOutput( + project=project, + runs=[r._run for r in runs], + ) + print(output.json()) def print_run_plan( diff --git a/src/dstack/_internal/core/models/gpus.py b/src/dstack/_internal/core/models/gpus.py new file mode 100644 index 0000000000..7f5b56b813 --- /dev/null +++ b/src/dstack/_internal/core/models/gpus.py @@ -0,0 +1,45 @@ +from typing import List, Literal, Optional + +import gpuhunt + +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import CoreModel +from dstack._internal.core.models.instances import InstanceAvailability +from dstack._internal.core.models.resources import Range + + +class BackendGpu(CoreModel): + """GPU specification from a backend offer.""" + + name: str + memory_mib: int + vendor: gpuhunt.AcceleratorVendor + availability: InstanceAvailability + spot: bool + count: int + price: float + region: str + + +class BackendGpus(CoreModel): + """Backend GPU specifications.""" + + backend_type: BackendType + gpus: List[BackendGpu] + regions: List[str] + + +class GpuGroup(CoreModel): + """GPU group that can handle all grouping scenarios.""" + + name: str + memory_mib: int + vendor: gpuhunt.AcceleratorVendor + availability: List[InstanceAvailability] + spot: List[Literal["spot", "on-demand"]] + count: Range[int] + price: Range[float] + backends: Optional[List[BackendType]] = None + backend: Optional[BackendType] = None + regions: Optional[List[str]] = None + region: Optional[str] = None diff --git a/src/dstack/_internal/server/schemas/gpus.py b/src/dstack/_internal/server/schemas/gpus.py index beda7e43ec..16d1191483 100644 --- a/src/dstack/_internal/server/schemas/gpus.py +++ b/src/dstack/_internal/server/schemas/gpus.py @@ -1,36 +1,12 @@ from typing import List, Literal, Optional -import gpuhunt from pydantic import Field -from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreModel -from dstack._internal.core.models.instances import InstanceAvailability -from dstack._internal.core.models.resources import Range +from dstack._internal.core.models.gpus import GpuGroup from dstack._internal.core.models.runs import RunSpec -class BackendGpu(CoreModel): - """GPU specification from a backend offer.""" - - name: str - memory_mib: int - vendor: gpuhunt.AcceleratorVendor - availability: InstanceAvailability - spot: bool - count: int - price: float - region: str - - -class BackendGpus(CoreModel): - """Backend GPU specifications.""" - - backend_type: BackendType - gpus: List[BackendGpu] - regions: List[str] - - class ListGpusRequest(CoreModel): """Request for listing GPUs with optional grouping.""" @@ -42,22 +18,6 @@ class ListGpusRequest(CoreModel): ) -class GpuGroup(CoreModel): - """GPU group that can handle all grouping scenarios.""" - - name: str - memory_mib: int - vendor: gpuhunt.AcceleratorVendor - availability: List[InstanceAvailability] - spot: List[Literal["spot", "on-demand"]] - count: Range[int] - price: Range[float] - backends: Optional[List[BackendType]] = None - backend: Optional[BackendType] = None - regions: Optional[List[str]] = None - region: Optional[str] = None - - class ListGpusResponse(CoreModel): """Response containing GPU specifications.""" diff --git a/src/dstack/_internal/server/services/gpus.py b/src/dstack/_internal/server/services/gpus.py index c5eee4b65c..9c52523d76 100644 --- a/src/dstack/_internal/server/services/gpus.py +++ b/src/dstack/_internal/server/services/gpus.py @@ -3,17 +3,13 @@ from dstack._internal.core.backends.base.backend import Backend from dstack._internal.core.errors import ServerClientError from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.gpus import BackendGpu, BackendGpus, GpuGroup from dstack._internal.core.models.instances import InstanceOfferWithAvailability from dstack._internal.core.models.profiles import SpotPolicy from dstack._internal.core.models.resources import Range from dstack._internal.core.models.runs import Requirements, RunSpec, get_policy_map from dstack._internal.server.models import ProjectModel -from dstack._internal.server.schemas.gpus import ( - BackendGpu, - BackendGpus, - GpuGroup, - ListGpusResponse, -) +from dstack._internal.server.schemas.gpus import ListGpusResponse from dstack._internal.server.services.offers import get_offers_by_requirements from dstack._internal.utils.common import get_or_error diff --git a/src/dstack/api/server/_gpus.py b/src/dstack/api/server/_gpus.py index f7bff4d49e..253410604d 100644 --- a/src/dstack/api/server/_gpus.py +++ b/src/dstack/api/server/_gpus.py @@ -1,10 +1,11 @@ -from typing import List, Optional +from typing import List, Literal, Optional, cast from pydantic import parse_obj_as from dstack._internal.core.compatibility.gpus import get_list_gpus_excludes +from dstack._internal.core.models.gpus import GpuGroup from dstack._internal.core.models.runs import RunSpec -from dstack._internal.server.schemas.gpus import GpuGroup, ListGpusRequest, ListGpusResponse +from dstack._internal.server.schemas.gpus import ListGpusRequest, ListGpusResponse from dstack.api.server._group import APIClientGroup @@ -15,7 +16,10 @@ def list_gpus( run_spec: RunSpec, group_by: Optional[List[str]] = None, ) -> List[GpuGroup]: - body = ListGpusRequest(run_spec=run_spec, group_by=group_by) + body = ListGpusRequest( + run_spec=run_spec, + group_by=cast(Optional[List[Literal["backend", "region", "count"]]], group_by), + ) resp = self._request( f"/api/project/{project_name}/gpus/list", body=body.json(exclude=get_list_gpus_excludes(body)),