diff --git a/src/discovery.rs b/src/discovery/mod.rs
similarity index 98%
rename from src/discovery.rs
rename to src/discovery/mod.rs
index f220ab7..cea2eec 100644
--- a/src/discovery.rs
+++ b/src/discovery/mod.rs
@@ -5,7 +5,9 @@ use socket2::{Domain, Protocol, Socket, Type};
 use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 use tokio::net::UdpSocket;
-use crate::udp_listen::NodeInfo;
+use udp_listen::NodeInfo;
+
+mod udp_listen;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct DiscoveryMessage {
diff --git a/src/udp_listen.rs b/src/discovery/udp_listen.rs
similarity index 98%
rename from src/udp_listen.rs
rename to src/discovery/udp_listen.rs
index 6cd7546..9fc5abe 100644
--- a/src/udp_listen.rs
+++ b/src/discovery/udp_listen.rs
@@ -1,5 +1,4 @@
 use crate::orchestration::PeerHandle;
-use crate::{discovery::DiscoveryMessage, node_service::HealthCheckRequest};
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::time::Duration;
@@ -9,6 +8,7 @@ use tokio::select;
 use tokio::sync::mpsc::UnboundedSender;
 use tonic::transport::Error;
 use tracing::{debug, error, info};
+use crate::discovery::DiscoveryMessage;
 use crate::topology::DeviceCapabilities;
 
 #[derive(Debug, Clone)]
diff --git a/src/main.rs b/src/main.rs
index 1130461..4804b32 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,7 +1,6 @@
 mod topology;
 mod orchestration;
 mod discovery;
-mod udp_listen;
 mod network;
 
 use serde::{Deserialize, Serialize};
@@ -13,7 +12,7 @@ use crate::node_service::{
     PromptRequest, SendOpaqueStatusRequest, SendResultRequest, Tensor, Topology as TopologyProto,
 };
 use node_service::node_service_server::{NodeService, NodeServiceServer};
-use node_service::{TensorRequest};
+use node_service::TensorRequest;
 use topology::Topology;
 
 pub mod node_service {