Fix concurrency

aigc-apps · Dec 9, 2024 · d20a1a4 · d20a1a4
1 parent ee1e7aa
commit d20a1a4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/pai_rag/tools/data_process/embed_workflow.py b/src/pai_rag/tools/data_process/embed_workflow.py
@@ -4,7 +4,7 @@
 import time
 from ray.data.datasource.filename_provider import _DefaultFilenameProvider
 from pai_rag.tools.data_process.actors.embed_actor import EmbedActor
-from pai_rag.tools.data_process.utils.ray_init import init_ray_env
+from pai_rag.tools.data_process.utils.ray_init import init_ray_env, get_concurrency
 
 
 def main(args):
@@ -18,7 +18,7 @@ def main(args):
             "config_file": args.config_file,
         },
         num_cpus=args.num_cpus,
-        concurrency=1,
+        concurrency=get_concurrency(),
         batch_size=args.batch_size,
     )
     logger.info("Embedding nodes completed.")

diff --git a/src/pai_rag/tools/data_process/utils/ray_init.py b/src/pai_rag/tools/data_process/utils/ray_init.py
@@ -25,6 +25,6 @@ def init_ray_env(working_dir, num_cpus_per_actor: int = 1):
 def get_concurrency(num_cpus_per_actor: int = 1):
     num_cpus_total = ray.cluster_resources()["CPU"] or 1
     if num_cpus_total > 1:
-        return int((num_cpus_total - 2) / num_cpus_per_actor)
+        return int((num_cpus_total - 1) / num_cpus_per_actor)
     else:
         return 1