Function GetTFGPUDeviceName

addons/triton/tensorflow_backend_tf.cc:310–332

Get the device name in model session given a non-negative device_id.

// Get the device name in model session given a non-negative device_id.
TRITONTF_Error*
GetTFGPUDeviceName(
    std::string* device_name, tensorflow::Session* session, const int device_id)
{
  if (device_id >= 0) {
    std::vector<tensorflow::DeviceAttributes> devices;
    RETURN_IF_TF_ERROR(session->ListDevices(&devices));
    for (const auto& d : devices) {
      if (d.device_type() == "GPU" || d.device_type() == "gpu") {
        // Session seems to be aware of all devices on the system,
        // thus need to filter out the correct full name for the device
        tensorflow::DeviceNameUtils::ParsedName parsed;
        if (tensorflow::DeviceNameUtils::ParseFullName(d.name(), &parsed)) {
          if (parsed.id == device_id) {
            *device_name = d.name();
            break;
          }
        }
      }
    }
  }
  return nullptr;
}

//
// TensorImpl