Repository: cubed-dev/cubed
Pull request #536 (2024-08-03 15:56:49)

Adding `compile_function` as execute option.
This allows users to use jit or aot compilation during the dag finalization process within executors. It should work straightforwardly on jax/numba style jit compilation. It's possible, but maybe ugly, to perform jax-aot-style compilation.

This should help make progress towards #304.

(Note to reviewer: I realized this could be merged ahead of #508, so I pulled it out into a non-dependent PR).

alxmrs: I think this PR is ready to merge. The next things to explore in this area are: 
- [ ] Testing AOT compilation
- [ ] Figuring out if this is the interface we need for exporting compiled functions (https://github.com/cubed-dev/cubed/issues/304#issuecomment-2258602859)

Patch:
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index e91070bf..9dc9f22b 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -1,4 +1,5 @@
 import atexit
+import dataclasses
 import inspect
 import shutil
 import tempfile
@@ -30,6 +31,8 @@ CONTEXT_ID = f"cubed-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4()}"
 # Delete local context dirs when Python exits
 CONTEXT_DIRS = set()
 
+Decorator = Callable
+
 
 def delete_on_exit(context_dir: str) -> None:
     if context_dir not in CONTEXT_DIRS and is_local_path(context_dir):
@@ -200,13 +203,45 @@ class Plan:
 
         return dag
 
+    def _compile_blockwise(self, dag, compile_function: Decorator) -> nx.MultiDiGraph:
+        """Compiles functions from all blockwise ops by mutating the input dag."""
+        # Recommended: make a copy of the dag before calling this function.
+
+        compile_with_config = 'config' in inspect.getfullargspec(compile_function).kwonlyargs
+
+        for n in dag.nodes:
+            node = dag.nodes[n]
+
+            if "primitive_op" not in node:
+                continue
+
+            if not isinstance(node["pipeline"].config, BlockwiseSpec):
+                continue
+
+            if compile_with_config:
+                compiled = compile_function(node["pipeline"].config.function, config=node["pipeline"].config)
+            else:
+                compiled = compile_function(node["pipeline"].config.function)
+
+            # node is a blockwise primitive_op.
+            # maybe we should investigate some sort of optics library for frozen dataclasses...
+            new_pipeline = dataclasses.replace(
+                node["pipeline"],
+                config=dataclasses.replace(node["pipeline"].config, function=compiled)
+            )
+            node["pipeline"] = new_pipeline
+
+        return dag
+
     @lru_cache
     def _finalize_dag(
-        self, optimize_graph: bool = True, optimize_function=None
+        self, optimize_graph: bool = True, optimize_function=None, compile_function: Optional[Decorator] = None,
     ) -> nx.MultiDiGraph:
         dag = self.optimize(optimize_function).dag if optimize_graph else self.dag
         # create a copy since _create_lazy_zarr_arrays mutates the dag
         dag = dag.copy()
+        if callable(compile_function):
+            dag = self._compile_blockwise(dag, compile_function)
         dag = self._create_lazy_zarr_arrays(dag)
         return nx.freeze(dag)
 
@@ -216,11 +251,12 @@ class Plan:
         callbacks=None,
         optimize_graph=True,
         optimize_function=None,
+        compile_function=None,
         resume=None,
         spec=None,
         **kwargs,
     ):
-        dag = self._finalize_dag(optimize_graph, optimize_function)
+        dag = self._finalize_dag(optimize_graph, optimize_function, compile_function)
 
         compute_id = f"compute-{datetime.now().strftime('%Y%m%dT%H%M%S.%f')}"
 
diff --git a/docs/configuration.md b/docs/configuration.md
new file mode 100644
index 00000000..e88831e6
--- /dev/null
+++ b/docs/configuration.md
@@ -0,0 +1,185 @@
+# Configuration
+
+Cubed uses [Donfig](https://donfig.readthedocs.io/en/latest/) for managing configuration of things like the directory used for temporary files, or for setting executor properties.
+
+This page covers how to specify configuration properties, and a reference with all the configuration options that you can use in Cubed.
+
+## Specification
+
+There are three main ways of specifying configuration in Cubed:
+1. by instantiating a `Spec` object,
+2. by using a YAML file and setting the `CUBED_CONFIG` environment variable, or
+3. by setting environment variables for individual properties.
+
+We look at each in turn.
+
+### `Spec` object
+
+This is how you configure Cubed directly from within a Python program - by instantiating a {py:class}`Spec <cubed.Spec>` object:
+
+```python
+import cubed
+
+spec = cubed.Spec(
+    work_dir="s3://cubed-tomwhite-temp",
+    allowed_mem="2GB",
+    executor_name="lithops",
+    executor_options=dict(use_backups=False, runtime="cubed-runtime", runtime_memory=2000)
+)
+```
+
+The `spec` instance is then passed to array creation functions as follows:
+
+```python
+import cubed.array_api as xp
+
+a = cubed.random.random((50000, 50000), chunks=(5000, 5000), spec=spec)
+b = cubed.random.random((50000, 50000), chunks=(5000, 5000), spec=spec)
+c = xp.add(a, b)
+```
+
+All arrays in any given computation must share the same `spec` instance.
+
+### YAML file
+
+A YAML file is a good way to encapsulate the configuration in a single file that lives outside the Python program.
+It's a useful way to package up the settings for running using a particular executor, so it can be reused.
+The Cubed [examples](https://github.com/cubed-dev/cubed/blob/main/examples/README.md) use YAML files for this reason.
+
+```yaml
+spec:
+  work_dir: "s3://cubed-$USER-temp"
+  allowed_mem: "2GB"
+  executor_name: "lithops"
+  executor_options:
+    use_backups: False
+    runtime: "cubed-runtime"
+    runtime_memory: 2000
+```
+
+Note that YAML files can use environment variables - in this example `$USER` will be expanded appropriately.
+
+To use the YAML file, set the `CUBED_CONFIG` environment variable to the file path before invoking your Python program:
+
+```shell
+CUBED_CONFIG=/path/to/cubed.yaml python ...
+```
+
+Donfig will actually look for YAML files in a variety of locations, see the [Donfig docs](https://donfig.readthedocs.io/en/latest/configuration.html#yaml-files) for details.
+
+### Environment variables
+
+You can also set [Donfig-style environment variables](https://donfig.readthedocs.io/en/latest/configuration.html#environment-variables) to set individual properties. Notice how double underscores are used to indicate nesting.
+
+```shell
+export CUBED_SPEC__WORK_DIR='s3://cubed-$USER-temp'
+export CUBED_SPEC__ALLOWED_MEM=2GB
+export CUBED_SPEC__EXECUTOR_NAME=lithops
+export CUBED_SPEC__EXECUTOR_OPTIONS__USE_BACKUPS=False
+export CUBED_SPEC__EXECUTOR_OPTIONS__RUNTIME=cubed-runtime
+export CUBED_SPEC__EXECUTOR_OPTIONS__RUNTIME_MEMORY=2000
+```
+
+This can be handy if you only have a couple of properties to set:
+
+```shell
+CUBED_SPEC__ALLOWED_MEM=2GB CUBED_SPEC__EXECUTOR_NAME=processes python ...
+```
+
+## Reference
+
+### Spec options
+
+These properties can be passed directly to the {py:class}`Spec <cubed.Spec>` constructor. Or, equivalently, they are directly under `spec` in a YAML file.
+
+| Property           | Default           | Description                                                                                                                             |
+|--------------------|-------------------|-----------------------------------------------------------------------------------------------------------------------------------------|
+| `work_dir`         | `None`            | The directory path (specified as an fsspec URL) used for storing intermediate data. If not set, the user's temporary directory is used. |
+| `allowed_mem`      | `200MB`           | The total memory available to a worker for running a task. This includes any `reserved_mem` that has been set.                          |
+| `reserved_mem`     | `100MB`           | The memory reserved on a worker for non-data use when running a task                                                                    |
+| `executor_name`    | `single-threaded` | The executor for running computations. One of `single-threaded`, `threads`, `processes`, `beam`, `coiled`, `dask`, `lithops`, `modal`.  |
+| `executor_options` | `None`            | Options to pass to the executor on construction. See below for possible options for each executor.                                      |
+
+
+### Executor options
+
+Different executors support different options. They are listed here for each executor.
+
+These properties are keys in the `executor_options` passed to the {py:class}`Spec <cubed.Spec>` constructor. Or, equivalently, they are directly under `spec.executor_options` in a YAML file.
+
+
+#### `single-threaded`
+
+The `single-threaded` executor is a simple executor mainly used for testing. It doesn't support any configuration options
+since it is deliberately designed not to have anything except the most basic features.
+
+#### `threads`
+
+| Property                     | Default | Description                                                                                        |
+|------------------------------|---------|----------------------------------------------------------------------------------------------------|
+| `retries`                    | 2       | The number of times to retry a task if it fails.                                                   |
+| `use_backups`                | `True`  | Whether to use backup tasks for mitigating stragglers.                                             |
+| `batch_size`                 | `None`  | Number of input tasks to submit to be run in parallel. The default is not to batch.                |
+| `compute_arrays_in_parallel` | `False` | Whether arrays are computed one at a time or in parallel.                                          |
+| `max_workers`                | `None`  | The maximum number of workers to use in the `ThreadPoolExecutor`. Defaults to number of CPU cores. |
+
+
+#### `processes`
+
+| Property                     | Default | Description                                                                                                                                |
+|------------------------------|---------|--------------------------------------------------------------------------------------------------------------------------------------------|
+| `use_backups`                | `True`  | Whether to use backup tasks for mitigating stragglers.                                                                                     |
+| `batch_size`                 | `None`  | Number of input tasks to submit to be run in parallel. `None` means don't batch.                                                           |
+| `compute_arrays_in_parallel` | `False` | Whether arrays are computed one at a time or in parallel.                                                                                  |
+| `max_workers`                | `None`  | The maximum number of workers to use in the `ProcessPoolExecutor`. Defaults to number of CPU cores.                                        |
+| `max_tasks_per_child`        | `None`  | The number of tasks to run in each child process. See the Python documentation for `concurrent.futures.ProcessPoolExecutor`. (Python 3.11) |
+
+Note that `retries` is not currently supported for the `processes` executor.
+
+#### `beam`
+
+The `beam` executor doesn't currently expose any configuration options.
+When running on Google Cloud Dataflow, [four retry attempts](https://cloud.google.com/dataflow/docs/pipeline-lifecycle#error_and_exception_handling) are made for failing tasks.
+
+#### `coiled`
+
+| Property        | Default | Description                                                                                             |
+|-----------------|---------|---------------------------------------------------------------------------------------------------------|
+| `coiled_kwargs` | `None`  | Keyword arguments to pass to [`coiled.function`](https://docs.coiled.io/user_guide/functions.html#api). |
+
+Note that there is currently no way to set retries or a timeout for the Coiled executor.
+
+#### `dask`
+
+| Property                     | Default | Description                                                                                                                     |
+|------------------------------|---------|---------------------------------------------------------------------------------------------------------------------------------|
+| `retries`                    | 2       | The number of times to retry a task if it fails.                                                                                |
+| `use_backups`                | `True`  | Whether to use backup tasks for mitigating stragglers.                                                                          |
+| `batch_size`                 | `None`  | Number of input tasks to submit to be run in parallel. The default is not to batch.                                             |
+| `compute_arrays_in_parallel` | `False` | Whether arrays are computed one at a time or in parallel.                                                                       |
+| `compute_kwargs`             | `None`  | Keyword arguments to pass to Dask's [`distributed.Client`](https://distributed.dask.org/en/latest/api.html#client) constructor. |
+
+Note that there is currently no way to set a timeout for the Dask executor.
+
+#### `lithops`
+
+| Property                     | Default | Description                                                                                                                                                                                                                                                                                                                                                       |
+|------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `retries`                    | 2       | The number of times to retry a task if it fails.                                                                                                                                                                                                                                                                                                                  |
+| `timeout`                    | `None`  | Tasks that take longer than the timeout will be automatically killed and retried. Defaults to the timeout specified when [deploying the lithops runtime image](https://lithops-cloud.github.io/docs/source/cli.html#lithops-runtime-deploy-runtime-name). This is 180 seconds in the [examples](https://github.com/cubed-dev/cubed/blob/main/examples/README.md). |
+| `use_backups`                | `True`  | Whether to use backup tasks for mitigating stragglers.                                                                                                                                                                                                                                                                                                            |
+| `compute_arrays_in_parallel` | `False` | Whether arrays are computed one at a time or in parallel.                                                                                                                                                                                                                                                                                                         |
+| Other properties             | N/A     | Other properties will be passed as keyword arguments to the [`lithops.executors.FunctionExecutor`](https://lithops-cloud.github.io/docs/source/api_futures.html#lithops.executors.FunctionExecutor) constructor.                                                                                                                                                  |
+
+Note that `batch_size` is not currently supported for Lithops.
+
+#### `modal`
+
+| Property                     | Default | Description                                                                         |
+|------------------------------|---------|-------------------------------------------------------------------------------------|
+| `cloud`                      | `aws`   | The cloud to run on. One of `aws` or `gcp`.                                         |
+| `use_backups`                | `True`  | Whether to use backup tasks for mitigating stragglers.                              |
+| `batch_size`                 | `None`  | Number of input tasks to submit to be run in parallel. The default is not to batch. |
+| `compute_arrays_in_parallel` | `False` | Whether arrays are computed one at a time or in parallel.                           |
+
+Currently the Modal executor in Cubed uses a hard-coded value of 2 for retries and 300 seconds for timeouts, neither of which can be changed through configuration.
diff --git a/docs/index.md b/docs/index.md
index 0a403f0d..c0549333 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -24,6 +24,7 @@ Intro slides <https://cubed-dev.github.io/cubed/cubed-intro.slides.html>
 Examples <https://github.com/tomwhite/cubed/tree/main/examples/README.md>
 api
 array-api
+configuration
 related-projects
 ```
 
diff --git a/examples/README.md b/examples/README.md
index fb45940a..ce3560a4 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -23,7 +23,7 @@ Follow the instructions for setting up Cubed to run on your executor runtime:
 
 | Executor  | Cloud  | Set up instructions                            |
 |-----------|--------|------------------------------------------------|
-| Processes | N/A    | N/A                                            |
+| Processes | N/A    | `pip install 'cubed[diagnostics]'`             |
 | Lithops   | AWS    | [lithops/aws/README.md](lithops/aws/README.md) |
 |           | Google | [lithops/gcp/README.md](lithops/gcp/README.md) |
 | Modal     | AWS    | [modal/aws/README.md](modal/aws/README.md)     |
diff --git a/examples/lithops/aws/README.md b/examples/lithops/aws/README.md
index 1a273669..7de189b5 100644
--- a/examples/lithops/aws/README.md
+++ b/examples/lithops/aws/README.md
@@ -12,7 +12,7 @@
 ```shell
 conda create --name cubed-lithops-aws-examples -y python=3.11
 conda activate cubed-lithops-aws-examples
-pip install -r requirements.txt  # use requirements file from same directory as this readme
+pip install 'cubed[lithops-aws]'
 ```
 
 2. Configure Lithops with an [AWS Lambda compute backend](https://lithops-cloud.github.io/docs/source/compute_config/aws_lambda.html), and an [AWS S3 storage backend](https://lithops-cloud.github.io/docs/source/storage_config/aws_s3.html).
diff --git a/examples/lithops/gcp/README.md b/examples/lithops/gcp/README.md
index fd7a4115..ef76cded 100644
--- a/examples/lithops/gcp/README.md
+++ b/examples/lithops/gcp/README.md
@@ -12,7 +12,7 @@
 # from this directory
 conda create --name cubed-lithops-gcp-examples -y python=3.11
 conda activate cubed-lithops-gcp-examples
-pip install -r requirements.txt  # use requirements file from same directory as this readme
+pip install 'cubed[lithops-gcp]'
 ```
 
 2. Configure Lithops with a [Google Cloud Functions compute backend](https://lithops-cloud.github.io/docs/source/compute_config/gcp_functions.html#configuration), and a [Google Cloud Storage backend](https://lithops-cloud.github.io/docs/source/storage_config/gcp_storage.html#configuration).
diff --git a/pyproject.toml b/pyproject.toml
index 97e1be3c..dad68184 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,6 +54,16 @@ beam = ["apache-beam", "gcsfs"]
 dask = ["dask"]
 dask-distributed = ["distributed"]
 lithops = ["lithops[aws] >= 2.7.0"]
+lithops-aws = [
+    "cubed[diagnostics]",
+    "lithops[aws]",
+    "s3fs",
+]
+lithops-gcp = [
+    "cubed[diagnostics]",
+    "lithops[gcp]",
+    "gcsfs",
+]
 modal = [
     "cubed[diagnostics]",
     "modal-client",


Test patch:
diff --git a/.github/workflows/jax-tests.yml b/.github/workflows/jax-tests.yml
index e8ca8662..c8e0aa9e 100644
--- a/.github/workflows/jax-tests.yml
+++ b/.github/workflows/jax-tests.yml
@@ -1,6 +1,7 @@
 name: JAX tests
 
 on:
+  pull_request:
   schedule:
     # Every weekday at 03:53 UTC, see https://crontab.guru/
     - cron: "53 3 * * 1-5"
diff --git a/cubed/tests/test_executor_features.py b/cubed/tests/test_executor_features.py
index f6bb5bd6..6a8dad86 100644
--- a/cubed/tests/test_executor_features.py
+++ b/cubed/tests/test_executor_features.py
@@ -315,3 +315,55 @@ def test_check_runtime_memory_processes(spec, executor):
 
     # OK if we use fewer workers
     c.compute(executor=executor, max_workers=max_workers // 2)
+
+
+COMPILE_FUNCTIONS = [lambda fn: fn]
+
+try:
+    from numba import jit as numba_jit
+    COMPILE_FUNCTIONS.append(numba_jit)
+except ModuleNotFoundError:
+    pass
+
+try:
+    if 'jax' in os.environ.get('CUBED_BACKEND_ARRAY_API_MODULE', ''):
+        from jax import jit as jax_jit
+        COMPILE_FUNCTIONS.append(jax_jit)
+except ModuleNotFoundError:
+    pass
+
+
+@pytest.mark.parametrize("compile_function", COMPILE_FUNCTIONS)
+def test_check_compilation(spec, executor, compile_function):
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    assert_array_equal(
+        c.compute(executor=executor, compile_function=compile_function), np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]])
+    )
+
+
+def test_compilation_can_fail(spec, executor):
+    def compile_function(func):
+        raise NotImplementedError(f"Cannot compile {func}")
+
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    with pytest.raises(NotImplementedError) as excinfo:
+        c.compute(executor=executor, compile_function=compile_function)
+    
+    assert "add" in str(excinfo.value), "Compile function was applied to add operation."
+
+
+def test_compilation_with_config_can_fail(spec, executor):
+    def compile_function(func, *, config=None):
+        raise NotImplementedError(f"Cannot compile {func} with {config}")
+
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    with pytest.raises(NotImplementedError) as excinfo:
+        c.compute(executor=executor, compile_function=compile_function)
+        
+    assert "BlockwiseSpec" in str(excinfo.value), "Compile function was applied with a config argument."