From ea21d983ca1f4e90094f0175930e6769415cfe2c Mon Sep 17 00:00:00 2001 From: kramstrom Date: Tue, 15 Oct 2024 17:02:18 +0200 Subject: [PATCH 1/3] init modal executor --- examples/modal_workflow/workflow.py | 106 ++++++++++++++++++++ redun/__init__.py | 1 + redun/executors/modal_executor.py | 147 ++++++++++++++++++++++++++++ 3 files changed, 254 insertions(+) create mode 100644 examples/modal_workflow/workflow.py create mode 100644 redun/executors/modal_executor.py diff --git a/examples/modal_workflow/workflow.py b/examples/modal_workflow/workflow.py new file mode 100644 index 00000000..b40fc8d4 --- /dev/null +++ b/examples/modal_workflow/workflow.py @@ -0,0 +1,106 @@ +from redun import task, File +from typing import List +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]}) + + # # a = task5() + + # # print(a) + + # # write to volume + # with open("/mnt/vol/test.txt", "w") as f: + # f.write("hello") + + # import os + + # print(os.listdir("/mnt/vol")) + + # # return a + + # return a + b + [ls("/root")] + + return [ls("./"), ls("/root"), ls("./redun")] 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..1c7d4972 --- /dev/null +++ b/redun/executors/modal_executor.py @@ -0,0 +1,147 @@ +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: + print("calling func") + res = func(*args, **kwargs) + print("putting success") + queue.put((id, "success")) + return res + except Exception: + print("putting error") + 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: + print("getting results") + res = self.queue.get_many(100, timeout=1) + print(f"got results {len(res)}") + 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): + print("stopping") + 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 From b2dfd678a8e9da69523ad26b626ea5b683fc44e9 Mon Sep 17 00:00:00 2001 From: kramstrom Date: Tue, 15 Oct 2024 17:09:08 +0200 Subject: [PATCH 2/3] clean --- examples/modal_workflow/workflow.py | 26 +++++++++----------------- redun/executors/modal_executor.py | 6 ------ 2 files changed, 9 insertions(+), 23 deletions(-) diff --git a/examples/modal_workflow/workflow.py b/examples/modal_workflow/workflow.py index b40fc8d4..adcaf27a 100644 --- a/examples/modal_workflow/workflow.py +++ b/examples/modal_workflow/workflow.py @@ -1,5 +1,4 @@ -from redun import task, File -from typing import List +from redun import task import modal redun_namespace = "redun.examples.modal_workflow" @@ -83,24 +82,17 @@ def ls(path: str): }, ) def main(x: str = "hi", y: str = "bye"): - # print("main") - # a = task4(x) - # b = task1({"list": [y], "first": a[0], "second": a[1]}) + print("main") + a = task4(x) + b = task1({"list": [y], "first": a[0], "second": a[1]}) - # # a = task5() - - # # print(a) - - # # write to volume - # with open("/mnt/vol/test.txt", "w") as f: - # f.write("hello") + with open("/mnt/vol/test.txt", "w") as f: + f.write("hello") - # import os + import os - # print(os.listdir("/mnt/vol")) + print(os.listdir("/mnt/vol")) - # # return a - # return a + b + [ls("/root")] + return a + b + [ls("/root")] - return [ls("./"), ls("/root"), ls("./redun")] diff --git a/redun/executors/modal_executor.py b/redun/executors/modal_executor.py index 1c7d4972..9702518a 100644 --- a/redun/executors/modal_executor.py +++ b/redun/executors/modal_executor.py @@ -22,13 +22,10 @@ def wrapper(*args, **kwargs): def queue_wrapper(func: Callable) -> Callable: def wrapper(id: str, queue: modal.Queue, *args, **kwargs): try: - print("calling func") res = func(*args, **kwargs) - print("putting success") queue.put((id, "success")) return res except Exception: - print("putting error") queue.put((id, "error")) raise @@ -93,9 +90,7 @@ def start(self) -> None: def result_thread_func(self): while not self.stopped: try: - print("getting results") res = self.queue.get_many(100, timeout=1) - print(f"got results {len(res)}") for job_id, status in res: job = self.jobs[job_id] try: @@ -111,7 +106,6 @@ def result_thread_func(self): continue def stop(self): - print("stopping") self.stopped = True self.result_thread.join() self.run_ctx.__exit__(None, None, None) From e6c99574904770ff72da835a1ef2dd10285faf67 Mon Sep 17 00:00:00 2001 From: Matt Rasmussen Date: Mon, 9 Jun 2025 12:49:18 -0700 Subject: [PATCH 3/3] add modal dep --- requirements-dev.txt | 1 + setup.py | 1 + 2 files changed, 2 insertions(+) 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: