Due to the distributed nature of JAX with Pathways, some operations might not
scale well due to communication overheads. While Pathways minimizes these
overheads with features like [asynchronous dispatch](https://jax.readthedocs.io/en/latest/async_dispatch.html), there are
some things you need to be aware of when you port JAX workloads to Pathways or
scale a JAX with Pathways workload to a large number of accelerators.

## Before you begin

Make sure you have:

- [Set up Cluster Toolkit](https://docs.cloud.google.com/cluster-toolkit/docs/setup/configure-environment.md.txt)
- [Installed Kubernetes tools](https://kubernetes.io/docs/tasks/tools/)
- [Installed the gcloud CLI](https://docs.cloud.google.com/sdk/docs/install)
- [Enabled the Google Kubernetes Engine API](https://docs.cloud.google.com/endpoints/docs/openapi/enable-api.md.txt)

## Process index

JAX with Pathways treats all devices across your Pathways cluster as local.
This simplifies device management and allows JAX to utilize all available
resources. In practice, this means:

- `jax.process_index()` is always 0 for all devices.
- `jax.devices()` and `jax.local_devices()` return all TPU devices across the entire job.

## Hardware type and colocation

For best performance, place all Pathways components and the user job in the
same Google Cloud cloud zone. Use a large CPU like the IFRT proxy and resource
manager. We recommend at least a dedicated `n2-standard-64` which comes with 64
vCPUs and 256 GB memory.

## PathwaysUtils

[Pathways-utils](https://github.com/AI-Hypercomputer/pathways-utils) is a
Python-based GitHub repository providing essential utilities and tools that let
you streamline the deployment and execution of JAX workloads on the Pathways on
Cloud architecture. This package handles the necessary adaptations for the
cloud environment, allowing JAX developers to focus on their core machine
learning workflows with minimal platform-specific configuration. Specifically,
it offers:

- A "proxy" JAX backend: this custom backend enables your JAX application to use the Pathways infrastructure by setting the `JAX_PLATFORMS=proxy` environment variable.
- Integrated Profiling Utilities: profiling capabilities that let you understand your application's performance. By using standard JAX profiling APIs like `jax.profiler.start_trace` and `jax.profiler.start_server`, you can profile not only your JAX code but also the underlying Pathways components, providing a holistic view of execution within the cloud environment.
- Distributed Checkpointing with Orbax: a custom Orbax checkpoint handler that lets you use distributed checkpoints and restore your checkpoints when using the Orbax library within the Pathways environment. This integration aims to work without requiring any changes to your existing Orbax checkpointing code as long as it imports `pathwaysutils`.
- Elastic Training Primitives: provides foundational elastic training primitives that you can use to build robust and scalable training workflows using Pathways. These primitives allow your training jobs to dynamically adapt to changes in available resources, improving efficiency and resilience in cloud environments.

## Checkpointing

[Orbax](https://github.com/google/orbax) is thoroughly tested with Pathways for
distributed checkpointing and restoring with Cloud Storage. When you
set the environment variable `ENABLE_PATHWAYS_PERSISTENCE=1` and call
`import pathwaysutils; pathwaysutils.initialize()` in your `train.py`, a custom
`ArrayHandler` is registered that efficiently handles checkpoint operations
through the [IFRT](https://github.com/openxla/xla/blob/main/xla/python/ifrt/README.md)
proxy, allowing Pathways workers on accelerators to directly save and restore data.

## Colocated Python

[Colocated Python](https://github.com/jax-ml/jax/blob/main/jax/experimental/colocated_python/api.py)
is an open-source JAX API that lets you run user-specified Python code
directly on the TPU or GPU hosts, which is more straightforward in multi-controller
JAX. By using this approach, you can avoid data transfer between the client and
TPU machines for compute-intensive tasks, such as data loading and
checkpointing. To configure your Pathways cluster to run the colocated Python
JAX API, build and push a sidecar container image by using the
[sample sidecar Dockerfile](https://github.com/AI-Hypercomputer/pathways-utils/blob/main/pathwaysutils/experimental/shared_pathways_service/dockerfiles/sample_sidecar.dockerfile).

## Data loading

During training we repeatedly load batches from a dataset to feed it into the
model. Having an efficient, async data loader which shards the batch across
hosts is important to avoid starving accelerators of work. When running training
with Pathways, the data loader runs on a CPU VM (unlike a TPU VM which is used
on multi-controller setups) and dispatches data to TPU VMs. This incurs a higher
latency in reading data but that is mitigated partially by reading ahead X
number of batches on the CPU host and dispatching the read data asynchronously
to the TPUs. This solution is sufficient when running at small to medium scale.

When you run a multiprocess data loader (such as
[Grain](https://github.com/google/grain)) in the head workload container on
GKE, the container shares the head node with the `pathways-proxy`
and `pathways-rm` sidecars. Because of this shared architecture, you must
configure the following:

- **Set an explicit worker count:** Inside a container, `os.cpu_count()` reports the host VM's total vCPUs rather than the container's CPU limit. Pass an explicit worker count (no more than the container's CPU limit) instead of deriving it from `os.cpu_count()`.
- **Mount a memory-backed `/dev/shm` volume:** Multiprocess data loaders use POSIX shared memory (`/dev/shm`) for inter-process communication. Ensure a memory-backed `emptyDir` (`medium: Memory`) volume is mounted at `/dev/shm` so that the prefetched batches do not exceed the default container shared-memory limit.

For optimal performance at scale, we strongly recommend colocating your input
data pipeline by using [colocated Python](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/porting-jax-workloads.md.txt#colocated-python) to run your data pipeline directly on
the accelerators. This eliminates the CPU bottleneck and leverages the TPU's
fast interconnects for data transfer.

You can find a reference implementation of migrating a TFDS based
input pipeline in the `RemoteIterator` implementation in
[multihost_dataloading.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/input_pipeline/multihost_dataloading.py).
This implementation works on both multi-controller JAX and Pathways in a
distributed manner using the colocated Python JAX API.

## Jax Versioning

Pathways releases are tightly coupled with JAX versions to ensure compatibility
and stability. To avoid potential issues, verify that your Pathways artifacts
and your JAX version are aligned. Each Pathways release clearly specifies the
compatible JAX versions through a tag of the form `jax-<version>`.

## Compilation Cache

Pathways persistent compilation cache is a feature that allows Pathways servers
to store compiled XLA executables in a persistent location, such as
Cloud Storage, to avoid redundant compilation. This feature is enabled by
default. The location of the cache is passed in as `--gcs_scratch_location`
flag to the resource manager and Pathways worker containers. To keep associated
storage costs to a minimum, the cache attaches a lifecycle policy to the
Cloud Storage location. There is a limit of 50 policies per
Cloud Storage bucket. Therefore, we recommend using a common
Cloud Storage location across all workloads.

If you omit the `--gcs_scratch_location` flag or leave it empty, Pathways automatically
falls back to an in-memory compilation cache. This cache does not persist compiled XLA
executables across Pathways deployments.

This cache is similar to the [JAX compilation cache](https://docs.jax.dev/en/latest/persistent_compilation_cache.html)
which is disabled by `pathwaysutils.initialize()` for Pathways workloads.

The following Cloud Storage permissions are required for the compilation cache:

- `storage.buckets.get`: To retrieve bucket metadata.
- `storage.buckets.update`: Essential for Pathways to set up object lifecycle policies to enforce TTL for cache eviction.
- `storage.objects.list`: To list existing cache objects within the bucket.
- `storage.objects.create`: To write new compiled executables to the cache.
- `storage.objects.get`: To read cached executables from the bucket.

## Profiling

You can use the JAX profiler to generate traces of a JAX program. There are two
common ways supported with Pathways:

- [Programmatic](https://jax.readthedocs.io/en/latest/profiling.html#programmatic-capture)
  - Programmatically capture profiles from your JAX code
- [Manual](https://jax.readthedocs.io/en/latest/profiling.html#manual-capture)
  - Capturing profiles on demand after starting the profiler server from your JAX code

In both cases, the profiles are written to a Cloud Storage bucket.
There will be multiple trace files created in the Cloud Storage bucket
potentially under different timestamp folders, for example:

- Main Python process which invoked the trace (typically your notebook VM): `<jax-client-vm-name>.xplane.pb`
- Pathways IFRT proxy: `client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb`
- Pathways Resource manager: `server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb`
- Pathways worker(s): `server.*<tpu-node-name>.xplane.pb`

These trace files can be analyzed with TensorBoard by running the following
command. For more information about TensorBoard and all of its profiling tools,
see [Optimize TensorFlow performance using the Profiler](https://www.tensorflow.org/guide/profiler).

```bash
# verify trace files are present
gcloud storage ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --load_fast=false --logdir=gs://BUCKET/PREFIX
```

Replace the following:

- `BUCKET` : a Cloud Storage bucket to store the trace files
- `PREFIX`: a path within your Cloud Storage bucket to store the trace files

### Programmatic profile capture

Capture a profile from inside your code. The profiles are saved inside
`gs://<bucket>/<prefix>` under a timestamp directory

```python
import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()
```

### Manual profile capture

To manually capture profile information, you must start the profiler server from
your Python code:

```python
import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_server(JAX_PROFILER_PORT)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op
```

While the profiler server is running, you can capture a profile and export the
data to the target Cloud Storage location:

```bash
export DURATION_IN_SECS=6
export OUTPUT_DIR="gs://BUCKET/PREFIX/"

curl -X POST http://localhost:JAX_PROFILER_PORT/profiling \
     -H "Content-Type: application/json" \
     -d "{\"duration_ms\": $((DURATION_IN_SECS * 1000)), \"repository_path\": \"${OUTPUT_DIR}\"}"
```

You can find timing information for IFRT proxy client methods like `Compile` and
`Execute` within your program's trace. These events, which detail the
interactions with the IFRT Proxy gRPC server during compilation and execution,
appear on the thread named `GrpcClientSessionUserFuturesWorkQueue`. By examining
this thread in your trace, you can gain insights into the performance of these
operations.

## XLA Flags

When you use Pathways, you must set the XLA flags in the `pathways-proxy`
container.

When you use Cluster Toolkit (the `gcluster job submit` command), set XLA flags by using `--pathways-proxy-args`:

```bash
--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"
```

## HLO Dump

To inspect High Level Optimizer (HLO) inputs that are given to
the XLA compiler, you can configure Pathways to dump the HLO to a specified
Cloud Storage location by using `--pathways-proxy-args`:

```bash
--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"
```

## What's next

- [Create a GKE cluster with Pathways](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster.md.txt)
- [Run a batch workload with Pathways](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/batch-workload.md.txt)
- [Pathways interactive mode](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-interactive-mode.md.txt)
- [Resilient training with Pathways](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/resilient-training.md.txt)
- [Troubleshoot Pathways on cloud](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/troubleshooting-pathways.md.txt)