Replace validators with dict() for serialization tweaks

Author: r4victor

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

@@ -1,8 +1,9 @@
 import re
 from enum import Enum
-from typing import Any, Union
+from typing import Any, Callable, Union
 
 import orjson
+from git import Optional
 from pydantic import Field
 from pydantic_duality import DualBaseModel
 from typing_extensions import Annotated
@@ -35,6 +36,37 @@ class Config:
         json_loads = orjson.loads
         json_dumps = _orjson_dumps
 
+    def json(
+        self,
+        *,
+        include: Optional[IncludeExcludeType] = None,
+        exclude: Optional[IncludeExcludeType] = None,
+        by_alias: bool = False,
+        skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        encoder: Optional[Callable[[Any], Any]] = None,
+        models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
+        **dumps_kwargs: Any,
+    ) -> str:
+        """
+        Override `json()` method so that it calls `dict()`.
+        Allows changing how models are serialized by overriding `dict()` only.
+        By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
+        """
+        data = self.dict(
+            by_alias=by_alias,
+            include=include,
+            exclude=exclude,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+        )
+        if self.__custom_root_type__:
+            data = data["__root__"]
+        return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
+
 
 class Duration(int):
     """

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

@@ -382,14 +382,6 @@ def schema_extra(schema: Dict[str, Any]):
     gpu: Annotated[Optional[GPUSpec], Field(description="The GPU requirements")] = None
     disk: Annotated[Optional[DiskSpec], Field(description="The disk resources")] = DEFAULT_DISK
 
-    # TODO: Remove in 0.20. Added for backward compatibility.
-    @root_validator
-    def _post_validate(cls, values):
-        cpu = values.get("cpu")
-        if isinstance(cpu, CPUSpec) and cpu.arch in [None, gpuhunt.CPUArchitecture.X86]:
-            values["cpu"] = cpu.count
-        return values
-
     def pretty_format(self) -> str:
         # TODO: Remove in 0.20. Use self.cpu directly
         cpu = parse_obj_as(CPUSpec, self.cpu)
@@ -407,3 +399,18 @@ def pretty_format(self) -> str:
             resources.update(disk_size=self.disk.size)
         res = pretty_resources(**resources)
         return res
+
+    def dict(self, *args, **kwargs) -> Dict:
+        # super() does not work with pydantic-duality
+        res = CoreModel.dict(self, *args, **kwargs)
+        self._update_serialized_cpu(res)
+        return res
+
+    # TODO: Remove in 0.20. Added for backward compatibility.
+    def _update_serialized_cpu(self, values: Dict):
+        cpu = values["cpu"]
+        if cpu:
+            arch = cpu.get("arch")
+            count = cpu.get("count")
+            if count and arch in [None, gpuhunt.CPUArchitecture.X86.value]:
+                values["cpu"] = count

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

@@ -325,56 +325,45 @@ def duration(self) -> timedelta:
             end_time = self.finished_at
         return end_time - self.submitted_at
 
-    @root_validator
-    def _status_message(cls, values) -> Dict:
-        try:
-            status = values["status"]
-            termination_reason = values["termination_reason"]
-            exit_code = values["exit_status"]
-        except KeyError:
-            return values
-        values["status_message"] = JobSubmission._get_status_message(
-            status=status,
-            termination_reason=termination_reason,
-            exit_status=exit_code,
-        )
-        return values
+    def dict(self, *args, **kwargs) -> Dict:
+        status_message = self._get_status_message()
+        error = self._get_error()
+        # super() does not work with pydantic-duality
+        res = CoreModel.dict(self, *args, **kwargs)
+        res["status_message"] = status_message
+        res["error"] = error
+        return res
 
-    @staticmethod
-    def _get_status_message(
-        status: JobStatus,
-        termination_reason: Optional[JobTerminationReason],
-        exit_status: Optional[int],
-    ) -> str:
-        if status == JobStatus.DONE:
+    def _get_status_message(self) -> Optional[str]:
+        if self.status == JobStatus.DONE:
             return "exited (0)"
-        elif status == JobStatus.FAILED:
-            if termination_reason == JobTerminationReason.CONTAINER_EXITED_WITH_ERROR:
-                return f"exited ({exit_status})"
-            elif termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY:
+        elif self.status == JobStatus.FAILED:
+            if self.termination_reason == JobTerminationReason.CONTAINER_EXITED_WITH_ERROR:
+                return f"exited ({self.exit_status})"
+            elif (
+                self.termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
+            ):
                 return "no offers"
-            elif termination_reason == JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY:
+            elif self.termination_reason == JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY:
                 return "interrupted"
             else:
                 return "error"
-        elif status == JobStatus.TERMINATED:
-            if termination_reason == JobTerminationReason.TERMINATED_BY_USER:
+        elif self.status == JobStatus.TERMINATED:
+            if self.termination_reason == JobTerminationReason.TERMINATED_BY_USER:
                 return "stopped"
-            elif termination_reason == JobTerminationReason.ABORTED_BY_USER:
+            elif self.termination_reason == JobTerminationReason.ABORTED_BY_USER:
                 return "aborted"
-        return status.value
+        return self.status.value
 
-    @root_validator
-    def _error(cls, values) -> Dict:
-        try:
-            termination_reason = values["termination_reason"]
-        except KeyError:
-            return values
-        values["error"] = JobSubmission._get_error(termination_reason=termination_reason)
-        return values
+    def _get_error(self) -> Optional[str]:
+        return JobSubmission._termination_reason_to_error(
+            termination_reason=self.termination_reason
+        )
 
     @staticmethod
-    def _get_error(termination_reason: Optional[JobTerminationReason]) -> Optional[str]:
+    def _termination_reason_to_error(
+        termination_reason: Optional[JobTerminationReason],
+    ) -> Optional[str]:
         error_mapping = {
             JobTerminationReason.INSTANCE_UNREACHABLE: "instance unreachable",
             JobTerminationReason.WAITING_INSTANCE_LIMIT_EXCEEDED: "waiting instance limit exceeded",
@@ -395,6 +384,12 @@ class Job(CoreModel):
     job_spec: JobSpec
     job_submissions: List[JobSubmission]
 
+    def get_last_termination_reason(self) -> Optional[JobTerminationReason]:
+        for submission in reversed(self.job_submissions):
+            if submission.termination_reason is not None:
+                return submission.termination_reason
+        return None
+
 
 class RunSpec(CoreModel):
     # TODO: run_name, working_dir are redundant here since they already passed in configuration
@@ -525,83 +520,66 @@ class Run(CoreModel):
     last_processed_at: datetime
     status: RunStatus
     status_message: Optional[str] = None
-    termination_reason: Optional[RunTerminationReason]
+    termination_reason: Optional[RunTerminationReason] = None
     run_spec: RunSpec
     jobs: List[Job]
-    latest_job_submission: Optional[JobSubmission]
+    latest_job_submission: Optional[JobSubmission] = None
     cost: float = 0
     service: Optional[ServiceSpec] = None
     deployment_num: int = 0  # default for compatibility with pre-0.19.14 servers
     # TODO: make error a computed field after migrating to pydanticV2
     error: Optional[str] = None
     deleted: Optional[bool] = None
 
-    @root_validator
-    def _error(cls, values) -> Dict:
-        try:
-            termination_reason = values["termination_reason"]
-        except KeyError:
-            return values
-        values["error"] = Run._get_error(termination_reason=termination_reason)
-        return values
+    def dict(self, *args, **kwargs) -> Dict:
+        status_message = self._get_status_message()
+        error = self._get_error()
+        # super() does not work with pydantic-duality
+        res = CoreModel.dict(self, *args, **kwargs)
+        res["status_message"] = status_message
+        res["error"] = error
+        return res
+
+    def _get_error(self) -> Optional[str]:
+        return Run._termination_reason_to_error(termination_reason=self.termination_reason)
 
     @staticmethod
-    def _get_error(termination_reason: Optional[RunTerminationReason]) -> Optional[str]:
+    def _termination_reason_to_error(
+        termination_reason: Optional[RunTerminationReason],
+    ) -> Optional[str]:
         if termination_reason == RunTerminationReason.RETRY_LIMIT_EXCEEDED:
             return "retry limit exceeded"
         elif termination_reason == RunTerminationReason.SERVER_ERROR:
             return "server error"
         else:
             return None
 
-    @root_validator
-    def _status_message(cls, values) -> Dict:
-        try:
-            status = values["status"]
-            jobs: List[Job] = values["jobs"]
-            retry_on_events = (
-                jobs[0].job_spec.retry.on_events if jobs and jobs[0].job_spec.retry else []
-            )
-            job_status = (
-                jobs[0].job_submissions[-1].status
-                if len(jobs) == 1 and jobs[0].job_submissions
-                else None
-            )
-            termination_reason = Run.get_last_termination_reason(jobs[0]) if jobs else None
-        except KeyError:
-            return values
-        values["status_message"] = Run._get_status_message(
-            status=status,
-            job_status=job_status,
-            retry_on_events=retry_on_events,
-            termination_reason=termination_reason,
-        )
-        return values
+    def _get_status_message(self) -> Optional[str]:
+        if len(self.jobs) == 0:
+            return self.status.value
 
-    @staticmethod
-    def get_last_termination_reason(job: "Job") -> Optional[JobTerminationReason]:
-        for submission in reversed(job.job_submissions):
-            if submission.termination_reason is not None:
-                return submission.termination_reason
-        return None
+        last_job = self.jobs[0]
+        last_job_termination_reason = last_job.get_last_termination_reason()
 
-    @staticmethod
-    def _get_status_message(
-        status: RunStatus,
-        job_status: Optional[JobStatus],
-        retry_on_events: List[RetryEvent],
-        termination_reason: Optional[JobTerminationReason],
-    ) -> str:
-        if job_status == JobStatus.PULLING:
-            return "pulling"
+        if len(self.jobs) == 1:
+            # FIXME: Clarify why show "pulling" only in case of one job
+            if (
+                last_job.job_submissions
+                and last_job.job_submissions[-1].status == JobStatus.PULLING
+            ):
+                return "pulling"
+
+        retry_on_events = last_job.job_spec.retry.on_events if last_job.job_spec.retry else []
         # Currently, `retrying` is shown only for `no-capacity` events
         if (
-            status in [RunStatus.SUBMITTED, RunStatus.PENDING]
-            and termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
+            self.status in [RunStatus.SUBMITTED, RunStatus.PENDING]
+            and last_job_termination_reason
+            == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
             and RetryEvent.NO_CAPACITY in retry_on_events
         ):
             return "retrying"
-        return status.value
+
+        return self.status.value
 
     def is_deployment_in_progress(self) -> bool:
         return any(

src/dstack/_internal/server/app.py

@@ -94,7 +94,6 @@ def create_app() -> FastAPI:
     app = FastAPI(
         docs_url="/api/docs",
         lifespan=lifespan,
-        default_response_class=CustomORJSONResponse,
     )
     app.state.proxy_dependency_injector = ServerProxyDependencyInjector()
     return app

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

@@ -13,9 +13,12 @@ class CustomORJSONResponse(Response):
     """
     Custom JSONResponse that uses orjson for serialization.
 
-    It's recommended to return this class from routers directly
-    to avoid the FastAPI's jsonable_encoder overhead.
+    It's recommended to return this class from routers directly instead of
+    returning pydantic models to avoid the FastAPI's jsonable_encoder overhead.
     See https://fastapi.tiangolo.com/advanced/custom-response/#use-orjsonresponse.
+
+    Beware that FastAPI skips model validation when responses are returned directly.
+    If serialization needs to be modified, override `dict()` instead of adding validators.
     """
 
     media_type = "application/json"

src/dstack/_internal/utils/json_utils.py

@@ -1,6 +1,12 @@
 import orjson
 from pydantic import BaseModel
 
+FREEZEGUN = True
+try:
+    from freezegun.api import FakeDatetime
+except ImportError:
+    FREEZEGUN = False
+
 
 def orjson_default(obj):
     if isinstance(obj, float):
@@ -10,6 +16,9 @@ def orjson_default(obj):
         # Allows calling orjson.dumps() on pydantic models
         # (e.g. to return from the API)
         return obj.dict()
+    if FREEZEGUN:
+        if isinstance(obj, FakeDatetime):
+            return obj.isoformat()
     raise TypeError
 
 

src/tests/_internal/core/models/test_runs.py

@@ -33,7 +33,7 @@ def test_job_termination_reason_to_retry_event_works_with_all_enum_variants():
         assert retry_event is None or isinstance(retry_event, RetryEvent)
 
 
-# Will fail if JobTerminationReason value is added without updaing JobSubmission._get_error
+# Will fail if JobTerminationReason value is added without updating JobSubmission._get_error
 def test_get_error_returns_expected_messages():
     no_error_reasons = [
         JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY,
@@ -47,7 +47,7 @@ def test_get_error_returns_expected_messages():
     ]
 
     for reason in JobTerminationReason:
-        if JobSubmission._get_error(reason) is None:
+        if JobSubmission._termination_reason_to_error(reason) is None:
             # Fail no-error reason is not in the list
             assert reason in no_error_reasons
 
@@ -62,6 +62,6 @@ def test_run_get_error_returns_none_for_specific_reasons():
     ]
 
     for reason in RunTerminationReason:
-        if Run._get_error(reason) is None:
+        if Run._termination_reason_to_error(reason) is None:
             # Fail no-error reason is not in the list
             assert reason in no_error_reasons

src/tests/_internal/server/routers/test_fleets.py

@@ -855,8 +855,8 @@ async def test_returns_plan(self, test_db, session: AsyncSession, client: AsyncC
         assert response.json() == {
             "project_name": project.name,
             "user": user.name,
-            "spec": spec.dict(),
-            "effective_spec": spec.dict(),
+            "spec": json.loads(spec.json()),
+            "effective_spec": json.loads(spec.json()),
             "current_resource": None,
             "offers": [json.loads(o.json()) for o in offers],
             "total_offers": len(offers),