Add cpu_device_tasks to tcp server

ml-energy · Oct 27, 2024 · 908ab1e · 908ab1e
1 parent b95d990
commit 908ab1e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/zeusd/src/main.rs b/zeusd/src/main.rs
@@ -50,7 +50,7 @@ async fn main() -> anyhow::Result<()> {
             let listener = TcpListener::bind(&config.tcp_bind_address)?;
             tracing::info!("Listening on {}", &listener.local_addr()?);
 
-            start_server_tcp(listener, device_tasks, num_workers)?.await?;
+            start_server_tcp(listener, device_tasks, cpu_device_tasks.clone(), num_workers)?.await?;
         }
     }
 

diff --git a/zeusd/src/startup.rs b/zeusd/src/startup.rs
@@ -116,6 +116,7 @@ pub fn start_server_uds(
 pub fn start_server_tcp(
     listener: TcpListener,
     device_tasks: GpuManagementTasks,
+    cpu_device_tasks: CpuManagementTasks,
     num_workers: usize,
 ) -> std::io::Result<Server> {
     let server = HttpServer::new(move || {
@@ -124,6 +125,7 @@ pub fn start_server_tcp(
             .service(web::scope("/gpu").configure(gpu_routes))
             .service(web::scope("/cpu").configure(cpu_routes))
             .app_data(web::Data::new(device_tasks.clone()))
+            .app_data(web::Data::new(cpu_device_tasks.clone()))
     })
     .workers(num_workers)
     .listen(listener)?