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

Review feedback

Author: peterschmidt85

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

@@ -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
 
 

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

@@ -3,11 +3,11 @@
 
 from rich.table import Table
 
-from dstack._internal.cli.models.offer import OfferCommandOutput
+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(
@@ -17,14 +17,14 @@ def print_gpu_json(
     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,
     )
 
-    output = OfferCommandOutput(
+    output = OfferCommandGroupByGpuOutput(
         project=project,
         requirements=req,
         group_by=group_by,

src/dstack/_internal/cli/utils/gpu.py

@@ -4,6 +4,7 @@
 from rich.markup import escape
 from rich.table import Table
 
+from dstack._internal.cli.models.offers import OfferCommandOutput, OfferRequirements
 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 +15,7 @@
 )
 from dstack._internal.core.models.profiles import (
     DEFAULT_RUN_TERMINATION_IDLE_TIME,
+    SpotPolicy,
     TerminationPolicy,
 )
 from dstack._internal.core.models.runs import (
@@ -24,6 +26,7 @@
     ProbeSpec,
     RunPlan,
     RunStatus,
+    get_policy_map,
 )
 from dstack._internal.core.models.runs import (
     Run as CoreRun,
@@ -43,33 +46,22 @@ 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,
-    }
-
-    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,
-            }
-        )
+    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,
+    )
 
-    import json
+    output = OfferCommandOutput(
+        project=run_plan.project_name,
+        user=run_plan.user,
+        requirements=requirements,
+        offers=job_plan.offers,
+        total_offers=job_plan.total_offers,
+    )
 
-    print(json.dumps(output, indent=2))
+    print(output.json())
 
 
 def print_run_plan(

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

src/dstack/_internal/core/models/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."""
 

src/dstack/_internal/server/schemas/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
 

src/dstack/_internal/server/services/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)),