From fe811625db37300f74064a52e80c130d7ae347ed Mon Sep 17 00:00:00 2001
From: pangengzheng <117730991+pangengzheng@users.noreply.github.com>
Date: Thu, 9 Feb 2023 11:20:07 +0800
Subject: [PATCH] Fix pscore test (#50349)

---
 paddle/fluid/distributed/ps/wrapper/fleet.cc          |  2 ++
 python/paddle/distributed/ps/the_one_ps.py            |  2 --
 python/paddle/fluid/tests/unittests/ps/CMakeLists.txt | 10 ++++++----
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/paddle/fluid/distributed/ps/wrapper/fleet.cc b/paddle/fluid/distributed/ps/wrapper/fleet.cc
index 93294bc84dc..7086f4ce95b 100644
--- a/paddle/fluid/distributed/ps/wrapper/fleet.cc
+++ b/paddle/fluid/distributed/ps/wrapper/fleet.cc
@@ -129,6 +129,8 @@ void FleetWrapper::InitWorker(const std::string& dist_desc,
           paddle::distributed::PSClientFactory::Create(ps_param));
       worker_ptr_->Configure(ps_param, dense_pull_regions, ps_env_, index);
     }
+    dist_desc_ = dist_desc;
+    is_initialized_ = true;
   } else {
     VLOG(3) << "Client can be initialized only once";
   }
diff --git a/python/paddle/distributed/ps/the_one_ps.py b/python/paddle/distributed/ps/the_one_ps.py
index 05d6ce9b78a..cca1d1552fb 100755
--- a/python/paddle/distributed/ps/the_one_ps.py
+++ b/python/paddle/distributed/ps/the_one_ps.py
@@ -1171,8 +1171,6 @@ class TheOnePSRuntime(RuntimeBase):
             gpus_env = os.getenv("FLAGS_selected_gpus")
             gpus_env = [int(s) for s in gpus_env.split(",")]
             main_program._fleet_opt["worker_places"] = gpus_env
-            PSGPU = core.PSGPU()
-            PSGPU.init_gpu_ps(gpus_env)
 
         def sync_strategy_envs():
             kwargs = {}
diff --git a/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt b/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt
index a41df079cbb..674d3b1baee 100755
--- a/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt
+++ b/python/paddle/fluid/tests/unittests/ps/CMakeLists.txt
@@ -4,7 +4,9 @@ file(
   "test_*.py")
 string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
 
-list(REMOVE_ITEM TEST_OPS "test_gpubox_ps")
+if(NOT WITH_HETERPS OR WITH_PSLIB)
+  list(REMOVE_ITEM TEST_OPS "test_gpubox_ps")
+endif()
 
 foreach(TEST_OP ${TEST_OPS})
   py_test_modules(${TEST_OP} MODULES ${TEST_OP})
@@ -12,6 +14,6 @@ foreach(TEST_OP ${TEST_OPS})
   set_tests_properties(${TEST_OP} PROPERTIES TIMEOUT 50)
 endforeach()
 
-#if(WITH_HETERPS)
-#  set_tests_properties(test_gpubox_ps PROPERTIES LABELS "RUN_TYPE=GPUPS")
-#endif()
+if(WITH_HETERPS AND NOT WITH_PSLIB)
+  set_tests_properties(test_gpubox_ps PROPERTIES LABELS "RUN_TYPE=GPUPS")
+endif()
-- 
GitLab