Fix generate-json-schema CI

Author: r4victor

.github/workflows/build.yml

@@ -237,8 +237,8 @@ jobs:
         run: uv sync
       - name: Generate json schema
         run: |
-          uv run python -c "from dstack._internal.core.models.configurations import DstackConfiguration; print(DstackConfiguration.schema_json(indent=2))" > configuration.json
-          uv run python -c "from dstack._internal.core.models.profiles import ProfilesConfig; print(ProfilesConfig.schema_json(indent=2))" > profiles.json
+          uv run python -c "from dstack._internal.core.models.configurations import DstackConfiguration; print(DstackConfiguration.schema_json())" > configuration.json
+          uv run python -c "from dstack._internal.core.models.profiles import ProfilesConfig; print(ProfilesConfig.schema_json())" > profiles.json
       - name: Upload json schema to S3
         run: |
           VERSION=$((${{ github.run_number }} + ${{ env.BUILD_INCREMENT }}))

.github/workflows/release.yml

@@ -311,8 +311,8 @@ jobs:
         run: uv sync
       - name: Generate json schema
         run: |
-          uv run python -c "from dstack._internal.core.models.configurations import DstackConfiguration; print(DstackConfiguration.schema_json(indent=2))" > configuration.json
-          uv run python -c "from dstack._internal.core.models.profiles import ProfilesConfig; print(ProfilesConfig.schema_json(indent=2))" > profiles.json
+          uv run python -c "from dstack._internal.core.models.configurations import DstackConfiguration; print(DstackConfiguration.schema_json())" > configuration.json
+          uv run python -c "from dstack._internal.core.models.profiles import ProfilesConfig; print(ProfilesConfig.schema_json())" > profiles.json
       - name: Upload json schema to S3
         run: |
           VERSION=${GITHUB_REF#refs/tags/}

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_default_options, orjson_default
+from dstack._internal.utils.json_utils import pydantic_orjson_dumps
 
 IncludeExcludeFieldType = Union[int, str]
 IncludeExcludeSetType = set[IncludeExcludeFieldType]
@@ -18,14 +18,6 @@
 IncludeExcludeType = Union[IncludeExcludeSetType, IncludeExcludeDictType]
 
 
-def _orjson_dumps(v: Any, *, default: Any) -> str:
-    return orjson.dumps(
-        v,
-        option=get_orjson_default_options(),
-        default=orjson_default,
-    ).decode()
-
-
 # DualBaseModel creates two classes for the model:
 # one with extra = "forbid" (CoreModel/CoreModel.__request__),
 # and another with extra = "ignore" (CoreModel.__response__).
@@ -34,7 +26,7 @@ def _orjson_dumps(v: Any, *, default: Any) -> str:
 class CoreModel(DualBaseModel):
     class Config:
         json_loads = orjson.loads
-        json_dumps = _orjson_dumps
+        json_dumps = pydantic_orjson_dumps
 
     def json(
         self,

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

@@ -19,7 +19,9 @@
 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
+from dstack._internal.utils.json_utils import (
+    pydantic_orjson_dumps_with_indent,
+)
 
 CommandsList = List[str]
 ValidPort = conint(gt=0, le=65536)
@@ -568,15 +570,6 @@ 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,
@@ -585,7 +578,7 @@ class DstackConfiguration(CoreModel):
 
     class Config:
         json_loads = orjson.loads
-        json_dumps = _orjson_dumps
+        json_dumps = pydantic_orjson_dumps_with_indent
 
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):

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

@@ -1,12 +1,14 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union, overload
 
+import orjson
 from pydantic import Field, root_validator, validator
 from typing_extensions import Annotated, Literal
 
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel, Duration
 from dstack._internal.utils.common import list_enum_values_for_annotation
+from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent
 from dstack._internal.utils.tags import tags_validator
 
 DEFAULT_RETRY_DURATION = 3600
@@ -343,6 +345,9 @@ class ProfilesConfig(CoreModel):
     profiles: List[Profile]
 
     class Config:
+        json_loads = orjson.loads
+        json_dumps = pydantic_orjson_dumps_with_indent
+
         schema_extra = {"$schema": "http://json-schema.org/draft-07/schema#"}
 
     def default(self) -> Optional[Profile]:

src/dstack/_internal/utils/json_utils.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import orjson
 from pydantic import BaseModel
 
@@ -15,6 +17,22 @@
     ASYNCPG = False
 
 
+def pydantic_orjson_dumps(v: Any, *, default: Any) -> str:
+    return orjson.dumps(
+        v,
+        option=get_orjson_default_options(),
+        default=orjson_default,
+    ).decode()
+
+
+def pydantic_orjson_dumps_with_indent(v: Any, *, default: Any) -> str:
+    return orjson.dumps(
+        v,
+        option=get_orjson_default_options() | orjson.OPT_INDENT_2,
+        default=orjson_default,
+    ).decode()
+
+
 def orjson_default(obj):
     if isinstance(obj, float):
         # orjson does not convert float subclasses be default