Integrate priorities into offloading

Author: jthomson04

lib/llm/src/block_manager/block.rs

@@ -82,6 +82,10 @@ pub trait BlockMetadata: Default + std::fmt::Debug + Clone + Ord + Send + Sync +
     /// Resets the metadata to the default value
     /// If called, the [BlockMetadata::is_reset()] should return true
     fn reset_metadata(&mut self);
+
+    /// The offload priority of the block. Higher priority blocks are offloaded first.
+    /// If the block should not be offloaded, return None.
+    fn offload_priority(&self) -> Option<u64>;
 }
 
 /// Marker trait for types that are mutable blocks
@@ -536,6 +540,10 @@ impl BlockMetadata for BasicMetadata {
     fn reset_metadata(&mut self) {
         self.priority = 0;
     }
+
+    fn offload_priority(&self) -> Option<u64> {
+        Some(self.priority as u64)
+    }
 }
 /// Collection that holds shared storage and layout
 #[derive(Debug)]

lib/llm/src/block_manager/pool/inactive.rs

@@ -518,6 +518,10 @@ pub(crate) mod tests {
         fn reset_metadata(&mut self) {
             self.priority = 0;
         }
+
+        fn offload_priority(&self) -> Option<u64> {
+            Some(self.priority as u64)
+        }
     }
 
     type TestPriorityKey = PriorityKey<TestMetadata>;

lib/llm/src/block_manager/pool/state.rs

@@ -179,9 +179,10 @@ impl<S: Storage, M: BlockMetadata> State<S, M> {
 
             let immutable = self.active.register(mutable)?;
 
-            // TODO: Make a way to set meaningful priority values, and maybe don't enqueue offloads for every registered block.
             if offload {
-                immutable.enqueue_offload(0).await.unwrap();
+                if let Some(priority) = immutable.metadata().offload_priority() {
+                    immutable.enqueue_offload(priority).await.unwrap();
+                }
             }
 
             immutable_blocks.push(immutable);