Browse Source

add policy manager

AlpinDale 1 year ago
parent
commit
ef4a42b05e
2 changed files with 43 additions and 1 deletions
  1. 1 1
      aphrodite/processing/block_manager.py
  2. 42 0
      aphrodite/processing/policy.py

+ 1 - 1
aphrodite/processing/block_manager.py

@@ -1,5 +1,5 @@
 """A block manager that manages token blocks."""
 """A block manager that manages token blocks."""
-from typing import Dict, List, Optional, Set, Tuple, Any
+from typing import Dict, List, Optional, Set, Tuple
 from aphrodite.block import PhysicalTokenBlock
 from aphrodite.block import PhysicalTokenBlock
 from aphrodite.sequence import Sequence, SequenceGroup, SequenceStatus
 from aphrodite.sequence import Sequence, SequenceGroup, SequenceStatus
 from aphrodite.utils import Device
 from aphrodite.utils import Device

+ 42 - 0
aphrodite/processing/policy.py

@@ -0,0 +1,42 @@
+from typing import List
+
+from aphrodite.sequence import SequenceGroup
+
+class Policy:
+
+    def get_priority(
+        self,
+        now: float,
+        seq_group: SequenceGroup,
+    ) -> float:
+        raise NotImplementedError
+
+    def sort_by_priority(
+        self,
+        now: float,
+        seq_groups: List[SequenceGroup],
+    ) -> List[SequenceGroup]:
+        return sorted(
+            seq_groups,
+            key=lambda seq_group: self.get_priority(now, seq_group),
+            reverse=True,
+        )
+
+class FCFS(Policy):
+
+    def get_priority(
+        self,
+        now: float,
+        seq_group: SequenceGroup,
+    ) -> float:
+        return now - seq_group.arrival_time
+
+class PolicyFactory:
+
+    _POLICY_REGISTRY = {
+        'fcfs': FCFS,
+    }
+
+    @classmethod
+    def get_policy(cls, policy_name: str, **kwargs) -> Policy:
+        return cls._POLICY_REGISTRY[policy_name](**kwargs)