From 523ec01dc97fe372429816183d7318c21e02c058 Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Fri, 20 Mar 2026 15:00:29 +0100 Subject: [PATCH] Do not request image config for default images Do not request the Docker image config for our default images when determining the user, since we know that our images do not set it. This helps avoid hitting Docker Hub's rate limits. --- .../services/jobs/configurators/base.py | 2 +- .../server/services/jobs/test_jobs.py | 69 +++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) create mode 100644 src/tests/_internal/server/services/jobs/test_jobs.py diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index b73c9bbe6..5bd7cad47 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -280,7 +280,7 @@ def _image_name(self) -> str: async def _user(self) -> Optional[UnixUser]: user = self.run_spec.configuration.user - if user is None: + if user is None and self.run_spec.configuration.image is not None: image_config = await self._get_image_config() user = image_config.user if user is None: diff --git a/src/tests/_internal/server/services/jobs/test_jobs.py b/src/tests/_internal/server/services/jobs/test_jobs.py new file mode 100644 index 000000000..53b427732 --- /dev/null +++ b/src/tests/_internal/server/services/jobs/test_jobs.py @@ -0,0 +1,69 @@ +from unittest.mock import patch + +import pytest + +from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.profiles import Profile +from dstack._internal.core.models.repos.local import LocalRunRepoData +from dstack._internal.core.models.runs import RunSpec +from dstack._internal.server.services.docker import ImageConfig +from dstack._internal.server.services.jobs import get_job_specs_from_run_spec + + +@pytest.mark.parametrize( + "configuration, expected_calls", + [ + pytest.param( + # No need to request the registry if our default image is used. + TaskConfiguration(commands=["sleep infinity"]), + 0, + id="default-dstack-image", + ), + pytest.param( + TaskConfiguration(image="ubuntu"), + 1, + id="custom-image", + ), + pytest.param( + TaskConfiguration(image="ubuntu", commands=["sleep infinity"]), + 1, + id="custom-image-with-commands", + ), + pytest.param( + TaskConfiguration(image="ubuntu", user="root"), + 1, + id="custom-image-with-user", + ), + pytest.param( + # Setting `commands` and `user` is a known hack that we advertised to some customers + # to avoid registry requests. + TaskConfiguration(image="ubuntu", commands=["sleep infinity"], user="root"), + 0, + id="custom-image-with-commands-and-user", + ), + ], +) +@pytest.mark.asyncio +async def test_get_job_specs_from_run_spec_image_config_calls( + configuration: TaskConfiguration, expected_calls: int +) -> None: + """ + Test the number of times we attempt to fetch the image config from the Docker registry. + + Whenever possible, we prefer not to request the registry to avoid hitting rate limits. + """ + + run_spec = RunSpec( + run_name="test-run", + repo_data=LocalRunRepoData(repo_dir="/"), + configuration=configuration, + profile=Profile(name="default"), + ssh_key_pub="user_ssh_key", + ) + fake_image_config = ImageConfig.parse_obj({"Entrypoint": ["/bin/bash"]}) + with patch( + "dstack._internal.server.services.jobs.configurators.base._get_image_config", + return_value=fake_image_config, + ) as mock_get_image_config: + await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) + assert mock_get_image_config.call_count == expected_calls