[Libfabric] Efa direct fallback and device mapping improvements (#817)

- Refactor libfabric device discovery and connection management
- Change getAvailableEfaDevices() to return provider-device mapping
- Update connection state handling in disconnect logic
- Modify rail manager to support provider-aware device selection
- Enhance topology detection for multi-provider environments
- Refactor libfabric memory registration to support device-specific operations
- Add device ID parameter to memory registration functions
- Update rail manager to handle device-aware memory registration
- Improve code formatting and line wrapping consistency
- Enhance topology management for multi-device scenarios
- Remove numa node mapping with EFA, instead select all rails in case of DRAM
- Clean up dead code to improve maintainability.
- Remove large blocks of commented-out code in nixlLibfabricEngine::disconnect() and simplify libfabric common utilities.

---------

Signed-off-by: Arun Karthik <[email protected]>

Author: akkart-aws

src/plugins/libfabric/libfabric_backend.cpp

@@ -513,86 +513,15 @@ nixlLibfabricEngine::disconnect(const std::string &remote_agent) {
                    << ", fi_addr: " << it->second->rail_remote_addr_list_[0];
         return NIXL_SUCCESS;
     }
-
-    NIXL_DEBUG << "Disconnecting from agent: " << remote_agent;
-
-    if (remote_agent != localAgent) {
-        // Send disconnect control message to remote peer - fire and forget semantics
-        NIXL_DEBUG << "Sending disconnect notification to remote agent: " << remote_agent;
-
-        // Use rail manager's serialization method with "src" prefix (we are sending our source
-        // endpoints)
-        std::string serialized_conn_info;
-        nixl_status_t serialize_status =
-            rail_manager.serializeConnectionInfo("src", serialized_conn_info);
-        if (serialize_status != NIXL_SUCCESS) {
-            NIXL_ERROR << "Rail manager serializeConnectionInfo failed";
-            return serialize_status;
-        }
-
-        // Allocate control request
-        const size_t control_rail_id = 0;
-        const size_t buffer_size =
-            (rail_manager.getNumDataRails() + rail_manager.getNumControlRails()) *
-                LF_EP_NAME_MAX_LEN +
-            1024;
-        nixlLibfabricReq *control_request =
-            rail_manager.getControlRail(control_rail_id).allocateControlRequest(buffer_size);
-        if (!control_request) {
-            NIXL_ERROR << "Failed to allocate control request for disconnect";
-            return NIXL_ERR_BACKEND;
-        }
-
-        memcpy(control_request->buffer, serialized_conn_info.data(), serialized_conn_info.length());
-
-        // Set the actual size of serialized data
-        control_request->buffer_size = serialized_conn_info.length();
-
-        nixl_status_t status = rail_manager.postControlMessage(
-            nixlLibfabricRailManager::ControlMessageType::DISCONNECT_REQ,
-            control_request,
-            it->second->control_rail_remote_addr_list_[0], // Use control rail 0
-            it->second->agent_index_);
-
-        if (status != NIXL_SUCCESS) {
-            NIXL_WARN << "Failed to send disconnect notification to " << remote_agent
-                      << ", proceeding with local cleanup anyway";
-            // Continue with cleanup even if notification failed
-        } else {
-            NIXL_DEBUG << "Disconnect notification sent successfully to " << remote_agent;
-        }
-    } else {
-        NIXL_DEBUG << "Skipping disconnect notification for self-connection";
-    }
-
-    // Clean up libfabric resources (AV entries) via rail manager
-    NIXL_DEBUG << "Cleaning up libfabric resources for agent: " << remote_agent;
-    // Clean up libfabric per-connection AV cleanup for both data and control rails
-    nixl_status_t data_cleanup_status = rail_manager.cleanupConnection(
-        nixlLibfabricRailManager::RailType::DATA, it->second->rail_remote_addr_list_);
-    if (data_cleanup_status != NIXL_SUCCESS) {
-        NIXL_ERROR << "Failed to clean up data rail resources for agent: " << remote_agent
-                   << " with status: " << data_cleanup_status;
-        return data_cleanup_status;
-    }
-
-    nixl_status_t control_cleanup_status = rail_manager.cleanupConnection(
-        nixlLibfabricRailManager::RailType::CONTROL, it->second->control_rail_remote_addr_list_);
-    if (control_cleanup_status != NIXL_SUCCESS) {
-        NIXL_ERROR << "Failed to clean up control rail resources for agent: " << remote_agent
-                   << " with status: " << control_cleanup_status;
-        return control_cleanup_status;
-    }
-
-    NIXL_DEBUG << "Successfully cleaned up libfabric resources for agent: " << remote_agent;
+    // TODO: Implement disconnect logic to cleanup the AV Address Entries from both local and remote
+    // AV.
 
     // Update connection state to DISCONNECTED before removing
     it->second->overall_state_ = ConnectionState::DISCONNECTED;
 
     // Remove connection from map
     connections_.erase(remote_agent);
     NIXL_DEBUG << "Connection erased from the connection map for agent: " << remote_agent;
-
     return NIXL_SUCCESS;
 }
 
@@ -837,6 +766,7 @@ nixlLibfabricEngine::registerMem(const nixlBlobDesc &mem,
     nixl_status_t status = rail_manager.registerMemory((void *)mem.addr,
                                                        mem.len,
                                                        nixl_mem,
+                                                       mem.devId,
                                                        priv->rail_mr_list_,
                                                        priv->rail_key_list_,
                                                        priv->selected_rails_);

src/utils/libfabric/libfabric_common.cpp

@@ -22,39 +22,74 @@
 #include <iomanip>
 #include <sstream>
 #include <atomic>
+#include <cstring>
 
 #include <rdma/fabric.h>
 #include <rdma/fi_domain.h>
 
 namespace LibfabricUtils {
 
-std::vector<std::string>
+
+std::pair<std::string, std::vector<std::string>>
 getAvailableEfaDevices() {
-    std::vector<std::string> devices;
+    std::unordered_map<std::string, std::vector<std::string>> provider_devices_map;
+    std::vector<std::string> all_efa_devices;
+    std::string fabric_name;
     struct fi_info *hints, *info;
     hints = fi_allocinfo();
     if (!hints) {
         NIXL_ERROR << "Failed to allocate fi_info for device discovery";
-        return devices;
+        return {fabric_name, all_efa_devices};
     }
 
+    // Important to initialize this to allow differentiation between EFA and EFA-Direct
+    hints->mode = ~0;
+
+    // Set required capabilities - let libfabric select the best provider
+    hints->caps = FI_READ | FI_WRITE | FI_RECV | FI_SEND | FI_REMOTE_READ | FI_REMOTE_WRITE |
+        FI_LOCAL_COMM | FI_REMOTE_COMM;
     hints->fabric_attr->prov_name = strdup("efa");
+    hints->ep_attr->type = FI_EP_RDM;
+
     int ret = fi_getinfo(FI_VERSION(1, 9), NULL, NULL, 0, hints, &info);
     if (ret) {
         NIXL_ERROR << "fi_getinfo failed during device discovery: " << fi_strerror(-ret);
         fi_freeinfo(hints);
-        return devices;
+        return {fabric_name, all_efa_devices};
     }
 
+    // Process providers and filter for EFA providers with RMA capabilities
     for (struct fi_info *cur = info; cur; cur = cur->next) {
-        if (cur->domain_attr && cur->domain_attr->name) {
-            devices.push_back(cur->domain_attr->name);
+        if (cur->domain_attr && cur->domain_attr->name && cur->fabric_attr &&
+            cur->fabric_attr->name) {
+
+            std::string device_name = cur->domain_attr->name;
+            std::string provider_name = cur->fabric_attr->name;
+
+            // Add device to the appropriate provider's vector
+            provider_devices_map[provider_name].push_back(device_name);
+
+            NIXL_TRACE << "Found EFA device: " << device_name << " with provider: " << provider_name
+                       << " (caps: 0x" << std::hex << cur->caps << std::dec << ")";
         }
     }
 
     fi_freeinfo(info);
     fi_freeinfo(hints);
-    return devices;
+
+    // Extract device names from the map, prioritizing efa-direct over efa
+    all_efa_devices.clear();
+    if (provider_devices_map.find("efa-direct") != provider_devices_map.end()) {
+        all_efa_devices = provider_devices_map["efa-direct"];
+        fabric_name = "efa-direct";
+        NIXL_TRACE << "Using efa-direct provider with " << all_efa_devices.size() << " devices";
+    } else if (provider_devices_map.find("efa") != provider_devices_map.end()) {
+        all_efa_devices = provider_devices_map["efa"];
+        fabric_name = "efa";
+        NIXL_TRACE << "Using efa provider with " << all_efa_devices.size() << " devices";
+    }
+
+    return {fabric_name, all_efa_devices};
 }
 
 std::string

src/utils/libfabric/libfabric_common.h

@@ -21,6 +21,7 @@
 #include <vector>
 #include <string>
 #include <unordered_set>
+#include <unordered_map>
 #include <cstring>
 
 #include "nixl.h"
@@ -151,7 +152,7 @@ preallocateXferIds(size_t count);
 // Utility functions
 namespace LibfabricUtils {
 // Device discovery
-std::vector<std::string>
+std::pair<std::string, std::vector<std::string>>
 getAvailableEfaDevices();
 // String utilities
 std::string

src/utils/libfabric/libfabric_rail.cpp

@@ -187,7 +187,9 @@ DataRequestPool::allocate(nixlLibfabricReq::OpType op_type) {
 
 // Rail Class Implementation
 
-nixlLibfabricRail::nixlLibfabricRail(const std::string &device, uint16_t id)
+nixlLibfabricRail::nixlLibfabricRail(const std::string &device,
+                                     const std::string &provider,
+                                     uint16_t id)
     : rail_id(id),
       device_name(device),
       blocking_cq_sread_supported(true),
@@ -203,7 +205,8 @@ nixlLibfabricRail::nixlLibfabricRail(const std::string &device, uint16_t id)
     memset(ep_name, 0, sizeof(ep_name));
 
     // Initialize all Libfabric resources for this rail
-    NIXL_TRACE << "Initializing rail " << rail_id << " with device: " << device_name;
+    NIXL_TRACE << "Initializing rail " << rail_id << " with device: " << device_name
+               << ", provider: " << provider;
 
     // Initialize hints for this rail
     struct fi_info *hints = fi_allocinfo();
@@ -212,9 +215,13 @@ nixlLibfabricRail::nixlLibfabricRail(const std::string &device, uint16_t id)
         throw std::runtime_error("Failed to allocate fi_info for rail " + std::to_string(rail_id));
     }
     hints->caps = 0;
-    hints->caps = FI_MSG | FI_RMA | FI_HMEM;
+    hints->caps = FI_MSG | FI_RMA;
     hints->caps |= FI_LOCAL_COMM | FI_REMOTE_COMM;
-    hints->mode = FI_CONTEXT | FI_CONTEXT2;
+    if (provider.c_str() == std::string("efa-direct")) {
+        hints->mode = FI_CONTEXT | FI_CONTEXT2;
+    } else {
+        hints->mode = FI_CONTEXT;
+    }
     hints->ep_attr->type = FI_EP_RDM;
     hints->domain_attr->mr_mode =
         FI_MR_LOCAL | FI_MR_HMEM | FI_MR_VIRT_ADDR | FI_MR_ALLOCATED | FI_MR_PROV_KEY;
@@ -250,8 +257,8 @@ nixlLibfabricRail::nixlLibfabricRail(const std::string &device, uint16_t id)
         cq_attr.size = 12288;
         ret = fi_cq_open(domain, &cq_attr, &cq, NULL);
         if (ret) {
-            NIXL_ERROR << "fi_cq_open failed for rail " << rail_id << ": " << fi_strerror(-ret)
-                       << " - trying FI_WAIT_NONE";
+            NIXL_WARN << "fi_cq_open failed for rail " << rail_id << ": " << fi_strerror(-ret)
+                      << " - trying FI_WAIT_NONE";
             if (ret == -FI_ENOSYS) {
                 NIXL_TRACE << "FI_WAIT_UNSPEC not supported, falling back to FI_WAIT_NONE for rail "
                            << rail_id;

src/utils/libfabric/libfabric_rail.h

@@ -214,7 +214,7 @@ class nixlLibfabricRail {
     struct fid_ep *endpoint; ///< Libfabric endpoint handle
 
     /** Initialize libfabric rail with all resources */
-    nixlLibfabricRail(const std::string &device, uint16_t id);
+    nixlLibfabricRail(const std::string &device, const std::string &provider, uint16_t id);
 
     /** Destroy rail and cleanup all libfabric resources */
     ~nixlLibfabricRail();