search

google / jetstream-pytorch

PyTorch/XLA integration with JetStream (https://github.com/google/JetStream) for LLM inference"
Apache License 2.0
33 stars 14 forks source link

Fix exception in ray_worker #144

Closed richardsliu closed 2 months ago

richardsliu commented 2 months ago

The cache manager recently added a parameter for env: https://github.com/google/jetstream-pytorch/blob/main/jetstream_pt/cache_manager.py#L94

But this is not passed in the Ray worker: https://github.com/google/jetstream-pytorch/blob/main/jetstream_pt/ray_worker.py#L353

This will cause the Ray process to crash when the cache manager tries to access self.env: https://github.com/google/jetstream-pytorch/blob/main/jetstream_pt/cache_manager.py#L106