Skip to content

Commit f336666

Browse files
roiedaninoroiedanino
authored
BINDINGS/RUST: adding missing Agent APIs - make_xfer_req, prepare_xfer_dlist (#684)
* SRC/BINDINGS/RUST: added missing Agent APIs - prep_xfer_dlist, make_xfer_req, releaseDlistH Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: added query_xfer_backend, get_local_partial_md, send_local_partial_md Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: added send_local_partial_md in agent.rs Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: minor fixes Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: fixed Rust compilation issues Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/SRC: minor fixes Signed-off-by: Roie Danino <[email protected]> * BINDINGS/RUST: fix clang formatting Signed-off-by: Roie Danino <[email protected]> * BINDINGS/RUST: fix clang formatting Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: disabled clang-format in the problematic part Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: disabled clang-format in the problematic part Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: fixed make_connection calls from tests Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: added tests to the new APIs Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: removed some of the code duplication Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: removed trailing whitespaces Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: fixed returned values lifetime using Result<,> Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUTS/TESTS: added <'a> to create_dlist Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: tests should return Result<,> Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/SRC: Added PartialEq trait to Backend struct Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: fixed vector Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: removed get_local_partial_md and send_local_partial_md (PR is too big) Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: having some issues in registration Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: fixed prepare xfer test Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST/TESTS: added and fixed tests for new APIs Signed-off-by: Roie Danino <[email protected]> * SRC/CORE: removed unnecessary changes and debug logging Signed-off-by: Roie Danino <[email protected]> * SRC/CORE: removed blank lines Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: removed prints Signed-off-by: Roie Danino <[email protected]> * SRC/BINDINGS/RUST: moving query_xfer_backend to the next PR (lines limit) Signed-off-by: Roie Danino <[email protected]> --------- Signed-off-by: Roie Danino <[email protected]>
1 parent 45204fa commit f336666

File tree

10 files changed

+444
-12
lines changed

10 files changed

+444
-12
lines changed

‎src/bindings/rust/src/agent.rs‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -374,15 +374,15 @@ impl Agent {
374374
}
375375
}
376376

377-
pub fn make_connection(&self, remote_agent: &str) -> Result<(), NixlError> {
377+
pub fn make_connection(&self, remote_agent: &str, opt_args: Option<&OptArgs>) -> Result<(), NixlError> {
378378
let remote_agent = CString::new(remote_agent)?;
379379
let inner_guard = self.inner.write().unwrap();
380380

381381
let status = unsafe {
382382
nixl_capi_agent_make_connection(
383383
inner_guard.handle.as_ptr(),
384384
remote_agent.as_ptr(),
385-
std::ptr::null_mut(),
385+
opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
386386
)
387387
};
388388

@@ -393,6 +393,64 @@ impl Agent {
393393
}
394394
}
395395

396+
pub fn prepare_xfer_dlist(
397+
&self,
398+
agent_name: &str,
399+
descs: &XferDescList,
400+
opt_args: Option<&OptArgs>,
401+
) -> Result<XferDlistHandle, NixlError> {
402+
let c_agent_name = CString::new(agent_name)?;
403+
let mut dlist_hndl = std::ptr::null_mut();
404+
let inner_guard = self.inner.read().unwrap();
405+
406+
let status = unsafe {
407+
nixl_capi_prep_xfer_dlist(
408+
inner_guard.handle.as_ptr(),
409+
c_agent_name.as_ptr(),
410+
descs.handle(),
411+
&mut dlist_hndl,
412+
opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr()),
413+
)
414+
};
415+
416+
match status {
417+
NIXL_CAPI_SUCCESS => Ok(XferDlistHandle::new(dlist_hndl, inner_guard.handle)),
418+
_ => Err(NixlError::BackendError),
419+
}
420+
}
421+
422+
pub fn make_xfer_req(&self, operation: XferOp,
423+
local_descs: &XferDlistHandle, local_indices: &[i32],
424+
remote_descs: &XferDlistHandle, remote_indices: &[i32],
425+
opt_args: Option<&OptArgs>) -> Result<XferRequest, NixlError> {
426+
let mut req = std::ptr::null_mut();
427+
let inner_guard = self.inner.read().unwrap();
428+
429+
let status = unsafe {
430+
nixl_capi_make_xfer_req(
431+
inner_guard.handle.as_ptr(),
432+
operation as bindings::nixl_capi_xfer_op_t,
433+
local_descs.handle(),
434+
local_indices.as_ptr(),
435+
local_indices.len() as usize,
436+
remote_descs.handle(),
437+
remote_indices.as_ptr(),
438+
remote_indices.len() as usize,
439+
&mut req,
440+
opt_args.map_or(std::ptr::null_mut(), |args| args.inner.as_ptr())
441+
)
442+
};
443+
444+
match status {
445+
NIXL_CAPI_SUCCESS => Ok(XferRequest::new(NonNull::new(req)
446+
.ok_or(NixlError::FailedToCreateXferRequest)?,
447+
self.inner.clone(),
448+
)),
449+
NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
450+
_ => Err(NixlError::BackendError),
451+
}
452+
}
453+
396454
/// Check if remote metadata for a specific agent is available
397455
///
398456
/// This function checks if the metadata for the specified remote agent has been

‎src/bindings/rust/src/descriptors.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@ use super::*;
1818
mod query;
1919
mod reg;
2020
mod xfer;
21+
mod xfer_dlist_handle;
2122

2223
pub use query::{QueryResponse, QueryResponseIterator, QueryResponseList};
2324
pub use reg::RegDescList;
2425
pub use xfer::XferDescList;
26+
pub use xfer_dlist_handle::XferDlistHandle;
2527

2628
/// Memory types supported by NIXL
2729
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]

‎src/bindings/rust/src/descriptors/reg.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -218,8 +218,7 @@ impl<'a> RegDescList<'a> {
218218
_ => Err(NixlError::BackendError),
219219
}?;
220220
if len > 0 {
221-
// TODO: Add API to get descriptor memory type
222-
MemType::Unknown
221+
self.get_type()?
223222
} else {
224223
desc_mem_type
225224
}

‎src/bindings/rust/src/descriptors/xfer.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -205,8 +205,7 @@ impl<'a> XferDescList<'a> {
205205
_ => Err(NixlError::BackendError),
206206
}?;
207207
if len > 0 {
208-
// TODO: Add API to get descriptor memory type
209-
MemType::Unknown
208+
self.get_type().unwrap()
210209
} else {
211210
desc_mem_type
212211
}

‎src/bindings/rust/src/descriptors/xfer_dlist_handle.rs‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
use super::*;
17+
18+
pub struct XferDlistHandle {
19+
inner: *mut bindings::nixl_capi_xfer_dlist_handle_s,
20+
agent: NonNull<bindings::nixl_capi_agent_s>
21+
}
22+
23+
impl XferDlistHandle {
24+
pub fn new(inner: *mut bindings::nixl_capi_xfer_dlist_handle_s,
25+
agent: NonNull<bindings::nixl_capi_agent_s>) -> Self {
26+
Self { inner, agent }
27+
}
28+
29+
pub fn handle(&self) -> *mut bindings::nixl_capi_xfer_dlist_handle_s {
30+
self.inner
31+
}
32+
}
33+
34+
impl Drop for XferDlistHandle {
35+
fn drop(&mut self) {
36+
unsafe {
37+
nixl_capi_release_xfer_dlist_handle(self.agent.as_ptr(),
38+
self.handle());
39+
}
40+
}
41+
}

‎src/bindings/rust/src/lib.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ use bindings::{
6969
nixl_capi_xfer_dlist_print, nixl_capi_reg_dlist_is_sorted, nixl_capi_gen_notif, nixl_capi_estimate_xfer_cost,
7070
nixl_capi_query_mem, nixl_capi_create_query_resp_list, nixl_capi_destroy_query_resp_list,
7171
nixl_capi_query_resp_list_size, nixl_capi_query_resp_list_has_value,
72-
nixl_capi_query_resp_list_get_params,
72+
nixl_capi_query_resp_list_get_params, nixl_capi_prep_xfer_dlist, nixl_capi_release_xfer_dlist_handle,
73+
nixl_capi_make_xfer_req
7374
};
7475

7576
// Re-export status codes
@@ -111,6 +112,10 @@ pub enum NixlError {
111112
RegDescListCreationFailed,
112113
#[error("Failed to add registration descriptor")]
113114
RegDescAddFailed,
115+
#[error("Failed to create XferDlistHandle")]
116+
FailedToCreateXferDlistHandle,
117+
#[error("Failed to create backend")]
118+
FailedToCreateBackend,
114119
}
115120

116121
/// A safe wrapper around NIXL memory list
@@ -192,6 +197,7 @@ pub struct Backend {
192197
unsafe impl Send for Backend {}
193198
unsafe impl Sync for Backend {}
194199

200+
195201
/// A safe wrapper around NIXL optional arguments
196202
pub struct OptArgs {
197203
inner: NonNull<bindings::nixl_capi_opt_args_s>,

‎src/bindings/rust/stubs.cpp‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
extern "C" {
2626

27+
// clang-format off
2728
// Internal struct definitions to match our opaque types
2829
// These are now stubs as their internal details are no longer used.
2930
struct nixl_capi_agent_s { /* empty */ };
@@ -37,6 +38,9 @@ struct nixl_capi_xfer_dlist_s { /* empty */ };
3738
struct nixl_capi_reg_dlist_s { /* empty */ };
3839
struct nixl_capi_xfer_req_s { /* empty */ };
3940
struct nixl_capi_notif_map_s { /* empty */ };
41+
struct nixl_capi_xfer_dlist_handle_s { /* empty */ };
42+
43+
// clang-format on
4044

4145
nixl_capi_status_t
4246
nixl_capi_stub_abort()
@@ -70,6 +74,35 @@ nixl_capi_load_remote_md(nixl_capi_agent_t agent, const void* data, size_t len,
7074
return nixl_capi_stub_abort();
7175
}
7276

77+
nixl_capi_status_t
78+
nixl_capi_prep_xfer_dlist(nixl_capi_agent_t agent,
79+
const char *agent_name,
80+
nixl_capi_xfer_dlist_t descs,
81+
nixl_capi_xfer_dlist_handle_t *dlist_handle,
82+
nixl_capi_opt_args_t opt_args) {
83+
return nixl_capi_stub_abort();
84+
}
85+
86+
nixl_capi_status_t
87+
nixl_capi_release_xfer_dlist_handle(nixl_capi_agent_t agent,
88+
nixl_capi_xfer_dlist_handle_t dlist_handle) {
89+
return nixl_capi_stub_abort();
90+
}
91+
92+
nixl_capi_status_t
93+
nixl_capi_make_xfer_req(nixl_capi_agent_t agent,
94+
nixl_capi_xfer_op_t operation,
95+
nixl_capi_xfer_dlist_handle_t local_descs,
96+
const int *local_indices,
97+
size_t local_indices_count,
98+
nixl_capi_xfer_dlist_handle_t remote_descs,
99+
const int *remote_indices,
100+
size_t remote_indices_count,
101+
nixl_capi_xfer_req_t *req_hndl,
102+
nixl_capi_opt_args_t opt_args) {
103+
return nixl_capi_stub_abort();
104+
}
105+
73106
nixl_capi_status_t
74107
nixl_capi_invalidate_remote_md(nixl_capi_agent_t agent, const char* remote_agent)
75108
{

0 commit comments

Comments
 (0)