Fix json schema reference and invalid properties errors (#2433)

* Replace *Schema models with manually specified extra types

* Allow json schema additionalProperties to fix invalid property errors

Author: r4victor

docs/docs/reference/dstack.yml/dev-environment.md

@@ -28,7 +28,7 @@ The `dev-environment` configuration type allows running [dev environments](../..
 
 ### `resources`
 
-#SCHEMA# dstack._internal.core.models.resources.ResourcesSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.ResourcesSpec
     overrides:
       show_root_heading: false
       type:
@@ -37,15 +37,15 @@ The `dev-environment` configuration type allows running [dev environments](../..
 
 #### `resources.gpu` { #resources-gpu data-toc-label="gpu" }
 
-#SCHEMA# dstack._internal.core.models.resources.GPUSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
 #### `resources.disk` { #resources-disk data-toc-label="disk" }
 
-#SCHEMA# dstack._internal.core.models.resources.DiskSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:
       show_root_heading: false
       type:

docs/docs/reference/dstack.yml/fleet.md

@@ -39,7 +39,7 @@ The `fleet` configuration type allows creating and updating fleets.
 
 ### `resources`
 
-#SCHEMA# dstack._internal.core.models.resources.ResourcesSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.ResourcesSpec
     overrides:
       show_root_heading: false
       type:
@@ -48,15 +48,15 @@ The `fleet` configuration type allows creating and updating fleets.
 
 #### `resouces.gpu` { #resources-gpu data-toc-label="gpu" } 
 
-#SCHEMA# dstack._internal.core.models.resources.GPUSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
 #### `resouces.disk` { #resources-disk data-toc-label="disk" }
 
-#SCHEMA# dstack._internal.core.models.resources.DiskSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:
       show_root_heading: false
       type:

docs/docs/reference/dstack.yml/service.md

@@ -90,7 +90,7 @@ The `service` configuration type allows running [services](../../concepts/servic
 
 ### `resources`
 
-#SCHEMA# dstack._internal.core.models.resources.ResourcesSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.ResourcesSpec
     overrides:
       show_root_heading: false
       type:
@@ -99,15 +99,15 @@ The `service` configuration type allows running [services](../../concepts/servic
 
 #### `resouces.gpu` { #resources-gpu data-toc-label="gpu" }
 
-#SCHEMA# dstack._internal.core.models.resources.GPUSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
 #### `resouces.disk` { #resources-disk data-toc-label="disk" }
 
-#SCHEMA# dstack._internal.core.models.resources.DiskSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:
       show_root_heading: false
       type:

docs/docs/reference/dstack.yml/task.md

@@ -28,7 +28,7 @@ The `task` configuration type allows running [tasks](../../concepts/tasks.md).
 
 ### `resources`
 
-#SCHEMA# dstack._internal.core.models.resources.ResourcesSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.ResourcesSpec
     overrides:
       show_root_heading: false
       type:
@@ -37,15 +37,15 @@ The `task` configuration type allows running [tasks](../../concepts/tasks.md).
 
 #### `resouces.gpu` { #resources-gpu data-toc-label="gpu" }
 
-#SCHEMA# dstack._internal.core.models.resources.GPUSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
 #### `resouces.disk` { #resources-disk data-toc-label="disk" }
 
-#SCHEMA# dstack._internal.core.models.resources.DiskSpecSchema
+#SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:
       show_root_heading: false
       type:

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

@@ -1,6 +1,6 @@
 import re
 from enum import Enum
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import Field, ValidationError, conint, constr, root_validator, validator
 from typing_extensions import Annotated, Literal
@@ -425,4 +425,9 @@ class DstackConfiguration(CoreModel):
     ]
 
     class Config:
-        schema_extra = {"$schema": "http://json-schema.org/draft-07/schema#"}
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any]):
+            schema["$schema"] = "http://json-schema.org/draft-07/schema#"
+            # Allow additionalProperties so that vscode and others not supporting
+            # top-level oneOf do not warn about properties being invalid.
+            schema["additionalProperties"] = True

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

@@ -19,6 +19,7 @@
     parse_idle_duration,
 )
 from dstack._internal.core.models.resources import Range, ResourcesSpec
+from dstack._internal.utils.json_schema import add_extra_schema_types
 
 
 class FleetStatus(str, Enum):
@@ -222,6 +223,14 @@ class InstanceGroupParams(CoreModel):
         ),
     ] = None
 
+    class Config:
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any], model: Type):
+            add_extra_schema_types(
+                schema["properties"]["nodes"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+
     _validate_idle_duration = validator("idle_duration", pre=True, allow_reuse=True)(
         parse_idle_duration
     )

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

@@ -8,6 +8,7 @@
 
 from dstack._internal.core.models.common import CoreModel
 from dstack._internal.utils.common import pretty_resources
+from dstack._internal.utils.json_schema import add_extra_schema_types
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -128,6 +129,22 @@ def __str__(self):
 
 
 class GPUSpec(CoreModel):
+    class Config:
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any]):
+            add_extra_schema_types(
+                schema["properties"]["count"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["memory"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["total_memory"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+
     vendor: Annotated[
         Optional[gpuhunt.AcceleratorVendor],
         Field(
@@ -233,6 +250,14 @@ def _vendor_from_string(cls, v: str) -> gpuhunt.AcceleratorVendor:
 
 
 class DiskSpec(CoreModel):
+    class Config:
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any]):
+            add_extra_schema_types(
+                schema["properties"]["size"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+
     size: Annotated[Range[Memory], Field(description="Disk size")]
 
     @classmethod
@@ -254,11 +279,26 @@ class ResourcesSpec(CoreModel):
     class Config:
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
-            schema.clear()
-            # replace strict schema with a more permissive one
-            ref_template = "#/definitions/ResourcesSpecRequest/definitions/{model}"
-            for field, value in ResourcesSpecSchema.schema(ref_template=ref_template).items():
-                schema[field] = value
+            add_extra_schema_types(
+                schema["properties"]["cpu"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["memory"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["shm_size"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["gpu"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["disk"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
 
     cpu: Annotated[Range[int], Field(description="The number of CPU cores")] = DEFAULT_CPU_COUNT
     memory: Annotated[Range[Memory], Field(description="The RAM size (e.g., `8GB`)")] = (
@@ -290,74 +330,3 @@ def pretty_format(self) -> str:
             resources.update(disk_size=self.disk.size)
         res = pretty_resources(**resources)
         return res
-
-
-IntRangeLike = Union[Range[Union[int, str]], int, str]
-MemoryRangeLike = Union[Range[Union[Memory, float, int, str]], float, int, str]
-MemoryLike = Union[Memory, float, int, str]
-GPULike = Union[GPUSpec, "GPUSpecSchema", int, str]
-DiskLike = Union[DiskSpec, "DiskSpecSchema", float, int, str]
-ComputeCapabilityLike = Union[ComputeCapability, float, str]
-
-
-class GPUSpecSchema(CoreModel):
-    vendor: Annotated[
-        Optional[gpuhunt.AcceleratorVendor],
-        Field(
-            description="The vendor of the GPU/accelerator, one of: `nvidia`, `amd`, `google` (alias: `tpu`), `intel`"
-        ),
-    ] = None
-    name: Annotated[
-        Optional[Union[List[str], str]], Field(description="The GPU name or list of names")
-    ] = None
-    count: Annotated[IntRangeLike, Field(description="The number of GPUs")] = DEFAULT_GPU_COUNT
-    memory: Annotated[
-        Optional[MemoryRangeLike],
-        Field(
-            description="The RAM size (e.g., `16GB`). Can be set to a range (e.g. `16GB..`, or `16GB..80GB`)"
-        ),
-    ] = None
-    total_memory: Annotated[
-        Optional[MemoryRangeLike],
-        Field(
-            description="The total RAM size (e.g., `32GB`). Can be set to a range (e.g. `16GB..`, or `16GB..80GB`)"
-        ),
-    ] = None
-    compute_capability: Annotated[
-        Optional[ComputeCapabilityLike],
-        Field(description="The minimum compute capability of the GPU (e.g., `7.5`)"),
-    ] = None
-
-
-class DiskSpecSchema(CoreModel):
-    size: Annotated[
-        MemoryRangeLike,
-        Field(
-            description="The disk size. Can be set to a range (e.g., `100GB..` or `100GB..200GB`)"
-        ),
-    ]
-
-
-class ResourcesSpecSchema(CoreModel):
-    cpu: Annotated[Optional[IntRangeLike], Field(description="The number of CPU cores")] = (
-        DEFAULT_CPU_COUNT
-    )
-    memory: Annotated[
-        Optional[MemoryRangeLike],
-        Field(description="The RAM size (e.g., `8GB`)"),
-    ] = DEFAULT_MEMORY_SIZE
-    shm_size: Annotated[
-        Optional[MemoryLike],
-        Field(
-            description="The size of shared memory (e.g., `8GB`). "
-            "If you are using parallel communicating processes (e.g., dataloaders in PyTorch), "
-            "you may need to configure this"
-        ),
-    ] = None
-    gpu: Annotated[
-        Optional[GPULike],
-        Field(
-            description="The GPU requirements. Can be set to a number, a string (e.g. `A100`, `80GB:2`, etc.), or an object"
-        ),
-    ] = None
-    disk: Annotated[Optional[DiskLike], Field(description="The disk resources")] = DEFAULT_DISK

src/dstack/_internal/utils/json_schema.py

@@ -0,0 +1,6 @@
+def add_extra_schema_types(schema_property: dict, extra_types: list[dict]):
+    if "allOf" in schema_property:
+        ref = schema_property.pop("allOf")[0]
+    else:
+        ref = {"type": schema_property.pop("type")}
+    schema_property["anyOf"] = [ref, *extra_types]