amazon-braket / amazon-braket/amazon-braket-containers

Make `braket_container.py` thread safe

Open
#326 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17
Forks
10
Avg merge
12h 33m
Merged PRs (30d)
11

Description

**Describe the feature you'd like**
The [`braket_container.py`](https://github.com/amazon-braket/amazon-braket-containers/blob/main/src/braket_container.py) script to launch the user-provided algorithm script is not thread safe, currently. When running a multi-node job with parallelization through MPI (the hyperparameter `sagemaker_mpi_enabled` makes SageMaker to invoke the `braket_container.py` with `mpirun`), this can create race conditions in paritcular in the step to download, extract and make available the user-provided code, when running a multi-node job.

The `braket_container.py` script should be made thread safe to account for jobs running on multiple instances or (GPU) cores with `sagemaker_mpi_enabled=True`.

**How would this feature be used? Please describe.**
The user shouldn't have to worry about this feature and, specifically, shouldn't have to change the `braket_container.py` script if they want to use MPI support for the jobs.

**Additional context**
There is an [example](https://github.com/amazon-braket/amazon-braket-examples/blob/dd3828c5db4cdcabae99b6535a630dacf97656c6/examples/nvidia_cuda_q/container/braket_container.py#L316) for a simple workaround for this issue in the amazon-braket-examples repository. I have created [an issue](https://github.com/amazon-braket/amazon-braket-examples/issues/679) there to document this doesn't ultimately solve the problem. But, actually, I think this should be addressed here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.