mirror of
https://github.com/vale981/ray
synced 2025-03-05 10:01:43 -05:00
d95dc2f2e5
This PR adds GPU support for pytorch and tensorflow predictor, as well as automatic setting `use_gpu` flag in `BatchPredictor`. Notable changes: - Added `use_gpu` flag in the constructor of `TorchPredictor` and `TensorflowPredictor` (note it's slightly different from our latest design doc that puts this flag at `predict()` call) - Added `use_gpu` flag to `SklearnPredictor` so its interface is compatible with `BatchPredictor` - Code to move both model weights and input tensor to default visible GPU at index 0 if flag is set - parametrized existing predictor tests to use GPU for both CPU & GPU coverage - Changed BUILD CI tests with an added `gpu` tag (I'm not 100% sure if that's a right way tho) Follow ups: https://github.com/ray-project/ray/issues/26249 is created in case our host has multiple GPU devices. It's a bit out of scope for this PR, but for GPU batch inference ideally we should be able to evenly use all GPU devices on host where CPU & DRAM are busy with pre-fetching + data movement to GPU. We might approximately do the same by scheduling same # of Predictor instances on the host, but that's worth verifying once benchmarks are set. |
||
|---|---|---|
| .. | ||
| hooks | ||
| copy_files.py | ||
| Dockerfile | ||
| Dockerfile.gpu | ||
| pipeline.gpu.large.yml | ||
| pipeline.gpu.yml | ||
| pipeline.macos.yml | ||
| pipeline.ml.yml | ||
| pipeline.windows.yml | ||
| pipeline.yml | ||