[Feature]: Add JSON output option for CLI commands (#3335)

* [Feature]: Add JSON output option for CLI commands #3322 (WIP)

This is a not-backward compatible change. It refactors existing `--json` code to make it eaasy to maintain and keep it extensible and backward-compatible in the future.

* [Feature]: Add JSON output option for CLI commands #3322 (WIP)

Review feedback

* [Feature]: Add JSON output option for CLI commands #3322 (WIP)

`dstack gateway` support

* [Feature]: Add JSON output option for CLI commands #3322 (WIP)

`dstack ps` support

Author: peterschmidt85

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:

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,

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:

src/dstack/_internal/cli/models/__init__.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]

src/dstack/_internal/cli/models/gateways.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]

src/dstack/_internal/cli/models/offers.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]

src/dstack/_internal/cli/models/runs.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,

src/dstack/_internal/cli/utils/gateway.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):

src/dstack/_internal/cli/utils/gpu.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(