NOTE: this keeps track of the state ACCORDING TO THE SCHEDULER - may not actually reflect the true state of requests. E.g., a "finished" sequence here may still have outstanding tokens to actually compute, and may not yet have had its request returned to the user. Also, we're
| 149 | |
| 150 | @dataclass |
| 151 | class SchedulingQueue: |
| 152 | """ |
| 153 | NOTE: this keeps track of the state ACCORDING TO THE SCHEDULER - may not |
| 154 | actually reflect the true state of requests. E.g., a "finished" sequence |
| 155 | here may still have outstanding tokens to actually compute, and may not |
| 156 | yet have had its request returned to the user. |
| 157 | |
| 158 | |
| 159 | Also, we're using the fact that Python dicts are ordered, so that these dicts |
| 160 | are insertion-order queues that we can also index by id. |
| 161 | """ |
| 162 | |
| 163 | decoding_seqs: dict[str, Sequence] = field(default_factory=dict) |
| 164 | prefilling_seqs: dict[str, Sequence] = field(default_factory=dict) |
| 165 | queued_seqs: dict[str, Sequence] = field(default_factory=dict) |
| 166 | |
| 167 | def get(self, sid: str): |
| 168 | out = ( |
| 169 | self.decoding_seqs.get(sid) |
| 170 | or self.prefilling_seqs.get(sid) |
| 171 | or self.queued_seqs.get(sid) |
| 172 | ) |
| 173 | if out is None: |
| 174 | raise ValueError(f"Request {sid} not found") |
| 175 | return out |
| 176 | |
| 177 | def __getitem__(self, sid: str): |
| 178 | return self.get(sid) |
| 179 | |
| 180 | def get_decoding(self, sid: str): |
| 181 | return self.decoding_seqs[sid] |
| 182 | |
| 183 | def get_prefilling(self, sid: str): |
| 184 | return self.prefilling_seqs[sid] |
| 185 | |
| 186 | def get_queued(self, sid: str): |
| 187 | return self.queued_seqs[sid] |
| 188 | |
| 189 | def add_decoding(self, seq: Sequence): |
| 190 | self.decoding_seqs[seq.id] = seq |
| 191 | |
| 192 | def add_prefilling(self, seq: Sequence): |
| 193 | self.prefilling_seqs[seq.id] = seq |
| 194 | |
| 195 | def add_queued(self, seq: Sequence): |
| 196 | self.queued_seqs[seq.id] = seq |
| 197 | |
| 198 | def remove_decoding(self, sid: str): |
| 199 | self.decoding_seqs.pop(sid) |
| 200 | |
| 201 | def remove_prefilling(self, sid: str): |
| 202 | self.prefilling_seqs.pop(sid) |
| 203 | |
| 204 | def remove_queued(self, sid: str): |
| 205 | self.queued_seqs.pop(sid) |
| 206 | |
| 207 | def remove(self, sid: str): |
| 208 | if self.in_decoding(sid): |