Interpolate registry_auth with secrets

Author: r4victor

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

@@ -34,6 +34,7 @@
     run_model_to_run,
     scale_run_replicas,
 )
+from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.services import update_service_desired_replica_count
 from dstack._internal.utils import common
 from dstack._internal.utils.logging import get_logger
@@ -385,7 +386,11 @@ async def _handle_run_replicas(
         )
         return
 
-    await _update_jobs_to_new_deployment_in_place(run_model, run_spec)
+    await _update_jobs_to_new_deployment_in_place(
+        session=session,
+        run_model=run_model,
+        run_spec=run_spec,
+    )
     if _has_out_of_date_replicas(run_model):
         non_terminated_replica_count = len(
             {j.replica_num for j in run_model.jobs if not j.status.is_finished()}
@@ -425,18 +430,25 @@ async def _handle_run_replicas(
             )
 
 
-async def _update_jobs_to_new_deployment_in_place(run_model: RunModel, run_spec: RunSpec) -> None:
+async def _update_jobs_to_new_deployment_in_place(
+    session: AsyncSession, run_model: RunModel, run_spec: RunSpec
+) -> None:
     """
     Bump deployment_num for jobs that do not require redeployment.
     """
-
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=run_model.project,
+    )
     for replica_num, job_models in group_jobs_by_replica_latest(run_model.jobs):
         if all(j.status.is_finished() for j in job_models):
             continue
         if all(j.deployment_num == run_model.deployment_num for j in job_models):
             continue
+        # FIXME: Handle getting image configuration errors or skip it.
         new_job_specs = await get_job_specs_from_run_spec(
             run_spec=run_spec,
+            secrets=secrets,
             replica_num=replica_num,
         )
         assert len(new_job_specs) == len(job_models), (

src/dstack/_internal/server/services/jobs/__init__.py

@@ -64,15 +64,23 @@
 logger = get_logger(__name__)
 
 
-async def get_jobs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[Job]:
+async def get_jobs_from_run_spec(
+    run_spec: RunSpec, secrets: Dict[str, str], replica_num: int
+) -> List[Job]:
     return [
         Job(job_spec=s, job_submissions=[])
-        for s in await get_job_specs_from_run_spec(run_spec, replica_num)
+        for s in await get_job_specs_from_run_spec(
+            run_spec=run_spec,
+            secrets=secrets,
+            replica_num=replica_num,
+        )
     ]
 
 
-async def get_job_specs_from_run_spec(run_spec: RunSpec, replica_num: int) -> List[JobSpec]:
-    job_configurator = _get_job_configurator(run_spec)
+async def get_job_specs_from_run_spec(
+    run_spec: RunSpec, secrets: Dict[str, str], replica_num: int
+) -> List[JobSpec]:
+    job_configurator = _get_job_configurator(run_spec=run_spec, secrets=secrets)
     job_specs = await job_configurator.get_job_specs(replica_num=replica_num)
     return job_specs
 
@@ -158,10 +166,10 @@ def delay_job_instance_termination(job_model: JobModel):
     job_model.remove_at = common.get_current_datetime() + timedelta(seconds=15)
 
 
-def _get_job_configurator(run_spec: RunSpec) -> JobConfigurator:
+def _get_job_configurator(run_spec: RunSpec, secrets: Dict[str, str]) -> JobConfigurator:
     configuration_type = RunConfigurationType(run_spec.configuration.type)
     configurator_class = _configuration_type_to_configurator_class_map[configuration_type]
-    return configurator_class(run_spec)
+    return configurator_class(run_spec=run_spec, secrets=secrets)
 
 
 _job_configurator_classes = [

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

@@ -68,8 +68,13 @@ class JobConfigurator(ABC):
     # JobSSHKey should be shared for all jobs in a replica for inter-node communication.
     _job_ssh_key: Optional[JobSSHKey] = None
 
-    def __init__(self, run_spec: RunSpec):
+    def __init__(
+        self,
+        run_spec: RunSpec,
+        secrets: Optional[Dict[str, str]] = None,
+    ):
         self.run_spec = run_spec
+        self.secrets = secrets or {}
 
     async def get_job_specs(self, replica_num: int) -> List[JobSpec]:
         job_spec = await self._get_job_spec(replica_num=replica_num, job_num=0, jobs_per_replica=1)
@@ -98,10 +103,20 @@ def _ports(self) -> List[PortMapping]:
     async def _get_image_config(self) -> ImageConfig:
         if self._image_config is not None:
             return self._image_config
+        interpolate = VariablesInterpolator({"secrets": self.secrets}).interpolate_or_error
+        registry_auth = self.run_spec.configuration.registry_auth
+        if registry_auth is not None:
+            try:
+                registry_auth = RegistryAuth(
+                    username=interpolate(registry_auth.username),
+                    password=interpolate(registry_auth.password),
+                )
+            except InterpolatorError as e:
+                raise ServerClientError(e.args[0])
         image_config = await run_async(
             _get_image_config,
             self._image_name(),
-            self.run_spec.configuration.registry_auth,
+            registry_auth,
         )
         self._image_config = image_config
         return image_config

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

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.configurations import PortMapping, RunConfigurationType
@@ -17,7 +17,7 @@
 class DevEnvironmentJobConfigurator(JobConfigurator):
     TYPE: RunConfigurationType = RunConfigurationType.DEV_ENVIRONMENT
 
-    def __init__(self, run_spec: RunSpec):
+    def __init__(self, run_spec: RunSpec, secrets: Dict[str, str]):
         if run_spec.configuration.ide == "vscode":
             __class = VSCodeDesktop
         elif run_spec.configuration.ide == "cursor":
@@ -29,7 +29,7 @@ def __init__(self, run_spec: RunSpec):
             version=run_spec.configuration.version,
             extensions=["ms-python.python", "ms-toolsai.jupyter"],
         )
-        super().__init__(run_spec)
+        super().__init__(run_spec=run_spec, secrets=secrets)
 
     def _shell_commands(self) -> List[str]:
         commands = self.ide.get_install_commands()

src/dstack/_internal/server/services/runs.py

@@ -82,6 +82,7 @@
 from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.services.projects import list_project_models, list_user_project_models
 from dstack._internal.server.services.resources import set_resources_defaults
+from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.users import get_user_model_by_name
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.random_names import generate_name
@@ -311,7 +312,12 @@ async def get_plan(
             ):
                 action = ApplyAction.UPDATE
 
-    jobs = await get_jobs_from_run_spec(effective_run_spec, replica_num=0)
+    secrets = await get_project_secrets_mapping(session=session, project=project)
+    jobs = await get_jobs_from_run_spec(
+        run_spec=effective_run_spec,
+        secrets=secrets,
+        replica_num=0,
+    )
 
     volumes = await get_job_configured_volumes(
         session=session,
@@ -462,6 +468,10 @@ async def submit_run(
         project=project,
         run_spec=run_spec,
     )
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=project,
+    )
 
     lock_namespace = f"run_names_{project.name}"
     if get_db().dialect_name == "sqlite":
@@ -513,7 +523,11 @@ async def submit_run(
             await services.register_service(session, run_model, run_spec)
 
         for replica_num in range(replicas):
-            jobs = await get_jobs_from_run_spec(run_spec, replica_num=replica_num)
+            jobs = await get_jobs_from_run_spec(
+                run_spec=run_spec,
+                secrets=secrets,
+                replica_num=replica_num,
+            )
             for job in jobs:
                 job_model = create_job_model_for_new_submission(
                     run_model=run_model,
@@ -1068,10 +1082,20 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
             await retry_run_replica_jobs(session, run_model, replica_jobs, only_failed=False)
             scheduled_replicas += 1
 
+        secrets = await get_project_secrets_mapping(
+            session=session,
+            project=run_model.project,
+        )
+
         for replica_num in range(
             len(active_replicas) + scheduled_replicas, len(active_replicas) + replicas_diff
         ):
-            jobs = await get_jobs_from_run_spec(run_spec, replica_num=replica_num)
+            # FIXME: Handle getting image configuration errors or skip it.
+            jobs = await get_jobs_from_run_spec(
+                run_spec=run_spec,
+                secrets=secrets,
+                replica_num=replica_num,
+            )
             for job in jobs:
                 job_model = create_job_model_for_new_submission(
                     run_model=run_model,
@@ -1084,8 +1108,14 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
 async def retry_run_replica_jobs(
     session: AsyncSession, run_model: RunModel, latest_jobs: List[JobModel], *, only_failed: bool
 ):
+    # FIXME: Handle getting image configuration errors or skip it.
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=run_model.project,
+    )
     new_jobs = await get_jobs_from_run_spec(
-        RunSpec.__response__.parse_raw(run_model.run_spec),
+        run_spec=RunSpec.__response__.parse_raw(run_model.run_spec),
+        secrets=secrets,
         replica_num=latest_jobs[0].replica_num,
     )
     assert len(new_jobs) == len(latest_jobs), (

src/dstack/_internal/server/testing/common.py

@@ -316,7 +316,9 @@ async def create_job(
     if deployment_num is None:
         deployment_num = run.deployment_num
     run_spec = RunSpec.parse_raw(run.run_spec)
-    job_spec = (await get_job_specs_from_run_spec(run_spec, replica_num=replica_num))[0]
+    job_spec = (
+        await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=replica_num)
+    )[0]
     job_spec.job_num = job_num
     job = JobModel(
         project_id=run.project_id,

src/tests/_internal/server/services/test_runs.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Union
 
 import pytest
 from pydantic import parse_obj_as
@@ -30,7 +30,7 @@ async def make_run(
     session: AsyncSession,
     replicas_statuses: List[JobStatus],
     status: RunStatus = RunStatus.RUNNING,
-    replicas: str = 1,
+    replicas: Union[str, int] = 1,
 ) -> RunModel:
     project = await create_project(session=session)
     user = await create_user(session=session)
@@ -70,7 +70,7 @@ async def make_run(
             status=job_status,
             replica_num=replica_num,
         )
-    await session.refresh(run)
+    await session.refresh(run, attribute_names=["project", "jobs"])
     return run