add worker fetch step

This commit is contained in:
Laura Abro
2025-04-24 11:52:26 -03:00
parent 08837a39c9
commit d2bbe4de95
2 changed files with 16 additions and 28 deletions

View File

@ -28,7 +28,7 @@ def prepare(runner, worker):
def execute(runner, worker, data):
"""Execute worker task step"""
url = f"{worker.url}/worker-task/{data['roundNumber']}"
url = f"{runner.config.middle_server_url}/summarizer/worker/fetch-todo"
response = requests.post(
url,
json={"signature": data["stakingSignature"], "stakingKey": data["stakingKey"]},
@ -41,32 +41,12 @@ def execute(runner, worker, data):
f"{result.get('message', 'No eligible todos')} for {worker.name} - continuing"
)
return {"success": True, "message": result.get("message")}
else:
response.raise_for_status()
if result.get("success") and "pr_url" in result:
round_key = str(runner.current_round)
round_state = runner.state["rounds"].setdefault(round_key, {})
# Initialize pr_urls if not exists
if "pr_urls" not in round_state:
round_state["pr_urls"] = {}
round_state["pr_urls"][worker.name] = result["pr_url"]
# Initialize submission_data if not exists
if "submission_data" not in round_state:
round_state["submission_data"] = {}
# Store submission data
round_state["submission_data"][worker.name] = {
"githubUsername": worker.env.get("GITHUB_USERNAME"),
"nodeType": "worker",
"prUrl": result["pr_url"],
"repoName": result.get("repoName"),
"repoOwner": result.get("repoOwner"),
"roundNumber": runner.current_round,
"taskId": runner.config.task_id,
"uuid": result.get("uuid"), # Should be provided by the worker
"stakingKey": worker.staking_public_key,
"pubKey": worker.public_key,
}
if result.get("success"):
runner.state["repo_url"] = (
f"https://github.com/{result['repo_owner']}/{result['repo_name']}"
)
return result

View File

@ -3,6 +3,7 @@
from prometheus_test import TestStep
from functools import partial
from .stages import (
worker_fetch,
worker_task,
worker_submission,
worker_audit,
@ -11,12 +12,19 @@ from .stages import (
steps = [
TestStep(
name="worker_fetch",
description="Fetch worker task",
prepare=worker_fetch.prepare,
execute=worker_fetch.execute,
worker="worker1",
),
TestStep(
name="worker_task",
description="Execute worker task",
prepare=worker_task.prepare,
execute=worker_task.execute,
worker="worker",
worker="worker1",
),
TestStep(
name="worker_submission",