Support CUBED_MODAL_REQUIREMENTS_FILE (#135)

cubed-dev · Sep 30, 2022 · 070bf81 · 070bf81
1 parent 7a2ec18
commit 070bf81
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 24 deletions.
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
@@ -1,3 +1,4 @@
+import os
 import time
 from asyncio.exceptions import TimeoutError
 
@@ -12,18 +13,23 @@
 
 stub = modal.Stub()
 
-image = modal.DebianSlim().pip_install(
-    [
-        "dask[array]",
-        "fsspec",
-        "networkx",
-        "pytest-mock",  # TODO: only needed for tests
-        "rechunker",
-        "s3fs",
-        "tenacity",
-        "zarr",
-    ]
-)
+requirements_file = os.getenv("CUBED_MODAL_REQUIREMENTS_FILE")
+
+if requirements_file:
+    image = modal.DebianSlim().pip_install_from_requirements(requirements_file)
+else:
+    image = modal.DebianSlim().pip_install(
+        [
+            "dask[array]",
+            "fsspec",
+            "networkx",
+            "pytest-mock",  # TODO: only needed for tests
+            "rechunker",
+            "s3fs",
+            "tenacity",
+            "zarr",
+        ]
+    )
 
 
 # Use a generator, since we want results to be returned as they finish and we don't care about order

diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
@@ -1,5 +1,6 @@
 import asyncio
 import copy
+import os
 import time
 from asyncio.exceptions import TimeoutError
 
@@ -16,18 +17,23 @@
 
 async_stub = modal.aio.AioStub()
 
-image = modal.DebianSlim().pip_install(
-    [
-        "dask[array]",
-        "fsspec",
-        "networkx",
-        "pytest-mock",  # TODO: only needed for tests
-        "rechunker",
-        "s3fs",
-        "tenacity",
-        "zarr",
-    ]
-)
+requirements_file = os.getenv("CUBED_MODAL_REQUIREMENTS_FILE")
+
+if requirements_file:
+    image = modal.DebianSlim().pip_install_from_requirements(requirements_file)
+else:
+    image = modal.DebianSlim().pip_install(
+        [
+            "dask[array]",
+            "fsspec",
+            "networkx",
+            "pytest-mock",  # TODO: only needed for tests
+            "rechunker",
+            "s3fs",
+            "tenacity",
+            "zarr",
+        ]
+    )
 
 
 @async_stub.generator(