Two fixes for StepMesh

fserver/csrc/public.hpp

@@ -178,7 +178,6 @@ void init() {
   q_signal_.store(0);;
 
   ps::StartPS(0, role_,  group_size_ * node_rank_ + gpu_, true);
-  Backend::Get()->SetDevice(gpu_);
   if (role_ == Node::WORKER) {
     fworker_ = new AFTensorWorker(instance_id_);
     barrier(true, true);

include/ps/ps.h

@@ -143,6 +143,13 @@ inline void StartPS(int customer_id, Node::Role role, int rank, bool do_barrier,
   Backend::Register("GPU", new GpuBackend());
 #endif
 
+  // scheduler do not need to attach to gpu
+  if (role != Node::SCHEDULER) {
+    int gpu = 0;
+    Environment::Get()->find("STEPMESH_GPU", &gpu, gpu);
+    Backend::Get()->SetDevice(gpu);
+  }
+
   int group_size = 1;
 
   Environment::Get()->find("DMLC_GROUP_SIZE", &group_size, group_size);

src/network_utils.h

@@ -335,20 +335,22 @@ static inline int GetAvailablePort(int num_ports, std::array<int, 32>* ports) {
  * \return 0 on failure or no cuda, 1 when getting the interface successfully
  */
 static inline int GetInterfaceAndIPByCurrentGpu(std::string* interface,
-                                                std::string* ip) {
+                                                std::string* ip, int* gpu) {
   interface->clear();
   ip->clear();
 
 #ifdef DMLC_USE_CUDA
-  int gpu = -1;
-  cudaGetDevice(&gpu);
-  if (gpu == -1) return 0;
+  cudaGetDevice(gpu);
+  if (*gpu == -1) return 0;
   char pciPath[512];
-  cudaDeviceProp deviceProp = {};
-  cudaGetDeviceProperties(&deviceProp, gpu);
-  snprintf(pciPath, sizeof(pciPath),
-           "/sys/class/pci_bus/0000:%02x/device/0000:%02x:%02x.0/device",
-           deviceProp.pciBusID, deviceProp.pciBusID, deviceProp.pciDeviceID);
+  char busId[16];
+  cudaError_t status;
+  status = cudaDeviceGetPCIBusId(busId, 16, *gpu);
+  PS_CHECK_EQ(status, cudaSuccess) << "cudaDeviceGetPCIBusId failed"
+                                   << " (" << cudaGetErrorString(status) << ")";
+  for (int i = 0; i < 16; i++) busId[i] = std::tolower(busId[i]);
+  snprintf(pciPath, sizeof(pciPath), "/sys/class/pci_bus/%.7s/device/%s/device",
+           busId, busId);
   char* path = realpath(pciPath, nullptr);
 
   if (path == nullptr) return 0;

src/van.cc

@@ -552,9 +552,10 @@ void Van::Start(int customer_id, bool standalone) {
       std::string interface;
       if (itf) interface = std::string(itf);
       if (interface == "auto" || interface == "AUTO") {
-        GetInterfaceAndIPByCurrentGpu(&interface, &ip);
-        PS_LOG(INFO) << "automatic detect interface and ip from gpu: "
-                     << interface << " (" << ip << ")";
+        int gpu = -1;
+        GetInterfaceAndIPByCurrentGpu(&interface, &ip, &gpu);
+        PS_LOG(INFO) << "automatic detect interface and ip from gpu(" << gpu
+                     << "): " << interface << " (" << ip << ")";
         Environment::Get()->set("DMLC_NODE_HOST", ip);
         Environment::Get()->set("DMLC_INTERFACE", interface);
       } else {