API: checkRemoteMD call (#203)

* API: checkRemoteMD call

Signed-off-by: Timothy Stamler <tstamler@nvidia.com>

* API: added descs to checkRemoteMD
* As well as minor reordering of the API

Signed-off-by: Timothy Stamler <tstamler@nvidia.com>
Signed-off-by: Moein Khazraee <moein@nvidia.com>
Co-authored-by: Moein Khazraee <moein@nvidia.com>

Author: tstamler

.gitlab/test_python.sh

@@ -43,7 +43,7 @@ pytest test/python
 
 echo "==== Running python example ===="
 cd examples/python
-python3 partial_md_example.py
 python3 blocking_send_recv_example.py --mode="target" --ip=127.0.0.1 --port=1234&
-sleep 1
+sleep 5
 python3 blocking_send_recv_example.py --mode="initiator" --ip=127.0.0.1 --port=1234
+python3 partial_md_example.py

examples/python/blocking_send_recv_example.py

@@ -20,7 +20,6 @@
 import torch
 
 from nixl._api import nixl_agent, nixl_agent_config
-from nixl._bindings import nixlNotFoundError
 
 
 def parse_args():
@@ -68,12 +67,11 @@ def parse_args():
 
         # Send desc list to initiator when metadata is ready
         while not ready:
-            try:
-                agent.send_notif("initiator", target_desc_str)
-            except nixlNotFoundError:
-                ready = False
-            else:
-                ready = True
+            ready = agent.check_remote_metadata("initiator")
+
+        agent.send_notif("initiator", target_desc_str)
+
+        print("Waiting for transfer")
 
         # Waiting for transfer
         # For now the notification is just UUID, could be any python bytes.
@@ -98,14 +96,13 @@ def parse_args():
         # Ensure remote metadata has arrived from fetch
         ready = False
         while not ready:
-            try:
-                xfer_handle = agent.initialize_xfer(
-                    "READ", initiator_descs, target_descs, "target", "UUID"
-                )
-            except nixlNotFoundError:
-                ready = False
-            else:
-                ready = True
+            ready = agent.check_remote_metadata("target")
+
+        print("Ready for transfer")
+
+        xfer_handle = agent.initialize_xfer(
+            "READ", initiator_descs, target_descs, "target", "UUID"
+        )
 
         if not xfer_handle:
             print("Creating transfer failed.")

examples/python/partial_md_example.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 
 import os
+import time
 
 import nixl._utils as nixl_utils
 from nixl._api import nixl_agent, nixl_agent_config
@@ -78,17 +79,13 @@
 
     # Wait for metadata to be loaded
     ready = False
-    xfer_handle_1 = 0
+
     while not ready:
-        try:
-            # initialize transfer mode
-            xfer_handle_1 = init_agent.initialize_xfer(
-                "READ", init_xfer_descs, target_xfer_descs1, "target", b"UUID1"
-            )
-        except nixlNotFoundError:
-            ready = False
-        else:
-            ready = True
+        ready = init_agent.check_remote_metadata("target", target_xfer_descs1)
+
+    xfer_handle_1 = init_agent.initialize_xfer(
+        "READ", init_xfer_descs, target_xfer_descs1, "target", b"UUID1"
+    )
 
     state = init_agent.transfer(xfer_handle_1)
     assert state != "ERR"
@@ -171,4 +168,6 @@
     for addr in malloc_addrs:
         nixl_utils.free_passthru(addr)
 
+    # Give sockets time to close
+    time.sleep(1)
     print("Test Complete.")

src/api/cpp/nixl.h

@@ -316,6 +316,59 @@ class nixlAgent {
                   const nixl_blob_t &msg,
                   const nixl_opt_args_t* extra_params = nullptr);
 
+        /*** Metadata handling through side channel ***/
+        /**
+         * @brief  Get metadata blob for this agent, to be given to other agents.
+         *
+         * @param  str [out]     The serialized metadata blob
+         * @return nixl_status_t Error code if call was not successful
+         */
+        nixl_status_t
+        getLocalMD (nixl_blob_t &str) const;
+
+        /**
+         * @brief  Get partial metadata blob for this agent, to be given to other agents.
+         *         If `descs` is empty, only backends' connection info is included in the metadata,
+         *         regardless of the value of `extra_params->includeConnInfo` and `descs` memory type.
+         *         If `descs` is non-empty, the metadata of the descriptors in the list are included,
+         *         and if `extra_params->includeConnInfo` is true, the connection info of the
+         *         backends supporting the memory type is also included.
+         *         If `extra_params->backends` is non-empty, only the descriptors supported by the
+         *         backends in the list and the backends' connection info are included in the metadata.
+         *
+         * @param  descs         [in]  Descriptor list to include in the metadata
+         * @param  str           [out] The serialized metadata blob
+         * @param  extra_params  [in]  Optional extra parameters used in getting partial metadata
+         * @return nixl_status_t       Error code if call was not successful
+         */
+        nixl_status_t
+        getLocalPartialMD(const nixl_reg_dlist_t &descs,
+                          nixl_blob_t &str,
+                          const nixl_opt_args_t* extra_params = nullptr) const;
+
+        /**
+         * @brief  Load other agent's metadata and unpack it internally. Now the local
+         *         agent can initiate transfers towards the remote agent.
+         *
+         * @param  remote_metadata  Serialized metadata blob to be loaded
+         * @param  agent_name [out] Agent name extracted from the loaded metadata blob
+         * @return nixl_status_t    Error code if call was not successful
+         */
+        nixl_status_t
+        loadRemoteMD (const nixl_blob_t &remote_metadata,
+                      std::string &agent_name);
+
+        /**
+         * @brief  Invalidate the remote agent metadata cached locally. This will
+         *         disconnect from that agent if already connected, and no more
+         *         transfers can be initiated towards that agent.
+         *
+         * @param  remote_agent  Remote agent name to invalidate its metadata blob
+         * @return nixl_status_t Error code if call was not successful
+         */
+        nixl_status_t
+        invalidateRemoteMD (const std::string &remote_agent);
+
         /*** Metadata handling through direct channels (p2p socket and ETCD) ***/
         /**
          * @brief  Send your own agent metadata to a remote location.
@@ -348,7 +401,7 @@ class nixlAgent {
          * @return nixl_status_t       Error code if call was not successful
          */
         nixl_status_t
-        sendLocalPartialMD(nixl_reg_dlist_t  &descs,
+        sendLocalPartialMD(const nixl_reg_dlist_t &descs,
                            const nixl_opt_args_t* extra_params = nullptr) const;
 
         /**
@@ -379,58 +432,18 @@ class nixlAgent {
         nixl_status_t
         invalidateLocalMD (const nixl_opt_args_t* extra_params = nullptr) const;
 
-        /*** Metadata handling through side channel ***/
         /**
-         * @brief  Get metadata blob for this agent, to be given to other agents.
+         * @brief  Check if metadata is available for a remote agent.
+         *         For partial metadata methods are used, the descriptor list in question
+         *         can be specified; otherwise, empty `descs` can be passed.
          *
-         * @param  str [out]     The serialized metadata blob
-         * @return nixl_status_t Error code if call was not successful
+         * @param  str           Remote agent to check for
+         * @return nixl_status_t Error code, NOT_FOUND if metadata not found
          */
         nixl_status_t
-        getLocalMD (nixl_blob_t &str) const;
+        checkRemoteMD (const std::string remote_name,
+                       const nixl_xfer_dlist_t &descs) const;
 
-        /**
-         * @brief  Get partial metadata blob for this agent, to be given to other agents.
-         *         If `descs` is empty, only backends' connection info is included in the metadata,
-         *         regardless of the value of `extra_params->includeConnInfo` and `descs` memory type.
-         *         If `descs` is non-empty, the metadata of the descriptors in the list are included,
-         *         and if `extra_params->includeConnInfo` is true, the connection info of the
-         *         backends supporting the memory type is also included.
-         *         If `extra_params->backends` is non-empty, only the descriptors supported by the
-         *         backends in the list and the backends' connection info are included in the metadata.
-         *
-         * @param  descs         [in]  Descriptor list to include in the metadata
-         * @param  str           [out] The serialized metadata blob
-         * @param  extra_params  [in]  Optional extra parameters used in getting partial metadata
-         * @return nixl_status_t       Error code if call was not successful
-         */
-        nixl_status_t
-        getLocalPartialMD(nixl_reg_dlist_t  &descs,
-                          nixl_blob_t &str,
-                          const nixl_opt_args_t* extra_params = nullptr) const;
-
-        /**
-         * @brief  Load other agent's metadata and unpack it internally. Now the local
-         *         agent can initiate transfers towards the remote agent.
-         *
-         * @param  remote_metadata  Serialized metadata blob to be loaded
-         * @param  agent_name [out] Agent name extracted from the loaded metadata blob
-         * @return nixl_status_t    Error code if call was not successful
-         */
-        nixl_status_t
-        loadRemoteMD (const nixl_blob_t &remote_metadata,
-                      std::string &agent_name);
-
-        /**
-         * @brief  Invalidate the remote agent metadata cached locally. This will
-         *         disconnect from that agent if already connected, and no more
-         *         transfers can be initiated towards that agent.
-         *
-         * @param  remote_agent  Remote agent name to invalidate its metadata blob
-         * @return nixl_status_t Error code if call was not successful
-         */
-        nixl_status_t
-        invalidateRemoteMD (const std::string &remote_agent);
 };
 
 #endif

src/api/cpp/nixl_params.h

@@ -70,9 +70,9 @@ class nixlAgentConfig {
         nixlAgentConfig (const bool use_prog_thread,
                          const bool use_listen_thread=false,
                          const int port=0,
+                         nixl_thread_sync_t sync_mode=nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT,
                          const uint64_t pthr_delay_us=0,
-                         const uint64_t lthr_delay_us = 100000,
-                         nixl_thread_sync_t sync_mode=nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT) :
+                         const uint64_t lthr_delay_us = 100000) :
                          useProgThread(use_prog_thread),
                          useListenThread(use_listen_thread),
                          listenPort(port),

src/api/python/_api.py

@@ -81,9 +81,18 @@ def __init__(
         if not nixl_conf:
             nixl_conf = nixl_agent_config()  # Using defaults set in nixl_agent_config
 
+        thread_config = (
+            nixlBind.NIXL_THREAD_SYNC_STRICT
+            if nixl_conf.enable_listen
+            else nixlBind.NIXL_THREAD_SYNC_NONE
+        )
+
         # Set agent config and instantiate an agent
         agent_config = nixlBind.nixlAgentConfig(
-            nixl_conf.enable_pthread, nixl_conf.enable_listen, nixl_conf.port
+            nixl_conf.enable_pthread,
+            nixl_conf.enable_listen,
+            nixl_conf.port,
+            thread_config,
         )
         self.agent = nixlBind.nixlAgent(agent_name, agent_config)
 
@@ -612,6 +621,29 @@ def get_partial_agent_metadata(
             handle_list.append(self.backends[backend_string])
         return self.agent.getLocalPartialMD(descs, inc_conn_info, handle_list)
 
+    """
+    @brief Add a remote agent using its metadata. After this call, current agent can
+            initiate transfers towards the remote agent.
+
+    @param metadata Metadata of the remote agent, received out-of-band in bytes.
+    @return Name of the added remote agent.
+    """
+
+    def add_remote_agent(self, metadata: bytes) -> str:
+        agent_name = self.agent.loadRemoteMD(metadata)
+        return agent_name
+
+    """
+    @brief Remove a remote agent. After this call, current agent cannot initiate
+            transfers towards the remote agent specified in the call anymore.
+            This call will also result in a disconnect between the two agents.
+
+    @param agent Name of the remote agent.
+    """
+
+    def remove_remote_agent(self, agent: str):
+        self.agent.invalidateRemoteMD(agent)
+
     """
     @brief Send all of your metadata to a peer or central metadata server.
 
@@ -674,27 +706,23 @@ def invalidate_local_metadata(
         self.agent.invalidateLocalMD(ip_addr, port)
 
     """
-    @brief Add a remote agent using its metadata. After this call, current agent can
-            initiate transfers towards the remote agent.
-
-    @param metadata Metadata of the remote agent, received out-of-band in bytes.
-    @return Name of the added remote agent.
-    """
-
-    def add_remote_agent(self, metadata: bytes) -> str:
-        agent_name = self.agent.loadRemoteMD(metadata)
-        return agent_name
-
-    """
-    @brief Remove a remote agent. After this call, current agent cannot initiate
-            transfers towards the remote agent specified in the call anymore.
-            This call will also result in a disconnect between the two agents.
+    @brief Check if the remote metadata for a specific agent is available.
+           When partial metadata methods are used, the descriptor list in question can be specified.
 
     @param agent Name of the remote agent.
+
+    @return True if available, False otherwise
     """
 
-    def remove_remote_agent(self, agent: str):
-        self.agent.invalidateRemoteMD(agent)
+    def check_remote_metadata(
+        self, agent: str, descs: nixlBind.nixlXferDList = None
+    ) -> bool:
+        if descs is None:  # Just empty list, mem_type not important
+            descs = nixlBind.nixlXferDList(nixlBind.DRAM_SEG)
+        if self.agent.checkRemoteMD(agent, descs) == nixlBind.NIXL_SUCCESS:
+            return True
+        else:
+            return False
 
     """
     @brief Get nixlXferDList from different input types:

src/bindings/python/nixl_bindings.cpp

@@ -123,6 +123,12 @@ PYBIND11_MODULE(_bindings, m) {
     m.attr("DEFAULT_COMM_PORT") = default_comm_port;
 
     //cast types
+    py::enum_<nixl_thread_sync_t>(m, "nixl_thread_sync_t")
+        .value("NIXL_THREAD_SYNC_NONE", nixl_thread_sync_t::NIXL_THREAD_SYNC_NONE)
+        .value("NIXL_THREAD_SYNC_STRICT", nixl_thread_sync_t::NIXL_THREAD_SYNC_STRICT)
+        .value("NIXL_THREAD_SYNC_DEFAULT", nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT)
+        .export_values();
+
     py::enum_<nixl_mem_t>(m, "nixl_mem_t")
         .value("DRAM_SEG", DRAM_SEG)
         .value("VRAM_SEG", VRAM_SEG)
@@ -282,7 +288,8 @@ PYBIND11_MODULE(_bindings, m) {
         //implicit constructor
         .def(py::init<bool>())
         .def(py::init<bool, bool>())
-        .def(py::init<bool, bool, int>());
+        .def(py::init<bool, bool, int>())
+        .def(py::init<bool, bool, int, nixl_thread_sync_t>());
 
     //note: pybind will automatically convert notif_map to python types:
     //so, a Dictionary of string: List<string>
@@ -504,6 +511,13 @@ PYBIND11_MODULE(_bindings, m) {
                     throw_nixl_exception(agent.getLocalPartialMD(descs, ret_str, &extra_params));
                     return py::bytes(ret_str);
                 }, py::arg("descs"), py::arg("inc_conn_info") = false, py::arg("backends") = std::vector<uintptr_t>({}))
+        .def("loadRemoteMD", [](nixlAgent &agent, const std::string &remote_metadata) -> py::bytes {
+                    //python can only interpret text strings
+                    std::string remote_name("");
+                    throw_nixl_exception(agent.loadRemoteMD(remote_metadata, remote_name));
+                    return py::bytes(remote_name);
+                })
+        .def("invalidateRemoteMD", &nixlAgent::invalidateRemoteMD)
         .def("sendLocalMD", [](nixlAgent &agent, std::string ip_addr, int port){
                     nixl_opt_args_t extra_params;
 
@@ -542,11 +556,5 @@ PYBIND11_MODULE(_bindings, m) {
 
                     throw_nixl_exception(agent.invalidateLocalMD(&extra_params));
                 }, py::arg("ip_addr") = std::string(""), py::arg("port") = 0 )
-        .def("loadRemoteMD", [](nixlAgent &agent, const std::string &remote_metadata) -> py::bytes {
-                    //python can only interpret text strings
-                    std::string remote_name("");
-                    throw_nixl_exception(agent.loadRemoteMD(remote_metadata, remote_name));
-                    return py::bytes(remote_name);
-                })
-        .def("invalidateRemoteMD", &nixlAgent::invalidateRemoteMD);
+        .def("checkRemoteMD", &nixlAgent::checkRemoteMD);
 }

src/core/meson.build

@@ -13,7 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-nixl_lib_deps = [nixl_infra, serdes_interface, stream_interface, dl_dep]
+# Add dependency on the common utility library which brings in logging deps
+nixl_lib_deps = [nixl_infra, serdes_interface, stream_interface, dl_dep, nixl_common_dep, thread_dep]
 
 if 'UCX' in static_plugins
     nixl_lib_deps += [ ucx_backend_interface, cuda_dep ]