Compilation cache doc

docs/persistent_compilation_cache.md

@@ -11,7 +11,7 @@ The compilation cache is enabled when the
 is set. This should be done prior to the first compilation. Set the location as
 follows:
 
-```
+```python
 import jax
 
 # Make sure this is called before jax runs any operations!
@@ -27,7 +27,7 @@ is an alternate way of setting `cache-location`.
 
 `cache-location` can be a directory on the local filesystem. For example:
 
-```
+```python
 import jax
 
 jax.config.update("jax_compilation_cache_dir", "/tmp/jax-cache")
@@ -67,8 +67,74 @@ Cloud Storage (GCS) bucket. We recommend the following configuration:
 Assuming that `gs://jax-cache` is the GCS bucket, set `cache-location` as
 follows:
 
-```
+```python
 import jax
 
 jax.config.update("jax_compilation_cache_dir", "gs://jax-cache")
 ```
+
+## How it works
+
+The JAX compilation cache works by hashing a number of parameters to create a signature for a compiled function these are:
+
+* The computation performed by the function captured by the non-optimized HLO of the JAX function being hashed
+* The Jaxlib version 
+* Relevant XLA compilation flags
+* Device configuration captured in general, by the number of devices and the topology of the devices. Currently for GPUs, the topology only contains a string representation of the GPU name
+* Compression algorithm used to compress the compiled executable
+* Any custom hooks
+
+When the signature for a function created using the parameters above matches that of a compiled function in the persistent cache, the function will not be compiled but just read and deserialized from the persistent cache. 
+
+### Compile time dependent caching
+
+JAX only caches executables that take a certain amount of time to compile. This threshold is controlled by the `jax_persistent_cache_min_compile_time_secs` configuration option. To cache every executable that is compiled regardless of compile time, set this value to zero:
+
+```python
+import jax
+
+jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+```
+
+## Different Runtimes
+
+Below we outline some observed behavior of the persistent compilation cache in a variety of different runtimes as it relates to which processes write to the cache.
+
+### Single node with single process and single device
+
+This is the simplest runtime and the only process that compiles and writes to the compilation cache is the singular process. The number of devices does not matter to the cache key in this setup, only the type of device does. 
+
+### Single node with single process and multiple devices
+
+Once again the only process that compiles and writes to the compilation cache is the singular proess. The difference between this setup and the previous is that now the number of devices matters in addition to the type of device. 
+
+### Multiple process and multiple devices (on either single or multiple nodes)
+
+In this runtime the first time a program is run (the persistent cache is cold / empty) all processes will compile, but only the process with rank 0 in the global communication group will write to the persistent cache. In subsequent runs, all processes will attempt to read from the persistent cache, so it is important for the persistent cache to be in a shared file system (eg: NFS) or remote storage (eg: GFS). If the persistent cache is local to rank 0, then all processes except rank 0 will once again compile in subsequent runs as a result of a compilation cache miss. 
+
+## Logging cache activity
+
+It can be helpful to examine what exactly is happening with the persistent compilation cache for debugging. While there is no singular canonical way of debugging and examining what's happening in the compilation cache, here are a few suggestions on how to begin.
+
+### Examining cache misses
+
+To merely examine and understand why there are cache misses JAX includes a configuration flag that enables the logging of all cache misses (including persistent compilation cache misses) with their explanations. This can be enabled by setting the following configuration.
+
+```python
+import jax
+
+jax.config.update("jax_explain_cache_misses", True)
+```
+
+## Pitfalls
+
+There are a couple of pitfalls that have currently been discovered:
+
+* Currently the persistent cache doesn't work with function that have host callbacks. In this situation, caching in completely avoided. 
+  - This is because the HLO contains a pointer to the callback and changes from run to run even if the computation and compute infrastructure is exactly the same. 
+
+* Currently the persistent cache doesn't work with a function that uses primitives that implement their own custom_partitioning. 
+  - The HLO of the function contains a pointer to the custom_partitioning callback, and leads to different cache keys for the same computation across runs. 
+  - In this situation, caching still proceeds, but a different key is produced every time, making the cache ineffective. 
+
+