Skip to content

Commit 2d80812

Browse files
committed
Worker version 20260626.01. Complete rewrite of database layer.
1 parent 6e767fa commit 2d80812

25 files changed

Lines changed: 3769 additions & 670 deletions

File tree

.gitignore

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ target/packer/out
1111
.stignore
1212
pyrightconfig.json
1313
TODO
14-
run.sh
14+
run*.sh
1515
.env
1616
__pycache__/
1717
*.py[co]
1818
client/warcupload/testdata
1919
upload-space/
20+
*.tar
21+
PIPELINES

client/worker/Dockerfile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ RUN apt-get update \
1212
&& apt-get clean
1313

1414
# 06fc825d3fbed9801110b2d3f562c44d72940862 is known to work.
15-
RUN pip3 install --break-system-packages rethinkdb git+https://github.com/internetarchive/brozzler@06fc825d3fbed9801110b2d3f562c44d72940862 websockets doublethink yt-dlp aiofiles
15+
RUN pip3 install --break-system-packages rethinkdb git+https://github.com/internetarchive/brozzler@06fc825d3fbed9801110b2d3f562c44d72940862 websockets doublethink yt-dlp aiofiles uuid_utils
1616
RUN pip3 install --break-system-packages --upgrade rethinkdb
1717

1818
RUN mkdir /app

client/worker/app.py

Lines changed: 36 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -105,20 +105,24 @@ async def warcprox_cleanup():
105105
except Exception:
106106
pass
107107

108-
async def run_job(ws: Websocket, full_job: dict, url: str, warc_prefix: str, ua: str, custom_js: typing.Optional[str], info_url: str):
109-
tries = full_job['_current_attempt']
110-
id = full_job['id']
108+
async def run_job(ws: Websocket, full_job: dict, info_url: str):
109+
job_id = full_job['job_id']
110+
page_id = full_job['page_id']
111+
attempt_id = full_job['attempt_id']
112+
assert "_" not in job_id
113+
warc_prefix = "mnbot-brozzler-" + job_id.replace("-", "_")
111114
#dedup_bucket = f"dedup-{id}-{tries}"
112115
dedup_bucket = ""
113-
stats_bucket = f"stats-{id}-{tries}"
116+
stats_bucket = f"stats-{job_id}"
114117
job = Job(
118+
attempt_id = attempt_id,
115119
full_job = full_job,
116-
url = url,
120+
url = full_job['payload'],
117121
warc_prefix = warc_prefix,
118122
dedup_bucket = dedup_bucket,
119123
stats_bucket = stats_bucket,
120-
ua = ua,
121-
custom_js = custom_js,
124+
ua = full_job['settings']['ua'],
125+
custom_js = full_job['settings']['custom_js'],
122126
cookie_jar = None,
123127
mnbot_info_url = info_url
124128
)
@@ -129,7 +133,10 @@ async def run_job(ws: Websocket, full_job: dict, url: str, warc_prefix: str, ua:
129133
PYTHON,
130134
os.path.join(os.path.dirname(sys.argv[0]), "browse.py"),
131135
str(pwrite),
132-
id, # useful for ps
136+
# These arguments are passed only because they are useful for ps
137+
job_id,
138+
page_id,
139+
attempt_id,
133140
stdin = subprocess.PIPE,
134141
stdout = subprocess.PIPE,
135142
stderr = subprocess.STDOUT,
@@ -174,16 +181,8 @@ async def run_job(ws: Websocket, full_job: dict, url: str, warc_prefix: str, ua:
174181
res = json.loads(res)
175182
type = res['type']
176183
payload = res['payload']
177-
if type in ("status_code", "outlinks", "final_url", "requisites", "custom_js"):
178-
await ws.store_result(id, type, tries, payload)
179-
elif type == "screenshot":
180-
# Don't tell the tracker to decode the thumbnail
181-
# if there isn't a thumbnail
182-
decode_fields = [k for k in ("full", "thumb") if payload[k]]
183-
await ws.store_result(id, type, tries, payload, decode_fields)
184-
elif type == "cjs_screenshot":
185-
decode_fields = ["full"]
186-
await ws.store_result(id, type, tries, payload, decode_fields)
184+
if type in ("status_code", "outlinks", "final_url", "requisites", "custom_js", "screenshot", "cjs_screenshot"):
185+
await ws.store_result(attempt_id, type, payload)
187186
elif type == "error":
188187
# Cancel tasks, since both stdout and pread are about to get closed.
189188
# Failing to do this results in a "Task exception was never retrieved"
@@ -233,6 +232,7 @@ async def ping_occasionally():
233232
pass
234233
await asyncio.sleep(15)
235234

235+
# Don't change this without changing the hardcoded slot=0 below and the hardcoded num_slots=1 in tracker.py.
236236
MAX_WORKERS = 1
237237
workers: dict[asyncio.Task, tuple[TaskType, dict | None]] = dict()
238238

@@ -266,41 +266,36 @@ def handle_sigint():
266266
logger.debug("not spinning up new item as we are pending a stop")
267267
continue
268268
logger.debug("spinning up worker")
269-
resp = await ws.claim_item()
269+
resp = await ws.claim_item(0)
270270
if resp:
271-
item, info_url = resp
272-
id = item['id']
273-
logger.info(f"Starting task {id}")
274-
url = item['item']
275-
assert "_" not in id
276-
prefix = "mnbot-brozzler-" + id.replace("-", "_")
271+
claim, info_url = resp
272+
print(claim, info_url)
273+
attempt_id = claim['attempt_id']
274+
job_id = claim['job_id']
275+
page_id = claim['page_id']
276+
logger.info(f"Starting claim {attempt_id} (for {job_id} : {page_id}")
277277
task = asyncio.create_task(run_job(
278278
ws,
279-
item,
280-
url,
281-
prefix,
282-
item['metadata']['ua'],
283-
item['metadata']['custom_js'],
284-
info_url
279+
claim,
280+
info_url,
285281
))
286-
task.set_name(id)
287-
workers[task] = (TaskType.ITEM, item)
282+
task.set_name(attempt_id)
283+
workers[task] = (TaskType.ITEM, claim)
288284
else:
289285
to_sleep = random.randint(10, 30)
290-
logger.info(f"No items found, blocking this worker for {to_sleep} seconds.")
286+
logger.info(f"No tasks found, blocking this worker for {to_sleep} seconds.")
291287
task = asyncio.create_task(asyncio.sleep(to_sleep))
292288
workers[task] = (TaskType.SLEEP, None)
293289
done: set[asyncio.Task] = (await asyncio.wait(workers, return_when = asyncio.FIRST_COMPLETED))[0]
294290
for finished_task in done:
295291
logger.debug(f"checking finished task {finished_task}")
296-
task_type, item = workers[finished_task]
292+
task_type, task_claim = workers[finished_task]
297293
del workers[finished_task]
298294
if task_type != TaskType.ITEM:
299295
logger.debug("nevermind, not an item")
300296
continue
301-
# item can't be None at this point
302-
id = item['id']
303-
tries = item['_current_attempt']
297+
# claim can't be None at this point
298+
attempt_id = task_claim['attempt_id']
304299
try:
305300
_dedup_bucket, _stats_bucket = finished_task.result()
306301
except Exception as e:
@@ -309,14 +304,14 @@ def handle_sigint():
309304
fatal = e.fatal
310305
else:
311306
fatal = False
312-
logger.exception(f"failed task {id}:")
307+
logger.exception(f"failed task {attempt_id}:")
313308
fmt = io.StringIO()
314309
finished_task.print_stack(file = fmt)
315310
message = f"Caught exception!\n{fmt.getvalue()}"
316-
await ws.fail_item(id, message, tries, fatal)
311+
await ws.fail_item(attempt_id, message, fatal)
317312
else:
318313
logger.info(f"task {id} was successful!")
319-
await ws.finish_item(id)
314+
await ws.finish_item(attempt_id)
320315
logger.debug("creating cleanup task")
321316
task = asyncio.create_task(warcprox_cleanup())
322317
workers[task] = (TaskType.CLEANUP, None)

client/worker/browse.py

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -342,12 +342,12 @@ def _brozzle(self) -> Result:
342342
ua = self._create_user_agent(version)
343343

344344
# Write job info to the WARC
345-
logger.debug("writing item info")
345+
logger.debug("writing job info")
346346
self._write_warcprox_record(
347-
"metadata:mnbot-job-metadata",
347+
f"metadata:mnbot-metadata/{self.job.attempt_id}",
348348
"application/json",
349349
json.dumps({
350-
"job": self.job.full_job,
350+
"claim": self.job.full_job,
351351
"version": VERSION,
352352
"browser": {
353353
"executable": self.chrome_exe,
@@ -403,7 +403,7 @@ def _brozzle(self) -> Result:
403403

404404
with self.websock_thread_lock:
405405
r = Result(
406-
id = self.job.full_job['id'],
406+
id = self.job.attempt_id,
407407
final_url = final_url,
408408
outlinks = list(outlinks),
409409
custom_js = custom_js_result,
@@ -413,10 +413,11 @@ def _brozzle(self) -> Result:
413413
)
414414
logger.debug("writing job result data")
415415
self._write_warcprox_record(
416-
"metadata:mnbot-job-result",
416+
f"metadata:mnbot-result/{self.job.attempt_id}",
417417
"application/json",
418418
json.dumps({
419-
"result": r.dict()
419+
"result": r.dict(),
420+
"attempt_id": self.job.attempt_id,
420421
}).encode(),
421422
self.job.warc_prefix
422423
)
@@ -445,17 +446,9 @@ def write_message(type, payload):
445446
write_message("requisites", [dataclasses.asdict(v) for v in result.requisites.values()])
446447
write_message("status_code", result.status_code)
447448

448-
screenshot = browser.screenshot
449-
thumbnail = browser.thumbnail
450-
if screenshot:
451-
screenshot = base64.b85encode(screenshot).decode()
452-
if thumbnail:
453-
thumbnail = base64.b85encode(thumbnail).decode()
454-
if screenshot or thumbnail:
455-
write_message("screenshot", {
456-
"full": screenshot,
457-
"thumb": thumbnail,
458-
})
449+
if browser.screenshot:
450+
screenshot = base64.b85encode(browser.screenshot).decode()
451+
write_message("screenshot", screenshot)
459452

460453
if result.custom_js_screenshot:
461454
write_message("cjs_screenshot", {"full": result.custom_js_screenshot})

client/worker/shared.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
# Update this whenever you make a change, cosmetic or not.
88
# During development you can ignore it, but when you actually
99
# push it to prod, it *must* be updated.
10-
VERSION = "20260412.01"
10+
VERSION = "20260626.01"
1111

1212
DEBUG = os.environ.get("DEBUG") == "1"
1313
if DEBUG:
@@ -19,6 +19,7 @@
1919
@dataclasses.dataclass
2020
class Job:
2121
full_job: dict
22+
attempt_id: str
2223
url: str
2324
warc_prefix: str
2425
dedup_bucket: str

client/worker/tracker.py

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
1-
### HEY YOU! Yeah, you!
2-
### If you are making any change to the client, please
3-
### update the version in meta.py.
4-
### No two versions in prod should have the same version number.
5-
61
############################# NOTE! ###############################
72
### If you are making any change to the client, please ###
83
### update the version in shared.py. ###
94
### No two versions in prod should have the same version number.###
105
############################ THANKS! ##############################
116

12-
import asyncio, logging, json
7+
import asyncio
8+
import json
9+
import logging
1310
import typing
1411

1512
from shared import VERSION
@@ -19,6 +16,14 @@
1916

2017
import websockets
2118
from websockets.asyncio.client import connect
19+
import uuid_utils.compat as uuid
20+
21+
class CustomEncoder(json.JSONEncoder):
22+
def default(self, obj):
23+
if isinstance(obj, uuid.UUID):
24+
return str(obj)
25+
return super().default(obj)
26+
encoder = CustomEncoder()
2227

2328
class Websocket:
2429
def __init__(self, url: str, advisory_handler):
@@ -32,13 +37,14 @@ async def _send_once(self, type, payload=None) -> tuple[int, dict]:
3237
async with self.lock:
3338
if not self.conn:
3439
logger.debug("creating connection")
35-
self.conn = await connect(self.url)
36-
await self.conn.send(json.dumps({"v": VERSION}))
40+
conn = await connect(self.url)
41+
await conn.send(encoder.encode({"v": VERSION, "p": 2, "num_slots": 1}))
42+
self.conn = conn
3743
self.seq += 1
3844
seq = self.seq
3945
message = {"type": type, "request": payload, "seq": seq}
40-
logger.debug(f"sending message (keys: {list(message.keys())})")
41-
await self.conn.send(json.dumps(message))
46+
logger.debug(f"sending {type} message (seq = {seq})")
47+
await self.conn.send(encoder.encode(message))
4248
while True:
4349
resp = json.loads(await self.conn.recv())
4450
if resp['seq'] is None: # Advisory
@@ -79,41 +85,39 @@ async def _send(self, type, payload=None):
7985
await asyncio.sleep(sleep)
8086
tries += 1
8187

82-
async def claim_item(self) -> typing.Optional[tuple[dict, str]]:
83-
status, resp = await self._send("Item:claim", {"pipeline_type": "brozzler"})
88+
async def claim_item(self, slot: int) -> typing.Optional[tuple[dict, str]]:
89+
status, resp = await self._send("Item:claim", {"slot": slot})
8490
if status != 200:
8591
raise RuntimeError(f"Bad response from server: {status} {resp}")
8692
if resp['item']:
8793
return resp['item'], resp['info_url']
8894
return None
8995

90-
async def fail_item(self, id: str, reason: str, tries: int, fatal: bool):
96+
async def fail_item(self, attempt_id: str, reason: str, fatal: bool):
9197
status, resp = await self._send(
9298
"Item:fail",
9399
{
94-
"id": id,
100+
"attempt_id": attempt_id,
95101
"message": reason,
96-
"attempt": tries,
97-
"fatal": fatal
102+
"fatal": fatal,
98103
}
99104
)
100105
if status != 204:
101106
raise RuntimeError(f"Bad response from server: {status} {resp}")
102107

103-
async def finish_item(self, id: str):
104-
status, resp = await self._send("Item:finish", {"id": id})
108+
async def finish_item(self, attempt_id: str):
109+
status, resp = await self._send("Item:finish", {"attempt_id": attempt_id})
105110
if status != 204:
106111
raise RuntimeError(f"Bad response from server: {status} {resp}")
107112

108-
async def store_result(self, id: str, result_type: str, tries: int, result, decode_fields = None):
113+
async def store_result(self, attempt_id: str, result_type: str, result):
114+
result_id = uuid.uuid7()
109115
pl = {
110-
"result_type": result_type,
111-
"attempt": tries,
112-
"result": result,
113-
"id": id
116+
"result_id": result_id,
117+
"attempt_id": attempt_id,
118+
"type": result_type,
119+
"payload": result,
114120
}
115-
if decode_fields:
116-
pl['decode_fields'] = decode_fields
117121
status, resp = await self._send("Item:store", pl)
118122
if status != 201:
119123
raise RuntimeError(f"Bad response from server: {status} {resp}")

0 commit comments

Comments
 (0)