[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.

Author: peterschmidt85

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
@@ -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/models/__init__.py

@@ -0,0 +1,19 @@
+from typing import List, Literal
+
+from dstack._internal.core.models.common import CoreConfig, generate_dual_core_model
+from dstack._internal.core.models.runs import Requirements
+from dstack._internal.server.schemas.gpus import GpuGroup
+from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent
+
+
+class OfferCommandOutputConfig(CoreConfig):
+    json_dumps = pydantic_orjson_dumps_with_indent
+
+
+class OfferCommandOutput(generate_dual_core_model(OfferCommandOutputConfig)):
+    """JSON output model for `dstack offer` command with GPU grouping."""
+
+    project: str
+    requirements: Requirements
+    group_by: List[Literal["gpu", "backend", "region", "count"]]
+    gpus: List[GpuGroup]

src/dstack/_internal/cli/models/offer.py

@@ -1,15 +1,21 @@
 import shutil
-from typing import List
+from typing import List, Literal
 
 from rich.table import Table
 
+from dstack._internal.cli.models.offer import OfferCommandOutput
 from dstack._internal.cli.utils.common import console
 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(
         resources=run_spec.configuration.resources,
@@ -18,49 +24,14 @@ def print_gpu_json(gpus, run_spec, group_by_cli, api_project):
         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 = OfferCommandOutput(
+        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):