Tako visibility refactoring

Author: spirali

crates/tako/src/control.rs

@@ -84,7 +84,7 @@ impl ServerRef {
 
     pub fn stop_worker(&self, worker_id: WorkerId) -> crate::Result<()> {
         let mut core = self.core_ref.get_mut();
-        if let Some(ref mut worker) = core.get_worker_mut(worker_id) {
+        if let Some(ref mut worker) = core.find_worker_mut(worker_id) {
             if !worker.is_stopping() {
                 worker.set_stop(LostWorkerReason::Stopped);
                 let mut comm = self.comm_ref.get_mut();

crates/tako/src/internal/common/resources/mod.rs

@@ -11,10 +11,12 @@ pub use descriptor::{
     ResourceDescriptor, ResourceDescriptorCoupling, ResourceDescriptorCouplingItem,
     ResourceDescriptorItem, ResourceDescriptorKind,
 };
+
 pub use map::{
     AMD_GPU_RESOURCE_NAME, CPU_RESOURCE_ID, CPU_RESOURCE_NAME, GlobalResourceMapping,
-    MEM_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME, ResourceRqMap,
+    MEM_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME,
 };
+
 pub use request::{
     AllocationRequest, ResourceAllocRequest, ResourceRequest, ResourceRequestEntries,
     ResourceRequestVariants, TimeRequest,

crates/tako/src/internal/common/stablemap.rs

@@ -86,6 +86,7 @@ where
     }
 
     #[inline]
+    #[allow(dead_code)]
     pub fn get<Q>(&self, key: &Q) -> &V
     where
         K: Borrow<Q>,

crates/tako/src/internal/scheduler/gap_cache.rs

@@ -160,11 +160,10 @@ fn compute_gap(
 #[cfg(test)]
 mod tests {
     use super::*;
-    
-    
+
     use crate::tests::utils::env::TestEnv;
     use crate::tests::utils::resources::ResBuilder;
-    
+
     use crate::tests::utils::worker::WorkerBuilder;
 
     #[test]

crates/tako/src/internal/scheduler/main.rs

@@ -53,11 +53,3 @@ pub(crate) fn run_scheduling(core: &mut Core, comm: &mut CommSender, now: Instan
     );
     comm.reset_scheduling_flag();
 }
-
-/*pub(crate) fn collect_assigned_not_running_tasks(core: &mut Core) -> Vec<TaskId> {
-    let mut result = Vec::new();
-    for worker in core.get_workers_mut() {
-        worker.collect_assigned_non_running_tasks(&mut result);
-    }
-    result
-}*/

crates/tako/src/internal/scheduler/mapping.rs

@@ -8,14 +8,14 @@ use std::cmp::Reverse;
 
 #[derive(Debug, Default)]
 #[cfg_attr(test, derive(Clone, PartialEq, PartialOrd, Eq, Ord))]
-pub struct WorkerTaskUpdate {
+pub(crate) struct WorkerTaskUpdate {
     pub(crate) assigned: Vec<(TaskId, ResourceVariantId)>,
     pub(crate) prefills: Vec<TaskId>,
     pub(crate) retracts: Vec<TaskId>,
 }
 
 #[derive(Debug, Default)]
-pub struct WorkerTaskMapping {
+pub(crate) struct WorkerTaskMapping {
     pub(crate) workers: Map<WorkerId, WorkerTaskUpdate>,
     pub(crate) mn_tasks_to_workers: Vec<TaskId>,
 }

crates/tako/src/internal/server/core.rs

@@ -39,7 +39,7 @@ pub(crate) struct CoreSplit<'a> {
 }
 
 #[derive(Default)]
-pub struct Core {
+pub(crate) struct Core {
     tasks: TaskMap,
     workers: WorkerMap,
     resource_map: GlobalResourceMapping,
@@ -109,12 +109,12 @@ impl Core {
         }
     }
 
-    pub fn new_worker_id(&mut self) -> WorkerId {
+    pub(crate) fn new_worker_id(&mut self) -> WorkerId {
         self.worker_id_counter += 1;
         WorkerId::new(self.worker_id_counter)
     }
 
-    pub fn worker_counter(&self) -> u32 {
+    pub(crate) fn worker_counter(&self) -> u32 {
         self.worker_id_counter
     }
 
@@ -175,26 +175,22 @@ impl Core {
     }
 
     #[inline]
-    pub fn get_worker_by_id(&self, id: WorkerId) -> Option<&Worker> {
-        self.workers.get(&id)
-    }
-
-    #[inline]
-    pub fn get_worker_by_id_or_panic(&self, id: WorkerId) -> &Worker {
+    pub fn get_worker(&self, id: WorkerId) -> &Worker {
         self.workers.get(&id).unwrap_or_else(|| {
             panic!("Asking for invalid worker id={id}");
         })
     }
 
     #[inline]
-    pub fn get_worker_mut_by_id_or_panic(&mut self, id: WorkerId) -> &mut Worker {
+    #[cfg(test)]
+    pub fn get_worker_mut(&mut self, id: WorkerId) -> &mut Worker {
         self.workers.get_mut(&id).unwrap_or_else(|| {
             panic!("Asking for invalid worker id={id}");
         })
     }
 
     #[inline]
-    pub fn get_worker_mut(&mut self, id: WorkerId) -> Option<&mut Worker> {
+    pub fn find_worker_mut(&mut self, id: WorkerId) -> Option<&mut Worker> {
         self.workers.get_mut(&id)
     }
 
@@ -203,21 +199,11 @@ impl Core {
         self.workers.values()
     }
 
-    #[inline]
-    pub fn get_workers_mut(&mut self) -> impl Iterator<Item = &mut Worker> {
-        self.workers.values_mut()
-    }
-
     #[inline]
     pub fn get_worker_map(&self) -> &WorkerMap {
         &self.workers
     }
 
-    #[inline]
-    pub fn has_workers(&self) -> bool {
-        !self.workers.is_empty()
-    }
-
     pub fn add_task(&mut self, task: Task, retracted: &mut Vec<TaskId>) {
         if task.is_ready() {
             self.task_queues.add_ready_task(&task, retracted);
@@ -260,11 +246,6 @@ impl Core {
         &self.tasks
     }
 
-    #[inline]
-    pub fn task_map_mut(&mut self) -> &mut TaskMap {
-        &mut self.tasks
-    }
-
     #[inline]
     pub fn get_task(&self, task_id: TaskId) -> &Task {
         self.tasks.get_task(task_id)
@@ -453,13 +434,6 @@ impl Core {
         self.resource_map.get_or_create_resource_id(name)
     }
 
-    pub fn convert_client_resource_rq(
-        &mut self,
-        resources: &crate::gateway::ResourceRequestVariants,
-    ) -> ResourceRequestVariants {
-        self.resource_map.convert_client_resource_rq(resources)
-    }
-
     #[inline]
     pub fn resource_map(&self) -> &GlobalResourceMapping {
         &self.resource_map
@@ -511,13 +485,10 @@ impl Core {
 mod tests {
     use crate::internal::server::core::Core;
     use crate::internal::server::task::Task;
-    
-    use crate::internal::server::worker::Worker;
-    use crate::internal::server::workergroup::WorkerGroup;
 
-    
+    use crate::internal::server::workergroup::WorkerGroup;
 
-    use crate::{TaskId, WorkerId};
+    use crate::TaskId;
 
     impl Core {
         pub fn worker_group(&self, group_name: &str) -> Option<&WorkerGroup> {
@@ -532,18 +503,6 @@ mod tests {
             }
         }
 
-        pub fn assert_worker_condition<F: Fn(&Worker) -> bool>(
-            &self,
-            worker_ids: &[WorkerId],
-            op: F,
-        ) {
-            for worker_id in worker_ids {
-                if !op(self.get_worker_by_id_or_panic(*worker_id)) {
-                    panic!("Worker {worker_id} does not satisfy the condition");
-                }
-            }
-        }
-
         pub fn assert_waiting(&self, task_ids: &[TaskId]) {
             self.assert_task_condition(task_ids, |t| t.is_waiting());
         }

crates/tako/src/internal/server/rpc.rs

@@ -269,7 +269,7 @@ async fn worker_rpc_loop(
     let mut core = core_ref.get_mut();
     let mut comm = comm_ref.get_mut();
     let reason = core
-        .get_worker_by_id_or_panic(worker_id)
+        .get_worker(worker_id)
         .stop_reason
         .map(|(r, _)| r)
         .unwrap_or(reason);
@@ -299,7 +299,7 @@ pub(crate) async fn worker_receive_loop<
                 on_retract_response(&mut core, &mut *comm, worker_id, &msg.retracted);
             }
             FromWorkerMessage::Heartbeat => {
-                if let Some(worker) = core.get_worker_mut(worker_id) {
+                if let Some(worker) = core.find_worker_mut(worker_id) {
                     log::trace!("Heartbeat received, worker={worker_id}");
                     worker.last_heartbeat = Instant::now();
                 };

crates/tako/src/internal/tests/test_reactor.rs

@@ -523,12 +523,7 @@ fn test_task_mn_fail() {
     comm.emptiness_check();
     assert!(rt.core().find_task(1.into()).is_none());
     for w in &ws {
-        assert!(
-            rt.core()
-                .get_worker_by_id_or_panic((*w).into())
-                .mn_assignment()
-                .is_none()
-        );
+        assert!(rt.core().get_worker((*w).into()).mn_assignment().is_none());
     }
 }
 
@@ -645,7 +640,7 @@ fn lost_worker_with_running_and_assign_tasks() {
 }
 
 fn check_worker_tasks_exact(core: &Core, worker_id: WorkerId, tasks: &[TaskId]) {
-    let worker = core.get_worker_by_id_or_panic(worker_id.into());
+    let worker = core.get_worker(worker_id.into());
     let sn = worker.sn_assignment().unwrap();
     assert_eq!(sn.assign_tasks.len(), tasks.len());
     for task in tasks {
@@ -654,7 +649,7 @@ fn check_worker_tasks_exact(core: &Core, worker_id: WorkerId, tasks: &[TaskId])
 }
 
 fn worker_has_task(core: &Core, worker_id: WorkerId, task_id: TaskId) -> bool {
-    core.get_worker_by_id_or_panic(worker_id.into())
+    core.get_worker(worker_id.into())
         .sn_assignment()
         .unwrap()
         .assign_tasks

crates/tako/src/internal/tests/test_scheduler_mn.rs

@@ -122,12 +122,7 @@ fn test_schedule_mn_simple() {
 
     assert!(!rt.task_exists(t3));
     for w in ws3 {
-        assert!(
-            rt.core()
-                .get_worker_by_id_or_panic(w)
-                .mn_assignment()
-                .is_none()
-        );
+        assert!(rt.core().get_worker(w).mn_assignment().is_none());
     }
 
     let mut comm = rt.schedule();