diff --git a/examples/modal_workflow/workflow.py b/examples/modal_workflow/workflow.py new file mode 100644 index 00000000..adcaf27a --- /dev/null +++ b/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")] + diff --git a/redun/__init__.py b/redun/__init__.py index fe0cd4ed..7abe4f65 100644 --- a/redun/__init__.py +++ b/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 diff --git a/redun/executors/modal_executor.py b/redun/executors/modal_executor.py new file mode 100644 index 00000000..9702518a --- /dev/null +++ b/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 diff --git a/requirements-dev.txt b/requirements-dev.txt index 300e1e77..d72440c2 100644 --- a/requirements-dev.txt +++ b/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 diff --git a/setup.py b/setup.py index b2a2b72e..b7362de2 100644 --- a/setup.py +++ b/setup.py @@ -43,6 +43,7 @@ "adlfs>=2024.2.0", "azure-ai-ml>=1.23.1", ], + "modal": "modal>=1.0", } if REQUIRE_POSTGRES: