270 lines
8.9 KiB
Python
270 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import time
|
|
import json
|
|
import threading
|
|
import http.server
|
|
from pathlib import Path
|
|
from urllib.parse import urlencode
|
|
from urllib.request import Request, urlopen
|
|
from urllib.error import URLError, HTTPError
|
|
from collections.abc import Mapping
|
|
from urllib.response import addinfourl
|
|
from typing import cast, Any
|
|
from datetime import datetime
|
|
|
|
class H(http.server.BaseHTTPRequestHandler):
|
|
def do_GET(self):
|
|
if self.path == "/healthz":
|
|
self.send_response(200)
|
|
self.end_headers()
|
|
self.wfile.write(b"ok")
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
def log_message(self, format, *args):
|
|
pass
|
|
|
|
def int_env(name: str, default: int) -> int:
|
|
try:
|
|
return int(os.getenv(name, str(default)))
|
|
except Exception:
|
|
return default
|
|
|
|
TASKS_PATH = Path("/app/tasks")
|
|
API_URL = os.getenv("API_URL", "https://api.sodew.ai/api/tasks").rstrip("/")
|
|
WORKER_UUID = os.getenv("WORKER_UUID", "worker-uuid")
|
|
WORKER_NAME = os.getenv("WORKER_NAME", "worker-name")
|
|
WORKER_POOL = os.getenv("WORKER_POOL", "")
|
|
WORKER_VERSION = "6.3.445"
|
|
GETTING_PAUSE = int_env("GETTING_PAUSE", 30)
|
|
SENDING_PAUSE = int_env("SENDING_PAUSE", 15)
|
|
|
|
HTTP_TIMEOUT = 15
|
|
DEFAULT_HEADERS = {"Accept": "application/json", "Content-Type": "application/json", "User-Agent": "kube-worker/1.1"}
|
|
|
|
def start_health():
|
|
t = threading.Thread(target=http.server.HTTPServer(('0.0.0.0', 8080), H).serve_forever, daemon=True)
|
|
t.start()
|
|
|
|
def ensure_dir() -> None:
|
|
TASKS_PATH.mkdir(parents=True, exist_ok=True)
|
|
|
|
def log_info(text: str) -> None:
|
|
print(datetime.now().strftime("[%Y.%m.%d %H:%M:%S]") + f" {text}")
|
|
|
|
def format_error_body(body: Any) -> str:
|
|
if body is None:
|
|
return "<no body>"
|
|
|
|
if isinstance(body, dict):
|
|
msg = body.get("message")
|
|
if isinstance(msg, str) and msg.strip():
|
|
return msg
|
|
|
|
try:
|
|
return json.dumps(body, ensure_ascii=False, sort_keys=True)
|
|
except Exception:
|
|
return repr(body)
|
|
|
|
if isinstance(body, list):
|
|
try:
|
|
return json.dumps(body, ensure_ascii=False, sort_keys=True)
|
|
except Exception:
|
|
return repr(body)
|
|
|
|
try:
|
|
return str(body)
|
|
except Exception:
|
|
return "<unprintable body>"
|
|
|
|
def task_read(path: Path) -> dict[str, str]:
|
|
result: dict[str, str] = {}
|
|
for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
k, v = line.split("=", 1)
|
|
result[k.strip()] = v.strip()
|
|
return result
|
|
|
|
def task_save(path: Path, data: Mapping[str, object]) -> None:
|
|
flat: dict[str, object] = dict(data)
|
|
lines: list[str] = []
|
|
for key, value in flat.items():
|
|
if not isinstance(key, str):
|
|
key = str(key)
|
|
|
|
if value is None:
|
|
value_str = ""
|
|
elif isinstance(value, (dict, list)):
|
|
try:
|
|
value_str = json.dumps(value, ensure_ascii=False)
|
|
except Exception:
|
|
value_str = str(value)
|
|
else:
|
|
value_str = str(value)
|
|
|
|
value_str = value_str.replace("\n", " ").replace("\r", " ")
|
|
lines.append(f"{key}={value_str}")
|
|
|
|
content = "\n".join(lines) + "\n"
|
|
path.write_text(content, encoding="utf-8")
|
|
|
|
def http_request_json(
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
params: Mapping[str, str] | None = None,
|
|
json_body: Mapping[str, object] | None = None,
|
|
headers: Mapping[str, str] | None = None,
|
|
timeout: int = HTTP_TIMEOUT,
|
|
) -> tuple[int, object | None]:
|
|
if params:
|
|
qs = urlencode(params)
|
|
url = f"{url}?{qs}"
|
|
body_bytes = None
|
|
req_headers = dict(DEFAULT_HEADERS)
|
|
if headers:
|
|
req_headers.update(headers)
|
|
if json_body is not None:
|
|
body_bytes = json.dumps(json_body).encode("utf-8")
|
|
req = Request(url, data=body_bytes, headers=req_headers, method=method)
|
|
try:
|
|
with urlopen(req, timeout=timeout) as resp:
|
|
resp_typed = cast(addinfourl, resp)
|
|
code = resp_typed.getcode() or 0
|
|
body = resp_typed.read()
|
|
|
|
except HTTPError as e:
|
|
try:
|
|
err_bytes = e.read()
|
|
except Exception:
|
|
err_bytes = b""
|
|
if not err_bytes:
|
|
return e.code, None
|
|
try:
|
|
return e.code, json.loads(err_bytes.decode("utf-8", errors="replace"))
|
|
except Exception:
|
|
return e.code, err_bytes.decode("utf-8", errors="replace")
|
|
except URLError as e:
|
|
return 0, {"error": "network", "reason": str(e)}
|
|
|
|
if not body:
|
|
return code, None
|
|
try:
|
|
return code, json.loads(body.decode("utf-8", errors="replace"))
|
|
except Exception:
|
|
return code, None
|
|
|
|
def task_receive() -> dict[str, Any] | None:
|
|
log_info(f"Receiving new tasks from API | Worker: {WORKER_NAME}")
|
|
url = f"{API_URL}/receive"
|
|
payload: dict[str, str] = {
|
|
"worker_name": WORKER_NAME,
|
|
"worker_uuid": WORKER_UUID,
|
|
"worker_version": WORKER_VERSION
|
|
}
|
|
if WORKER_POOL and WORKER_POOL.strip():
|
|
payload["worker_pool"] = WORKER_POOL.strip()
|
|
|
|
code, body = http_request_json("POST", url, json_body=payload)
|
|
|
|
if code == 204:
|
|
log_info("Code 204 | No tasks from API")
|
|
return None
|
|
|
|
if code == 200:
|
|
try:
|
|
log_info("Code: 200 | Raw data: " + (json.dumps(body, ensure_ascii=False, sort_keys=True) if isinstance(body, (dict, list)) else repr(body)))
|
|
except Exception:
|
|
log_info("Code: 200 | Raw data: <unserializable>")
|
|
|
|
if isinstance(body, dict):
|
|
d = cast(dict[str, object], body)
|
|
try:
|
|
log_info("Task: received | " + json.dumps(d, ensure_ascii=False, sort_keys=True))
|
|
except Exception:
|
|
log_info("Task: received | <unserializable dict>")
|
|
|
|
if "uuid" in d and "action" in d:
|
|
return d
|
|
|
|
log_info("Task: missing required fields 'uuid' or 'action'")
|
|
else:
|
|
log_info("Task: not recognized (body is not a dict)")
|
|
else:
|
|
message_error = format_error_body(body)
|
|
log_info(f"Code: {code} | Error: {message_error}")
|
|
|
|
return None
|
|
|
|
def main() -> None:
|
|
start_health()
|
|
ensure_dir()
|
|
|
|
while True:
|
|
task_files = list(TASKS_PATH.glob("*.task"))
|
|
task_count = len(task_files)
|
|
log_info(f"Task files found: {task_count}")
|
|
|
|
if not task_files:
|
|
task = task_receive()
|
|
if task:
|
|
uuid = str(task.get("uuid", "")).strip()
|
|
action = str(task.get("action", "")).strip()
|
|
task.pop("uuid", None)
|
|
task["status"] = "waiting"
|
|
if uuid and action:
|
|
log_info(f"Task from API: {uuid}")
|
|
path = TASKS_PATH / f"{uuid}.task"
|
|
if not path.exists():
|
|
task_save(path=path, data=task)
|
|
else:
|
|
log_info(f"Task: {uuid} | Error: uuid or action is empty, skip")
|
|
time.sleep(GETTING_PAUSE)
|
|
else:
|
|
for path in task_files:
|
|
uuid = path.stem
|
|
try:
|
|
data = task_read(path)
|
|
log_info(f"Task: {uuid} | Parse task success")
|
|
except Exception:
|
|
log_info(f"Task: {uuid} | Task not recognized")
|
|
continue
|
|
action = (data.get("action") or "unknown").strip().lower()
|
|
if action == "pause":
|
|
continue
|
|
status = (data.get("status") or "unknown").strip().lower()
|
|
|
|
payload: dict[str, str] = dict(data)
|
|
payload["worker_uuid"] = WORKER_UUID
|
|
payload["status"] = status
|
|
log_info(f"Task: {uuid} | Payload: {json.dumps(payload, ensure_ascii=False)}")
|
|
url = f"{API_URL}/{uuid}/status"
|
|
code, body = http_request_json("PATCH", url, json_body=payload)
|
|
if not (200 <= code < 300):
|
|
message_error = format_error_body(body)
|
|
log_info(f"Task: {uuid} | Code: {code} | Error: {message_error}")
|
|
continue
|
|
|
|
if status not in {"new", "waiting", "execution"}:
|
|
log_info(f"Task: {uuid} | Remove...")
|
|
try:
|
|
path.unlink(missing_ok=True)
|
|
except Exception:
|
|
pass
|
|
|
|
time.sleep(SENDING_PAUSE)
|
|
|
|
if __name__ == "__main__":
|
|
print(f"[ Worker Agent {WORKER_VERSION} | {datetime.now():%Y-%m-%d %H-%M-%S} ______________ ]")
|
|
print(f"🔹Worker: {WORKER_NAME} | {WORKER_UUID}")
|
|
print(f"🔹Tasks path: {TASKS_PATH}")
|
|
print(f"🔹API URL: {API_URL}")
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
print("🔥Interrupted by user.")
|