A local in-memory and on-disk WorkQueue implementation, which uses a local workspace directory to store the queue index, lock files, and completed results for persistent resumption across process restarts.
| 161 | |
| 162 | |
| 163 | class LocalWorkQueue(WorkQueue): |
| 164 | """ |
| 165 | A local in-memory and on-disk WorkQueue implementation, which uses |
| 166 | a local workspace directory to store the queue index, lock files, |
| 167 | and completed results for persistent resumption across process restarts. |
| 168 | """ |
| 169 | |
| 170 | def __init__(self, workspace_path: str): |
| 171 | """ |
| 172 | Initialize the local work queue. |
| 173 | |
| 174 | Args: |
| 175 | workspace_path: Local directory path where the queue index, |
| 176 | results, and locks are stored. |
| 177 | """ |
| 178 | self.workspace_path = os.path.abspath(workspace_path) |
| 179 | os.makedirs(self.workspace_path, exist_ok=True) |
| 180 | |
| 181 | # Local index file (compressed) |
| 182 | self._index_path = os.path.join(self.workspace_path, "work_index_list.csv.zstd") |
| 183 | |
| 184 | # Output directory for completed tasks |
| 185 | self._results_dir = os.path.join(self.workspace_path, "results") |
| 186 | os.makedirs(self._results_dir, exist_ok=True) |
| 187 | |
| 188 | # Directory for lock files |
| 189 | self._locks_dir = os.path.join(self.workspace_path, "worker_locks") |
| 190 | os.makedirs(self._locks_dir, exist_ok=True) |
| 191 | |
| 192 | # Internal queue |
| 193 | self._queue: Queue[Any] = Queue() |
| 194 | |
| 195 | async def populate_queue(self, work_paths: List[str], items_per_group: int) -> None: |
| 196 | """ |
| 197 | Add new items to the work queue (local version). |
| 198 | |
| 199 | Args: |
| 200 | work_paths: Each individual path (local in this context) |
| 201 | that we will process over |
| 202 | items_per_group: Number of items to group together in a single work item |
| 203 | """ |
| 204 | # Treat them as local paths, but keep variable name for consistency |
| 205 | all_paths = set(work_paths) |
| 206 | logger.info(f"Found {len(all_paths):,} total paths") |
| 207 | |
| 208 | # Load existing work groups from local index |
| 209 | existing_lines = await asyncio.to_thread(download_zstd_csv_local, self._index_path) |
| 210 | existing_groups = {} |
| 211 | for line in existing_lines: |
| 212 | if line.strip(): |
| 213 | parts = line.strip().split(",") |
| 214 | group_hash = parts[0] |
| 215 | group_paths = parts[1:] |
| 216 | existing_groups[group_hash] = group_paths |
| 217 | |
| 218 | existing_path_set = {p for paths in existing_groups.values() for p in paths} |
| 219 | new_paths = all_paths - existing_path_set |
| 220 | logger.info(f"{len(new_paths):,} new paths to add to the workspace") |