Review and TODO fixes

- Docs
- Compatibility
- More tests
- Optimized DB queries
- Limit max probes and probe timeout
- Defaults out of `ProbeConfig`
- Stricter validation
- Migration rebase
- Locking fix

Author: jvstme

docs/docs/concepts/services.md

@@ -187,6 +187,53 @@ port: 8000
 
 </div>
 
+### Probes
+
+Configure one or more HTTP probes to periodically check the health of the service.
+
+<div editor-title="service.dstack.yml">
+
+```yaml
+type: service
+name: my-service
+port: 80
+image: my-app:latest
+probes:
+- type: http
+  url: /health
+  interval: 15s
+```
+
+</div>
+
+You can track probe statuses in `dstack ps --verbose`.
+
+<div class="termy">
+
+```shell
+$ dstack ps --verbose
+
+ NAME                            BACKEND          STATUS   PROBES  SUBMITTED
+ my-service deployment=1                          running          11 mins ago
+   replica=0 job=0 deployment=0  aws (us-west-2)  running  ✓       11 mins ago
+   replica=1 job=0 deployment=1  aws (us-west-2)  running  ×       1 min ago
+```
+
+</div>
+
+??? info "Probe statuses"
+    The following symbols are used for probe statuses:
+
+    - `×` &mdash; the last probe execution failed.
+    - `~` &mdash; the last probe execution succeeded, but the [`ready_after`](../reference/dstack.yml/service.md#ready_after) threshold is not yet reached.
+    - `✓` &mdash; the last `ready_after` probe executions succeeded.
+
+    If multiple probes are configured for the service, their statuses are displayed in the order in which the probes appear in the configuration.
+
+Probes are executed for each service replica while the replica is `running`. Probe statuses do not affect how `dstack` handles replicas, except during [rolling deployments](#rolling-deployment).
+
+See the [reference](../reference/dstack.yml/service.md#probes) for more probe configuration options.
+
 ### Path prefix { #path-prefix }
 
 If your `dstack` project doesn't have a [gateway](gateways.md), services are hosted with the
@@ -758,7 +805,7 @@ Update the run? [y/n]:
 
 </div>
 
-If approved, `dstack` gradually updates the service replicas. To update a replica, `dstack` starts a new replica, waits for it to become `running`, then terminates the old replica. This process is repeated for each replica, one at a time.
+If approved, `dstack` gradually updates the service replicas. To update a replica, `dstack` starts a new replica, waits for it to become `running` and for all of its [probes](#probes) to pass, then terminates the old replica. This process is repeated for each replica, one at a time.
 
 You can track the progress of rolling deployment in both `dstack apply` or `dstack ps`. 
 Older replicas have lower `deployment` numbers; newer ones have higher.
@@ -792,8 +839,6 @@ The rolling deployment stops when all replicas are updated or when a new deploym
     To trigger a rolling deployment when no properties have changed (e.g., after updating [secrets](secrets.md) or to restart all replicas),  
     make a minor config change, such as adding a dummy [environment variable](#environment-variables).
 
-<!-- TODO: probes -->
-
 --8<-- "docs/concepts/snippets/manage-runs.ext"
 
 !!! info "What's next?"

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

@@ -106,6 +106,16 @@ The `service` configuration type allows running [services](../../concepts/servic
           type:
             required: true
 
+### `probes`
+
+#### `probes[n]`
+
+#SCHEMA# dstack._internal.core.models.configurations.ProbeConfig
+    overrides:
+      show_root_heading: false
+      type:
+        required: true
+
 ### `retry`
 
 #SCHEMA# dstack._internal.core.models.profiles.ProfileRetry

docs/docs/reference/environment-variables.md

@@ -124,6 +124,8 @@ For more details on the options below, refer to the [server deployment](../guide
 - `DSTACK_DB_MAX_OVERFLOW`{ #DSTACK_DB_MAX_OVERFLOW } - The client DB connections pool allowed overflow. Defaults to `20`.
 - `DSTACK_SERVER_BACKGROUND_PROCESSING_FACTOR`{ #DSTACK_SERVER_BACKGROUND_PROCESSING_FACTOR } - The number of background jobs for processing server resources. Increase if you need to process more resources per server replica quickly. Defaults to `1`.
 - `DSTACK_SERVER_BACKGROUND_PROCESSING_DISABLED`{ #DSTACK_SERVER_BACKGROUND_PROCESSING_DISABLED } - Disables background processing if set to any value. Useful to run only web frontend and API server.
+- `DSTACK_SERVER_MAX_PROBES_PER_JOB`{ #DSTACK_SERVER_MAX_PROBES_PER_JOB } - Maximum number of probes allowed in a run configuration. Validated at apply time.
+- `DSTACK_SERVER_MAX_PROBE_TIMEOUT`{ #DSTACK_SERVER_MAX_PROBE_TIMEOUT } - Maximum allowed timeout for a probe. Validated at apply time.
 
 ??? info "Internal environment variables"
      The following environment variables are intended for development purposes:

src/dstack/_internal/cli/utils/run.py

@@ -5,7 +5,7 @@
 from rich.table import Table
 
 from dstack._internal.cli.utils.common import NO_OFFERS_WARNING, add_row_from_dict, console
-from dstack._internal.core.models.configurations import DevEnvironmentConfiguration, ProbeConfig
+from dstack._internal.core.models.configurations import DevEnvironmentConfiguration
 from dstack._internal.core.models.instances import InstanceAvailability
 from dstack._internal.core.models.profiles import (
     DEFAULT_RUN_TERMINATION_IDLE_TIME,
@@ -14,6 +14,7 @@
 from dstack._internal.core.models.runs import (
     JobStatus,
     Probe,
+    ProbeSpec,
     RunPlan,
 )
 from dstack._internal.core.services.profiles import get_termination
@@ -241,13 +242,14 @@ def get_runs_table(
 
 
 def _format_job_probes(
-    probe_configs: list[ProbeConfig], probes: list[Probe], job_status: JobStatus
+    probe_specs: list[ProbeSpec], probes: list[Probe], job_status: JobStatus
 ) -> str:
     if not probes or job_status != JobStatus.RUNNING:
         return ""
     statuses = []
-    for probe_config, probe in zip(probe_configs, probes):
-        if probe.success_streak >= probe_config.ready_after:
+    for probe_spec, probe in zip(probe_specs, probes):
+        # NOTE: the symbols are documented in concepts/services.md, keep in sync.
+        if probe.success_streak >= probe_spec.ready_after:
             status = "[code]✓[/]"
         elif probe.success_streak > 0:
             status = "[warning]~[/]"

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

@@ -53,6 +53,8 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD
             job_submissions_excludes["exit_status"] = True
         if all(js.deployment_num == 0 for js in job_submissions):
             job_submissions_excludes["deployment_num"] = True
+        if all(not js.probes for js in job_submissions):
+            job_submissions_excludes["probes"] = True
         latest_job_submission = current_resource.latest_job_submission
         if latest_job_submission is not None:
             latest_job_submission_excludes: IncludeExcludeDictType = {}
@@ -69,6 +71,8 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD
                 latest_job_submission_excludes["exit_status"] = True
             if latest_job_submission.deployment_num == 0:
                 latest_job_submission_excludes["deployment_num"] = True
+            if not latest_job_submission.probes:
+                latest_job_submission_excludes["probes"] = True
     return {"plan": apply_plan_excludes}
 
 
@@ -120,7 +124,8 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
         profile_excludes.add("startup_order")
     if configuration.stop_criteria is None:
         configuration_excludes["stop_criteria"] = True
-    # TODO: probes
+    if isinstance(configuration, ServiceConfiguration) and not configuration.probes:
+        configuration_excludes["probes"] = True
     if profile is not None and profile.stop_criteria is None:
         profile_excludes.add("stop_criteria")
     if not configuration.files:
@@ -155,7 +160,8 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType:
         spec_excludes["file_archives"] = True
     if all(s.service_port is None for s in job_specs):
         spec_excludes["service_port"] = True
-    # TODO: probes
+    if all(not s.probes for s in job_specs):
+        spec_excludes["probes"] = True
 
     return spec_excludes
 

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

@@ -34,6 +34,11 @@
 DEFAULT_REPO_DIR = "/workflow"
 MIN_PROBE_TIMEOUT = 1
 MIN_PROBE_INTERVAL = 1
+DEFAULT_PROBE_URL = "/"
+DEFAULT_PROBE_TIMEOUT = 10
+DEFAULT_PROBE_INTERVAL = 15
+DEFAULT_PROBE_READY_AFTER = 1
+MAX_PROBE_URL_LEN = 2048
 
 
 class RunConfigurationType(str, Enum):
@@ -166,53 +171,69 @@ class RateLimit(CoreModel):
 
 class ProbeConfig(CoreModel):
     type: Literal["http"]  # expect other probe types in the future, namely `exec`
-    url: Annotated[str, Field(description="The URL to request")] = "/"
+    url: Annotated[
+        Optional[str], Field(description=f"The URL to request. Defaults to `{DEFAULT_PROBE_URL}`")
+    ] = None
     timeout: Annotated[
-        Union[int, str],
-        Field(description=("Maximum amount of time the HTTP request is allowed to take")),
-    ] = "10s"
+        Optional[Union[int, str]],
+        Field(
+            description=(
+                f"Maximum amount of time the HTTP request is allowed to take. Defaults to `{DEFAULT_PROBE_TIMEOUT}s`"
+            )
+        ),
+    ] = None
     interval: Annotated[
-        Union[int, str],
+        Optional[Union[int, str]],
         Field(
             description=(
                 "Minimum amount of time between the end of one probe execution"
-                " and the start of the next"
+                f" and the start of the next. Defaults to `{DEFAULT_PROBE_INTERVAL}s`"
             )
         ),
-    ] = "15s"
+    ] = None
     ready_after: Annotated[
-        int,
+        Optional[int],
         Field(
             ge=1,
             description=(
-                "The number of consecutive successful probe executions required for the job"
-                " to be considered ready. Used during rolling deployments"
+                "The number of consecutive successful probe executions required for the replica"
+                " to be considered ready. Used during rolling deployments."
+                f" Defaults to `{DEFAULT_PROBE_READY_AFTER}`"
             ),
         ),
-    ] = 1
+    ] = None
 
     class Config:
         frozen = True
 
     @validator("timeout")
-    def parse_timeout(cls, v: Union[int, str]) -> int:
+    def parse_timeout(cls, v: Optional[Union[int, str]]) -> Optional[int]:
+        if v is None:
+            return v
         parsed = parse_duration(v)
         if parsed < MIN_PROBE_TIMEOUT:
             raise ValueError(f"Probe timeout cannot be shorter than {MIN_PROBE_TIMEOUT}s")
         return parsed
 
     @validator("interval")
-    def parse_interval(cls, v: Union[int, str]) -> int:
+    def parse_interval(cls, v: Optional[Union[int, str]]) -> Optional[int]:
+        if v is None:
+            return v
         parsed = parse_duration(v)
         if parsed < MIN_PROBE_INTERVAL:
             raise ValueError(f"Probe interval cannot be shorter than {MIN_PROBE_INTERVAL}s")
         return parsed
 
     @validator("url")
-    def validate_url(cls, v: str) -> str:
-        # TODO: stricter constraints to avoid HTTPX URL parsing errors
+    def validate_url(cls, v: Optional[str]) -> Optional[str]:
+        if v is None:
+            return v
         if not v.startswith("/"):
             raise ValueError("Must start with `/`")
+        if len(v) > MAX_PROBE_URL_LEN:
+            raise ValueError(f"Cannot be longer than {MAX_PROBE_URL_LEN} characters")
+        if not v.isprintable():
+            raise ValueError("Cannot contain non-printable characters")
         return v
 
 
@@ -504,7 +525,7 @@ class ServiceConfigurationParams(CoreModel):
     rate_limits: Annotated[list[RateLimit], Field(description="Rate limiting rules")] = []
     probes: Annotated[
         list[ProbeConfig],
-        Field(unique_items=True, description="List of probes used to determine job health"),
+        Field(description="List of probes used to determine job health"),
     ] = []
 
     @validator("port")
@@ -569,6 +590,16 @@ def validate_rate_limits(cls, v: list[RateLimit]) -> list[RateLimit]:
             )
         return v
 
+    @validator("probes")
+    def validate_probes(cls, v: list[ProbeConfig]) -> list[ProbeConfig]:
+        if len(v) != len(set(v)):
+            # Using a custom validator instead of Field(unique_items=True) to avoid Pydantic bug:
+            # https://github.com/pydantic/pydantic/issues/3765
+            # Because of the bug, our gen_schema_reference.py fails to determine the type of
+            # ServiceConfiguration.probes and insert the correct hyperlink.
+            raise ValueError("Probes must be unique")
+        return v
+
 
 class ServiceConfiguration(
     ProfileParams, BaseRunConfigurationWithCommands, ServiceConfigurationParams

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

@@ -1,6 +1,6 @@
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Literal, Optional, Type
 
 from pydantic import UUID4, Field, root_validator
 from typing_extensions import Annotated
@@ -10,7 +10,6 @@
 from dstack._internal.core.models.configurations import (
     DEFAULT_REPO_DIR,
     AnyRunConfiguration,
-    ProbeConfig,
     RunConfiguration,
     ServiceConfiguration,
 )
@@ -224,6 +223,14 @@ class JobSSHKey(CoreModel):
     public: str
 
 
+class ProbeSpec(CoreModel):
+    type: Literal["http"]  # expect other probe types in the future, namely `exec`
+    url: str
+    timeout: int
+    interval: int
+    ready_after: int
+
+
 class JobSpec(CoreModel):
     replica_num: int = 0  # default value for backward compatibility
     job_num: int
@@ -257,7 +264,7 @@ class JobSpec(CoreModel):
     file_archives: list[FileArchiveMapping] = []
     # None for non-services and pre-0.19.19 services. See `get_service_port`
     service_port: Optional[int] = None
-    probes: list[ProbeConfig] = []
+    probes: list[ProbeSpec] = []
 
 
 class JobProvisioningData(CoreModel):