Add more probe request configuration options

docs/docs/concepts/services.md

@@ -230,7 +230,34 @@ $ dstack ps --verbose
 
     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).
+Probes are executed for each service replica while the replica is `running`. A probe execution is considered successful if the replica responds with a `2xx` status code. Probe statuses do not affect how `dstack` handles replicas, except during [rolling deployments](#rolling-deployment).
+
+??? info "HTTP request configuration"
+    You can configure the HTTP request method, headers, and other properties. To include secret values in probe requests, use environment variable interpolation, which is enabled for the `url`, `headers[i].value`, and `body` properties.
+
+    <div editor-title="service.dstack.yml">
+
+    ```yaml
+    type: service
+    name: my-service
+    port: 80
+    image: my-app:latest
+    env:
+    - PROBES_API_KEY
+    probes:
+    - type: http
+      method: post
+      url: /check-health
+      headers:
+      - name: X-API-Key
+        value: ${{ env.PROBES_API_KEY }}
+      - name: Content-Type
+        value: application/json
+      body: '{"level": 2}'
+      timeout: 20s
+    ```
+
+    </div>
 
 See the [reference](../reference/dstack.yml/service.md#probes) for more probe configuration options.
 

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

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

src/dstack/_internal/cli/services/configurators/run.py

@@ -341,6 +341,14 @@ def interpolate_env(self, conf: BaseRunConfiguration):
                     username=interpolator.interpolate_or_error(conf.registry_auth.username),
                     password=interpolator.interpolate_or_error(conf.registry_auth.password),
                 )
+            if isinstance(conf, ServiceConfiguration):
+                for probe in conf.probes:
+                    for header in probe.headers:
+                        header.value = interpolator.interpolate_or_error(header.value)
+                    if probe.url:
+                        probe.url = interpolator.interpolate_or_error(probe.url)
+                    if probe.body:
+                        probe.body = interpolator.interpolate_or_error(probe.body)
         except InterpolatorError as e:
             raise ConfigurationError(e.args[0])
 

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

@@ -19,6 +19,7 @@
 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.common import has_duplicates
 from dstack._internal.utils.json_utils import (
     pydantic_orjson_dumps_with_indent,
 )
@@ -38,6 +39,7 @@
 DEFAULT_PROBE_TIMEOUT = 10
 DEFAULT_PROBE_INTERVAL = 15
 DEFAULT_PROBE_READY_AFTER = 1
+DEFAULT_PROBE_METHOD = "get"
 MAX_PROBE_URL_LEN = 2048
 
 
@@ -169,11 +171,54 @@ class RateLimit(CoreModel):
     ] = 0
 
 
+HTTPMethod = Literal["get", "post", "put", "delete", "patch", "head"]
+
+
+class HTTPHeaderSpec(CoreModel):
+    name: Annotated[
+        str,
+        Field(
+            description="The name of the HTTP header",
+            min_length=1,
+            max_length=256,
+        ),
+    ]
+    value: Annotated[
+        str,
+        Field(
+            description="The value of the HTTP header",
+            min_length=1,
+            max_length=2048,
+        ),
+    ]
+
+
 class ProbeConfig(CoreModel):
     type: Literal["http"]  # expect other probe types in the future, namely `exec`
     url: Annotated[
         Optional[str], Field(description=f"The URL to request. Defaults to `{DEFAULT_PROBE_URL}`")
     ] = None
+    method: Annotated[
+        Optional[HTTPMethod],
+        Field(
+            description=(
+                "The HTTP method to use for the probe (e.g., `get`, `post`, etc.)."
+                f" Defaults to `{DEFAULT_PROBE_METHOD}`"
+            )
+        ),
+    ] = None
+    headers: Annotated[
+        list[HTTPHeaderSpec],
+        Field(description="A list of HTTP headers to include in the request", max_items=16),
+    ] = []
+    body: Annotated[
+        Optional[str],
+        Field(
+            description="The HTTP request body to send with the probe",
+            min_length=1,
+            max_length=2048,
+        ),
+    ] = None
     timeout: Annotated[
         Optional[Union[int, str]],
         Field(
@@ -203,9 +248,6 @@ class ProbeConfig(CoreModel):
         ),
     ] = None
 
-    class Config:
-        frozen = True
-
     @validator("timeout")
     def parse_timeout(cls, v: Optional[Union[int, str]]) -> Optional[int]:
         if v is None:
@@ -236,6 +278,13 @@ def validate_url(cls, v: Optional[str]) -> Optional[str]:
             raise ValueError("Cannot contain non-printable characters")
         return v
 
+    @root_validator
+    def validate_body_matches_method(cls, values):
+        method: HTTPMethod = values["method"]
+        if values["body"] is not None and method in ["get", "head"]:
+            raise ValueError(f"Cannot set request body for the `{method}` method")
+        return values
+
 
 class BaseRunConfiguration(CoreModel):
     type: Literal["none"]
@@ -592,7 +641,7 @@ def validate_rate_limits(cls, v: list[RateLimit]) -> list[RateLimit]:
 
     @validator("probes")
     def validate_probes(cls, v: list[ProbeConfig]) -> list[ProbeConfig]:
-        if len(v) != len(set(v)):
+        if has_duplicates(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

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

@@ -8,8 +8,11 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import ApplyAction, CoreModel, NetworkMode, RegistryAuth
 from dstack._internal.core.models.configurations import (
+    DEFAULT_PROBE_METHOD,
     DEFAULT_REPO_DIR,
     AnyRunConfiguration,
+    HTTPHeaderSpec,
+    HTTPMethod,
     RunConfiguration,
     ServiceConfiguration,
 )
@@ -226,6 +229,9 @@ class JobSSHKey(CoreModel):
 class ProbeSpec(CoreModel):
     type: Literal["http"]  # expect other probe types in the future, namely `exec`
     url: str
+    method: HTTPMethod = DEFAULT_PROBE_METHOD
+    headers: list[HTTPHeaderSpec] = []
+    body: Optional[str] = None
     timeout: int
     interval: int
     ready_after: int

src/dstack/_internal/server/background/tasks/process_probes.py

@@ -116,9 +116,13 @@ async def _execute_probe(probe: ProbeModel, probe_spec: ProbeSpec) -> bool:
 
     try:
         async with _get_service_replica_client(probe.job) as client:
-            resp = await client.get(
-                "http://dstack" + probe_spec.url,
+            resp = await client.request(
+                method=probe_spec.method,
+                url="http://dstack" + probe_spec.url,
+                headers=[(h.name, h.value) for h in probe_spec.headers],
+                data=probe_spec.body,
                 timeout=probe_spec.timeout,
+                follow_redirects=False,
             )
             logger.debug("%s: probe status code: %s", fmt(probe), resp.status_code)
             return resp.is_success

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -12,6 +12,7 @@
 from dstack._internal.core.models.common import RegistryAuth
 from dstack._internal.core.models.configurations import (
     DEFAULT_PROBE_INTERVAL,
+    DEFAULT_PROBE_METHOD,
     DEFAULT_PROBE_READY_AFTER,
     DEFAULT_PROBE_TIMEOUT,
     DEFAULT_PROBE_URL,
@@ -372,6 +373,9 @@ def _probe_config_to_spec(c: ProbeConfig) -> ProbeSpec:
         timeout=c.timeout if c.timeout is not None else DEFAULT_PROBE_TIMEOUT,
         interval=c.interval if c.interval is not None else DEFAULT_PROBE_INTERVAL,
         ready_after=c.ready_after if c.ready_after is not None else DEFAULT_PROBE_READY_AFTER,
+        method=c.method if c.method is not None else DEFAULT_PROBE_METHOD,
+        headers=c.headers,
+        body=c.body,
     )
 
 

src/dstack/_internal/utils/common.py

@@ -222,6 +222,21 @@ def remove_prefix(text: str, prefix: str) -> str:
     return text
 
 
+def has_duplicates(iterable: Iterable[Any]) -> bool:
+    """
+    Checks if there are any duplicate items in the given iterable.
+
+    O(n^2) implementation, but works with iterables with unhashable items.
+    For iterables with hashable items, prefer len(set(iterable)) != len(iterable).
+    """
+    seen = []
+    for item in iterable:
+        if item in seen:
+            return True
+        seen.append(item)
+    return False
+
+
 T = TypeVar("T")
 
 

src/tests/_internal/utils/test_common.py

@@ -8,6 +8,7 @@
     batched,
     concat_url_path,
     format_duration_multiunit,
+    has_duplicates,
     local_time,
     make_proxy_url,
     parse_memory,
@@ -124,6 +125,30 @@ def test_forbids_negative(self) -> None:
             format_duration_multiunit(-1)
 
 
+@pytest.mark.parametrize(
+    "iterable, expected",
+    [
+        ([1, 2, 3, 4], False),
+        ([1, 2, 3, 4, 2], True),
+        (iter([1, 2, 3, 4]), False),
+        (iter([1, 2, 3, 4, 2]), True),
+        ("abcde", False),
+        ("hello", True),
+        ([1, "a"], False),
+        ([1, "a", 1], True),
+        ([[1, 2], [3, 4]], False),
+        ([[1, 2], [1, 2]], True),
+        ([{"a": "b"}, {"a": "c"}], False),
+        ([{"a": "b"}, {"a": "b"}], True),
+        ([{}, {}], True),
+        ([], False),
+        ([1], False),
+    ],
+)
+def test_has_duplicates(iterable, expected):
+    assert has_duplicates(iterable) == expected
+
+
 class TestParseMemory:
     @pytest.mark.parametrize(
         "memory,as_units,expected",