Ray framework: distributed train ML บน python ได้ ด้วย Ray framework

Bank Eakasit
4 min readAug 29, 2020

--

ระหว่างไล่ศึกษา technology distributed computing ไปเจอตัวนี้เข้า ซึ่งจริง ๆ ก็ไม่ใหม่เท่าไร (2017) มันชื่อว่า Ray ครับ สั้น ๆ คือ Python framework ที่จัดการสร้างและรันระบบ distributed application ให้เรา เหมาะมากสำหรับ cluster และก็อาจจะง่ายสำหรับ desktop grid computing ด้วย (เช่น คอมบ้านหลายเครื่องว่างอยู่ แชร์งานไปทำหน่อย) ผมยังไม่ได้นำไปใช้งานอะไร แต่อยากมาแบ่งปันก่อน

จินตนาการง่าย ๆ ถ้าเรารัน ML บนเครื่องเราเครื่องเดียว มันก็ใช้ได้มากสุดแค่ตามจำนวน core บนเครื่องเรา ถ้าเรามีเครื่องเหลือล่ะ เราอยากใช้งาน core/GPU พวกนั้นล่ะ มันจะทำให้งานของเราเร็วขึ้น 2–3–4++… เท่าได้เลยนะ แต่เราจะทำยังไงดีล่ะ จะให้เขียนระบบจัดการส่งโค้ดไปรันเครื่องอื่นเอง? เพ้อเจ้อมาก 555 แต่นั่นแหละ คือสิ่งที่ Ray Framework จะจัดการให้ เราแค่ระบุ annotation และโค้ดเพิ่ม ส่วน complexity ด้านล่างเดี๋ยวมันไปจัดการให้เอง

https://docs.ray.io/en/latest/walkthrough.html

ตัวจัดการหลักคือ Ray Core สำหรับคนเคยเขียน OpenMP หรือ MPI มันใกล้เคียงมากเลยแหละ แต่ทั้งนี้ framework เขาก็รู้ว่าถ้าให้ทุกคนมาเขียนพวกนี้เองเนี่ย ถึงแม้จะฟังดูเป็นไปได้ แต่ก็ยังซับซ้อนอยู่ดี และมันก็มีงานหลายประเภทที่หลาย ๆ คนใช้บ่อย ๆ เช่น ML training ดังนั้น Ray จึงมี library เพิ่มให้อีกสำหรับแนว ML training โดยเฉพาะ เรียกได้ว่า สั่ง distributed ได้ในไม่กี่บรรทัดเลย

ยกตัวอย่างเช่น Ray Tune ที่ใช้ในการหา hyperparameter (learning rate, momentum, etc) เช่น grid search ซึ่งปกติเสียเวลามาก จะเห็นว่าโค้ดนั้นก็ดูปกติเหมือนที่เราเขียนทั่วไป เพิ่มมาแค่ 2 บรรทัดเอง คือ tune.report สำหรับรวมผลลัพธ์กลับไปที่เครื่องหลัก (Reduce) และ tune.run สำหรับเริ่มการทำงาน (Map)

https://docs.ray.io/en/latest/tune/index.html

นอกจาก Tune ยังมี library ที่ทำงานส่วนอื่น เช่น rllib (สำหรับ Reinforcement learning), RaySGD (สำหรับ gradient descent) และ RayServe (สำหรับ serve model) และยัง support การทำงานร่วมกับ library ชื่อดัง เช่น Torch, Keras, Tensorflow, Scikit และอื่น ๆ อีกมาก เข้าไปดูใน document เขา มี demo ให้เต็มไปหมดเลย บางอันก็ดูง่ายมาก ปรับจาก model ที่เราเขียนปกตินิดเดียว บางอันก็ดูซับซ้อนหน่อย

ผมไม่ได้ทำงาน ML เป็นหลัก เลยขอหยิบ demo tune (ซึ่งผมคุ้นเคยที่สุด) มาดู โค้ดเรียบง่าย และลอง demo เขาก็ดูเร็วขึ้นจริง เริ่มจาก tune ผมหยิบโค้ดเขามาเลย เป็นการ tune Pytorch https://docs.ray.io/en/latest/tune/tutorials/tune-tutorial.html แล้วปรับ search_space นิดหน่อยให้เป็น grid search ไป

Install ลงทุกเครื่อง https://docs.ray.io/en/latest/cluster/index.html#starting-ray-on-each-machine

ลง dependency ทั้งเครื่องหลักและเครื่องอื่น ๆ ที่จะใช้ให้เรียบร้อยครับ (ผม install ด้วย root เพื่อให้มันสร้าง bin file ใน /usr/local/bin)

sudo pip install -U ray[tune] # หรือ lib ที่ต้องใช้อื่น ๆ
pip install torch torchvision
ray start --head --port=6379

เครื่องอื่น ๆ ต้องลง dependency ให้เหมือนกัน ต่างกันที่ ray start

sudo pip install -U ray[tune] # หรือ lib ที่ต้องใช้อื่น ๆ
pip install torch torchvision
ray start --address='ip:port' --redis_password='xxx'

ข้อเสียคือ version ทุกเครื่องต้องตรงกันนะครับ ไม่งั้นโดนด่า T..T

RuntimeError: Version mismatch: The cluster was started with:
Ray: 0.8.7
Python: 3.6.9
This process on node 165.22.xxx.xxx was started with:
Ray: 0.8.7
Python: 3.8.2

อันนี้ผมสร้าง droplet digital ocean มา 3 เครื่อง: CPU 4 cores x3, memory 8GB x3

ตอน init ray ใน python

ray.init(address="auto") # กรณีอยู่บน head node
ray.init(address='ip:port', redis_password='xxx') # กรณีรันบนเครื่องอื่น

จากนั้นก็สั่งทำงานตามโค้ดเลย จะเห็นว่า Ray detects เจอทั้งหมด 12 CPU cores และ memory ให้ใช้เยอะเลย (แต่ก็น้อยกว่า 8GB x3)

เวลาทำงาน มี status อัพเดทให้ด้วย งานที่เสร็จแล้วจะเป็น terminated ครับ

ได้เวลามาตามนี้ครับ

กรณี Ray 1 core (ปกติคงไม่เป็นอันนี้ แต่ทำให้ดู overhead)
CPU times: user 13.6 s, sys: 2.02 s, total: 15.6 s
Wall time: 4min 11s
กรณี วน loop python (1 core)
CPU times: user 6min, sys: 6.52 s, total: 6min 7s
Wall time: 3min 12s
จะเห็นว่ามี overhead พอสมควรเลย แต่พอเป็นหลาย core/cluster แล้ว overhead นี้ก็พอมองข้ามได้ครับกรณี Ray 1 droplet - 4 cores
CPU times: user 8.08 s, sys: 1.19 s, total: 9.26 s
Wall time: 1min 15s
กรณี Ray 3 droplet เชื่อมกัน - 12 cores
CPU times: user 6.99 s, sys: 1.11 s, total: 8.1 s
Wall time: 32.3 s

หรือแบบซับซ้อนขึ้นมาหน่อย เขาก็มีตัวอย่างเป็น google colab ให้นะครับ

อยากทดลองแบบ GPU ให้เหมือนกัน แต่ไม่มี resource ครับ บรัย T..T

ตรวจจำนวน resources, nodes กับ cluster ได้

print(ray.nodes())
print(ray.cluster_resources())

กรณี error มันจะทำงานต่อให้ครบโดยใช้เครื่องที่เหลือ แล้ว report ตัวที่ error เก็บไว้ครับ ยังรู้สึกว่า resilience ยังไม่ perfect น่าจะพัฒนาต่อได้อีก

อีกเรื่องคือ ผมลองสร้าง server แบบ global (Singapore, Germany) แล้วพบว่า Ray มันเชื่อมกันไม่ได้ครับ คิดว่า latency สูงแล้ว scheduler เลยไม่เลือกรวม cluster กัน แต่ยังหา documentation อธิบายตรงนี้ไม่ได้

ภาพรวม ผมมองว่า library ของมันออกแบบมาให้ใช้ได้ค่อนข้างง่ายทีเดียว การ install ก็ผ่าน pip แถมตอนเชื่อมเป็น cluster ก็แค่สั่ง init บรรทัดเดียว มีเงื่อนไขเรื่อง version บ้าง มันเลยยังเป็น perfect desktop grid ไม่ได้ ออกแนว home/lab cluster มากกว่า

ถ้าถามผมว่า ต่อจากนี้ ผมควรเขียนใส่ dependency Ray เข้าไปเลยไหม คำตอบของผมคือ คงไม่ครับ ถ้ายังไม่ได้ train นานจริง ๆ (เป็นชั่วโมง หรือเป็นวัน) แต่ต่อไปก็คงพยายามเขียนให้เป็น class / immutable / pipeline ซึ่งเป็น practice ที่ดีอยู่แล้ว (แต่ไม่ค่อยจะทำ อ๊ากกกก 555) และเวลาถ้าต้องการเพิ่ม Ray ก็ปรับแก้ได้เร็ว ไม่กี่บรรทัดด้วย

- Aug 29, 2020 -

--

--

Bank Eakasit
Bank Eakasit

No responses yet