Skip to content

climate_ref_celery.executor #

Executor for running diagnostics asynchronously using Celery

CeleryExecutor #

Bases: Executor

Run a diagnostic asynchronously

Celery is an asynchronous task queue/job queue based on distributed message passing. Celery uses a message broker to distribute tasks across a cluster of worker nodes. The worker nodes are responsible for executing the tasks. The message broker used in this case is Redis. The worker node may be running on the same machine as the client or on a different machine, either natively or via a docker container.

We cannot resume tasks that are in progress if the process terminates. That should be possible tracking some additional state in the database.

Source code in packages/climate-ref-celery/src/climate_ref_celery/executor.py 
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
class CeleryExecutor(Executor):
    """
    Run a diagnostic asynchronously

    Celery is an asynchronous task queue/job queue based on distributed message passing.
    Celery uses a message broker to distribute tasks across a cluster of worker nodes.
    The worker nodes are responsible for executing the tasks.
    The message broker used in this case is [Redis](https://github.com/redis/redis).
    The worker node may be running on the same machine as the client or on a different machine,
    either natively or via a docker container.

    We cannot resume tasks that are in progress if the process terminates.
    That should be possible tracking some additional state in the database.
    """

    name = "celery"

    # Results are persisted by the worker via the ``handle_result`` task link,
    # so ``join`` only waits for completion.
    # Skipping it (queue-and-exit) does not lose work.
    # The workers keep processing and ingesting results in the background.
    collects_results_on_join = False

    def __init__(self, *, config: Config, **kwargs: Any) -> None:
        self.config = config
        self._database: Database | None = kwargs.get("database")
        super().__init__(**kwargs)  # type: ignore
        self._results: list[celery.result.AsyncResult[ExecutionResult]] = []
        self._execution_ids: list[int] = []
        self._queue_counts: Counter[str] = Counter()
        # Resolving provider slugs imports the provider packages, so only do it when a table is set
        known_providers = _configured_provider_slugs(config) if os.environ.get(ROUTES_ENV_VAR) else None
        self._routing_table = load_routing_table(known_providers=known_providers)

    def run(
        self,
        definition: ExecutionDefinition,
        execution: Execution | None = None,
    ) -> None:
        """
        Run a diagnostic calculation

        This will queue the diagnostic to be run by a Celery worker.
        The executions will be stored in the database when the task completes if `execution`
        is specified.
        No result will be returned from this function.
        Instead, you can periodically check the status of the task in the database.

        Tasks may not execute immediately if the correct workers are not available.
        The task will be queued and executed when a worker becomes available.

        Parameters
        ----------
        definition
            A description of the information needed for this execution of the diagnostic
            This includes relative paths to the data files,
            which will be converted to absolute paths when being executed
        execution
            Result of the diagnostic execution
            This is a database object that contains the executions of the execution.
            If provided, it will be updated with the executions of the execution.
            This may happen asynchronously, so the executions may not be immediately available.
        """
        from climate_ref_celery.worker_tasks import handle_failure, handle_result  # noqa: PLC0415

        diagnostic = definition.diagnostic

        name = generate_task_name(diagnostic.provider, diagnostic)
        queue = self._routing_table.queue_for(diagnostic.provider.slug, diagnostic.slug)

        async_result = app.send_task(
            name,
            args=[definition, self.config.log_level],
            # ``exclusive`` is left at its default.
            # The worker's concurrency is set on the worker,
            # not here, so this process cannot promise the worker's cgroup holds one execution.
            # Peak memory is measured over the task's process tree instead.
            kwargs={"measure": self.config.executor.measure_resources},
            queue=queue,
            link=handle_result.s(execution_id=execution.id).set(queue="celery") if execution else None,
            link_error=handle_failure.s(execution_id=execution.id).set(queue="celery") if execution else None,
        )
        logger.debug(f"Celery task {async_result.id} for {diagnostic.slug} submitted to queue {queue}")
        self._queue_counts[queue] += 1
        self._results.append(async_result)
        if execution is not None:
            self._execution_ids.append(execution.id)

    def log_submission_summary(self) -> None:
        """
        Log a per-queue count of the tasks submitted since the last summary
        """
        for queue, count in sorted(self._queue_counts.items()):
            logger.info(f"Submitted {count} executions to queue {queue}")
        self._queue_counts.clear()

    def join(self, timeout: float) -> None:
        """
        Wait for all executions to finish

        This will block until all executions have finished running,
        and display a progress bar while waiting.

        Any tasks still running after the timeout will continue to run in the background.

        Parameters
        ----------
        timeout
            Maximum time to wait in seconds before raising a TimeoutError.
            A non-positive value (``<= 0``) waits indefinitely.

        Raises
        ------
        TimeoutError
            If all executions aren't completed within a positive timeout
        """
        start_time = time.time()
        refresh_time = 0.5  # Time to wait between checking for completed tasks in seconds

        results = self._results
        t = LoggedProgress(total=len(results), desc="Waiting for executions to complete", unit="execution")

        try:
            while results:
                # Wait for a short time before checking for completed executions
                time.sleep(refresh_time)

                elapsed_time = time.time() - start_time

                # ``timeout <= 0`` means wait indefinitely (no overall deadline).
                if timeout > 0 and elapsed_time > timeout:
                    raise TimeoutError("Not all tasks completed within the specified timeout")

                # Iterate over a copy of the list and remove finished tasks
                for result in results[:]:
                    if result.ready():
                        t.update(n=1)
                        results.remove(result)
        finally:
            t.close()

        deadline = start_time + timeout if timeout > 0 else None
        self._wait_for_recording(deadline)

    def _wait_for_recording(self, deadline: float | None) -> None:
        """
        Wait for the worker callbacks to record the outcome of each submitted execution

        A task is ready before its ``handle_result`` or ``handle_failure`` callback has run,
        so the database can still show it in progress.
        """
        ids, self._execution_ids = self._execution_ids, []
        if self._database is None or not ids:
            return

        end = time.time() + RECORDING_GRACE_SECONDS
        if deadline is not None:
            end = min(end, deadline)
        while True:
            with self._database.session_scope() as session:
                pending = (
                    session.query(func.count(Execution.id))
                    .filter(Execution.id.in_(ids), Execution.successful.is_(None))
                    .scalar()
                )
            if not pending or time.time() >= end:
                return
            time.sleep(0.5)

join(timeout) #

Wait for all executions to finish

This will block until all executions have finished running, and display a progress bar while waiting.

Any tasks still running after the timeout will continue to run in the background.

Parameters:

Name Type Description Default
timeout float

Maximum time to wait in seconds before raising a TimeoutError. A non-positive value (<= 0) waits indefinitely.

 required

Raises:

Type Description
TimeoutError

If all executions aren't completed within a positive timeout
Source code in packages/climate-ref-celery/src/climate_ref_celery/executor.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def join(self, timeout: float) -> None:
    """
    Wait for all executions to finish

    This will block until all executions have finished running,
    and display a progress bar while waiting.

    Any tasks still running after the timeout will continue to run in the background.

    Parameters
    ----------
    timeout
        Maximum time to wait in seconds before raising a TimeoutError.
        A non-positive value (``<= 0``) waits indefinitely.

    Raises
    ------
    TimeoutError
        If all executions aren't completed within a positive timeout
    """
    start_time = time.time()
    refresh_time = 0.5  # Time to wait between checking for completed tasks in seconds

    results = self._results
    t = LoggedProgress(total=len(results), desc="Waiting for executions to complete", unit="execution")

    try:
        while results:
            # Wait for a short time before checking for completed executions
            time.sleep(refresh_time)

            elapsed_time = time.time() - start_time

            # ``timeout <= 0`` means wait indefinitely (no overall deadline).
            if timeout > 0 and elapsed_time > timeout:
                raise TimeoutError("Not all tasks completed within the specified timeout")

            # Iterate over a copy of the list and remove finished tasks
            for result in results[:]:
                if result.ready():
                    t.update(n=1)
                    results.remove(result)
    finally:
        t.close()

    deadline = start_time + timeout if timeout > 0 else None
    self._wait_for_recording(deadline)

log_submission_summary() #

Log a per-queue count of the tasks submitted since the last summary

Source code in packages/climate-ref-celery/src/climate_ref_celery/executor.py 
140
141
142
143
144
145
146
def log_submission_summary(self) -> None:
    """
    Log a per-queue count of the tasks submitted since the last summary
    """
    for queue, count in sorted(self._queue_counts.items()):
        logger.info(f"Submitted {count} executions to queue {queue}")
    self._queue_counts.clear()

run(definition, execution=None) #

Run a diagnostic calculation

This will queue the diagnostic to be run by a Celery worker. The executions will be stored in the database when the task completes if execution is specified. No result will be returned from this function. Instead, you can periodically check the status of the task in the database.

Tasks may not execute immediately if the correct workers are not available. The task will be queued and executed when a worker becomes available.

Parameters:

Name Type Description Default
definition ExecutionDefinition

A description of the information needed for this execution of the diagnostic This includes relative paths to the data files, which will be converted to absolute paths when being executed

 required
execution Execution | None

Result of the diagnostic execution This is a database object that contains the executions of the execution. If provided, it will be updated with the executions of the execution. This may happen asynchronously, so the executions may not be immediately available.

 None
Source code in packages/climate-ref-celery/src/climate_ref_celery/executor.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def run(
    self,
    definition: ExecutionDefinition,
    execution: Execution | None = None,
) -> None:
    """
    Run a diagnostic calculation

    This will queue the diagnostic to be run by a Celery worker.
    The executions will be stored in the database when the task completes if `execution`
    is specified.
    No result will be returned from this function.
    Instead, you can periodically check the status of the task in the database.

    Tasks may not execute immediately if the correct workers are not available.
    The task will be queued and executed when a worker becomes available.

    Parameters
    ----------
    definition
        A description of the information needed for this execution of the diagnostic
        This includes relative paths to the data files,
        which will be converted to absolute paths when being executed
    execution
        Result of the diagnostic execution
        This is a database object that contains the executions of the execution.
        If provided, it will be updated with the executions of the execution.
        This may happen asynchronously, so the executions may not be immediately available.
    """
    from climate_ref_celery.worker_tasks import handle_failure, handle_result  # noqa: PLC0415

    diagnostic = definition.diagnostic

    name = generate_task_name(diagnostic.provider, diagnostic)
    queue = self._routing_table.queue_for(diagnostic.provider.slug, diagnostic.slug)

    async_result = app.send_task(
        name,
        args=[definition, self.config.log_level],
        # ``exclusive`` is left at its default.
        # The worker's concurrency is set on the worker,
        # not here, so this process cannot promise the worker's cgroup holds one execution.
        # Peak memory is measured over the task's process tree instead.
        kwargs={"measure": self.config.executor.measure_resources},
        queue=queue,
        link=handle_result.s(execution_id=execution.id).set(queue="celery") if execution else None,
        link_error=handle_failure.s(execution_id=execution.id).set(queue="celery") if execution else None,
    )
    logger.debug(f"Celery task {async_result.id} for {diagnostic.slug} submitted to queue {queue}")
    self._queue_counts[queue] += 1
    self._results.append(async_result)
    if execution is not None:
        self._execution_ids.append(execution.id)