Use orjson with indentation for json schema

Author: r4victor

src/dstack/_internal/core/models/common.py

@@ -8,7 +8,7 @@
 from pydantic_duality import DualBaseModel
 from typing_extensions import Annotated
 
-from dstack._internal.utils.json_utils import get_orjson_options, orjson_default
+from dstack._internal.utils.json_utils import get_orjson_default_options, orjson_default
 
 IncludeExcludeFieldType = Union[int, str]
 IncludeExcludeSetType = set[IncludeExcludeFieldType]
@@ -21,7 +21,7 @@
 def _orjson_dumps(v: Any, *, default: Any) -> str:
     return orjson.dumps(
         v,
-        option=get_orjson_options(),
+        option=get_orjson_default_options(),
         default=orjson_default,
     ).decode()
 

src/dstack/_internal/core/models/configurations.py

@@ -4,6 +4,7 @@
 from pathlib import PurePosixPath
 from typing import Any, Dict, List, Optional, Union
 
+import orjson
 from pydantic import Field, ValidationError, conint, constr, root_validator, validator
 from typing_extensions import Annotated, Literal
 
@@ -18,6 +19,7 @@
 from dstack._internal.core.models.services import AnyModel, OpenAIChatModel
 from dstack._internal.core.models.unix import UnixUser
 from dstack._internal.core.models.volumes import MountPoint, VolumeConfiguration, parse_mount_point
+from dstack._internal.utils.json_utils import get_orjson_default_options, orjson_default
 
 CommandsList = List[str]
 ValidPort = conint(gt=0, le=65536)
@@ -566,13 +568,25 @@ def parse_apply_configuration(data: dict) -> AnyApplyConfiguration:
 AnyDstackConfiguration = AnyApplyConfiguration
 
 
+# Custom _orjson_dumps for DstackConfiguration with indentation
+def _orjson_dumps(v: Any, *, default: Any) -> str:
+    return orjson.dumps(
+        v,
+        option=get_orjson_default_options() | orjson.OPT_INDENT_2,
+        default=orjson_default,
+    ).decode()
+
+
 class DstackConfiguration(CoreModel):
     __root__: Annotated[
         AnyDstackConfiguration,
         Field(discriminator="type"),
     ]
 
     class Config:
+        json_loads = orjson.loads
+        json_dumps = _orjson_dumps
+
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             schema["$schema"] = "http://json-schema.org/draft-07/schema#"

src/dstack/_internal/server/utils/routers.py

@@ -6,7 +6,7 @@
 
 from dstack._internal.core.errors import ServerClientError, ServerClientErrorCode
 from dstack._internal.core.models.common import CoreModel
-from dstack._internal.utils.json_utils import get_orjson_options, orjson_default
+from dstack._internal.utils.json_utils import get_orjson_default_options, orjson_default
 
 
 class CustomORJSONResponse(Response):
@@ -26,7 +26,7 @@ class CustomORJSONResponse(Response):
     def render(self, content: Any) -> bytes:
         return orjson.dumps(
             content,
-            option=get_orjson_options(),
+            option=get_orjson_default_options(),
             default=orjson_default,
         )
 

src/dstack/_internal/utils/json_utils.py

@@ -32,5 +32,5 @@ def orjson_default(obj):
     raise TypeError
 
 
-def get_orjson_options() -> int:
+def get_orjson_default_options() -> int:
     return orjson.OPT_NON_STR_KEYS