1313
1414import mkdocs_gen_files
1515import yaml
16- from pydantic .main import BaseModel
16+ from pydantic import BaseModel , RootModel
17+ from pydantic_core import PydanticUndefined
1718from typing_extensions import Annotated , Any , Dict , Literal , Type , Union , get_args , get_origin
1819
1920from dstack ._internal .core .models .resources import Range
2526logger .info ("Generating schema reference..." )
2627
2728
28- def _is_linkable_type (annotation : Any ) -> bool :
29- """Check if a type annotation contains a BaseModel subclass (excluding Range)."""
29+ def _unwrap_optional (annotation : Any ) -> Any :
30+ """The non-`None` member of an `Optional[...]`, or the annotation unchanged."""
31+ if get_origin (annotation ) is Union :
32+ args = [a for a in get_args (annotation ) if a is not type (None )]
33+ if len (args ) == 1 :
34+ return args [0 ]
35+ return annotation
36+
37+
38+ def _linkable_model (annotation : Any ) -> Optional [type ]:
39+ """
40+ The `BaseModel` subclass a field links to in the reference, if any.
41+
42+ pydantic v2 strips `Annotated` off `FieldInfo.annotation`, so the shape this used to unwrap by
43+ hand (`Annotated[Optional[SSHParams], Field(...)]`) now arrives as plain `Optional[SSHParams]`.
44+ Recursing over the annotation covers both, and also catches a bare model field, which the old
45+ `get_args(...)[0]` approach silently missed.
46+ """
3047 origin = get_origin (annotation )
48+ # The container cases come first: `get_origin(Annotated[X, ...])` is `Annotated`, which is
49+ # itself a class, so testing `inspect.isclass` first would stop before ever unwrapping it.
50+ if origin in (Annotated , Union , list ):
51+ for arg in get_args (annotation ):
52+ if arg is type (None ):
53+ continue
54+ found = _linkable_model (arg )
55+ if found is not None :
56+ return found
57+ return None
3158 type_ = origin if origin is not None else annotation
32- if inspect .isclass (type_ ):
33- return issubclass (type_ , BaseModel ) and not issubclass (type_ , Range )
34- if origin is Annotated :
35- return _is_linkable_type (get_args (annotation )[0 ])
36- if origin is Union :
37- return any (_is_linkable_type (arg ) for arg in get_args (annotation ))
38- if origin is list :
39- args = get_args (annotation )
40- return bool (args ) and _is_linkable_type (args [0 ])
41- return False
59+ if inspect .isclass (type_ ) and issubclass (type_ , BaseModel ) and not issubclass (type_ , Range ):
60+ return type_
61+ return None
62+
63+
64+ # Scalar JSON Schema types, mapped to how the docs spell them. Deliberately an allowlist rather
65+ # than a full mapping: `array` and `object` would only restate what the annotation already renders
66+ # more precisely (`list[str]` gaining a bare `list`, `dict` gaining `object`), and merging anything
67+ # into a bracketed type corrupts it, since `_enrich_type_from_schema` splits the rendered type on
68+ # `" | "` — which `list["no-capacity" | "interruption"]` contains.
69+ _ENRICHABLE = {"string" : "str" , "integer" : "int" , "boolean" : "bool" , "number" : "float" }
70+
71+
72+ def _shorthand_primitives (model : Type ) -> list :
73+ """
74+ The primitive types a model accepts in place of its object form, e.g. `8` or `arm:8` for
75+ `CPUSpec`. Taken from the model's own validation JSON Schema, which is the same declaration
76+ that produces the published `configuration.json`.
77+ """
78+ try :
79+ schema = model .model_json_schema (mode = "validation" )
80+ except Exception :
81+ return []
82+ found = {
83+ _ENRICHABLE [entry ["type" ]]
84+ for entry in schema .get ("anyOf" , [])
85+ if entry .get ("type" ) in _ENRICHABLE
86+ }
87+ return sorted (found , key = _type_sort_key )
4288
4389
4490def _type_sort_key (t : str ) -> tuple :
@@ -57,7 +103,7 @@ def _type_sort_key(t: str) -> tuple:
57103 return (5 , t )
58104
59105
60- def get_friendly_type (annotation : Type ) -> str :
106+ def get_friendly_type (annotation : Any ) -> str :
61107 """Get a user-friendly type string for documentation.
62108
63109 Produces types like: ``int | str``, ``"vscode" | "cursor"``, ``list[object]``.
@@ -112,10 +158,11 @@ def get_friendly_type(annotation: Type) -> str:
112158
113159 # Range — depends on inner type parameter
114160 if issubclass (annotation , Range ):
115- min_field = annotation .__fields__ .get ("min" )
116- if min_field and inspect .isclass (min_field .type_ ):
161+ min_field = annotation .model_fields .get ("min" )
162+ inner = _unwrap_optional (min_field .annotation ) if min_field else None
163+ if inspect .isclass (inner ):
117164 # Range[Memory] → str, Range[int] → int | str
118- if issubclass (min_field . type_ , float ):
165+ if issubclass (inner , float ):
119166 return "str"
120167 return "int | str"
121168
@@ -127,13 +174,16 @@ def get_friendly_type(annotation: Type) -> str:
127174
128175 # BaseModel subclass (not Range)
129176 if issubclass (annotation , BaseModel ) and not issubclass (annotation , Range ):
130- # Root models (with __root__ field) — resolve from the root type
131- if "__root__" in annotation .__fields__ :
132- return get_friendly_type (annotation .__fields__ ["__root__" ].annotation )
133- # Models with custom __get_validators__ accept primitive input (int, str)
134- # in addition to the full object form (e.g., GPUSpec, CPUSpec, DiskSpec)
135- if "__get_validators__" in annotation .__dict__ :
136- return "int | str | object"
177+ # Root models — resolve from the root type
178+ if issubclass (annotation , RootModel ):
179+ return get_friendly_type (annotation .model_fields ["root" ].annotation )
180+ # Models that define their own core schema also accept a shorthand. Read which
181+ # primitives from the model's own JSON Schema rather than assuming `int | str`:
182+ # `CPUSpec` takes both, but `FilePathMapping` and `RepoSpec` take only a string.
183+ if "__get_pydantic_core_schema__" in annotation .__dict__ :
184+ shorthand = _shorthand_primitives (annotation )
185+ if shorthand :
186+ return " | " .join ([* shorthand , "object" ])
137187 return "object"
138188
139189 # ComputeCapability (tuple subclass that parses "7.5" strings)
@@ -163,33 +213,24 @@ def get_friendly_type(annotation: Type) -> str:
163213 return str (annotation )
164214
165215
166- _JSON_SCHEMA_TYPE_MAP = {
167- "string" : "str" ,
168- "integer" : "int" ,
169- "number" : "float" ,
170- "boolean" : "bool" ,
171- "array" : "list" ,
172- "object" : "object" ,
173- }
174-
175-
176216def _enrich_type_from_schema (friendly_type : str , prop_schema : Dict [str , Any ]) -> str :
177217 """Enrich the friendly type with extra accepted types from the JSON schema.
178218
179- Models may define ``schema_extra`` that adds ``anyOf`` entries for fields
180- that accept alternative input types (e.g., duration fields typed as ``int``
181- but also accepting ``str`` like ``"5m"``).
219+ A field's annotation is its *post-validation* type, so it does not show what a before-validator
220+ also accepts — a duration typed ``int`` takes ``"5m"``, ``false`` and ``"off"`` as well. Those
221+ come from the type's ``json_schema_input_type``, i.e. the same declaration that produces the
222+ published schema.
182223 """
183224 any_of = prop_schema .get ("anyOf" )
184225 if not any_of :
185226 return friendly_type
186- # Only consider string/integer — the most common alternative input types.
187- # Skip boolean (typically a backward-compat artifact) and object/array.
188- _ENRICHABLE = {"string" : "str" , "integer" : "int" }
189227 schema_types = set ()
190228 for entry in any_of :
191- # Skip entries with enum constraints — those are already captured as literal values
192- if "enum" in entry :
229+ # A single accepted value (`Literal["off"]`) is more useful spelled out than as `str`.
230+ # Duplicates are removed below, so an annotation that already shows it is unaffected.
231+ literals = [entry ["const" ]] if "const" in entry else entry .get ("enum" , [])
232+ if literals :
233+ schema_types .update (f'"{ v } for v in literals if isinstance (v , str ))
193234 continue
194235 mapped = _ENRICHABLE .get (entry .get ("type" , "" ))
195236 if mapped :
@@ -200,9 +241,6 @@ def _enrich_type_from_schema(friendly_type: str, prop_schema: Dict[str, Any]) ->
200241 if not new_parts :
201242 return friendly_type
202243 all_parts = list (set (current_parts ) | new_parts )
203- # If str is now present, single-value literals are redundant
204- if "str" in all_parts :
205- all_parts = [p for p in all_parts if not p .startswith ('"' ) or p in all_parts ]
206244 all_parts .sort (key = _type_sort_key )
207245 return " | " .join (all_parts )
208246
@@ -228,15 +266,17 @@ def generate_schema_reference(
228266 "" ,
229267 ]
230268 )
231- # Get JSON schema to detect extra accepted types from schema_extra
269+ # The schema says what a field *accepts*, which is wider than its annotation wherever a
270+ # before-validator coerces. `mode="validation"` is pydantic's default, but state it: the
271+ # serialization schema carries the narrow type and would defeat the whole point.
232272 try :
233- schema_props = cls .schema ( ).get ("properties" , {})
273+ schema_props = cls .model_json_schema ( mode = "validation" ).get ("properties" , {})
234274 except Exception :
235275 schema_props = {}
236- for name , field in cls .__fields__ .items ():
276+ for name , field in cls .model_fields .items ():
237277 default = field .default
238278 default_repr : Optional [str ]
239- if default is None :
279+ if default is None or default is PydanticUndefined :
240280 default_repr = None
241281 elif isinstance (default , (list , tuple , dict )) and len (default ) == 0 :
242282 default_repr = None
@@ -252,24 +292,17 @@ def generate_schema_reference(
252292 friendly_type = _enrich_type_from_schema (friendly_type , schema_props .get (name , {}))
253293 values = dict (
254294 name = name ,
255- description = field .field_info . description ,
295+ description = field .description ,
256296 type = friendly_type ,
257297 default = default_repr ,
258- required = field .required ,
298+ required = field .is_required () ,
259299 )
260300 # TODO: If the field doesn't have description (e.g. BaseConfiguration.type), we could fallback to docstring
261301 if values ["description" ]:
262302 if overrides and name in overrides :
263303 values .update (overrides [name ])
264- field_type = next (iter (get_args (field .annotation )), None )
265- # TODO: This is a dirty workaround
266- if field_type :
267- if field .annotation .__name__ == "Annotated" :
268- if field_type .__name__ in ["Optional" , "List" , "list" , "Union" ]:
269- field_type = get_args (field_type )[0 ]
270- base_model = _is_linkable_type (field_type )
271- else :
272- base_model = False
304+ field_type = _linkable_model (field .annotation )
305+ base_model = field_type is not None
273306 _defaults = (
274307 f"Defaults to `{ values ['default' ]}
275308 if not base_model and values .get ("default" )
0 commit comments