Fix Modal deprecations

cubed-dev · Dec 11, 2022 · f90d09c · f90d09c
1 parent deb6d2b
commit f90d09c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 11 deletions.
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
@@ -16,9 +16,9 @@
 requirements_file = os.getenv("CUBED_MODAL_REQUIREMENTS_FILE")
 
 if requirements_file:
-    image = modal.DebianSlim().pip_install_from_requirements(requirements_file)
+    image = modal.Image.debian_slim().pip_install_from_requirements(requirements_file)
 else:
-    image = modal.DebianSlim().pip_install(
+    image = modal.Image.debian_slim().pip_install(
         [
             "dask[array]",
             "fsspec",
@@ -33,7 +33,9 @@
 
 
 # Use a generator, since we want results to be returned as they finish and we don't care about order
-@stub.generator(image=image, secret=modal.ref("my-aws-secret"), memory=2000, retries=2)
+@stub.generator(
+    image=image, secret=modal.Secret.from_name("my-aws-secret"), memory=2000, retries=2
+)
 def run_remotely(input, func=None, config=None):
     print(f"running remotely on {input}")
     peak_memory_start = peak_memory()

diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
@@ -20,9 +20,9 @@
 requirements_file = os.getenv("CUBED_MODAL_REQUIREMENTS_FILE")
 
 if requirements_file:
-    image = modal.DebianSlim().pip_install_from_requirements(requirements_file)
+    image = modal.Image.debian_slim().pip_install_from_requirements(requirements_file)
 else:
-    image = modal.DebianSlim().pip_install(
+    image = modal.Image.debian_slim().pip_install(
         [
             "dask[array]",
             "fsspec",
@@ -37,7 +37,7 @@
 
 
 @async_stub.generator(
-    image=image, secret=modal.ref("my-aws-secret"), memory=2000, retries=2
+    image=image, secret=modal.Secret.from_name("my-aws-secret"), memory=2000, retries=2
 )
 async def async_run_remotely(input, func=None, config=None):
     print(f"running remotely on {input}")

diff --git a/cubed/tests/runtime/test_modal_async.py b/cubed/tests/runtime/test_modal_async.py
@@ -26,7 +26,7 @@ def write_int_to_file(path, i):
 
 stub = modal.aio.AioStub()
 
-image = modal.DebianSlim().pip_install(
+image = modal.Image.debian_slim().pip_install(
     [
         "dask[array]",
         "fsspec",
@@ -40,14 +40,14 @@ def write_int_to_file(path, i):
 )
 
 
-@stub.function(image=image, secret=modal.ref("my-aws-secret"))
+@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"))
 def never_fail(i):
     invocation_count_file = join_path(tmp_path, f"{i}")
     write_int_to_file(invocation_count_file, 1)
     return i
 
 
-@stub.function(image=image, secret=modal.ref("my-aws-secret"), retries=2)
+@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2)
 async def fail_on_first_invocation(i):
     invocation_count_file = join_path(tmp_path, f"{i}")
     fs = fsspec.open(invocation_count_file).fs
@@ -60,7 +60,7 @@ async def fail_on_first_invocation(i):
     return i
 
 
-@stub.function(image=image, secret=modal.ref("my-aws-secret"))
+@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"))
 async def fail_on_first_invocation_no_retry(i):
     invocation_count_file = join_path(tmp_path, f"{i}")
     fs = fsspec.open(invocation_count_file).fs
@@ -73,7 +73,7 @@ async def fail_on_first_invocation_no_retry(i):
     return i
 
 
-@stub.function(image=image, secret=modal.ref("my-aws-secret"))
+@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"))
 async def sleep_on_first_invocation(i):
     print(f"sleep_on_first_invocation {i}")
     invocation_count_file = join_path(tmp_path, f"{i}")