import logging
import math
import pickle
import subprocess
import threading
import typing
import warnings
from collections import defaultdict
from concurrent.futures import Future
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union

import typeguard

from parsl import curvezmq
from parsl.addresses import get_all_addresses
from parsl.app.errors import RemoteExceptionWrapper
from parsl.data_provider.staging import Staging
from parsl.executors.errors import (
    BadMessage,
    InvalidResourceSpecification,
    ScalingFailed,
)
from parsl.executors.high_throughput import zmq_pipes
from parsl.executors.high_throughput.errors import CommandClientTimeoutError
from parsl.executors.high_throughput.manager_selector import (
    ManagerSelector,
    RandomManagerSelector,
)
from parsl.executors.status_handling import BlockProviderExecutor
from parsl.jobs.states import TERMINAL_STATES, JobState, JobStatus
from parsl.process_loggers import wrap_with_logs
from parsl.providers import LocalProvider
from parsl.providers.base import ExecutionProvider
from parsl.serialize import deserialize, pack_res_spec_apply_message
from parsl.serialize.errors import DeserializationError, SerializationError
from parsl.usage_tracking.api import UsageInformation
from parsl.utils import RepresentationMixin

logger = logging.getLogger(__name__)

DEFAULT_LAUNCH_CMD = ("process_worker_pool.py {debug} {max_workers_per_node} "
                      "-a {addresses} "
                      "-p {prefetch_capacity} "
                      "-c {cores_per_worker} "
                      "-m {mem_per_worker} "
                      "--poll {poll_period} "
                      "--task_port={task_port} "
                      "--result_port={result_port} "
                      "--cert_dir {cert_dir} "
                      "--logdir={logdir} "
                      "--block_id={{block_id}} "
                      "--hb_period={heartbeat_period} "
                      "{address_probe_timeout_string} "
                      "--hb_threshold={heartbeat_threshold} "
                      "--drain_period={drain_period} "
                      "--cpu-affinity {cpu_affinity} "
                      "{enable_mpi_mode} "
                      "--mpi-launcher={mpi_launcher} "
                      "--available-accelerators {accelerators}")

DEFAULT_INTERCHANGE_LAUNCH_CMD = ["interchange.py"]

GENERAL_HTEX_PARAM_DOCS = """provider : :class:`~parsl.providers.base.ExecutionProvider`
       Provider to access computation resources. Can be one of :class:`~parsl.providers.aws.aws.EC2Provider`,
        :class:`~parsl.providers.condor.condor.Condor`,
        :class:`~parsl.providers.googlecloud.googlecloud.GoogleCloud`,
        :class:`~parsl.providers.gridEngine.gridEngine.GridEngine`,
        :class:`~parsl.providers.local.local.Local`,
        :class:`~parsl.providers.sge.sge.GridEngine`,
        :class:`~parsl.providers.slurm.slurm.Slurm`, or
        :class:`~parsl.providers.torque.torque.Torque`.

    label : str
        Label for this executor instance.

    launch_cmd : str
        Command line string to launch the process_worker_pool from the provider. The command line string
        will be formatted with appropriate values for the following values (debug, task_url, result_url,
        cores_per_worker, nodes_per_block, heartbeat_period ,heartbeat_threshold, logdir). For example:
        launch_cmd="process_worker_pool.py {debug} -c {cores_per_worker} --task_url={task_url} --result_url={result_url}"

    interchange_launch_cmd : Sequence[str]
        Custom sequence of command line tokens to launch the interchange process from the executor. If
        undefined, the executor will use the default "interchange.py" command.

    address : string
        An address to connect to the main Parsl process which is reachable from the network in which
        workers will be running. This field expects an IPv4 or IPv6 address.
        Most login nodes on clusters have several network interfaces available, only some of which
        can be reached from the compute nodes. This field can be used to limit the executor to listen
        only on a specific interface, and limiting connections to the internal network.
        By default, the executor will attempt to enumerate and connect through all possible addresses.
        Setting an address here overrides the default behavior.
        default=None

    loopback_address: string
        Specify address used for internal communication between executor and interchange.
        Supports IPv4 and IPv6 addresses
        default=127.0.0.1

    worker_ports : (int, int)
        Specify the ports to be used by workers to connect to Parsl. If this option is specified,
        worker_port_range will not be honored.

    worker_port_range : (int, int)
        Worker ports will be chosen between the two integers provided.

    interchange_port_range : (int, int)
        Port range used by Parsl to communicate with the Interchange.

    working_dir : str
        Working dir to be used by the executor.

    worker_debug : Bool
        Enables worker debug logging.

    prefetch_capacity : int
        Number of tasks that could be prefetched over available worker capacity.
        When there are a few tasks (<100) or when tasks are long running, this option should
        be set to 0 for better load balancing. Default is 0.

    address_probe_timeout : int | None
        Managers attempt connecting over many different addresses to determine a viable address.
        This option sets a time limit in seconds on the connection attempt.
        Default of None implies 30s timeout set on worker.

    heartbeat_threshold : int
        Seconds since the last message from the counterpart in the communication pair:
        (interchange, manager) after which the counterpart is assumed to be un-available. Default: 120s

    heartbeat_period : int
        Number of seconds after which a heartbeat message indicating liveness is sent to the
        counterpart (interchange, manager). Default: 30s

    poll_period : int
        Timeout period to be used by the executor components in milliseconds. Increasing poll_periods
        trades performance for cpu efficiency. Default: 10ms

    drain_period : int
        The number of seconds after start when workers will begin to drain
        and then exit. Set this to a time that is slightly less than the
        maximum walltime of batch jobs to avoid killing tasks while they
        execute. For example, you could set this to the walltime minus a grace
        period for the batch job to start the workers, minus the expected
        maximum length of an individual task.

    worker_logdir_root : string
        In case of a remote file system, specify the path to where logs will be kept.

    encrypted : bool
        Flag to enable/disable encryption (CurveZMQ). Default is False.

    manager_selector: ManagerSelector
        Determines what strategy the interchange uses to select managers during task distribution.
        See API reference under "Manager Selectors" regarding the various manager selectors.
        Default: 'RandomManagerSelector'
"""  # Documentation for params used by both HTEx and MPIEx


class HighThroughputExecutor(BlockProviderExecutor, RepresentationMixin, UsageInformation):
    __doc__ = f"""Executor designed for cluster-scale

    The HighThroughputExecutor system has the following components:
      1. The HighThroughputExecutor instance which is run as part of the Parsl script.
      2. The Interchange which acts as a load-balancing proxy between workers and Parsl
      3. The multiprocessing based worker pool which coordinates task execution over several
         cores on a node.
      4. ZeroMQ pipes connect the HighThroughputExecutor, Interchange and the process_worker_pool

    Here is a diagram

    .. code:: python


                        |  Data   |  Executor   |  Interchange  | External Process(es)
                        |  Flow   |             |               |
                   Task | Kernel  |             |               |
                 +----->|-------->|------------>|->outgoing_q---|-> process_worker_pool
                 |      |         |             | batching      |    |         |
           Parsl<---Fut-|         |             | load-balancing|  result   exception
                     ^  |         |             | watchdogs     |    |         |
                     |  |         |    Result   |               |    |         |
                     |  |         |    Queue    |               |    V         V
                     |  |         |    Thread<--|-incoming_q<---|--- +---------+
                     |  |         |      |      |               |
                     |  |         |      |      |               |
                     +----update_fut-----+


    Each of the workers in each process_worker_pool has access to its local rank through
    an environmental variable, ``PARSL_WORKER_RANK``. The local rank is unique for each process
    and is an integer in the range from 0 to the number of workers per in the pool minus 1.
    The workers also have access to the ID of the worker pool as ``PARSL_WORKER_POOL_ID``
    and the size of the worker pool as ``PARSL_WORKER_COUNT``.


    Parameters
    ----------

    {GENERAL_HTEX_PARAM_DOCS}

    cores_per_worker : float
        cores to be assigned to each worker. Oversubscription is possible
        by setting cores_per_worker < 1.0. Default=1

    mem_per_worker : float
        GB of memory required per worker. If this option is specified, the node manager
        will check the available memory at startup and limit the number of workers such that
        the there's sufficient memory for each worker. Default: None

    max_workers_per_node : int
        Caps the number of workers launched per node. Default: None

    cpu_affinity: string
        Whether or how each worker process sets thread affinity. Options include "none" to forgo
        any CPU affinity configuration, "block" to assign adjacent cores to workers
        (ex: assign 0-1 to worker 0, 2-3 to worker 1), and
        "alternating" to assign cores to workers in round-robin
        (ex: assign 0,2 to worker 0, 1,3 to worker 1).
        The "block-reverse" option assigns adjacent cores to workers, but assigns
        the CPUs with large indices to low index workers (ex: assign 2-3 to worker 1, 0,1 to worker 2)

    available_accelerators: int | list
        Accelerators available for workers to use. Each worker will be pinned to exactly one of the provided
        accelerators, and no more workers will be launched than the number of accelerators.

        Either provide the list of accelerator names or the number available. If a number is provided,
        Parsl will create names as integers starting with 0.

        default: empty list

    """

    @typeguard.typechecked
    def __init__(self,
                 label: str = 'HighThroughputExecutor',
                 provider: ExecutionProvider = LocalProvider(),
                 launch_cmd: Optional[str] = None,
                 interchange_launch_cmd: Optional[Sequence[str]] = None,
                 address: Optional[str] = None,
                 loopback_address: str = "127.0.0.1",
                 worker_ports: Optional[Tuple[int, int]] = None,
                 worker_port_range: Optional[Tuple[int, int]] = (54000, 55000),
                 interchange_port_range: Optional[Tuple[int, int]] = (55000, 56000),
                 storage_access: Optional[List[Staging]] = None,
                 working_dir: Optional[str] = None,
                 worker_debug: bool = False,
                 cores_per_worker: float = 1.0,
                 mem_per_worker: Optional[float] = None,
                 max_workers_per_node: Optional[Union[int, float]] = None,
                 cpu_affinity: str = 'none',
                 available_accelerators: Union[int, Sequence[str]] = (),
                 prefetch_capacity: int = 0,
                 heartbeat_threshold: int = 120,
                 heartbeat_period: int = 30,
                 drain_period: Optional[int] = None,
                 poll_period: int = 10,
                 address_probe_timeout: Optional[int] = None,
                 worker_logdir_root: Optional[str] = None,
                 manager_selector: ManagerSelector = RandomManagerSelector(),
                 block_error_handler: Union[bool, Callable[[BlockProviderExecutor, Dict[str, JobStatus]], None]] = True,
                 encrypted: bool = False):

        logger.debug("Initializing HighThroughputExecutor")

        BlockProviderExecutor.__init__(self, provider=provider, block_error_handler=block_error_handler)
        self.label = label
        self.worker_debug = worker_debug
        self.storage_access = storage_access
        self.working_dir = working_dir
        self.cores_per_worker = cores_per_worker
        self.mem_per_worker = mem_per_worker
        self.prefetch_capacity = prefetch_capacity
        self.address = address
        self.address_probe_timeout = address_probe_timeout
        self.manager_selector = manager_selector
        self.loopback_address = loopback_address

        if self.address:
            self.all_addresses = address
        else:
            self.all_addresses = ','.join(get_all_addresses())

        self.max_workers_per_node = max_workers_per_node or float("inf")

        mem_slots = self.max_workers_per_node
        cpu_slots = self.max_workers_per_node
        if hasattr(self.provider, 'mem_per_node') and \
                self.provider.mem_per_node is not None and \
                mem_per_worker is not None and \
                mem_per_worker > 0:
            mem_slots = math.floor(self.provider.mem_per_node / mem_per_worker)
        if hasattr(self.provider, 'cores_per_node') and \
                self.provider.cores_per_node is not None:
            cpu_slots = math.floor(self.provider.cores_per_node / cores_per_worker)

        # Set the list of available accelerators
        if isinstance(available_accelerators, int):
            # If the user provide an integer, create some names for them
            available_accelerators = list(map(str, range(available_accelerators)))
        self.available_accelerators = list(available_accelerators)

        # Determine the number of workers per node
        self._workers_per_node = min(self.max_workers_per_node, mem_slots, cpu_slots)
        if len(self.available_accelerators) > 0:
            self._workers_per_node = min(self._workers_per_node, len(available_accelerators))
        if self._workers_per_node == float('inf'):
            self._workers_per_node = 1  # our best guess-- we do not have any provider hints

        self._task_counter = 0
        self.worker_ports = worker_ports
        self.worker_port_range = worker_port_range
        self.interchange_proc: Optional[subprocess.Popen] = None
        self.interchange_port_range = interchange_port_range
        self.heartbeat_threshold = heartbeat_threshold
        self.heartbeat_period = heartbeat_period
        self.drain_period = drain_period
        self.poll_period = poll_period
        self.run_dir = '.'
        self.worker_logdir_root = worker_logdir_root
        self.cpu_affinity = cpu_affinity
        self.encrypted = encrypted
        self.cert_dir = None

        if not launch_cmd:
            launch_cmd = DEFAULT_LAUNCH_CMD
        self.launch_cmd = launch_cmd

        if not interchange_launch_cmd:
            interchange_launch_cmd = DEFAULT_INTERCHANGE_LAUNCH_CMD
        self.interchange_launch_cmd = interchange_launch_cmd

        self._result_queue_thread_exit = threading.Event()
        self._result_queue_thread: Optional[threading.Thread] = None

    radio_mode = "htex"
    enable_mpi_mode: bool = False
    mpi_launcher: str = "mpiexec"

    def _warn_deprecated(self, old: str, new: str):
        warnings.warn(
            f"{old} is deprecated and will be removed in a future release. "
            f"Please use {new} instead.",
            DeprecationWarning,
            stacklevel=2
        )

    @property
    def logdir(self):
        return "{}/{}".format(self.run_dir, self.label)

    @property
    def worker_logdir(self):
        if self.worker_logdir_root is not None:
            return "{}/{}".format(self.worker_logdir_root, self.label)
        return self.logdir

    def validate_resource_spec(self, resource_specification: dict):
        """HTEX supports the following *Optional* resource specifications:
        priority: lower value is higher priority"""
        if resource_specification:
            acceptable_fields = {'priority'}
            keys = set(resource_specification.keys())
            invalid_keys = keys - acceptable_fields
            if invalid_keys:
                message = "Task resource specification only accepts these types of resources: {}".format(
                    ', '.join(acceptable_fields))
                logger.error(message)
                raise InvalidResourceSpecification(set(invalid_keys), message)
        return

    def initialize_scaling(self):
        """Compose the launch command and scale out the initial blocks.
        """
        debug_opts = "--debug" if self.worker_debug else ""
        max_workers_per_node = "" if self.max_workers_per_node == float('inf') else "--max_workers_per_node={}".format(self.max_workers_per_node)
        enable_mpi_opts = "--enable_mpi_mode " if self.enable_mpi_mode else ""

        address_probe_timeout_string = ""
        if self.address_probe_timeout:
            address_probe_timeout_string = "--address_probe_timeout={}".format(self.address_probe_timeout)

        l_cmd = self.launch_cmd.format(debug=debug_opts,
                                       prefetch_capacity=self.prefetch_capacity,
                                       address_probe_timeout_string=address_probe_timeout_string,
                                       addresses=self.all_addresses,
                                       task_port=self.worker_task_port,
                                       result_port=self.worker_result_port,
                                       cores_per_worker=self.cores_per_worker,
                                       mem_per_worker=self.mem_per_worker,
                                       max_workers_per_node=max_workers_per_node,
                                       nodes_per_block=self.provider.nodes_per_block,
                                       heartbeat_period=self.heartbeat_period,
                                       heartbeat_threshold=self.heartbeat_threshold,
                                       drain_period=self.drain_period,
                                       poll_period=self.poll_period,
                                       cert_dir=self.cert_dir,
                                       logdir=self.worker_logdir,
                                       cpu_affinity=self.cpu_affinity,
                                       enable_mpi_mode=enable_mpi_opts,
                                       mpi_launcher=self.mpi_launcher,
                                       accelerators=" ".join(self.available_accelerators))
        self.launch_cmd = l_cmd
        logger.debug("Launch command: {}".format(self.launch_cmd))

        logger.debug("Starting HighThroughputExecutor with provider:\n%s", self.provider)

    def start(self):
        """Create the Interchange process and connect to it.
        """
        if self.encrypted and self.cert_dir is None:
            logger.debug("Creating CurveZMQ certificates")
            self.cert_dir = curvezmq.create_certificates(self.logdir)
        elif not self.encrypted and self.cert_dir:
            raise AttributeError(
                "The certificates directory path attribute (cert_dir) is defined, but the "
                "encrypted attribute is set to False. You must either change cert_dir to "
                "None or encrypted to True."
            )

        self.outgoing_q = zmq_pipes.TasksOutgoing(
            self.loopback_address, self.interchange_port_range, self.cert_dir
        )
        self.incoming_q = zmq_pipes.ResultsIncoming(
            self.loopback_address, self.interchange_port_range, self.cert_dir
        )
        self.command_client = zmq_pipes.CommandClient(
            self.loopback_address, self.interchange_port_range, self.cert_dir
        )

        self._result_queue_thread = None
        self._start_result_queue_thread()
        self._start_local_interchange_process()

        logger.debug("Created result queue thread: %s", self._result_queue_thread)

        self.initialize_scaling()

    @wrap_with_logs
    def _result_queue_worker(self):
        """Listen to the queue for task result messages and handle them.

        Depending on the message, tasks will be updated with results or exceptions.

        .. code:: python

            {
               "task_id" : <task_id>
               "result"  : serialized result object, if task succeeded
               ... more tags could be added later
            }

            {
               "task_id" : <task_id>
               "exception" : serialized exception object, on failure
            }
        """
        logger.debug("Result queue worker starting")

        while not self.bad_state_is_set and not self._result_queue_thread_exit.is_set():
            try:
                msgs = self.incoming_q.get(timeout_ms=self.poll_period)
                if msgs is None:  # timeout
                    continue

            except IOError as e:
                logger.exception("Caught broken queue with exception code {}: {}".format(e.errno, e))
                return

            except Exception as e:
                logger.exception("Caught unknown exception: {}".format(e))
                return

            else:

                for serialized_msg in msgs:
                    try:
                        msg = pickle.loads(serialized_msg)
                    except pickle.UnpicklingError:
                        raise BadMessage("Message received could not be unpickled")

                    if msg['type'] == 'result':
                        try:
                            tid = msg['task_id']
                        except Exception:
                            raise BadMessage("Message received does not contain 'task_id' field")

                        if tid == -1 and 'exception' in msg:
                            logger.warning("Executor shutting down due to exception from interchange")
                            exception = deserialize(msg['exception'])
                            self.set_bad_state_and_fail_all(exception)
                            break

                        task_fut = self.tasks.pop(tid)

                        if 'result' in msg:
                            result = deserialize(msg['result'])
                            task_fut.set_result(result)

                        elif 'exception' in msg:
                            try:
                                s = deserialize(msg['exception'])
                                # s should be a RemoteExceptionWrapper... so we can reraise it
                                if isinstance(s, RemoteExceptionWrapper):
                                    try:
                                        s.reraise()
                                    except Exception as e:
                                        task_fut.set_exception(e)
                                elif isinstance(s, Exception):
                                    task_fut.set_exception(s)
                                else:
                                    raise ValueError("Unknown exception-like type received: {}".format(type(s)))
                            except Exception as e:
                                # TODO could be a proper wrapped exception?
                                task_fut.set_exception(
                                    DeserializationError("Received exception, but handling also threw an exception: {}".format(e)))
                        else:
                            raise BadMessage("Message received is neither result or exception")
                    else:
                        raise BadMessage("Message received with unknown type {}".format(msg['type']))

        logger.info("Closing result ZMQ pipe")
        self.incoming_q.close()
        logger.info("Result queue worker finished")

    def _start_local_interchange_process(self) -> None:
        """ Starts the interchange process locally

        Starts the interchange process locally and uses the command queue to
        get the worker task and result ports that the interchange has bound to.
        """

        interchange_config = {"client_address": self.loopback_address,
                              "client_ports": (self.outgoing_q.port,
                                               self.incoming_q.port,
                                               self.command_client.port),
                              "interchange_address": self.address,
                              "worker_ports": self.worker_ports,
                              "worker_port_range": self.worker_port_range,
                              "hub_address": self.hub_address,
                              "hub_zmq_port": self.hub_zmq_port,
                              "logdir": self.logdir,
                              "heartbeat_threshold": self.heartbeat_threshold,
                              "poll_period": self.poll_period,
                              "logging_level": logging.DEBUG if self.worker_debug else logging.INFO,
                              "cert_dir": self.cert_dir,
                              "manager_selector": self.manager_selector,
                              "run_id": self.run_id,
                              }

        config_pickle = pickle.dumps(interchange_config)

        self.interchange_proc = subprocess.Popen(self.interchange_launch_cmd, stdin=subprocess.PIPE)
        stdin = self.interchange_proc.stdin
        assert stdin is not None, "Popen should have created an IO object (vs default None) because of PIPE mode"

        logger.debug("Popened interchange process. Writing config object")
        stdin.write(config_pickle)
        stdin.flush()
        stdin.close()
        logger.debug("Sent config object. Requesting worker ports")
        try:
            (self.worker_task_port, self.worker_result_port) = self.command_client.run("WORKER_PORTS", timeout_s=120)
        except CommandClientTimeoutError:
            logger.error("Interchange has not completed initialization. Aborting")
            raise Exception("Interchange failed to start")
        logger.debug("Got worker ports")

    def _start_result_queue_thread(self):
        """Method to start the result queue thread as a daemon.

        Checks if a thread already exists, then starts it.
        Could be used later as a restart if the result queue thread dies.
        """
        if self._result_queue_thread is None:
            logger.debug("Starting result queue thread")
            self._result_queue_thread = threading.Thread(target=self._result_queue_worker, name="HTEX-Result-Queue-Thread")
            self._result_queue_thread.daemon = True
            self._result_queue_thread.start()
            logger.debug("Started result queue thread")

        else:
            logger.error("Result queue thread already exists, returning")

    def hold_worker(self, worker_id: str) -> None:
        """Puts a worker on hold, preventing scheduling of additional tasks to it.

        This is called "hold" mostly because this only stops scheduling of tasks,
        and does not actually kill the worker.

        Parameters
        ----------

        worker_id : str
            Worker id to be put on hold
        """
        self.command_client.run("HOLD_WORKER;{}".format(worker_id))
        logger.debug("Sent hold request to manager: {}".format(worker_id))

    @property
    def outstanding(self) -> int:
        """Returns the count of tasks outstanding across the interchange
        and managers"""
        return len(self.tasks)

    @property
    def connected_workers(self) -> int:
        """Returns the count of workers across all connected managers"""
        return self.command_client.run("WORKERS")

    def connected_managers(self) -> List[Dict[str, typing.Any]]:
        """Returns a list of dicts one for each connected managers.
        The dict contains info on manager(str:manager_id), block_id,
        worker_count, tasks(int), idle_durations(float), active(bool)
        """
        return self.command_client.run("MANAGERS")

    def connected_blocks(self) -> List[str]:
        """List of connected block ids"""
        return self.command_client.run("CONNECTED_BLOCKS")

    def _hold_block(self, block_id):
        """ Sends hold command to all managers which are in a specific block

        Parameters
        ----------
        block_id : str
             Block identifier of the block to be put on hold
        """

        managers = self.connected_managers()

        for manager in managers:
            if manager['block_id'] == block_id:
                logger.debug("Sending hold to manager: {}".format(manager['manager']))
                self.hold_worker(manager['manager'])

    def submit(self, func, resource_specification, *args, **kwargs):
        """Submits work to the outgoing_q.

        The outgoing_q is an external process listens on this
        queue for new work. This method behaves like a submit call as described here `Python docs: <https://docs.python.org/3/
        library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor>`_

        Args:
            - func (callable) : Callable function
            - resource_specification (dict): Dictionary containing relevant info about task that is needed by underlying executors.
            - args (list) : List of arbitrary positional arguments.

        Kwargs:
            - kwargs (dict) : A dictionary of arbitrary keyword args for func.

        Returns:
              Future
        """

        self.validate_resource_spec(resource_specification)

        if self.bad_state_is_set:
            raise self.executor_exception

        self._task_counter += 1
        task_id = self._task_counter

        # handle people sending blobs gracefully
        if logger.getEffectiveLevel() <= logging.DEBUG:
            args_to_print = tuple([ar if len(ar := repr(arg)) < 100 else (ar[:100] + '...') for arg in args])
            logger.debug("Pushing function {} to queue with args {}".format(func, args_to_print))

        fut = Future()
        fut.parsl_executor_task_id = task_id
        self.tasks[task_id] = fut

        try:
            fn_buf = pack_res_spec_apply_message(func, args, kwargs,
                                                 resource_specification=resource_specification,
                                                 buffer_threshold=1024 * 1024)
        except TypeError:
            raise SerializationError(func.__name__)

        msg = {"task_id": task_id, "resource_spec": resource_specification, "buffer": fn_buf}

        # Post task to the outgoing queue
        self.outgoing_q.put(msg)

        # Return the future
        return fut

    @property
    def workers_per_node(self) -> Union[int, float]:
        return self._workers_per_node

    def scale_in(self, blocks: int, max_idletime: Optional[float] = None) -> List[str]:
        """Scale in the number of active blocks by specified amount.

        The scale in method here is very rude. It doesn't give the workers
        the opportunity to finish current tasks or cleanup. This is tracked
        in issue #530

        Parameters
        ----------

        blocks : int
             Number of blocks to terminate and scale_in by

        max_idletime: float
             A time to indicate how long a block should be idle to be a
             candidate for scaling in.

             If None then blocks will be force scaled in even if they are busy.

             If a float, then only idle blocks will be terminated, which may be less than
             the requested number.

        Returns
        -------
        List of block IDs scaled in
        """
        logger.debug(f"Scale in called, blocks={blocks}")

        @dataclass
        class BlockInfo:
            tasks: int  # sum of tasks in this block
            idle: float  # shortest idle time of any manager in this block

        # block_info will be populated from two sources:
        # the Job Status Poller mutable block list, and the list of blocks
        # which have connected to the interchange.

        def new_block_info():
            return BlockInfo(tasks=0, idle=float('inf'))

        block_info: Dict[str, BlockInfo] = defaultdict(new_block_info)

        for block_id, job_status in self._status.items():
            if job_status.state not in TERMINAL_STATES:
                block_info[block_id] = new_block_info()

        managers = self.connected_managers()
        for manager in managers:
            if not manager['active']:
                continue
            b_id = manager['block_id']
            block_info[b_id].tasks += manager['tasks']
            block_info[b_id].idle = min(block_info[b_id].idle, manager['idle_duration'])

        # The scaling policy is that longest idle blocks should be scaled down
        # in preference to least idle (most recently used) blocks.
        # Other policies could be implemented here.

        sorted_blocks = sorted(block_info.items(), key=lambda item: (-item[1].idle, item[1].tasks))

        logger.debug(f"Scale in selecting from {len(sorted_blocks)} blocks")
        if max_idletime is None:
            block_ids_to_kill = [x[0] for x in sorted_blocks[:blocks]]
        else:
            block_ids_to_kill = []
            for x in sorted_blocks:
                if x[1].idle > max_idletime and x[1].tasks == 0:
                    block_ids_to_kill.append(x[0])
                    if len(block_ids_to_kill) == blocks:
                        break

            logger.debug("Selected idle block ids to kill: {}".format(
                block_ids_to_kill))
            if len(block_ids_to_kill) < blocks:
                logger.warning(f"Could not find enough blocks to kill: wanted {blocks} but only selected {len(block_ids_to_kill)}")

        # Hold the block
        for block_id in block_ids_to_kill:
            self._hold_block(block_id)

        # Now kill via provider
        # Potential issue with multiple threads trying to remove the same blocks
        to_kill = [self.blocks_to_job_id[bid] for bid in block_ids_to_kill if bid in self.blocks_to_job_id]

        r = self.provider.cancel(to_kill)
        job_ids = self._filter_scale_in_ids(to_kill, r)

        # to_kill block_ids are fetched from self.blocks_to_job_id
        # If a block_id is in self.blocks_to_job_id, it must exist in self.job_ids_to_block
        block_ids_killed = [self.job_ids_to_block[jid] for jid in job_ids]

        return block_ids_killed

    def _get_launch_command(self, block_id: str) -> str:
        if self.launch_cmd is None:
            raise ScalingFailed(self, "No launch command")
        launch_cmd = self.launch_cmd.format(block_id=block_id)
        return launch_cmd

    def status(self) -> Dict[str, JobStatus]:
        job_status = super().status()
        connected_blocks = self.connected_blocks()
        for job_id in job_status:
            job_info = job_status[job_id]
            if job_info.terminal and job_id not in connected_blocks and job_info.state != JobState.SCALED_IN:
                logger.debug("Rewriting job %s from status %s to MISSING", job_id, job_info)
                job_status[job_id].state = JobState.MISSING
                if job_status[job_id].message is None:
                    job_status[job_id].message = (
                        "Job is marked as MISSING since the workers failed to register "
                        "to the executor. Check the stdout/stderr logs in the submit_scripts "
                        "directory for more debug information"
                    )
        return job_status

    def shutdown(self, timeout: float = 10.0):
        """Shutdown the executor, including the interchange. This does not
        shut down any workers directly - workers should be terminated by the
        scaling mechanism or by heartbeat timeout.

        Parameters
        ----------

        timeout : float
            Amount of time to wait for the Interchange process to terminate before
            we forcefully kill it.
        """
        if self.interchange_proc is None:
            logger.info("HighThroughputExecutor has not started; skipping shutdown")
            return

        logger.info("Attempting HighThroughputExecutor shutdown")

        logger.info("Terminating interchange and result queue thread")
        self._result_queue_thread_exit.set()
        self.interchange_proc.terminate()
        try:
            self.interchange_proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            logger.warning("Unable to terminate Interchange process; sending SIGKILL")
            self.interchange_proc.kill()

        logger.info("Closing ZMQ pipes")

        # These pipes are used in a thread unsafe manner. If you have traced a
        # problem to this block of code, you might consider what is happening
        # with other threads that access these.

        # incoming_q is not closed here because it is used by the results queue
        # worker which is not shut down at this point.

        if hasattr(self, 'outgoing_q'):
            logger.info("Closing outgoing_q")
            self.outgoing_q.close()

        if hasattr(self, 'command_client'):
            logger.info("Closing command client")
            self.command_client.close()

        logger.info("Waiting for result queue thread exit")
        if self._result_queue_thread:
            self._result_queue_thread.join()

        logger.info("Finished HighThroughputExecutor shutdown attempt")

    def get_usage_information(self):
        return {"mpi": self.enable_mpi_mode}