Add Modal Executor [WIP]

examples/modal_workflow/workflow.py

@@ -0,0 +1,98 @@
+from redun import task
+import modal
+
+redun_namespace = "redun.examples.modal_workflow"
+
+
+@task()
+def task3(x):
+    print("task3")
+    return [x["first"], x["second"]]
+
+
+@task(
+    executor="modal",
+    modal_args={
+        "cpu": 10,
+        "mounts": [
+            modal.Mount.from_local_dir(
+                "./examples/01_hello_world", remote_path="/root/hello_world"
+            )
+        ],
+        "network_file_systems": {
+            "/mnt/nfs": modal.NetworkFileSystem.from_name("my-redun-nfs", create_if_missing=True)
+        },
+    },
+)
+def task2(x):
+    print("task2")
+
+    import os
+
+    print(os.listdir("/root/hello_world"))
+
+    with open("/mnt/nfs/test.txt", "w") as f:
+        f.write("hello")
+
+    print(os.listdir("/mnt/nfs"))
+
+    return task3(x) + x["list"]
+
+
+@task(
+    executor="modal",
+    modal_args={"cloud": "gcp", "secrets": [modal.Secret.from_name("my-huggingface-secret")]},
+)
+def task1(x):
+    print("task1")
+    return task2(x)
+
+
+@task(
+    executor="modal",
+    modal_args={"region": "eu", "image": modal.Image.debian_slim().pip_install("redun", "pandas")},
+)
+def task4(x):
+    import pandas as pd
+
+    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    print(df.head())
+
+    print("task4")
+    return [x, x]
+
+
+@task(executor="modal", script=True)
+def ls(path: str):
+    # Calls grep to find pattern in a file.
+
+    import time
+
+    time.sleep(1)
+    return f"""
+    ls {path}
+    """
+
+
+@task(
+    executor="modal",
+    modal_args={
+        "gpu": "A10G",
+        "volumes": {"/mnt/vol": modal.Volume.from_name("my-redun-volume", create_if_missing=True)},
+    },
+)
+def main(x: str = "hi", y: str = "bye"):
+    print("main")
+    a = task4(x)
+    b = task1({"list": [y], "first": a[0], "second": a[1]})
+
+    with open("/mnt/vol/test.txt", "w") as f:
+        f.write("hello")
+
+    import os
+
+    print(os.listdir("/mnt/vol"))
+
+
+    return a + b + [ls("/root")]
+

redun/__init__.py

@@ -3,6 +3,7 @@
 from redun.executors.aws_batch import AWSBatchExecutor
 from redun.executors.aws_glue import AWSGlueExecutor
 from redun.executors.docker import DockerExecutor
+from redun.executors.modal_executor import ModalExecutor
 
 try:
     from redun.executors.k8s import K8SExecutor

redun/executors/modal_executor.py

@@ -0,0 +1,141 @@
+import typing
+import queue
+import modal
+from typing import Optional, Callable
+import threading
+from redun.config import create_config_section
+from redun.executors.base import Executor, register_executor
+from redun.scripting import exec_script
+from redun.task import get_task_registry
+
+if typing.TYPE_CHECKING:
+    from redun.scheduler import Job, Scheduler
+
+
+def exec_script_func(func: Callable) -> Callable:
+    def wrapper(*args, **kwargs):
+        return exec_script(func(*args, **kwargs))
+
+    return wrapper
+
+
+def queue_wrapper(func: Callable) -> Callable:
+    def wrapper(id: str, queue: modal.Queue, *args, **kwargs):
+        try:
+            res = func(*args, **kwargs)
+            queue.put((id, "success"))
+            return res
+        except Exception:
+            queue.put((id, "error"))
+            raise
+
+    return wrapper
+
+
+@register_executor("modal")
+class ModalExecutor(Executor):
+    """
+    A redun Executor for running jobs on Modal.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        scheduler: Optional["Scheduler"] = None,
+        config=None,
+    ):
+        super().__init__(name, scheduler=scheduler)
+
+        # Parse config.
+        if not config:
+            config = create_config_section()
+
+        self.finished_setup = False
+
+        self.modal_function_calls = {}
+
+        self.jobs = {}
+        self.stopped = False
+
+    def start(self) -> None:
+        self.image = modal.Image.debian_slim().pip_install("redun")
+        self.modal_app = modal.App(f"redun-modal-job-{self.name}", image=self.image)
+
+        self.functions = {}
+        for t in get_task_registry():
+            task_options = t.get_task_options()
+
+            if task_options.get("executor") == "modal":
+                modal_args = task_options.get("modal_args", {})
+                modal_args["image"] = modal_args.get("image", self.image)
+                modal_args["name"] = modal_args.get("name", t.fullname)
+                modal_args["serialized"] = True
+
+                func = queue_wrapper(exec_script_func(t.func) if t.script else t.func)
+
+                self.functions[t.fullname] = self.modal_app.function(**modal_args)(func)
+
+        self.output_ctx = modal.enable_output()
+        self.output_ctx.__enter__()
+        self.run_ctx = self.modal_app.run()
+        self.run_ctx.__enter__()
+
+        self.queue_ctx = modal.Queue.ephemeral()
+        self.queue = self.queue_ctx.__enter__()
+
+        self.result_thread = threading.Thread(target=self.result_thread_func)
+        self.result_thread.start()
+        self.finished_setup = True
+
+    def result_thread_func(self):
+        while not self.stopped:
+            try:
+                res = self.queue.get_many(100, timeout=1)
+                for job_id, status in res:
+                    job = self.jobs[job_id]
+                    try:
+                        return_value = self.modal_function_calls[job_id].get()
+                    except Exception as e:
+                        return_value = e
+
+                    if status == "success":
+                        self._scheduler.done_job(job, return_value)
+                    else:
+                        self._scheduler.reject_job(job, return_value)
+            except queue.Empty:
+                continue
+
+    def stop(self):
+        self.stopped = True
+        self.result_thread.join()
+        self.run_ctx.__exit__(None, None, None)
+        self.output_ctx.__exit__(None, None, None)
+        self.queue_ctx.__exit__(None, None, None)
+
+    def _submit(self, job: "Job") -> None:
+        assert job.args
+        if not self.finished_setup:
+            self.start()
+
+        args, kwargs = job.args
+
+        try:
+            self.jobs[job.id] = job
+            modal_exec_func = self.functions[job.task.fullname]
+        except KeyError:
+            raise ValueError(f"No modal function found for task {job.task.fullname}")
+
+        self.modal_function_calls[job.id] = modal_exec_func.spawn(
+            job.id, self.queue, *args, **kwargs
+        )
+
+    def submit(self, job: "Job") -> None:
+        assert not job.task.script
+        self._submit(job)
+
+    def submit_script(self, job: "Job") -> None:
+        assert job.task.script
+        self._submit(job)
+
+    def scratch_root(self) -> str:
+        pass

requirements-dev.txt

@@ -17,6 +17,7 @@ google-cloud-compute==1.11.0
 urllib3<2
 adlfs==2024.2.0
 azure-ai-ml==1.23.1
+modal==1.0.3
 
 # Pin boto for testing. We can unpinned if moto is updated.
 aiobotocore==2.12.3

setup.py

@@ -43,6 +43,7 @@
         "adlfs>=2024.2.0",
         "azure-ai-ml>=1.23.1",
     ],
+    "modal": "modal>=1.0",
 }
 
 if REQUIRE_POSTGRES: