Update CI image with develop docker image

[lsy323]

Update CI image with the develop docker image.

Plan to take the following steps to update CI image:

• Enable PyTorch/XLA build in GHA workflows with new docker image.
• Enable PyTorch/XLA build & tests in GHA workflows with new docker image.
• Create PR in upstream to update the CI image used in upstream CI for PyTorch/XLA tests: Bump xla_base version tag to v1.1 pytorch#109757

Summary of the change in this PR:

• The new docker image is built from .circleci/docker/Dockerfile. The image is using PyTorch/XLA develop container as a base image, plus some additional dependencies for CI and PyTorch CI. The new image is at gcr.io/tpu-pytorch/xla_base:dev-3.8_cuda_12.1
• The new docker image is python3.8 + Debian:Bullseye + CUDA 12.1
• CI image uses Conda for Python environment (Python 3.8), see comment for detail.

[lsy323]

Build and CPU tests passed with CUDA dev container.

However, the GPU test TestPjRtDistributedDataParallel.test_ddp_correctness_env_init in test/pjrt/test_ddp.py is failing. It seems that it turned into a bad state, causing the other test got stuck.

Try skipping ddp test to see if any other test is failing in GPU dev container.

[yeounoh]

also leave this as-is

[yeounoh]

nit. Let's just use the image tag as prefix to cache-silo-ci-dev-, so it would look like cache-silo-ci-dev-3.8_cuda_12.1.

[lsy323]

sg, let me update the cache name

[yeounoh]

Left some comments, will approve after testing is complete.

[lsy323]

Synced with @yeounoh offline, since we have additional dependencies for CI (e.g. sccache), which are not required by development. So we create an other CI image based on the dev image, installing the necessary dependencies in the CI image. The CI Dockerfile will reuse the .circleci/docker/Dockerfile.

I updated the .circleci/docker/Dockerfile to use dev container as base image, now the CI is passing on CPU and GPU.

I had to make the following changes to Dockerfile to make it work

• Remove the jenkins user, otherwise I'm seeing permission error when executing the ./install_valgrind.sh. Tried to move jenkins creation at the begining but somehow didn't help. In this case, valgrind installation will be done by root user (need to remove sudo)
• Removed conda installation, it's not needed anymore

[lsy323]

Also we cannot create jenkins user at the end of the Dockerfile, otherwise, we need sudo to install pytorch, pytorch/xla

[lsy323]

After offline discussion with @yeounoh, the following changes are made since last checkpoint:

1. Preserve the jenkins user(As a normal user instead of root), as it's used in upstream CI template
2. Because of change 1, we need to add --user to python setup.py command, this wasn't needed before, since we were using conda and conda env is private for the jenkins user.
3. In Dockerfile, change exec permission for bazel and install sudo package
4. Recover the changes to the paths in previous commit, as we are switching back to jenkins user

[yeounoh]

LGTM

[vanbasten23]

This is awesome. Thanks Siyuan!