fal 1.45.2__py3-none-any.whl → 1.46.1__py3-none-any.whl

fal/distributed/worker.py

@@ -0,0 +1,791 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import os
+import pickle
+import queue
+import threading
+import time
+import traceback
+import warnings
+from collections.abc import AsyncIterator, Callable, Coroutine
+from concurrent.futures import Future
+from functools import partial
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from fal.distributed.utils import (
+    KeepAliveTimer,
+    distributed_deserialize,
+    distributed_serialize,
+    encode_text_event,
+    launch_distributed_processes,
+)
+
+if TYPE_CHECKING:
+    import torch
+    import torch.multiprocessing as mp
+    from zmq.sugar.socket import Socket
+
+
+class DistributedWorker:
+    """
+    A base class for distributed workers.
+    """
+
+    queue: queue.Queue[bytes]
+    loop: asyncio.AbstractEventLoop
+    thread: threading.Thread
+
+    def __init__(
+        self,
+        rank: int = 0,
+        world_size: int = 1,
+    ) -> None:
+        self.rank = rank
+        self.world_size = world_size
+        self.queue = queue.Queue()
+
+        try:
+            import uvloop
+
+            self.loop = uvloop.new_event_loop()
+        except ImportError:
+            self.loop = asyncio.new_event_loop()
+
+        self._start_thread()
+
+    def _start_thread(self) -> None:
+        """
+        Start the thread.
+        """
+        self.thread = threading.Thread(target=self._run_forever, daemon=True)
+        self.thread.start()
+
+    def _run_forever(self) -> None:
+        """
+        Run the event loop forever.
+        """
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_forever()
+
+    # Public API
+
+    @property
+    def device(self) -> torch.device:
+        """
+        :return: The device for the current worker.
+        """
+        import torch
+
+        return torch.device(f"cuda:{self.rank}")
+
+    @property
+    def running(self) -> bool:
+        """
+        :return: Whether the event loop is running.
+        """
+        return self.thread.is_alive()
+
+    def initialize(self, **kwargs: Any) -> None:
+        """
+        Initialize the worker.
+        """
+        import torch
+
+        torch.cuda.set_device(self.device)
+        self.rank_print(f"Initializing worker on device {self.device}")
+
+        setup_start = time.time()
+        future = self.run_in_worker(self.setup, **kwargs)
+        future.result()
+        setup_duration = time.time() - setup_start
+        self.rank_print(f"Setup took {setup_duration:.2f} seconds")
+
+    def add_streaming_result(
+        self,
+        result: Any,
+        image_format: str = "jpeg",
+        as_text_event: bool = False,
+    ) -> None:
+        """
+        Add a streaming result to the queue.
+        :param result: The result to add to the queue.
+        """
+        if as_text_event:
+            serialized = encode_text_event(
+                result, is_final=False, image_format=image_format
+            )
+        else:
+            serialized = distributed_serialize(
+                result, is_final=False, image_format=image_format
+            )
+
+        self.queue.put_nowait(serialized)
+
+    def add_streaming_error(self, error: Exception) -> None:
+        """
+        Add an error to the queue.
+        :param error: The error to add to the queue.
+        """
+        self.queue.put_nowait(
+            distributed_serialize({"error": str(error)}, is_final=False)
+        )
+
+    def rank_print(self, message: str, debug: bool = False) -> None:
+        """
+        Print a message with the rank of the current worker.
+        :param message: The message to print.
+        :param debug: Whether to print the message as a debug message.
+        """
+        prefix = "[debug] " if debug else ""
+        print(f"{prefix}[rank {self.rank}] {message}")
+
+    def submit(self, coro: Coroutine[Any, Any, Any]) -> Future[Any]:
+        """
+        Submit a coroutine to the event loop.
+        :param coro: The coroutine to submit to the event loop.
+        :return: A future that will resolve to the result of the coroutine.
+        """
+        if not self.running:
+            raise RuntimeError("Event loop is not running.")
+        return asyncio.run_coroutine_threadsafe(coro, self.loop)
+
+    def shutdown(self, timeout: Optional[Union[int, float]] = None) -> None:
+        """
+        Shutdown the event loop.
+        :param timeout: The timeout for the shutdown.
+        """
+        try:
+            self.run_in_worker(self.teardown).result()
+        except Exception as e:
+            self.rank_print(f"Error during teardown: {e}\n{traceback.format_exc()}")
+
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join(timeout=timeout)
+
+    async def _run_sync_in_executor(
+        self,
+        func: Callable[..., Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """Run a synchronous function in the executor."""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, partial(func, *args, **kwargs))
+
+    def run_in_worker(
+        self,
+        func: Callable[..., Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Future[Any]:
+        """
+        Run a function in the worker.
+        """
+        if inspect.iscoroutinefunction(func):
+            coro = func(*args, **kwargs)
+        else:
+            # Using in place of asyncio.to_thread
+            # since it's not available in Python 3.8
+            coro = self._run_sync_in_executor(func, *args, **kwargs)
+
+        return self.submit(coro)
+
+    # Overrideables
+
+    def setup(self, **kwargs: Any) -> None:
+        """
+        Override this method to set up the worker.
+        This method is called once per worker.
+        """
+        return
+
+    def teardown(self) -> None:
+        """
+        Override this method to tear down the worker.
+        This method is called once per worker.
+        """
+        return
+
+    def __call__(self, streaming: bool = False, **kwargs: Any) -> Any:
+        """
+        Override this method to run the worker.
+        """
+        return {}
+
+
+class DistributedRunner:
+    """
+    A class to launch and manage distributed workers.
+    """
+
+    zmq_socket: Optional[Socket[Any]]
+    context: Optional[mp.ProcessContext]
+    keepalive_timer: Optional[KeepAliveTimer]
+
+    def __init__(
+        self,
+        worker_cls: type[DistributedWorker] = DistributedWorker,
+        world_size: int = 1,
+        master_addr: str = "127.0.0.1",
+        master_port: int = 29500,
+        worker_addr: str = "127.0.0.1",
+        worker_port: int = 54923,
+        timeout: int = 86400,  # 24 hours
+        keepalive_payload: dict[str, Any] = {},
+        keepalive_interval: Optional[Union[int, float]] = None,
+        cwd: Optional[Union[str, Path]] = None,
+        set_device: Optional[bool] = None,  # deprecated
+    ) -> None:
+        self.worker_cls = worker_cls
+        self.world_size = world_size
+        self.master_addr = master_addr
+        self.master_port = master_port
+        self.worker_addr = worker_addr
+        self.worker_port = worker_port
+        self.timeout = timeout
+        self.cwd = cwd
+        self.zmq_socket = None
+        self.context = None
+        self.keepalive_payload = keepalive_payload
+        self.keepalive_interval = keepalive_interval
+        self.keepalive_timer = None
+
+        if set_device is not None:
+            warnings.warn("set_device is deprecated and will be removed in the future.")
+
+    def is_alive(self) -> bool:
+        """
+        Check if the distributed worker processes are alive.
+        :return: True if the distributed processes are alive, False otherwise.
+        """
+        if self.context is None:
+            return False
+        for process in self.context.processes:
+            if not process.is_alive():
+                return False
+        return True
+
+    def terminate(self, timeout: Union[int, float] = 10) -> None:
+        """
+        Terminates the distributed worker processes.
+        This method should be called to clean up the worker processes.
+        """
+        if self.context is not None:
+            for process in self.context.processes:
+                if process.is_alive():
+                    process.terminate()
+                    process.join(timeout=timeout)
+
+    def gather_errors(self) -> list[Exception]:
+        """
+        Gathers errors from the distributed worker processes.
+
+        This method should be called to collect any errors that occurred
+        during execution.
+
+        :return: A list of exceptions raised by the worker processes.
+        """
+        errors = []
+
+        if self.context is not None:
+            for error_file in self.context.error_files:
+                if os.path.exists(error_file):
+                    with open(error_file, "rb") as f:
+                        error = pickle.loads(f.read())
+                        errors.append(error)
+
+                    os.remove(error_file)
+
+        return errors
+
+    def ensure_alive(self) -> None:
+        """
+        Ensures that the distributed worker processes are alive.
+        If the processes are not alive, it raises an error.
+        """
+        if not self.is_alive():
+            raise RuntimeError(
+                f"Distributed processes are not running. Errors: {self.gather_errors()}"
+            )
+
+    def get_zmq_socket(self) -> Socket[Any]:
+        """
+        Returns a ZeroMQ socket of the specified type.
+        :param socket_type: The type of the ZeroMQ socket.
+        :return: A ZeroMQ socket.
+        """
+        if self.zmq_socket is not None:
+            return self.zmq_socket
+
+        import zmq
+        import zmq.asyncio
+
+        context = zmq.asyncio.Context()
+        socket = context.socket(zmq.ROUTER)
+        socket.bind(f"tcp://{self.worker_addr}:{self.worker_port}")
+        self.zmq_socket = socket
+        return socket
+
+    def close_zmq_socket(self) -> None:
+        """
+        Closes the ZeroMQ socket.
+        """
+        if self.zmq_socket is not None:
+            try:
+                self.zmq_socket.close()
+            except Exception as e:
+                print(
+                    f"[debug] Error closing ZeroMQ socket: {e}\n"
+                    f"{traceback.format_exc()}"
+                )
+            self.zmq_socket = None
+
+    def run(self, **kwargs: Any) -> None:
+        """
+        The main function to run the distributed worker.
+
+        This function is called by each worker process spawned by
+        `torch.multiprocessing.spawn`. This method must be synchronous.
+
+        :param kwargs: The arguments to pass to the worker.
+        """
+        import torch.distributed as dist
+        import zmq
+
+        # Set up communication
+        rank = int(os.environ["RANK"])
+        context = zmq.Context()
+        socket = context.socket(zmq.DEALER)
+        socket.setsockopt(zmq.IDENTITY, str(rank).encode("utf-8"))
+        socket.connect(f"tcp://{self.worker_addr}:{self.worker_port}")
+
+        # Create and setup the worker
+        worker = self.worker_cls(rank, self.world_size)
+        try:
+            worker.initialize(**kwargs)
+        except Exception as e:
+            worker.rank_print(
+                f"Error during initialization: {e}\n{traceback.format_exc()}"
+            )
+            socket.send(b"EXIT")
+            socket.close()
+            return
+
+        # Wait until all workers are ready
+        socket.send(b"READY")
+        dist.barrier()
+
+        # Define execution methods to invoke from workers
+        def execute(payload: bytes) -> Any:
+            """
+            Execute the worker function with the given payload synchronously.
+            :param payload: The payload to send to the worker.
+            :return: The result from the worker.
+            """
+            payload_dict = distributed_deserialize(payload)
+            assert isinstance(payload_dict, dict)
+            payload_dict["streaming"] = False
+
+            try:
+                future = worker.run_in_worker(worker.__call__, **payload_dict)
+                result = future.result()
+            except Exception as e:
+                error_output = {"error": str(e)}
+                worker.rank_print(
+                    f"Error in execution: {error_output}\n{traceback.format_exc()}"
+                )
+                result = error_output
+
+            dist.barrier()
+            if worker.rank != 0:
+                return
+
+            socket.send(distributed_serialize(result, is_final=True))
+
+        def stream(payload: bytes, as_text_events: bool) -> None:
+            """
+            Stream the result from the worker function with the given payload.
+            :param payload: The payload to send to the worker.
+            :return: An async iterator that yields the result from the worker.
+            """
+            payload_dict = distributed_deserialize(payload)
+            assert isinstance(payload_dict, dict)
+            payload_dict["streaming"] = True
+            image_format = payload_dict.get("image_format", "jpeg")
+            encoded_response: Optional[bytes] = None
+
+            try:
+                future = worker.run_in_worker(worker.__call__, **payload_dict)
+                while not future.done():
+                    try:
+                        intermediate = worker.queue.get(timeout=0.1)
+                        if intermediate is not None and worker.rank == 0:
+                            socket.send(intermediate)  # already serialized
+                    except queue.Empty:
+                        pass
+                result = future.result()
+            except Exception as e:
+                error_output = {"error": str(e)}
+                worker.rank_print(
+                    f"Error in streaming: {error_output}\n{traceback.format_exc()}"
+                )
+                if worker.rank == 0:
+                    if as_text_events:
+                        encoded_response = encode_text_event(error_output)
+                    else:
+                        encoded_response = distributed_serialize(error_output)
+            else:
+                if worker.rank == 0:
+                    if as_text_events:
+                        encoded_response = encode_text_event(
+                            result, is_final=True, image_format=image_format
+                        )
+                    else:
+                        encoded_response = distributed_serialize(
+                            result, is_final=True, image_format=image_format
+                        )
+
+            dist.barrier()
+            if worker.rank != 0:
+                return
+
+            if encoded_response is not None:
+                socket.send(encoded_response)
+            socket.send(b"DONE")
+
+        # Runtime code
+        if rank == 0:
+            worker.rank_print("Master worker is ready to receive tasks.")
+            while True:
+                serialized_data = socket.recv()
+                streaming = serialized_data[0] == ord("1")
+                as_text_events = serialized_data[1] == ord("1")
+                serialized_data = serialized_data[2:]
+                params = [serialized_data, streaming, as_text_events]
+                dist.broadcast_object_list(params, src=0)
+
+                if serialized_data == b"EXIT":
+                    worker.rank_print("Received exit payload, exiting.")
+                    break
+
+                if streaming:
+                    stream(serialized_data, as_text_events)
+                else:
+                    execute(serialized_data)
+        else:
+            worker.rank_print("Worker waiting for tasks.")
+            while True:
+                try:
+                    params = [None, None, None]
+                    dist.broadcast_object_list(params, src=0)
+                    payload, streaming, as_text_events = params  # type: ignore[assignment]
+                    if payload == b"EXIT":
+                        worker.rank_print("Received exit payload, exiting.")
+                        break
+
+                    if streaming:
+                        stream(payload, as_text_events)  # type: ignore[arg-type]
+                    else:
+                        execute(payload)  # type: ignore[arg-type]
+                except Exception as e:
+                    worker.rank_print(f"Error in worker: {e}\n{traceback.format_exc()}")
+
+        # Teardown
+        worker.rank_print("Worker is tearing down.")
+        try:
+            worker.shutdown()
+            worker.rank_print("Worker torn down successfully.")
+        except Exception as e:
+            worker.rank_print(f"Error during teardown: {e}\n{traceback.format_exc()}")
+
+        socket.send(b"EXIT")
+        socket.close()
+
+    async def start(self, timeout: int = 1800, **kwargs: Any) -> None:
+        """
+        Starts the distributed worker processes.
+        :param timeout: The timeout for the distributed processes.
+        """
+        import zmq
+
+        if self.is_alive():
+            raise RuntimeError("Distributed processes are already running.")
+
+        self.context = launch_distributed_processes(
+            self.run,
+            world_size=self.world_size,
+            master_addr=self.master_addr,
+            master_port=self.master_port,
+            timeout=self.timeout,
+            cwd=self.cwd,
+            **kwargs,
+        )
+
+        try:
+            ready_workers: set[int] = set()
+            socket = self.get_zmq_socket()
+            start_time = time.perf_counter()
+
+            while len(ready_workers) < self.world_size:
+                try:
+                    ident, msg = await socket.recv_multipart(flags=zmq.NOBLOCK)  # type: ignore[misc]
+
+                    if msg != b"READY":
+                        worker_id = ident.decode("utf-8")
+                        worker_msg = msg.decode("utf-8")
+                        raise RuntimeError(
+                            f"Unexpected message from worker {worker_id}: {worker_msg}"
+                        )
+
+                    print(f"[debug] Worker {ident.decode('utf-8')} is ready.")
+                    ready_workers.add(ident)
+                except zmq.Again:
+                    total_wait_time = time.perf_counter() - start_time
+                    if total_wait_time > timeout:
+                        raise TimeoutError(
+                            f"Timeout reached after {timeout} seconds while "
+                            f"waiting for workers to be ready."
+                        )
+                    await asyncio.sleep(0.5)
+                self.ensure_alive()
+        except Exception as e:
+            print(f"[debug] Error during startup: {e}\n{traceback.format_exc()}")
+            self.terminate(timeout=timeout)
+            raise RuntimeError("Failed to start distributed processes.") from e
+
+        print("[debug] All workers are ready and running.")
+
+        # Start the keepalive timer
+        self.maybe_start_keepalive()
+
+    def keepalive(self, timeout: Optional[Union[int, float]] = 60.0) -> None:
+        """
+        Sends the keepalive payload to the worker.
+        """
+        # Cancel the keepalive timer
+        self.maybe_cancel_keepalive()
+        loop_thread = None
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            loop_thread = threading.Thread(target=loop.run_forever, daemon=True)
+            loop_thread.start()
+
+        future = asyncio.run_coroutine_threadsafe(
+            self.invoke(self.keepalive_payload), loop
+        )
+
+        try:
+            future.result(timeout=timeout)
+        except Exception as e:
+            print(f"[debug] Error during keepalive: {e}\n{traceback.format_exc()}")
+            raise RuntimeError("Failed to run keepalive.") from e
+        finally:
+            if loop_thread is not None:
+                loop.call_soon_threadsafe(loop.stop)
+                loop_thread.join(timeout=timeout)
+            # Restart the keepalive timer
+            self.maybe_start_keepalive()
+
+    def maybe_start_keepalive(self) -> None:
+        """
+        Starts the keepalive timer if it is set.
+        """
+        if self.keepalive_timer is None and self.keepalive_interval is not None:
+            self.keepalive_timer = KeepAliveTimer(
+                self.keepalive, self.keepalive_interval, start=True
+            )
+
+    def maybe_reset_keepalive(self) -> None:
+        """
+        Resets the keepalive timer if it is set.
+        """
+        if self.keepalive_timer is not None:
+            self.keepalive_timer.reset()
+
+    def maybe_cancel_keepalive(self) -> None:
+        """
+        Cancels the keepalive timer if it is set.
+        """
+        if self.keepalive_timer is not None:
+            self.keepalive_timer.cancel()
+            self.keepalive_timer = None
+
+    async def stop(self, timeout: int = 10) -> None:
+        """
+        Stops the distributed worker processes.
+        :param timeout: The timeout for the distributed processes to stop.
+        """
+        import zmq
+
+        if not self.is_alive():
+            raise RuntimeError("Distributed processes are not running.")
+
+        self.maybe_cancel_keepalive()
+        worker_exits: set[int] = set()
+        socket = self.get_zmq_socket()
+        await socket.send_multipart([b"0", b"00EXIT"])
+
+        wait_start = time.perf_counter()
+        while len(worker_exits) < self.world_size:
+            try:
+                ident, msg = await socket.recv_multipart(flags=zmq.NOBLOCK)  # type: ignore[misc]
+                if msg == b"EXIT":
+                    worker_exits.add(ident)
+                    if len(worker_exits) == self.world_size:
+                        print("[debug] All workers have exited.")
+                        await asyncio.sleep(1)  # Allow time for cleanup
+                        break
+
+            except zmq.Again:
+                # No messages available, continue waiting
+                await asyncio.sleep(0.1)
+
+            if time.perf_counter() - wait_start > timeout:
+                print(
+                    f"[debug] Timeout reached after {timeout} seconds, "
+                    f"stopping waiting."
+                )
+                break
+
+            if not self.is_alive():
+                print("[debug] All workers have exited prematurely.")
+                break
+
+        if self.is_alive():
+            print("[debug] Some workers did not exit cleanly, terminating them.")
+            self.terminate(timeout=timeout)
+
+        self.close_zmq_socket()
+
+    async def stream(
+        self,
+        payload: dict[str, Any] = {},
+        timeout: Optional[int] = None,
+        streaming_timeout: Optional[int] = None,
+        as_text_events: bool = False,
+    ) -> AsyncIterator[Any]:
+        """
+        Streams the result from the distributed worker.
+        :param payload: The payload to send to the worker.
+        :param timeout: The timeout for the overall operation.
+        :param streaming_timeout: The timeout in-between streamed results.
+        :param as_text_events: Whether to yield results as text events.
+        :return: An async iterator that yields the result from the worker.
+        """
+        import zmq
+
+        self.ensure_alive()
+        self.maybe_cancel_keepalive()  # Cancel until the streaming is done
+        socket = self.get_zmq_socket()
+        payload_serialized = distributed_serialize(payload, is_final=True)
+        await socket.send_multipart(
+            [b"0", b"1" + (b"1" if as_text_events else b"0") + payload_serialized]
+        )
+
+        start_time = time.perf_counter()
+        last_yield_time = start_time
+        yielded_once = False
+
+        while True:
+            iter_start_time = time.perf_counter()
+            try:
+                rank, response = await socket.recv_multipart(flags=zmq.NOBLOCK)  # type: ignore[misc]
+            except zmq.Again:
+                if timeout is not None and iter_start_time - start_time > timeout:
+                    raise TimeoutError(f"Streaming timed out after {timeout} seconds.")
+                if (
+                    streaming_timeout is not None
+                    and iter_start_time - last_yield_time > streaming_timeout
+                ):
+                    raise TimeoutError(
+                        f"Streaming timed out after {streaming_timeout} "
+                        f"seconds of inactivity."
+                    )
+
+                await asyncio.sleep(0.1)
+                continue
+
+            assert rank == b"0", "Expected response from worker with rank 0"
+
+            if response == b"DONE":
+                if not yielded_once:
+                    raise RuntimeError("No data was yielded from the worker.")
+                break
+
+            if as_text_events:
+                yield response
+            else:
+                yield distributed_deserialize(response)
+
+            yielded_once = True
+            last_yield_time = iter_start_time
+
+        self.maybe_start_keepalive()  # Restart the keepalive timer
+
+    async def invoke(
+        self,
+        payload: dict[str, Any] = {},
+        timeout: Optional[int] = None,
+    ) -> Any:
+        """
+        Invokes the distributed worker with the given payload.
+        :param payload: The payload to send to the worker.
+        :param timeout: The timeout for the overall operation.
+        :return: The result from the worker.
+        """
+        import zmq
+
+        self.ensure_alive()
+        self.maybe_cancel_keepalive()  # Cancel until the invocation is done
+        socket = self.get_zmq_socket()
+        payload_serialized = distributed_serialize(payload, is_final=True)
+
+        await socket.send_multipart([b"0", b"00" + payload_serialized])
+
+        # Wait for the response from the worker
+        start_time = time.perf_counter()
+        while True:
+            try:
+                rank, response = await socket.recv_multipart(flags=zmq.NOBLOCK)  # type: ignore[misc]
+                break  # Exit the loop if we received a response
+            except zmq.Again:
+                elapsed = time.perf_counter() - start_time
+                if timeout is not None and elapsed > timeout:
+                    raise TimeoutError(f"Invocation timed out after {timeout} seconds.")
+
+                await asyncio.sleep(0.1)
+                self.ensure_alive()
+
+        self.maybe_start_keepalive()  # Restart the keepalive timer
+        assert rank == b"0", "Expected response from worker with rank 0"
+        return distributed_deserialize(response)
+
+    async def __aenter__(self) -> DistributedRunner:
+        """
+        Enter the context manager.
+        :return: The DistributedRunner instance.
+        """
+        await self.start()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        exc_traceback: Optional[traceback.StackSummary],
+    ) -> None:
+        """
+        Exit the context manager.
+        :param exc_type: The type of the exception raised, if any.
+        :param exc_value: The value of the exception raised, if any.
+        :param traceback: The traceback of the exception raised, if any.
+        """
+        try:
+            await self.stop()
+        except Exception as e:
+            print(f"[debug] Error during cleanup: {e}\n{traceback.format_exc()}")