diff --git a/docs/docs/concepts/services.md b/docs/docs/concepts/services.md index b33f561bc8..368e9ad8e9 100644 --- a/docs/docs/concepts/services.md +++ b/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. + +
+ + ```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 + ``` + +
See the [reference](../reference/dstack.yml/service.md#probes) for more probe configuration options. diff --git a/docs/docs/reference/dstack.yml/service.md b/docs/docs/reference/dstack.yml/service.md index 92b7944844..85612d862c 100644 --- a/docs/docs/reference/dstack.yml/service.md +++ b/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 diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 4a554a9d72..363497e01f 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/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]) diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 2078904640..39530696ab 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/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 diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index 1e7754e5b2..c3635d3181 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/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 diff --git a/src/dstack/_internal/server/background/tasks/process_probes.py b/src/dstack/_internal/server/background/tasks/process_probes.py index 1ce618b12b..5ed9375d13 100644 --- a/src/dstack/_internal/server/background/tasks/process_probes.py +++ b/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 diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 92fadf3447..1a67ad3cf7 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/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, ) diff --git a/src/dstack/_internal/utils/common.py b/src/dstack/_internal/utils/common.py index 56382abd38..2b2cfa4aad 100644 --- a/src/dstack/_internal/utils/common.py +++ b/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") diff --git a/src/tests/_internal/utils/test_common.py b/src/tests/_internal/utils/test_common.py index 1764fedcff..140627580f 100644 --- a/src/tests/_internal/utils/test_common.py +++ b/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",