Fix missing jrd update

Author: r4victor

src/dstack/_internal/server/background/scheduled_tasks/running_jobs.py

@@ -470,7 +470,7 @@ async def _process_running_job_pulling_state(
         context.job_submission.age,
     )
     shim_state = await common_utils.run_async(
-        _get_shim_pulling_state,
+        _sync_shim_pulling_state,
         server_ssh_private_keys,
         job_provisioning_data,
         None,
@@ -481,11 +481,12 @@ async def _process_running_job_pulling_state(
         return
 
     if shim_state == _ShimPullingState.READY:
+        job_runtime_data = get_job_runtime_data(context.job_model)
         runner_availability = await common_utils.run_async(
             _get_runner_availability,
             server_ssh_private_keys,
             job_provisioning_data,
-            None,
+            job_runtime_data,
         )
         if runner_availability == _RunnerAvailability.UNAVAILABLE:
             _reset_disconnected_at(session, context.job_model)
@@ -507,7 +508,7 @@ async def _process_running_job_pulling_state(
                 _submit_job_to_runner,
                 server_ssh_private_keys,
                 job_provisioning_data,
-                None,
+                job_runtime_data,
                 session=session,
                 run=context.run,
                 job_model=context.job_model,
@@ -856,7 +857,7 @@ def _get_runner_availability(ports: Dict[int, int]) -> _RunnerAvailability:
 
 
 @runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT])
-def _get_shim_pulling_state(
+def _sync_shim_pulling_state(
     ports: Dict[int, int],
     job_model: JobModel,
 ) -> Union[Literal[False], _ShimPullingState]:

src/tests/_internal/server/background/scheduled_tasks/test_running_jobs.py

@@ -39,6 +39,7 @@
 from dstack._internal.server import settings as server_settings
 from dstack._internal.server.background.scheduled_tasks.running_jobs import (
     _patch_base_image_for_aws_efa,
+    _RunnerAvailability,
     process_running_jobs,
 )
 from dstack._internal.server.models import JobModel
@@ -579,6 +580,91 @@ async def test_pulling_shim_runner_not_ready(
         assert job is not None
         assert job.status == JobStatus.PULLING
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_pulling_shim_uses_runtime_port_mapping_for_runner_calls(
+        self,
+        test_db,
+        session: AsyncSession,
+        ssh_tunnel_mock: Mock,
+        shim_client_mock: Mock,
+    ):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        repo = await create_repo(session=session, project_id=project.id)
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+        )
+        instance = await create_instance(
+            session=session,
+            project=project,
+            status=InstanceStatus.BUSY,
+        )
+        job = await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.PULLING,
+            job_provisioning_data=get_job_provisioning_data(dockerized=True),
+            job_runtime_data=get_job_runtime_data(network_mode="bridge", ports=None),
+            instance=instance,
+            instance_assigned=True,
+        )
+        shim_client_mock.get_task.return_value.status = TaskStatus.RUNNING
+        shim_client_mock.get_task.return_value.ports = [
+            PortMapping(container=10022, host=32771),
+            PortMapping(container=10999, host=32772),
+        ]
+
+        expected_ports = {
+            10022: 32771,
+            10999: 32772,
+        }
+
+        def assert_runner_availability(_, __, job_runtime_data):
+            assert job_runtime_data is not None
+            assert job_runtime_data.ports == expected_ports
+            return _RunnerAvailability.AVAILABLE
+
+        def assert_submit_job_to_runner(_, __, job_runtime_data, **kwargs):
+            assert job_runtime_data is not None
+            assert job_runtime_data.ports == expected_ports
+            return True
+
+        with (
+            patch(
+                "dstack._internal.server.background.scheduled_tasks.running_jobs._get_runner_availability",
+                side_effect=assert_runner_availability,
+            ) as get_runner_availability_mock,
+            patch(
+                "dstack._internal.server.background.scheduled_tasks.running_jobs._submit_job_to_runner",
+                side_effect=assert_submit_job_to_runner,
+            ) as submit_job_to_runner_mock,
+            patch(
+                "dstack._internal.server.background.scheduled_tasks.running_jobs._get_job_file_archives",
+                new_callable=AsyncMock,
+                return_value=[],
+            ),
+            patch(
+                "dstack._internal.server.background.scheduled_tasks.running_jobs._get_job_code",
+                new_callable=AsyncMock,
+                return_value=b"",
+            ),
+        ):
+            await process_running_jobs()
+
+            ssh_tunnel_mock.assert_called_once()
+            get_runner_availability_mock.assert_called_once()
+            submit_job_to_runner_mock.assert_called_once()
+
+        await session.refresh(job)
+        assert job is not None
+        assert job.status == JobStatus.PULLING
+        jrd = JobRuntimeData.__response__.parse_raw(job.job_runtime_data)
+        assert jrd.ports == expected_ports
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_pulling_shim_failed(self, test_db, session: AsyncSession):