worker ¶

@abstractmethod
def clear_log_queue(self) -> None:
    """Stop forwarding worker logs to the queue. Call after generate_video."""
    self.collective_rpc("clear_log_queue")

@abstractmethod
def collective_rpc(self,
                   method: str | Callable[..., _R],
                   timeout: float | None = None,
                   args: tuple = (),
                   kwargs: dict[str, Any] | None = None) -> list[_R]:
    """
    Execute an RPC call on all workers.

    Args:
        method: Name of the worker method to execute, or a callable that
            is serialized and sent to all workers to execute.

            If the method is a callable, it should accept an additional
            `self` argument, in addition to the arguments passed in `args`
            and `kwargs`. The `self` argument will be the worker object.
        timeout: Maximum time in seconds to wait for execution. Raises a
            :exc:`TimeoutError` on timeout. `None` means wait indefinitely.
        args: Positional arguments to pass to the worker method.
        kwargs: Keyword arguments to pass to the worker method.

    Returns:
        A list containing the results from each worker.

    Note:
        It is recommended to use this API to only pass control messages,
        and set up data-plane communication to pass data.
    """
    raise NotImplementedError

@abstractmethod
def merge_lora_weights(self) -> None:
    """
    Merge the LoRA weights for the workers.
    """
    raise NotImplementedError

@abstractmethod
def set_log_queue(self, log_queue: Queue | None) -> None:
    """Forward worker logs to the given queue. Call before generate_video."""
    self.collective_rpc("set_log_queue", kwargs={"log_queue": log_queue})

@abstractmethod
def set_lora_adapter(self,
                     lora_nickname: str,
                     lora_path: str | None = None,
                     strength: float = 1.0,
                     accumulate: bool = False) -> None:
    """
    Set the LoRA adapter for the workers.
    """
    raise NotImplementedError

@abstractmethod
def shutdown(self) -> None:
    """
    Shutdown the executor.
    """
    raise NotImplementedError

@abstractmethod
def unmerge_lora_weights(self) -> None:
    """
    Unmerge the LoRA weights for the workers.
    """
    raise NotImplementedError

def __del__(self):
    """Ensure cleanup on garbage collection"""
    self.shutdown()

def __enter__(self):
    """Support for context manager protocol"""
    return self

def __exit__(self, exc_type, exc_val, exc_tb):
    """Ensure cleanup when exiting context"""
    self.shutdown()

def clear_log_queue(self) -> None:
    """Stop forwarding worker logs to the queue. Call after generate_video."""
    self.collective_rpc("clear_log_queue")

def set_log_queue(self, log_queue: Queue | None) -> None:
    """Forward worker logs to the given queue. Call before generate_video."""
    self.collective_rpc("set_log_queue", kwargs={"log_queue": log_queue})

def shutdown(self) -> None:
    """Properly shut down the executor and its workers"""
    if hasattr(self, 'shutting_down') and self.shutting_down:
        return  # Prevent multiple shutdown calls

    logger.info("Shutting down MultiprocExecutor...")

    # Check if workers were initialized (they might not be if initialization failed)
    if not hasattr(self, 'workers') or not self.workers:
        logger.info("No workers to shut down.")
        return

    self.shutting_down = True

    # First try gentle termination
    try:
        # Send termination message to all workers
        for worker in self.workers:
            with contextlib.suppress(Exception):
                worker.pipe.send({"method": "shutdown", "args": (), "kwargs": {}})

        # Give workers some time to exit gracefully
        start_time = time.perf_counter()
        while time.perf_counter() - start_time < 5.0:  # 5 seconds timeout
            if all(not worker.proc.is_alive() for worker in self.workers):
                break
            time.sleep(0.1)

        # Force terminate any remaining workers
        for worker in self.workers:
            if worker.proc.is_alive():
                worker.proc.terminate()

        # Final timeout for terminate
        start_time = time.perf_counter()
        while time.perf_counter() - start_time < 2.0:  # 2 seconds timeout
            if all(not worker.proc.is_alive() for worker in self.workers):
                break
            time.sleep(0.1)

        # Kill if still alive
        for worker in self.workers:
            if worker.proc.is_alive():
                worker.proc.kill()
            worker.proc.join(timeout=1.0)

    except Exception as e:
        logger.error("Error during shutdown: %s", e)
        # Last resort, try to kill all workers
        for worker in self.workers:
            with contextlib.suppress(Exception):
                if worker.proc.is_alive():
                    worker.proc.kill()

    # Clean up pipes
    for worker in self.workers:
        with contextlib.suppress(Exception):
            worker.pipe.close()

    self.workers = []
    logger.info("MultiprocExecutor shutdown complete")