change deprecated modal Stub to App (#2038)

axolotl-ai-cloud · Nov 11, 2024 · 3931a42 · 3931a42
1 parent dc8f905
commit 3931a42
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/cicd/multigpu.py b/cicd/multigpu.py
@@ -10,7 +10,7 @@
 import jinja2
 import modal
 from jinja2 import select_autoescape
-from modal import Image, Stub
+from modal import App, Image
 
 cicd_path = pathlib.Path(__file__).parent.resolve()
 
@@ -46,7 +46,7 @@
     .pip_install("fastapi==0.110.0", "pydantic==2.6.3")
 )
 
-stub = Stub("Axolotl CI/CD", secrets=[])
+app = App("Axolotl CI/CD", secrets=[])
 
 
 N_GPUS = int(os.environ.get("N_GPUS", 2))
@@ -61,7 +61,7 @@ def run_cmd(cmd: str, run_folder: str):
         exit(exit_code)  # pylint: disable=consider-using-sys-exit
 
 
-@stub.function(
+@app.function(
     image=cicd_image,
     gpu=GPU_CONFIG,
     timeout=60 * 60,
@@ -72,6 +72,6 @@ def cicd_pytest():
     run_cmd("./cicd/multigpu.sh", "/workspace/axolotl")
 
 
-@stub.local_entrypoint()
+@app.local_entrypoint()
 def main():
     cicd_pytest.remote()
diff --git a/cicd/tests.py b/cicd/tests.py
@@ -10,7 +10,7 @@
 import jinja2
 import modal
 from jinja2 import select_autoescape
-from modal import Image, Stub
+from modal import App, Image
 
 cicd_path = pathlib.Path(__file__).parent.resolve()
 
@@ -47,7 +47,7 @@
     .pip_install("fastapi==0.110.0", "pydantic==2.6.3")
 )
 
-stub = Stub("Axolotl CI/CD", secrets=[])
+app = App("Axolotl CI/CD", secrets=[])
 
 
 N_GPUS = int(os.environ.get("N_GPUS", 1))
@@ -62,7 +62,7 @@ def run_cmd(cmd: str, run_folder: str):
         exit(exit_code)  # pylint: disable=consider-using-sys-exit
 
 
-@stub.function(
+@app.function(
     image=cicd_image,
     gpu=GPU_CONFIG,
     timeout=60 * 60,
@@ -73,6 +73,6 @@ def cicd_pytest():
     run_cmd("./cicd/cicd.sh", "/workspace/axolotl")
 
 
-@stub.local_entrypoint()
+@app.local_entrypoint()
 def main():
     cicd_pytest.remote()