prov/efa: Update UT to make sure SHM isn't used for Neuron

Add assert statements to existing Neuron ut's to make sure that the SHM provider is not used. Signed-off-by: Seth Zegelstein <[email protected]>
ofiwg · Dec 6, 2023 · 4d0d297 · 4d0d297
1 parent c0a7f41
commit 4d0d297
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 1 deletion.
diff --git a/prov/efa/test/efa_unit_test_hmem.c b/prov/efa/test/efa_unit_test_hmem.c
@@ -2,13 +2,17 @@
 
 
 #if HAVE_NEURON
+#include <dlfcn.h>
+#include "nrt/nrt.h"
+#include "nrt/nrt_experimental.h"
+
 /**
  * @brief Verify when neuron_alloc failed (return null),
  * efa_domain_open, which call efa_hmem_info_update_neuron
  * when HAVE_NEURON=1, will still return 0 but leave
  * efa_hmem_info[FI_HMEM_NEURON].initialized and
  * efa_hmem_info[FI_HMEM_NEURON].p2p_supported_by_device as false.
- * 
+ *
  * @param[in]	state		struct efa_resource that is managed by the framework
  */
 void test_efa_hmem_info_update_neuron(struct efa_resource **state)
@@ -46,9 +50,74 @@ void test_efa_hmem_info_update_neuron(struct efa_resource **state)
         assert_false(efa_domain->hmem_info[FI_HMEM_NEURON].initialized);
         assert_false(efa_domain->hmem_info[FI_HMEM_NEURON].p2p_supported_by_device);
 }
+
+
+int neuron_init(void)
+{
+	NRT_STATUS (*nrt_init)(nrt_framework_type_t framework, const char *fw_version, const char *fal_version);
+	NRT_STATUS ret;
+	void *neuron_handle = NULL;
+
+	neuron_handle = dlopen("libnrt.so.1", RTLD_NOW);
+	if (!neuron_handle) {
+		return -1;
+	}
+
+	nrt_init = dlsym(neuron_handle, "nrt_init");
+	if (!nrt_init) {
+		return -1;
+	}
+
+	ret = nrt_init(NRT_FRAMEWORK_TYPE_NO_FW, "2.0", "");
+	if (ret != NRT_SUCCESS) {
+		return -1;
+	}
+
+	return 0;
+}
+
+
+void test_efa_hmem_neuron_no_shm(struct efa_resource **state)
+{
+	int ret;
+	struct efa_resource *resource = *state;
+	struct efa_domain *efa_domain;
+
+	int no_neuron = neuron_init();
+	if (no_neuron) {
+		/* This test requires that it be run on a neuron instance */
+		skip();
+	}
+
+	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM);
+	assert_non_null(resource->hints);
+
+	ret = fi_getinfo(FI_VERSION(1, 14), NULL, NULL, 0ULL, resource->hints, &resource->info);
+	assert_int_equal(ret, 0);
+
+	ret = fi_fabric(resource->info->fabric_attr, &resource->fabric, NULL);
+	assert_int_equal(ret, 0);
+
+	ret = fi_domain(resource->fabric, resource->info, &resource->domain, NULL);
+
+	efa_domain = container_of(resource->domain, struct efa_domain,
+				  util_domain.domain_fid.fid);
+
+	/* Make sure that SHM is not being used with Neuron */
+	assert_null(efa_domain->shm_domain);
+	assert_null(efa_domain->shm_info);
+	assert_null(efa_domain->fabric->shm_fabric);
+
+	/* TODO: When we fix inject size for neuron, we should test for it */
+}
 #else
 void test_efa_hmem_info_update_neuron()
 {
         skip();
 }
+
+void test_efa_hmem_neuron_no_shm(struct efa_resource **state)
+{
+	skip();
+}
 #endif /* HAVE_NEURON */
diff --git a/prov/efa/test/efa_unit_tests.c b/prov/efa/test/efa_unit_tests.c
@@ -133,6 +133,7 @@ int main(void)
 		cmocka_unit_test_setup_teardown(test_efa_rdm_msg_send_to_local_peer_with_null_desc, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		cmocka_unit_test_setup_teardown(test_efa_fork_support_request_initialize_when_ibv_fork_support_is_needed, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		cmocka_unit_test_setup_teardown(test_efa_fork_support_request_initialize_when_ibv_fork_support_is_unneeded, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+		cmocka_unit_test_setup_teardown(test_efa_hmem_neuron_no_shm, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown), // Needs to be last b/c this will shut down shm on neuron
 	};
 
 	cmocka_set_message_output(CM_OUTPUT_XML);

diff --git a/prov/efa/test/efa_unit_tests.h b/prov/efa/test/efa_unit_tests.h
@@ -119,6 +119,7 @@ void test_info_check_hmem_cuda_support_on_api_lt_1_18();
 void test_info_check_hmem_cuda_support_on_api_ge_1_18();
 void test_info_check_no_hmem_support_when_not_requested();
 void test_efa_hmem_info_update_neuron();
+void test_efa_hmem_neuron_no_shm();
 void test_efa_use_device_rdma_env1_opt1();
 void test_efa_use_device_rdma_env0_opt0();
 void test_efa_use_device_rdma_env1_opt0();