API: added descs to checkRemoteMD

* As well as minor reordering of the API

Signed-off-by: Moein Khazraee <moein@nvidia.com>

Author: mkhazraee

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():

examples/python/partial_md_example.py

@@ -80,7 +80,7 @@
     ready = False
 
     while not ready:
-        ready = init_agent.check_remote_metadata("target")
+        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"

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;
 
         /**
@@ -381,65 +434,16 @@ class nixlAgent {
 
         /**
          * @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           Remote agent to check for
          * @return nixl_status_t Error code, NOT_FOUND if metadata not found
          */
         nixl_status_t
-        checkRemoteMD (const std::string remote_name) const;
+        checkRemoteMD (const std::string remote_name,
+                       const nixl_xfer_dlist_t &descs) const;
 
-        /*** 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(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/python/_api.py

@@ -612,6 +612,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,42 +697,24 @@ def invalidate_local_metadata(
         self.agent.invalidateLocalMD(ip_addr, port)
 
     """
-    @brief Check if the remote metadata for a specific agent is available
+    @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 check_remote_metadata(self, agent: str) -> bool:
-        if self.agent.checkRemoteMD(agent) == nixlBind.NIXL_SUCCESS:
+    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 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 Get nixlXferDList from different input types:
             a) list of 3 element tuples (address, len, device ID) alongside a mandatory memory type

src/bindings/python/nixl_bindings.cpp

@@ -504,6 +504,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,12 +549,5 @@ PYBIND11_MODULE(_bindings, m) {
 
                     throw_nixl_exception(agent.invalidateLocalMD(&extra_params));
                 }, py::arg("ip_addr") = std::string(""), py::arg("port") = 0 )
-        .def("checkRemoteMD", &nixlAgent::checkRemoteMD)
-        .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);
 }