From 163298e39948e000a482f05a14f4b7235a5f15f1 Mon Sep 17 00:00:00 2001
From: Dmitry Meyer <me@undef.im>
Date: Mon, 6 Jan 2025 11:11:16 +0000
Subject: [PATCH] Fix volume attachment with dstack backend (#2175)

Fixes: https://github.com/dstackai/dstack/issues/2172
See: https://github.com/dstackai/dstack/issues/2172
See: https://github.com/dstackai/dstack/pull/1682
---
 .../_internal/server/background/tasks/process_submitted_jobs.py | 2 +-
 src/dstack/_internal/server/services/jobs/__init__.py           | 2 ++
 src/dstack/_internal/server/services/runs.py                    | 2 +-
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
index 581e8ef1c..393538fe8 100644
--- a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
+++ b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
@@ -546,7 +546,7 @@ async def _attach_volumes(
             volume = volume_model_to_volume(volume_model)
             try:
                 if (
-                    job_provisioning_data.backend != volume.configuration.backend
+                    job_provisioning_data.get_base_backend() != volume.configuration.backend
                     or job_provisioning_data.region != volume.configuration.region
                 ):
                     continue
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
index 96c7c7657..c2cf5ee80 100644
--- a/src/dstack/_internal/server/services/jobs/__init__.py
+++ b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -95,6 +95,8 @@ def job_model_to_job_submission(job_model: JobModel) -> JobSubmission:
         job_provisioning_data.instance_type.resources.description = (
             job_provisioning_data.instance_type.resources.pretty_format()
         )
+        # TODO do we really still need this magic? See https://github.com/dstackai/dstack/pull/1682
+        # i.e., replacing `jpd.backend` with `jpd.get_base_backend()` should give the same result
         if (
             job_provisioning_data.backend == BackendType.DSTACK
             and job_provisioning_data.backend_data is not None
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
index 0fdb5c066..f66e36c7c 100644
--- a/src/dstack/_internal/server/services/runs.py
+++ b/src/dstack/_internal/server/services/runs.py
@@ -844,7 +844,7 @@ def get_job_mount_point_volume(
     """
     for volume in volumes:
         if (
-            volume.configuration.backend != job_provisioning_data.backend
+            volume.configuration.backend != job_provisioning_data.get_base_backend()
             or volume.configuration.region != job_provisioning_data.region
         ):
             continue