在GCP上,基于Ray/Python的分布式计算完全没有并行性

人工智能 2026-07-12

我正在尝试搭建一个用于NLP处理的分布式环境。我在GCP上使用Ray和 Python。我有一个主节点和若干工作节点。无论运行1 个工作节点还是8 个工作节点,得到10个结果所需的时间都一样,完全没有并行性。我已经为此苦苦折腾了好几个小时,仍然没有头绪。

以下是我的调度器代码:

import ray
from pymongo import MongoClient
import worker  # your worker.py will define the Ray actor
import logging, threading, time
from datetime import datetime

def start_orchestration():
    logging.info("--- Starting orchestration ---")
    ray.init(address="auto")  # connect to Ray cluster (or just ray.init() locally)

    cpus_per_actor = 2
    nlp_actors = []

    def scale_actors():
        resources = ray.cluster_resources()
        total_cpus = int(resources.get("CPU", 0))
        desired = total_cpus // cpus_per_actor

        print(f"Desired actors={desired}, current={len(nlp_actors)}")

        # Add actors if needed
        while len(nlp_actors) < desired:
            new_actor = worker.NLPActor.options(num_cpus=cpus_per_actor).remote(
                "gs://babyllm-data/embeddings.pkl"
            )
            nlp_actors.append(new_actor)

        # Remove actors if cluster shrank
        while len(nlp_actors) > desired:
            actor_to_remove = nlp_actors.pop()
            ray.kill(actor_to_remove)

    def scaler_loop(interval=30):
        while True:
            try:
                scale_actors()
            except Exception as e:
                logging.error(f"[Scaler] Error {e}")
            time.sleep(interval)

    # Initial scale
    scale_actors()

    # Run scaler in background
    threading.Thread(target=scaler_loop, daemon=True).start()

    logging.info("--- NLPActor started ---")

    # Mongo Connection
    mongo = MongoClient("mongodb+srv://....../")
    col = mongo["BabyLLM"]["simple_texts_wiki_v2"]

    MAX_IN_FLIGHT = len(nlp_actors) * 2 

    logging.info("--- Processing Docs ---")
    active_futures = []
    cursor = col.find({"text_l1": {"$exists": False}})#{"text_mini": {"$exists": False}})
    counter = 0

    start_time = datetime.now()
    for doc in cursor:
        # 1. Submit the task and keep moving
        actor = nlp_actors[counter % len(nlp_actors)]
        f = actor.extract_and_replace.remote(doc["_id"], doc["title"], doc["text"][:1000])
        active_futures.append(f)

        # 2. Only wait once we've filled our "pipeline" buffer
        if len(active_futures) >= MAX_IN_FLIGHT:
            # Get whichever task finishes FIRST (not necessarily the one we just sent)
            done, active_futures = ray.wait(active_futures, num_returns=1)

            for finished in done:
                try:
                    res = ray.get(finished)
                    if res:
                        # Save immediately as requested
                        # col.update_one(
                        #     {"_id": res["_id"]},
                        #     {"$set": {
                        #         "text_l1": res["text_mini_l1"]
                        #         # "text_mini": res["text_mini"],
                        #         # "text_mini_l1": res["text_mini_l1"]
                        #     }}
                        # )
                        counter += 1
                        if counter % 10 == 0:
                            print(f"Saved {counter} documents total. Took {(datetime.now() - start_time).total_seconds()} seconds.")
                            start_time = datetime.now()
                except Exception as e:
                    logging.error(f"Task failed: {e}")

    # 3. Clean up the final tasks left in the pipeline
    while active_futures:
        done, active_futures = ray.wait(active_futures, num_returns=1)
        for finished in done:
            res = ray.get(finished)
            if res:
                col.update_one({"_id": res["_id"]}, {"$set": {"text_mini": res["text_mini"]}})
                counter += 1

    logging.info("--- Job Complete ---")

if __name__ == "__main__":
    start_orchestration()

解决方案

我认为 MAX_IN_FLIGHT 的计算时机太早,因此如果在 scale_actors() 之后再计算,可能并非所有执行单元都已经就绪;把它放到循环内部会更好。

if len(active_futures) >= len(nlp_actors) * 2:

我还注意到 counter 只有在任务完成时才会自增,因此多个任务最终会分配到同一个执行单元。在这种情况下,有一个更好的做法,就是简单地轮询分配(round robin)。

actor = nlp_actors[len(active_futures) % len(nlp_actors)]

最后给出一个建议:如果你在 NLPActor 内使用spaCy或 transformers,它们可能会阻塞全局解释器锁(GIL),从而尽管设置了并行,执行单元仍然串行运行。

站内所有文章版权归属LeftHeroAI导航站,无授权禁止任何主体转载、抄袭、复制内容,亦不得私自架设镜像站点。一经侵权,本站将通过法律途径追责。

相关文章