Repository: alpa-projects/alpa

Issue:
AttributeError: 'DistributedArray' object has no attribute 'uuid'
**Please describe the bug**
When call `DistributedArray.block_until_ready` just like jax DeviceArray, alpa raise AttributeError:
```
2023-03-22 23:41:18,342	ERROR tensor_computing_v1.py:87 -- 'DistributedArray' object has no attribute 'uuid'
Traceback (most recent call last):
  File "/home/admin/ray-pack/tmp/job/05000080/package/tensor_computing_v1.py", line 85, in main
    alpa_compute()
  File "/home/admin/ray-pack/tmp/job/05000080/package/tensor_computing_v1.py", line 60, in alpa_compute
    results = timeit.repeat(
  File "/home/admin/micromamba/envs/alpa/lib/python3.8/timeit.py", line 238, in repeat
    return Timer(stmt, setup, timer, globals).repeat(repeat, number)
  File "/home/admin/micromamba/envs/alpa/lib/python3.8/timeit.py", line 205, in repeat
    t = self.timeit(number)
  File "/home/admin/micromamba/envs/alpa/lib/python3.8/timeit.py", line 177, in timeit
    timing = self.inner(it, self.timer)
  File "<timeit-src>", line 6, in inner
  File "/home/admin/ray-pack/tmp/job/05000080/package/tensor_computing_v1.py", line 56, in execute
    result.block_until_ready()
  File "/home/admin/ray-pack/tmp/job/05000080/pyenv/lib/python3.8/site-packages/alpa/device_mesh.py", line 1536, in block_until_ready
    self.device_mesh.block_until_ready_remote_buffers(self.uuid)
AttributeError: 'DistributedArray' object has no attribute 'uuid'
```
**Please describe the expected behavior**

**System information and environment**
**System information and environment**
- OS Platform and Distribution: Linux REPL7 docker
- Python version: 3.8.16
- CUDA version: 11.2
- NCCL version: nccl-2.14.3.1
- cupy version: 11.6.0
- GPU model and memory: A10
- Alpa version: master
- TensorFlow version: not installed
- JAX version: jaxlib-0.3.22.cuda112.cudnn810-cp38-cp38

**To Reproduce**


**Screenshots**
![image](https://user-images.githubusercontent.com/12445254/227132689-971478f1-5173-4d2c-a111-a6a07dcb27e5.png)
**Code snippet to reproduce the problem**
```python
def matrix_compute(t):
    import jax.numpy as np
    t_transposed = np.transpose(t)
    dot_matrix = np.dot(t, t_transposed)
    v1_row_norm = np.linalg.norm(t, axis=1).reshape(-1, 1)
    v2_col_norm = np.linalg.norm(t_transposed, axis=0).reshape(1, -1)
    norm_matrix = np.dot(v1_row_norm, v2_col_norm)
    res = dot_matrix / norm_matrix
    res = np.where(np.isneginf(res), 0, res)
    return res

matrix_compute_jit = alpa.parallelize(matrix_compute)
result = matrix_compute_jit(vector_matrix)
result.block_until_ready()
```
**Additional information**
Add any other context about the problem here or include any logs that would be helpful to diagnose the problem.


Discussion:

Resolving patch:
diff --git a/alpa/device_mesh.py b/alpa/device_mesh.py
--- a/alpa/device_mesh.py
+++ b/alpa/device_mesh.py
@@ -1528,7 +1528,7 @@ def prefetch(self):
 
     def block_until_ready(self):
         """Block until all remote buffers of this array are ready."""
-        self.device_mesh.block_until_ready_remote_buffers(self.uuid)
+        self.device_mesh.block_until_ready_remote_buffers([self.remote_ref])
 
     def delete(self):
         self.remote_ref = None


Test patch: