opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/client/agent/client_agent.py

@@ -0,0 +1,218 @@
+import asyncio
+import logging
+import sys
+import threading
+from concurrent.futures import Future
+from typing import Optional
+
+import zmq.asyncio
+
+from scaler.client.agent.disconnect_manager import ClientDisconnectManager
+from scaler.client.agent.future_manager import ClientFutureManager
+from scaler.client.agent.heartbeat_manager import ClientHeartbeatManager
+from scaler.client.agent.object_manager import ClientObjectManager
+from scaler.client.agent.task_manager import ClientTaskManager
+from scaler.client.serializer.mixins import Serializer
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.async_connector import ZMQAsyncConnector
+from scaler.io.mixins import AsyncConnector
+from scaler.protocol.python.common import ObjectStorageAddress
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    ClientHeartbeatEcho,
+    ClientShutdownResponse,
+    GraphTask,
+    ObjectInstruction,
+    Task,
+    TaskCancel,
+    TaskCancelConfirm,
+    TaskLog,
+    TaskResult,
+)
+from scaler.protocol.python.mixins import Message
+from scaler.utility.event_loop import create_async_loop_routine
+from scaler.utility.exceptions import ClientCancelledException, ClientQuitException, ClientShutdownException
+from scaler.utility.identifiers import ClientID
+
+
+class ClientAgent(threading.Thread):
+    def __init__(
+        self,
+        identity: ClientID,
+        client_agent_address: ZMQConfig,
+        scheduler_address: ZMQConfig,
+        context: zmq.Context,
+        future_manager: ClientFutureManager,
+        stop_event: threading.Event,
+        timeout_seconds: int,
+        heartbeat_interval_seconds: int,
+        serializer: Serializer,
+        object_storage_address: Optional[str] = None,
+    ):
+        threading.Thread.__init__(self, daemon=True)
+
+        self._stop_event = stop_event
+        self._timeout_seconds = timeout_seconds
+        self._heartbeat_interval_seconds = heartbeat_interval_seconds
+        self._serializer = serializer
+
+        self._identity = identity
+        self._client_agent_address = client_agent_address
+        self._scheduler_address = scheduler_address
+        self._context = context
+        self._object_storage_address: Future[ObjectStorageAddress] = Future()
+        if object_storage_address is not None:
+            manual_config = ZMQConfig.from_string(object_storage_address)
+            self._object_storage_address_override = ObjectStorageAddress.new_msg(manual_config.host, manual_config.port)
+        else:
+            self._object_storage_address_override = None
+
+        self._future_manager = future_manager
+
+        self._connector_internal: AsyncConnector = ZMQAsyncConnector(
+            context=zmq.asyncio.Context.shadow(self._context),
+            name="client_agent_internal",
+            socket_type=zmq.PAIR,
+            bind_or_connect="bind",
+            address=self._client_agent_address,
+            callback=self.__on_receive_from_client,
+            identity=None,
+        )
+        self._connector_external: AsyncConnector = ZMQAsyncConnector(
+            context=zmq.asyncio.Context.shadow(self._context),
+            name="client_agent_external",
+            socket_type=zmq.DEALER,
+            address=self._scheduler_address,
+            bind_or_connect="connect",
+            callback=self.__on_receive_from_scheduler,
+            identity=self._identity,
+        )
+
+        self._disconnect_manager: Optional[ClientDisconnectManager] = None
+        self._heartbeat_manager: Optional[ClientHeartbeatManager] = None
+        self._task_manager: Optional[ClientTaskManager] = None
+
+    def __initialize(self):
+        self._disconnect_manager = ClientDisconnectManager()
+        self._heartbeat_manager = ClientHeartbeatManager(
+            death_timeout_seconds=self._timeout_seconds, storage_address_future=self._object_storage_address
+        )
+        self._object_manager = ClientObjectManager(identity=self._identity)
+        self._task_manager = ClientTaskManager()
+
+        # register all managers
+        self._disconnect_manager.register(
+            connector_internal=self._connector_internal, connector_external=self._connector_external
+        )
+        self._object_manager.register(
+            connector_internal=self._connector_internal, connector_external=self._connector_external
+        )
+        self._task_manager.register(
+            connector_external=self._connector_external,
+            object_manager=self._object_manager,
+            future_manager=self._future_manager,
+        )
+        self._heartbeat_manager.register(connector_external=self._connector_external)
+
+    def __run_loop(self):
+        self._loop = asyncio.new_event_loop()
+        self._task = self._loop.create_task(self.__get_loops())
+        self._loop.run_until_complete(self._task)
+        self._loop.close()
+
+    def run(self):
+        self.__initialize()
+        self.__run_loop()
+
+    def get_object_storage_address(self) -> ObjectStorageAddress:
+        """Returns the object storage address, or block until it receives it."""
+        if self._object_storage_address_override is not None:
+            return self._object_storage_address_override
+        return self._object_storage_address.result()
+
+    async def __on_receive_from_client(self, message: Message):
+        if isinstance(message, ClientDisconnect):
+            await self._disconnect_manager.on_client_disconnect(message)
+            return
+
+        if isinstance(message, ObjectInstruction):
+            await self._object_manager.on_object_instruction(message)
+            return
+
+        if isinstance(message, Task):
+            await self._task_manager.on_new_task(message)
+            return
+
+        if isinstance(message, TaskCancel):
+            await self._task_manager.on_cancel_task(message)
+            return
+
+        if isinstance(message, GraphTask):
+            await self._task_manager.on_new_graph_task(message)
+            return
+
+        raise TypeError(f"Unknown {message=}")
+
+    async def __on_receive_from_scheduler(self, message: Message):
+        if isinstance(message, ClientShutdownResponse):
+            await self._disconnect_manager.on_client_shutdown_response(message)
+            return
+
+        if isinstance(message, ClientHeartbeatEcho):
+            await self._heartbeat_manager.on_heartbeat_echo(message)
+            return
+
+        if isinstance(message, TaskLog):
+            log_type = sys.stdout if message.log_type == TaskLog.LogType.Stdout else sys.stderr
+            print(message.content, file=log_type, end="")
+            return
+
+        if isinstance(message, TaskResult):
+            await self._task_manager.on_task_result(message)
+            return
+
+        if isinstance(message, TaskCancelConfirm):
+            await self._task_manager.on_task_cancel_confirm(message)
+            return
+
+        raise TypeError(f"Unknown {message=}")
+
+    async def __get_loops(self):
+        await self._heartbeat_manager.send_heartbeat()
+
+        loops = [
+            create_async_loop_routine(self._connector_external.routine, 0),
+            create_async_loop_routine(self._connector_internal.routine, 0),
+            create_async_loop_routine(self._heartbeat_manager.routine, self._heartbeat_interval_seconds),
+        ]
+
+        exception = None
+        try:
+            await asyncio.gather(*loops)
+        except BaseException as e:
+            exception = e
+        finally:
+            self._stop_event.set()  # always set the stop event before setting futures' exceptions
+
+            await self._object_manager.clear_all_objects(clear_serializer=True)
+
+            self._connector_external.destroy()
+            self._connector_internal.destroy()
+
+        if exception is None:
+            return
+
+        if not self._object_storage_address.done():
+            self._object_storage_address.set_exception(exception)
+
+        if isinstance(exception, asyncio.CancelledError):
+            logging.error("ClientAgent: async. loop cancelled")
+            self._future_manager.set_all_futures_with_exception(ClientCancelledException("client cancelled"))
+        elif isinstance(exception, (ClientQuitException, ClientShutdownException)):
+            logging.info("ClientAgent: client quitting")
+            self._future_manager.set_all_futures_with_exception(exception)
+        elif isinstance(exception, TimeoutError):
+            logging.error(f"ClientAgent: client timeout when connecting to {self._scheduler_address.to_address()}")
+            self._future_manager.set_all_futures_with_exception(exception)
+        else:
+            raise exception

scaler/client/agent/disconnect_manager.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+from scaler.client.agent.mixins import DisconnectManager
+from scaler.io.mixins import AsyncConnector
+from scaler.protocol.python.message import ClientDisconnect, ClientShutdownResponse
+from scaler.utility.exceptions import ClientQuitException, ClientShutdownException
+
+
+class ClientDisconnectManager(DisconnectManager):
+    def __init__(self):
+        self._connector_internal: Optional[AsyncConnector] = None
+        self._connector_external: Optional[AsyncConnector] = None
+
+    def register(self, connector_internal: AsyncConnector, connector_external: AsyncConnector):
+        self._connector_internal = connector_internal
+        self._connector_external = connector_external
+
+    async def on_client_disconnect(self, disconnect: ClientDisconnect):
+        await self._connector_external.send(disconnect)
+
+        if disconnect.disconnect_type == ClientDisconnect.DisconnectType.Disconnect:
+            raise ClientQuitException("client disconnecting")
+
+    async def on_client_shutdown_response(self, response: ClientShutdownResponse):
+        await self._connector_internal.send(response)
+
+        raise ClientShutdownException("cluster shutting down")

scaler/client/agent/future_manager.py

@@ -0,0 +1,112 @@
+import logging
+import threading
+from concurrent.futures import Future, InvalidStateError
+from typing import Dict, Optional
+
+from scaler.client.agent.mixins import FutureManager
+from scaler.client.future import ScalerFuture
+from scaler.client.serializer.mixins import Serializer
+from scaler.protocol.python.common import TaskCancelConfirmType, TaskResultType, TaskState
+from scaler.protocol.python.message import TaskCancelConfirm, TaskResult
+from scaler.utility.exceptions import WorkerDiedError
+from scaler.utility.identifiers import ObjectID, TaskID
+from scaler.utility.metadata.profile_result import retrieve_profiling_result_from_task_result
+
+
+class ClientFutureManager(FutureManager):
+    def __init__(self, serializer: Serializer):
+        self._lock = threading.RLock()
+        self._serializer = serializer
+
+        self._task_id_to_future: Dict[TaskID, ScalerFuture] = dict()
+
+    def add_future(self, future: Future):
+        assert isinstance(future, ScalerFuture)
+        with self._lock:
+            future.set_running_or_notify_cancel()
+            self._task_id_to_future[future.task_id] = future
+
+    def cancel_all_futures(self):
+        with self._lock:
+            futures_to_cancel = list(self._task_id_to_future.values())
+
+        # Actually cancelling the futures should occur without holding the future manager's lock. That's because
+        # `cancel()` is blocking, and requires the manager to process result and cancel confirm messages.
+
+        logging.info(f"canceling {len(futures_to_cancel)} task(s)")
+        for future in futures_to_cancel:
+            future.cancel()
+
+    def set_all_futures_with_exception(self, exception: Exception):
+        with self._lock:
+            for future in self._task_id_to_future.values():
+                try:
+                    future.set_exception(exception)
+                except InvalidStateError:
+                    continue  # Future got canceled
+
+            self._task_id_to_future.clear()
+
+    def on_task_result(self, result: TaskResult):
+        with self._lock:
+            task_id = result.task_id
+            if task_id not in self._task_id_to_future:
+                return
+
+            future = self._task_id_to_future.pop(task_id)
+            assert result.task_id == future.task_id
+
+            profile_result = retrieve_profiling_result_from_task_result(result)
+
+            if result.result_type == TaskResultType.FailedWorkerDied:
+                future.set_exception(
+                    WorkerDiedError(f"worker died when processing task: {task_id.hex()}"), profile_result
+                )
+
+            elif result.result_type == TaskResultType.Success:
+                result_object_id = self.__get_result_object_id(result)
+                future.set_result_ready(result_object_id, TaskState.Success, profile_result)
+
+            elif result.result_type == TaskResultType.Failed:
+                result_object_id = self.__get_result_object_id(result)
+                future.set_result_ready(result_object_id, TaskState.Failed, profile_result)
+
+            else:
+                raise TypeError(f"{result.task_id.hex()}: Unknown task status: {result.result_type}")
+
+    def on_task_cancel_confirm(self, cancel_confirm: TaskCancelConfirm):
+        with self._lock:
+            task_id = cancel_confirm.task_id
+            if task_id not in self._task_id_to_future:
+                return
+
+            future = self._task_id_to_future.pop(task_id)
+            assert cancel_confirm.task_id == future.task_id
+
+            if cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.Canceled:
+                future.set_canceled()
+
+            elif cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.CancelNotFound:
+                logging.error(f"{task_id!r}: task to cancel not found")
+                future.set_canceled()
+
+            elif cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.CancelFailed:
+                logging.error(f"{task_id!r}: task cancel failed")
+                self._task_id_to_future[task_id] = future
+
+            else:
+                raise TypeError(
+                    f"{task_id}: unknown task cancel confirm type:" f" {cancel_confirm.cancel_confirm_type}"
+                )
+
+    @staticmethod
+    def __get_result_object_id(result: TaskResult) -> Optional[ObjectID]:
+        if len(result.results) == 1:
+            result_object_id = ObjectID(result.results[0])
+        elif len(result.results) == 0:
+            # this will happen only if umbrella task is done
+            result_object_id = None
+        else:
+            raise ValueError(f"{result.task_id!r}: received multiple objects for the results: {len(result.results)=}")
+
+        return result_object_id

scaler/client/agent/heartbeat_manager.py

@@ -0,0 +1,74 @@
+import time
+from concurrent.futures import Future
+from typing import Optional
+
+import psutil
+
+from scaler.client.agent.mixins import HeartbeatManager, ObjectManager
+from scaler.io.mixins import AsyncConnector
+from scaler.protocol.python.common import ObjectStorageAddress
+from scaler.protocol.python.message import ClientHeartbeat, ClientHeartbeatEcho
+from scaler.protocol.python.status import Resource
+from scaler.utility.mixins import Looper
+
+
+class ClientHeartbeatManager(Looper, HeartbeatManager):
+    def __init__(self, death_timeout_seconds: int, storage_address_future: Future):
+        self._death_timeout_seconds = death_timeout_seconds
+        self._object_storage_address = storage_address_future
+
+        self._process = psutil.Process()
+
+        self._last_scheduler_contact = time.time()
+        self._start_timestamp_ns = 0
+        self._latency_us = 0
+        self._connected = False
+
+        self._connector_external: Optional[AsyncConnector] = None
+        self._object_manager: Optional[ObjectManager] = None
+
+    def register(self, connector_external: AsyncConnector):
+        self._connector_external = connector_external
+
+    async def send_heartbeat(self):
+        await self._connector_external.send(
+            ClientHeartbeat.new_msg(
+                Resource.new_msg(int(self._process.cpu_percent() * 10), self._process.memory_info().rss),
+                self._latency_us,
+            )
+        )
+
+    async def on_heartbeat_echo(self, heartbeat: ClientHeartbeatEcho):
+        if not self._connected:
+            self._connected = True
+
+        self._last_scheduler_contact = time.time()
+        if self._start_timestamp_ns == 0:
+            # not handling echo if we didn't send out heartbeat
+            return
+
+        self._latency_us = int(((time.time_ns() - self._start_timestamp_ns) / 2) // 1_000)
+        self._start_timestamp_ns = 0
+
+        if self._object_storage_address.done():
+            return
+
+        self._object_storage_address.set_result(heartbeat.object_storage_address())
+
+    async def routine(self):
+        if time.time() - self._last_scheduler_contact > self._death_timeout_seconds:
+            raise TimeoutError(
+                f"Timeout when connecting to scheduler {self._connector_external.address} "
+                f"in {self._death_timeout_seconds} seconds"
+            )
+
+        if self._start_timestamp_ns != 0:
+            # already sent heartbeat, expecting heartbeat echo, so not sending
+            return
+
+        await self.send_heartbeat()
+        self._start_timestamp_ns = time.time_ns()
+
+    def get_object_storage_address(self) -> ObjectStorageAddress:
+        """Returns the object storage configuration, or block until it receives it."""
+        return self._object_storage_address.result()

scaler/client/agent/mixins.py

@@ -0,0 +1,89 @@
+import abc
+from concurrent.futures import Future
+
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    ClientHeartbeatEcho,
+    ClientShutdownResponse,
+    GraphTask,
+    ObjectInstruction,
+    Task,
+    TaskCancelConfirm,
+    TaskResult,
+)
+
+
+class HeartbeatManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def send_heartbeat(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_heartbeat_echo(self, heartbeat: ClientHeartbeatEcho):
+        raise NotImplementedError()
+
+
+class TimeoutManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def update_last_seen_time(self):
+        raise NotImplementedError()
+
+
+class ObjectManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def on_object_instruction(self, object_instruction: ObjectInstruction):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_result(self, result: TaskResult):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def clear_all_objects(self, clear_serializer: bool):
+        raise NotImplementedError()
+
+
+class TaskManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def on_new_task(self, task: Task):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_new_graph_task(self, task: GraphTask):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_result(self, result: TaskResult):
+        raise NotImplementedError()
+
+
+class FutureManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def add_future(self, future: Future):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def cancel_all_futures(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_all_futures_with_exception(self, exception: Exception):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_result(self, result: TaskResult):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_cancel_confirm(self, cancel_confirm: TaskCancelConfirm):
+        raise NotImplementedError()
+
+
+class DisconnectManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def on_client_disconnect(self, disconnect: ClientDisconnect):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_client_shutdown_response(self, response: ClientShutdownResponse):
+        raise NotImplementedError()

scaler/client/agent/object_manager.py

@@ -0,0 +1,98 @@
+from typing import Optional, Set
+
+from scaler.client.agent.mixins import ObjectManager
+from scaler.io.mixins import AsyncConnector
+from scaler.protocol.python.common import ObjectMetadata
+from scaler.protocol.python.message import ObjectInstruction, TaskResult
+from scaler.utility.identifiers import ClientID, ObjectID
+
+
+class ClientObjectManager(ObjectManager):
+    def __init__(self, identity: ClientID):
+        self._sent_object_ids: Set[ObjectID] = set()
+        self._sent_serializer_id: Optional[ObjectID] = None
+
+        self._identity = identity
+
+        self._connector_internal: Optional[AsyncConnector] = None
+        self._connector_external: Optional[AsyncConnector] = None
+
+    def register(self, connector_internal: AsyncConnector, connector_external: AsyncConnector):
+        self._connector_internal = connector_internal
+        self._connector_external = connector_external
+
+    async def on_object_instruction(self, instruction: ObjectInstruction):
+        if instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Create:
+            await self.__send_object_creation(instruction)
+        elif instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Delete:
+            await self.__delete_objects(instruction)
+        elif instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Clear:
+            await self.clear_all_objects(clear_serializer=False)
+
+    def on_task_result(self, task_result: TaskResult):
+        self._sent_object_ids.update((ObjectID(object_id_bytes) for object_id_bytes in task_result.results))
+
+    async def clear_all_objects(self, clear_serializer):
+        cleared_object_ids = self._sent_object_ids.copy()
+
+        if clear_serializer:
+            self._sent_serializer_id = None
+        elif self._sent_serializer_id is not None:
+            cleared_object_ids.remove(self._sent_serializer_id)
+
+        self._sent_object_ids.difference_update(cleared_object_ids)
+
+        await self._connector_external.send(
+            ObjectInstruction.new_msg(
+                ObjectInstruction.ObjectInstructionType.Delete,
+                self._identity,
+                ObjectMetadata.new_msg(tuple(cleared_object_ids)),
+            )
+        )
+
+    async def __send_object_creation(self, instruction: ObjectInstruction):
+        assert instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Create
+
+        new_object_ids = set(instruction.object_metadata.object_ids) - self._sent_object_ids
+        if not new_object_ids:
+            return
+
+        if ObjectMetadata.ObjectContentType.Serializer in instruction.object_metadata.object_types:
+            if self._sent_serializer_id is not None:
+                raise ValueError("trying to send multiple serializers.")
+
+            serializer_index = instruction.object_metadata.object_types.index(
+                ObjectMetadata.ObjectContentType.Serializer
+            )
+            self._sent_serializer_id = instruction.object_metadata.object_ids[serializer_index]
+
+        new_object_content = ObjectMetadata.new_msg(
+            *zip(
+                *filter(
+                    lambda object_pack: object_pack[0] in new_object_ids,
+                    zip(
+                        instruction.object_metadata.object_ids,
+                        instruction.object_metadata.object_types,
+                        instruction.object_metadata.object_names,
+                    ),
+                )
+            )
+        )
+
+        self._sent_object_ids.update(set(new_object_content.object_ids))
+
+        await self._connector_external.send(
+            ObjectInstruction.new_msg(
+                ObjectInstruction.ObjectInstructionType.Create, instruction.object_user, new_object_content
+            )
+        )
+
+    async def __delete_objects(self, instruction: ObjectInstruction):
+        assert instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Delete
+
+        if self._sent_serializer_id in instruction.object_metadata.object_ids:
+            raise ValueError("trying to delete serializer.")
+
+        self._sent_object_ids.difference_update(instruction.object_metadata.object_ids)
+
+        await self._connector_external.send(instruction)

scaler/client/agent/task_manager.py

@@ -0,0 +1,64 @@
+from typing import Optional, Set
+
+from scaler.client.agent.future_manager import ClientFutureManager
+from scaler.client.agent.mixins import ObjectManager, TaskManager
+from scaler.io.mixins import AsyncConnector
+from scaler.protocol.python.message import GraphTask, Task, TaskCancel, TaskCancelConfirm, TaskResult
+
+
+class ClientTaskManager(TaskManager):
+    def __init__(self):
+        self._task_ids: Set[bytes] = set()
+
+        self._connector_external: Optional[AsyncConnector] = None
+        self._object_manager: Optional[ObjectManager] = None
+        self._future_manager: Optional[ClientFutureManager] = None
+
+    def register(
+        self, connector_external: AsyncConnector, object_manager: ObjectManager, future_manager: ClientFutureManager
+    ):
+        self._connector_external = connector_external
+        self._object_manager = object_manager
+        self._future_manager = future_manager
+
+    async def on_new_task(self, task: Task):
+        self._task_ids.add(task.task_id)
+        await self._connector_external.send(task)
+
+    async def on_cancel_task(self, task_cancel: TaskCancel):
+        # We might receive a cancel task event on a previously finished task if:
+        # -	The scheduler sends a TaskResult message to the client agent
+        # - The client sends a TaskCancel message to the client agent, as it's not yet aware the task finished.
+        # - The client agent processes the TaskResult message and removes the task from self._task_ids
+        # - The client agent processes the TaskCancel message (that was already queued before processing the
+        #   TaskResult), and fails on self._task_ids.remove() as the task_id no longer exists.
+
+        if task_cancel.task_id not in self._task_ids:
+            return
+
+        await self._connector_external.send(task_cancel)
+
+    async def on_new_graph_task(self, task: GraphTask):
+        self._task_ids.add(task.task_id)
+        self._task_ids.update(set(task.targets))
+        await self._connector_external.send(task)
+
+    async def on_task_result(self, result: TaskResult):
+        # All task result objects must be propagated to the object manager, even if we do not track the task anymore
+        # (e.g. if it got cancelled). If we don't, we might lose track of these result objects and not properly clear
+        # them.
+        self._object_manager.on_task_result(result)
+
+        if result.task_id not in self._task_ids:
+            return
+
+        self._task_ids.remove(result.task_id)
+
+        self._future_manager.on_task_result(result)
+
+    async def on_task_cancel_confirm(self, task_cancel_confirm: TaskCancelConfirm):
+        if task_cancel_confirm.task_id not in self._task_ids:
+            return
+
+        self._task_ids.remove(task_cancel_confirm.task_id)
+        self._future_manager.on_task_cancel_confirm(task_cancel_confirm)