From a1760e20cd660e793aec2c5cbd5ad8463424370a Mon Sep 17 00:00:00 2001
From: Elias Freider <freider@users.noreply.github.com>
Date: Wed, 27 Nov 2024 13:19:01 +0100
Subject: [PATCH] Fixes issue where Volume.from_name eagerly assigns event loop
 (#2581)

* Fixes issue where Volume.from_name eagerly assigns the main thread event loop to internal lock
---
 modal/image.py      |  2 +-
 modal/object.py     |  2 +-
 modal/volume.py     | 19 +++++++++++++------
 test/volume_test.py | 20 ++++++++++++++++++++
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/modal/image.py b/modal/image.py
index dfcbc540c..24538d437 100644
--- a/modal/image.py
+++ b/modal/image.py
@@ -1714,7 +1714,7 @@ def workdir(self, path: Union[str, PurePosixPath]) -> "_Image":
         """
 
         def build_dockerfile(version: ImageBuilderVersion) -> DockerfileSpec:
-            commands = ["FROM base", f"WORKDIR {shlex.quote(path)}"]
+            commands = ["FROM base", f"WORKDIR {shlex.quote(str(path))}"]
             return DockerfileSpec(commands=commands, context_files={})
 
         return _Image._from_args(
diff --git a/modal/object.py b/modal/object.py
index 4d6d02a8c..22c6cfc0f 100644
--- a/modal/object.py
+++ b/modal/object.py
@@ -17,7 +17,7 @@
 
 _BLOCKING_O = synchronize_api(O)
 
-EPHEMERAL_OBJECT_HEARTBEAT_SLEEP = 300
+EPHEMERAL_OBJECT_HEARTBEAT_SLEEP: int = 300
 
 
 def _get_environment_name(environment_name: Optional[str] = None, resolver: Optional[Resolver] = None) -> Optional[str]:
diff --git a/modal/volume.py b/modal/volume.py
index 015f77a09..cc67aee12 100644
--- a/modal/volume.py
+++ b/modal/volume.py
@@ -122,14 +122,21 @@ def g():
     ```
     """
 
-    _lock: asyncio.Lock
+    _lock: Optional[asyncio.Lock] = None
 
-    def _initialize_from_empty(self):
+    async def _get_lock(self):
         # To (mostly*) prevent multiple concurrent operations on the same volume, which can cause problems under
         # some unlikely circumstances.
         # *: You can bypass this by creating multiple handles to the same volume, e.g. via lookup. But this
         # covers the typical case = good enough.
-        self._lock = asyncio.Lock()
+
+        # Note: this function runs no async code but is marked as async to ensure it's
+        # being run inside the synchronicity event loop and binds the lock to the
+        # correct event loop on Python 3.9 which eagerly assigns event loops on
+        # constructions of locks
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        return self._lock
 
     @staticmethod
     def new():
@@ -188,7 +195,7 @@ async def ephemeral(
         environment_name: Optional[str] = None,
         version: "typing.Optional[modal_proto.api_pb2.VolumeFsVersion.ValueType]" = None,
         _heartbeat_sleep: float = EPHEMERAL_OBJECT_HEARTBEAT_SLEEP,
-    ) -> AsyncIterator["_Volume"]:
+    ) -> AsyncGenerator["_Volume", None]:
         """Creates a new ephemeral volume within a context manager:
 
         Usage:
@@ -269,7 +276,7 @@ async def create_deployed(
 
     @live_method
     async def _do_reload(self, lock=True):
-        async with self._lock if lock else asyncnullcontext():
+        async with (await self._get_lock()) if lock else asyncnullcontext():
             req = api_pb2.VolumeReloadRequest(volume_id=self.object_id)
             _ = await retry_transient_errors(self._client.stub.VolumeReload, req)
 
@@ -280,7 +287,7 @@ async def commit(self):
         If successful, the changes made are now persisted in durable storage and available to other containers accessing
         the volume.
         """
-        async with self._lock:
+        async with await self._get_lock():
             req = api_pb2.VolumeCommitRequest(volume_id=self.object_id)
             try:
                 # TODO(gongy): only apply indefinite retries on 504 status.
diff --git a/test/volume_test.py b/test/volume_test.py
index 6a00862f8..4e449adca 100644
--- a/test/volume_test.py
+++ b/test/volume_test.py
@@ -381,3 +381,23 @@ async def test_open_files_error_annotation(tmp_path):
 def test_invalid_name(servicer, client, name):
     with pytest.raises(InvalidError, match="Invalid Volume name"):
         modal.Volume.lookup(name)
+
+
+@pytest.fixture()
+def unset_main_thread_event_loop():
+    try:
+        event_loop = asyncio.get_event_loop()
+    except RuntimeError:
+        event_loop = None
+
+    asyncio.set_event_loop(None)
+    try:
+        yield
+    finally:
+        asyncio.set_event_loop(event_loop)  # reset so we don't break other tests
+
+
+@pytest.mark.usefixtures("unset_main_thread_event_loop")
+def test_lock_is_py39_safe(set_env_client):
+    vol = modal.Volume.from_name("my_vol", create_if_missing=True)
+    vol.reload()