TheRock/rocm-systems/emulation/mirage/rocjitsu/tests/daemon.rs

//! Integration tests for the shared in-process RocJitsu daemon API.
//!
//! Stands up the `librocjitsu` daemon through the Rust lifecycle wrapper,
//! connects a client speaking the daemon RPC protocol, performs the
//! handshake the KMD interposer would, and verifies the daemon serves a
//! live simulated device. The whole test is skipped when no rocjitsu KMD
//! library is discoverable on this machine (install rocjitsu under
//! `$ROCM_HOME` or a sibling monorepo build to exercise it).

use std::io::{Read, Write};
use std::os::unix::net::UnixStream;

use mirage_core::common::MaybeRef;
use mirage_core::emulator::{EmulatorDaemon, EmulatorDef, ExecMode};
use mirage_rocjitsu::daemon::Daemon;
use mirage_rocjitsu::{kmd_config, kmd_preload};
use rocjitsu_sys::RjDaemonStatus;

/// Build the 16-byte RPC header the wire protocol uses.
fn header(opcode: u16, request_id: u32, payload_bytes: u32, result: i32) -> [u8; 16] {
    let mut h = [0u8; 16];
    h[0..2].copy_from_slice(&opcode.to_ne_bytes());
    h[4..8].copy_from_slice(&request_id.to_ne_bytes());
    h[8..12].copy_from_slice(&payload_bytes.to_ne_bytes());
    h[12..16].copy_from_slice(&result.to_ne_bytes());
    h
}

fn read_exact(stream: &mut UnixStream, buf: &mut [u8]) {
    stream.read_exact(buf).expect("read response");
}

#[test]
fn daemon_serves_handshake() {
    let _g = mirage_core::paths::test_env_lock();
    let tmp = tempfile::tempdir().unwrap();
    mirage_core::paths::set_test_root(tmp.path());

    // The daemon needs the KMD library (it exports the rj_vm_* API);
    // skip cleanly when rocjitsu is not installed.
    let Some(lib) = kmd_preload() else {
        eprintln!("rocjitsu KMD library not found; skipping daemon test");
        return;
    };

    // Synthesise a real sim config from a builtin agent.
    let agent_report = mirage_builtin::ensure_agents(false).unwrap();
    let agent_name = agent_report
        .iter()
        .map(|(n, _)| n.clone())
        .next()
        .expect("at least one builtin agent");
    let def = EmulatorDef {
        emulator: "rocjitsu".to_string(),
        plugins: Default::default(),
        exec_mode: ExecMode::Functional,
        options: Default::default(),
        topology: MaybeRef::Owned(mirage_core::topology::TopologyDef {
            num_nodes: 1,
            gpus_per_node: 1,
            agent: MaybeRef::Ref(agent_name),
        }),
    };
    let config = kmd_config(&def, None).expect("sim config should materialise");

    let runtime_dir = tmp.path().join("rt");
    let daemon = Daemon::start(&lib, &config, &runtime_dir).expect("daemon should start");
    assert_eq!(daemon.status(), RjDaemonStatus::Running);
    assert!(
        daemon.socket_path().exists(),
        "daemon socket should be bound"
    );

    // Connect and perform the handshake the interposer would.
    let mut stream = UnixStream::connect(daemon.socket_path()).expect("connect to daemon");
    stream
        .write_all(&header(0 /* HANDSHAKE */, 1, 0, 0))
        .expect("send handshake");

    let mut resp = [0u8; 16];
    read_exact(&mut stream, &mut resp);
    let payload_bytes = u32::from_ne_bytes(resp[8..12].try_into().unwrap()) as usize;
    let result = i32::from_ne_bytes(resp[12..16].try_into().unwrap());
    assert_eq!(result, 0, "handshake should succeed");
    assert!(
        payload_bytes >= 16,
        "handshake payload should carry the response struct"
    );

    let mut payload = vec![0u8; payload_bytes];
    read_exact(&mut stream, &mut payload);
    let version = u32::from_ne_bytes(payload[0..4].try_into().unwrap());
    let topo_len = u32::from_ne_bytes(payload[8..12].try_into().unwrap()) as usize;
    let drm_len = u32::from_ne_bytes(payload[12..16].try_into().unwrap()) as usize;
    assert_eq!(version, 3, "protocol version should match the interposer");
    // RpcHandshakeResponse is 328 bytes (16 fixed fields + 312-byte
    // RpcGpuInfo), then the topology and DRM path strings.
    assert_eq!(
        payload_bytes,
        328 + topo_len + drm_len,
        "handshake payload framing should be self-consistent"
    );
    assert!(topo_len > 0, "daemon should report a sysfs topology path");

    // Cleanly close the client session.
    stream
        .write_all(&header(2 /* CLOSE */, 2, 0, 0))
        .expect("send close");
    let mut close_resp = [0u8; 16];
    read_exact(&mut stream, &mut close_resp);

    // Shutting the daemon down removes its socket.
    let socket_path = daemon.socket_path().to_path_buf();
    Box::new(daemon).stop();
    assert!(
        !socket_path.exists(),
        "daemon socket should be removed on shutdown"
    );
}

/// A second handshake on a fresh connection must also succeed: the
/// daemon serves many clients over its lifetime (one per workload
/// open of `/dev/kfd`).
#[test]
fn daemon_serves_multiple_clients() {
    let _g = mirage_core::paths::test_env_lock();
    let tmp = tempfile::tempdir().unwrap();
    mirage_core::paths::set_test_root(tmp.path());

    let Some(lib) = kmd_preload() else {
        eprintln!("rocjitsu KMD library not found; skipping daemon test");
        return;
    };

    let agent_report = mirage_builtin::ensure_agents(false).unwrap();
    let agent_name = agent_report.iter().map(|(n, _)| n.clone()).next().unwrap();
    let def = EmulatorDef {
        emulator: "rocjitsu".to_string(),
        plugins: Default::default(),
        exec_mode: ExecMode::Functional,
        options: Default::default(),
        topology: MaybeRef::Owned(mirage_core::topology::TopologyDef {
            num_nodes: 1,
            gpus_per_node: 1,
            agent: MaybeRef::Ref(agent_name),
        }),
    };
    let config = kmd_config(&def, None).unwrap();
    let runtime_dir = tmp.path().join("rt");
    let daemon = Daemon::start(&lib, &config, &runtime_dir).expect("daemon should start");
    assert_eq!(daemon.status(), RjDaemonStatus::Running);

    for req_id in 0..3u32 {
        let mut stream = UnixStream::connect(daemon.socket_path()).expect("connect");
        stream.write_all(&header(0, req_id, 0, 0)).unwrap();
        let mut resp = [0u8; 16];
        read_exact(&mut stream, &mut resp);
        let payload_bytes = u32::from_ne_bytes(resp[8..12].try_into().unwrap()) as usize;
        assert!(payload_bytes >= 16, "client {req_id} handshake framed");
        let mut payload = vec![0u8; payload_bytes];
        read_exact(&mut stream, &mut payload);
        stream.write_all(&header(2, req_id, 0, 0)).unwrap();
        let mut close_resp = [0u8; 16];
        read_exact(&mut stream, &mut close_resp);
    }
}