Executor

Module with executor logic.

Executor

Executor class definition

Source code in executor/app/main.py

class Executor(metaclass=LoggableMeta):
    """Executor class definition"""
    _LOG = logging.getLogger("geokube.Executor")

    def __init__(self, broker, store_path):
        self._store = store_path
        broker_conn = pika.BlockingConnection(
            pika.ConnectionParameters(host=broker, heartbeat=10),
        )
        self._conn = broker_conn
        self._channel = broker_conn.channel()
        self._db = DBManager()

    def create_dask_cluster(self, dask_cluster_opts: dict = None):
        """Create a Dask cluster with given options.

        Parameters
        ----------
        dask_cluster_opts : optional `dict`
            A dictionary with cluster parameters.
        """
        if dask_cluster_opts is None:
            dask_cluster_opts = {}
            dask_cluster_opts["scheduler_port"] = int(
                os.getenv("DASK_SCHEDULER_PORT", 8188)
            )
            dask_cluster_opts["processes"] = True
            port = int(os.getenv("DASK_DASHBOARD_PORT", 8787))
            dask_cluster_opts["dashboard_address"] = f":{port}"
            dask_cluster_opts["n_workers"] = None
            dask_cluster_opts["memory_limit"] = "auto"
        self._worker_id = self._db.create_worker(
            status="enabled",
            dask_scheduler_port=dask_cluster_opts["scheduler_port"],
            dask_dashboard_address=dask_cluster_opts["dashboard_address"],
        )
        self._LOG.info(
            "creating Dask Cluster with options: `%s`",
            dask_cluster_opts,
            extra={"track_id": self._worker_id},
        )
        dask_cluster = LocalCluster(
            n_workers=dask_cluster_opts["n_workers"],
            scheduler_port=dask_cluster_opts["scheduler_port"],
            dashboard_address=dask_cluster_opts["dashboard_address"],
            memory_limit=dask_cluster_opts["memory_limit"],
        )
        self._LOG.info(
            "creating Dask Client...", extra={"track_id": self._worker_id}
        )
        self._dask_client = Client(dask_cluster)
        self._nanny = Nanny(self._dask_client.cluster.scheduler.address)

    def maybe_restart_cluster(self, status: RequestStatus):
        """Restart the run Dask cluster when needed.

        Resturt the cluster if request status was set to `TIMEOUT` or
        the cluster died.

        Parameters
        ----------
        statis : `RequestStatus`
            A status of a request being processed by the cluster.
        """
        if status is RequestStatus.TIMEOUT:
            self._LOG.info("recreating the cluster due to timeout")
            self._dask_client.cluster.close()
            self.create_dask_cluster()
        if self._dask_client.cluster.status is Status.failed:
            self._LOG.info("attempt to restart the cluster...")
            try:
                asyncio.run(self._nanny.restart())
            except Exception as err:
                self._LOG.error(
                    "couldn't restart the cluster due to an error: %s", err
                )
                self._LOG.info("closing the cluster")
                self._dask_client.cluster.close()
        if self._dask_client.cluster.status is Status.closed:
            self._LOG.info("recreating the cluster")
            self.create_dask_cluster()

    def ack_message(self, channel, delivery_tag):
        """Acknowledge the broker message."""
        if channel.is_open:
            channel.basic_ack(delivery_tag)
        else:
            self._LOG.info(
                "cannot acknowledge the message. channel is closed!"
            )
            pass

    def retry_until_timeout(
        self,
        future,
        message: Message,
        retries: int = 30,
        sleep_time: int = 10,
    ):
        """Retry processing the `future` object.

        Parameters
        ----------
        future : `Future`
            A future object to being computed
        message : `Message`
            A message object
        retries : `int`, default `30`
            A number of trials
        sleep_time : `int`, default `10`
            A number of seconds to sleep between trials

        Returns
        -------
        result : `tuple` of (location_path, status, fail_reason)
        """
        assert retries is not None, "`retries` cannot be `None`"
        assert sleep_time is not None, "`sleep_time` cannot be `None`"
        status = fail_reason = location_path = None
        try:
            self._LOG.debug(
                "attempt to get result for the request",
                extra={"track_id": message.request_id},
            )
            for _ in range(retries):
                if future.done():
                    self._LOG.debug(
                        "result is done",
                        extra={"track_id": message.request_id},
                    )
                    location_path = future.result()
                    status = RequestStatus.DONE
                    self._LOG.debug(
                        "result save under: %s",
                        location_path,
                        extra={"track_id": message.request_id},
                    )
                    break
                self._LOG.debug(
                    f"result is not ready yet. sleeping {sleep_time} sec",
                    extra={"track_id": message.request_id},
                )
                time.sleep(sleep_time)
            else:
                self._LOG.info(
                    "processing timout",
                    extra={"track_id": message.request_id},
                )
                future.cancel()
                status = RequestStatus.TIMEOUT
                fail_reason = "Processing timeout"
        except Exception as e:
            self._LOG.error(
                "failed to get result due to an error: %s",
                e,
                exc_info=True,
                stack_info=True,
                extra={"track_id": message.request_id},
            )
            status = RequestStatus.FAILED
            fail_reason = f"{type(e).__name__}: {str(e)}"
        return (location_path, status, fail_reason)

    def handle_message(self, connection, channel, delivery_tag, body):
        message: Message = Message(body)
        self._LOG.debug(
            "executing query: `%s`",
            message.content,
            extra={"track_id": message.request_id},
        )

        # TODO: estimation size should be updated, too
        self._db.update_request(
            request_id=message.request_id,
            worker_id=self._worker_id,
            status=RequestStatus.RUNNING,
        )

        self._LOG.debug(
            "submitting job for workflow request",
            extra={"track_id": message.request_id},
        )
        future = self._dask_client.submit(
            process,
            message=message,
            compute=False,
        )
        location_path, status, fail_reason = self.retry_until_timeout(
            future,
            message=message,
            retries=int(os.environ.get("RESULT_CHECK_RETRIES")),
        )
        self._db.update_request(
            request_id=message.request_id,
            worker_id=self._worker_id,
            status=status,
            location_path=location_path,
            size_bytes=self.get_size(location_path),
            fail_reason=fail_reason,
        )
        self._LOG.debug(
            "acknowledging request", extra={"track_id": message.request_id}
        )
        cb = functools.partial(self.ack_message, channel, delivery_tag)
        connection.add_callback_threadsafe(cb)

        self.maybe_restart_cluster(status)
        self._LOG.debug(
            "request acknowledged", extra={"track_id": message.request_id}
        )

    def on_message(self, channel, method_frame, header_frame, body, args):
        (connection, threads) = args
        delivery_tag = method_frame.delivery_tag
        t = threading.Thread(
            target=self.handle_message,
            args=(connection, channel, delivery_tag, body),
        )
        t.start()
        threads.append(t)

    def subscribe(self, etype):
        self._LOG.debug(
            "subscribe channel: %s_queue", etype, extra={"track_id": "N/A"}
        )
        self._channel.queue_declare(queue=f"{etype}_queue", durable=True)
        self._channel.basic_qos(prefetch_count=1)

        threads = []
        on_message_callback = functools.partial(
            self.on_message, args=(self._conn, threads)
        )

        self._channel.basic_consume(
            queue=f"{etype}_queue", on_message_callback=on_message_callback
        )

    def listen(self):
        while True:
            self._channel.start_consuming()

    def get_size(self, location_path):
        if location_path and os.path.exists(location_path):
            return os.path.getsize(location_path)
        return None

ack_message(channel, delivery_tag)

Acknowledge the broker message.

Source code in executor/app/main.py

def ack_message(self, channel, delivery_tag):
    """Acknowledge the broker message."""
    if channel.is_open:
        channel.basic_ack(delivery_tag)
    else:
        self._LOG.info(
            "cannot acknowledge the message. channel is closed!"
        )
        pass

create_dask_cluster(dask_cluster_opts=None)

Create a Dask cluster with given options.

Parameters

dask_cluster_opts : optional dict A dictionary with cluster parameters.

Source code in executor/app/main.py

def create_dask_cluster(self, dask_cluster_opts: dict = None):
    """Create a Dask cluster with given options.

    Parameters
    ----------
    dask_cluster_opts : optional `dict`
        A dictionary with cluster parameters.
    """
    if dask_cluster_opts is None:
        dask_cluster_opts = {}
        dask_cluster_opts["scheduler_port"] = int(
            os.getenv("DASK_SCHEDULER_PORT", 8188)
        )
        dask_cluster_opts["processes"] = True
        port = int(os.getenv("DASK_DASHBOARD_PORT", 8787))
        dask_cluster_opts["dashboard_address"] = f":{port}"
        dask_cluster_opts["n_workers"] = None
        dask_cluster_opts["memory_limit"] = "auto"
    self._worker_id = self._db.create_worker(
        status="enabled",
        dask_scheduler_port=dask_cluster_opts["scheduler_port"],
        dask_dashboard_address=dask_cluster_opts["dashboard_address"],
    )
    self._LOG.info(
        "creating Dask Cluster with options: `%s`",
        dask_cluster_opts,
        extra={"track_id": self._worker_id},
    )
    dask_cluster = LocalCluster(
        n_workers=dask_cluster_opts["n_workers"],
        scheduler_port=dask_cluster_opts["scheduler_port"],
        dashboard_address=dask_cluster_opts["dashboard_address"],
        memory_limit=dask_cluster_opts["memory_limit"],
    )
    self._LOG.info(
        "creating Dask Client...", extra={"track_id": self._worker_id}
    )
    self._dask_client = Client(dask_cluster)
    self._nanny = Nanny(self._dask_client.cluster.scheduler.address)

maybe_restart_cluster(status)

Restart the run Dask cluster when needed.

Resturt the cluster if request status was set to TIMEOUT or the cluster died.

Parameters

statis : RequestStatus A status of a request being processed by the cluster.

Source code in executor/app/main.py

def maybe_restart_cluster(self, status: RequestStatus):
    """Restart the run Dask cluster when needed.

    Resturt the cluster if request status was set to `TIMEOUT` or
    the cluster died.

    Parameters
    ----------
    statis : `RequestStatus`
        A status of a request being processed by the cluster.
    """
    if status is RequestStatus.TIMEOUT:
        self._LOG.info("recreating the cluster due to timeout")
        self._dask_client.cluster.close()
        self.create_dask_cluster()
    if self._dask_client.cluster.status is Status.failed:
        self._LOG.info("attempt to restart the cluster...")
        try:
            asyncio.run(self._nanny.restart())
        except Exception as err:
            self._LOG.error(
                "couldn't restart the cluster due to an error: %s", err
            )
            self._LOG.info("closing the cluster")
            self._dask_client.cluster.close()
    if self._dask_client.cluster.status is Status.closed:
        self._LOG.info("recreating the cluster")
        self.create_dask_cluster()

retry_until_timeout(future, message, retries=30, sleep_time=10)

Retry processing the future object.

Parameters

future : Future A future object to being computed message : Message A message object retries : int, default 30 A number of trials sleep_time : int, default 10 A number of seconds to sleep between trials

Returns

result : tuple of (location_path, status, fail_reason)

Source code in executor/app/main.py

def retry_until_timeout(
    self,
    future,
    message: Message,
    retries: int = 30,
    sleep_time: int = 10,
):
    """Retry processing the `future` object.

    Parameters
    ----------
    future : `Future`
        A future object to being computed
    message : `Message`
        A message object
    retries : `int`, default `30`
        A number of trials
    sleep_time : `int`, default `10`
        A number of seconds to sleep between trials

    Returns
    -------
    result : `tuple` of (location_path, status, fail_reason)
    """
    assert retries is not None, "`retries` cannot be `None`"
    assert sleep_time is not None, "`sleep_time` cannot be `None`"
    status = fail_reason = location_path = None
    try:
        self._LOG.debug(
            "attempt to get result for the request",
            extra={"track_id": message.request_id},
        )
        for _ in range(retries):
            if future.done():
                self._LOG.debug(
                    "result is done",
                    extra={"track_id": message.request_id},
                )
                location_path = future.result()
                status = RequestStatus.DONE
                self._LOG.debug(
                    "result save under: %s",
                    location_path,
                    extra={"track_id": message.request_id},
                )
                break
            self._LOG.debug(
                f"result is not ready yet. sleeping {sleep_time} sec",
                extra={"track_id": message.request_id},
            )
            time.sleep(sleep_time)
        else:
            self._LOG.info(
                "processing timout",
                extra={"track_id": message.request_id},
            )
            future.cancel()
            status = RequestStatus.TIMEOUT
            fail_reason = "Processing timeout"
    except Exception as e:
        self._LOG.error(
            "failed to get result due to an error: %s",
            e,
            exc_info=True,
            stack_info=True,
            extra={"track_id": message.request_id},
        )
        status = RequestStatus.FAILED
        fail_reason = f"{type(e).__name__}: {str(e)}"
    return (location_path, status, fail_reason)

persist_datacube(kube, message, base_path)

Save geokube.DataCube given the message and base_path.

Parameters

kube : geokube.DataCube A data cube to save message : Message A message with details like dataset or product ID base_path : str Base path to save

Source code in executor/app/main.py

def persist_datacube(
    kube: DataCube,
    message: Message,
    base_path: str | os.PathLike,
) -> str | os.PathLike:
    """Save `geokube.DataCube` given the `message` and `base_path`.

    Parameters
    ----------
    kube : `geokube.DataCube`
        A data cube to save
    message : `Message`
        A message with details like dataset or product ID
    base_path : `str`
        Base path to save
    """
    if rcp85_filename_condition(kube, message):
        path = get_file_name_for_climate_downscaled(kube, message)
    else:
        var_names = list(kube.fields.keys())
        if len(kube) == 1:
            path = "_".join(
                [
                    var_names[0],
                    message.dataset_id,
                    message.product_id,
                    message.request_id,
                ]
            )
        else:
            path = "_".join(
                [message.dataset_id, message.product_id, message.request_id]
            )
    kube._properties["history"] = get_history_message()
    if isinstance(message.content, GeoQuery):
        format = message.content.format
        format_args = message.content.format_args
    else:
        format = "netcdf"
    match format:
        case "netcdf":
            full_path = os.path.join(base_path, f"{path}.nc")
            kube.to_netcdf(full_path)
        case "geojson":
            full_path = os.path.join(base_path, f"{path}.json")
            kube.to_geojson(full_path)
        case "png":
            full_path = os.path.join(base_path, f"{path}.png")
            kube.to_image(full_path, **format_args)
        case "jpeg":
            full_path = os.path.join(base_path, f"{path}.jpg")
            kube.to_image(full_path, **format_args)
        case _:
            raise ValueError(f"format `{format}` is not supported")
    return full_path

persist_dataset(dset, message, base_path)

Save geokube.Dataset given the message and base_path.

Under the hood uses persist_datacube function.

Parameters

kube : geokube.Dataset A data cube to save message : Message A message with details like dataset or product ID base_path : str Base path to save

Source code in executor/app/main.py

def persist_dataset(
    dset: Dataset,
    message: Message,
    base_path: str | os.PathLike,
):
    """Save `geokube.Dataset` given the `message` and `base_path`.

    Under the hood uses `persist_datacube` function.

    Parameters
    ----------
    kube : `geokube.Dataset`
        A data cube to save
    message : `Message`
        A message with details like dataset or product ID
    base_path : `str`
        Base path to save
    """    
    def _get_attr_comb(dataframe_item, attrs):
        return "_".join([dataframe_item[attr_name] for attr_name in attrs])

    def _persist_single_datacube(dataframe_item, base_path, format, format_args=None):
        if not format_args:
            format_args = {}
        dcube = dataframe_item[dset.DATACUBE_COL]
        if isinstance(dcube, Delayed):
            dcube = dcube.compute()
        if len(dcube) == 0:
            return None
        for field in dcube.fields.values():
            if 0 in field.shape:
                return None
        attr_str = _get_attr_comb(dataframe_item, dset._Dataset__attrs)
        var_names = list(dcube.fields.keys())
        if len(dcube) == 1:
            path = "_".join(
                [
                    var_names[0],
                    message.dataset_id,
                    message.product_id,
                    attr_str,
                    message.request_id,
                ]
            )
        else:
            path = "_".join(
                [
                    message.dataset_id,
                    message.product_id,
                    attr_str,
                    message.request_id,
                ]
            )
        match format:
            case "netcdf":
                full_path = os.path.join(base_path, f"{path}.nc")
                dcube.to_netcdf(full_path)
            case "geojson":
                full_path = os.path.join(base_path, f"{path}.json")
                dcube.to_geojson(full_path)
            case "png":
                full_path = os.path.join(base_path, f"{path}.png")
                dcube.to_image(full_path, **format_args)
            case "jpeg":
                full_path = os.path.join(base_path, f"{path}.jpg")
                dcube.to_image(full_path, **format_args)
            case _:
                raise ValueError(f"format: {format} is not supported!")
        return full_path

    if isinstance(message.content, GeoQuery):
        format = message.content.format
        format_args = message.content.format_args
    else:
        format = "netcdf"
    datacubes_paths = dset.data.apply(
        _persist_single_datacube, base_path=base_path, format=format, format_args=format_args, axis=1
    )
    paths = datacubes_paths[~datacubes_paths.isna()]
    if len(paths) == 0:
        return None
    elif len(paths) == 1:
        return paths.iloc[0]
    zip_name = "_".join(
        [message.dataset_id, message.product_id, message.request_id]
    )
    path = os.path.join(base_path, f"{zip_name}.zip")
    with ZipFile(path, "w") as archive:
        for file in paths:
            archive.write(file, arcname=os.path.basename(file))
    for file in paths:
        os.remove(file)
    return path

process(message, compute)

Process a message and compute (if needed).

Parameters

message : Message A message to process compute : bool A flag to indicate if result should be computed

Source code in executor/app/main.py

def process(message: Message, compute: bool):
    """Process a message and compute (if needed).

    Parameters
    ----------
    message : `Message`
        A message to process
    compute : bool
        A flag to indicate if result should be computed
    """
    res_path = os.path.join(_BASE_DOWNLOAD_PATH, message.request_id)
    os.makedirs(res_path, exist_ok=True)
    match message.type:
        case MessageType.QUERY:
            kube = Datastore().query(
                message.dataset_id,
                message.product_id,
                message.content,
                compute,
            )
        case MessageType.WORKFLOW:
            kube = Workflow.from_tasklist(message.content).compute()
        case _:
            raise ValueError("unsupported message type")
    if isinstance(kube, Field):
        kube = DataCube(
            fields=[kube],
            properties=kube.properties,
            encoding=kube.encoding,
        )
    match kube:
        case DataCube():
            return persist_datacube(kube, message, base_path=res_path)
        case Dataset():
            return persist_dataset(kube, message, base_path=res_path)
        case _:
            raise TypeError(
                "expected geokube.DataCube or geokube.Dataset, but passed"
                f" {type(kube).__name__}"
            )

Module contains definitions of messages in the executor.

Message

Message class definition.

Source code in executor/app/messaging.py

class Message:
    """Message class definition."""
    _LOG = logging.getLogger("geokube.Message")

    request_id: int
    dataset_id: str = "<unknown>"
    product_id: str = "<unknown>"
    type: MessageType
    content: GeoQuery | Workflow

    def __init__(self, load: bytes) -> None:
        """Create `Message` instances.

        Parameters
        ----------
        load : `bytes`
            Bytes containing message load
        """
        self.request_id, msg_type, *query = load.decode().split(
            MESSAGE_SEPARATOR
        )
        match MessageType(msg_type):
            case MessageType.QUERY:
                self._LOG.debug("processing content of `query` type")
                assert len(query) == 3, "improper content for query message"
                self.dataset_id, self.product_id, self.content = query
                self.content: GeoQuery = GeoQuery.parse(self.content)
                self.type = MessageType.QUERY
            case MessageType.WORKFLOW:
                self._LOG.debug("processing content of `workflow` type")
                assert len(query) == 1, "improper content for workflow message"
                self.content: Workflow = Workflow.parse(query[0])
                self.dataset_id = self.content.dataset_id
                self.product_id = self.content.product_id
                self.type = MessageType.WORKFLOW
            case _:
                self._LOG.error("type `%s` is not supported", msg_type)
                raise ValueError(f"type `{msg_type}` is not supported!")

__init__(load)

Create Message instances.

Parameters

load : bytes Bytes containing message load

Source code in executor/app/messaging.py

def __init__(self, load: bytes) -> None:
    """Create `Message` instances.

    Parameters
    ----------
    load : `bytes`
        Bytes containing message load
    """
    self.request_id, msg_type, *query = load.decode().split(
        MESSAGE_SEPARATOR
    )
    match MessageType(msg_type):
        case MessageType.QUERY:
            self._LOG.debug("processing content of `query` type")
            assert len(query) == 3, "improper content for query message"
            self.dataset_id, self.product_id, self.content = query
            self.content: GeoQuery = GeoQuery.parse(self.content)
            self.type = MessageType.QUERY
        case MessageType.WORKFLOW:
            self._LOG.debug("processing content of `workflow` type")
            assert len(query) == 1, "improper content for workflow message"
            self.content: Workflow = Workflow.parse(query[0])
            self.dataset_id = self.content.dataset_id
            self.product_id = self.content.product_id
            self.type = MessageType.WORKFLOW
        case _:
            self._LOG.error("type `%s` is not supported", msg_type)
            raise ValueError(f"type `{msg_type}` is not supported!")

Module with LoggableMeta metaclass

LoggableMeta

Bases: type

Metaclass for dealing with logger levels and handlers

Source code in executor/app/meta.py

class LoggableMeta(type):
    """Metaclass for dealing with logger levels and handlers"""

    def __new__(cls, child_cls, bases, namespace):
        # NOTE: method is called while creating a class, not an instance!
        res = super().__new__(cls, child_cls, bases, namespace)
        if hasattr(res, "_LOG"):
            format_ = os.environ.get(
                "LOGGING_FORMAT",
                "%(asctime)s %(name)s %(levelname)s %(lineno)d"
                " %(track_id)s %(message)s",
            )
            formatter = logging.Formatter(format_)
            logging_level = os.environ.get("LOGGING_LEVEL", "INFO")
            res._LOG.setLevel(logging_level)
            stream_handler = logging.StreamHandler()
            stream_handler.setFormatter(formatter)
            stream_handler.setLevel(logging_level)
            res._LOG.addHandler(stream_handler)
            for handler in logging.getLogger("geokube").handlers:
                handler.setFormatter(formatter)
        return res