|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Reference dispatch command for the oz-agent-worker "command" backend. |
| 3 | +
|
| 4 | +It reads the task ``DispatchPayload`` (JSON) on stdin, transforms it into a |
| 5 | +hypothetical runtime's REST API shape, and POSTs it to ``OZ_DISPATCH_URL``. |
| 6 | +
|
| 7 | +Use it as a template for delegating task execution to a self-hosted runtime that |
| 8 | +is already configured to run oz agents on demand. The part you customize is |
| 9 | +``transform()`` — adapt it to your runtime's request schema. Everything else |
| 10 | +(reading stdin, auth, timeouts, exit-code contract) can stay as-is. |
| 11 | +
|
| 12 | +Only the Python standard library is used, so there are no dependencies to |
| 13 | +install on the worker host. |
| 14 | +
|
| 15 | +Required environment: |
| 16 | + OZ_DISPATCH_URL REST endpoint to POST the transformed body to. |
| 17 | +
|
| 18 | +Optional environment: |
| 19 | + OZ_DISPATCH_AUTH_HEADER Authorization header value (e.g. "Bearer ..."). |
| 20 | + OZ_DISPATCH_TIMEOUT_SECS Request timeout in seconds (default 30). |
| 21 | +
|
| 22 | +Also provided by the worker (no need to set these yourself): |
| 23 | + OZ_TASK_ID, OZ_EXECUTION_ID, OZ_WORKER_BACKEND, OZ_SERVER_ROOT_URL, OZ_DOCKER_IMAGE |
| 24 | +
|
| 25 | +Exit semantics (the contract the command backend relies on): |
| 26 | + exit 0 => task accepted for dispatch; the remote runtime now owns it and |
| 27 | + the agent reports terminal state to Warp itself. |
| 28 | + exit != 0 => dispatch failed; the worker marks the task failed. |
| 29 | +""" |
| 30 | + |
| 31 | +import json |
| 32 | +import os |
| 33 | +import sys |
| 34 | +import urllib.error |
| 35 | +import urllib.request |
| 36 | + |
| 37 | + |
| 38 | +def transform(payload): |
| 39 | + """Map the worker's DispatchPayload onto our runtime's API shape. |
| 40 | +
|
| 41 | + This is an example of an arbitrary, not-hard transformation: fields are |
| 42 | + renamed and nested under a ``run`` object, sidecar mounts are restructured |
| 43 | + (``mount_path`` -> ``path``), and a few values are lifted into a metadata |
| 44 | + block. Replace the body of this function to match your own API. |
| 45 | + """ |
| 46 | + task = payload.get("task") or {} |
| 47 | + definition = task.get("task_definition") or {} |
| 48 | + |
| 49 | + return { |
| 50 | + "run": { |
| 51 | + "task_id": payload["task_id"], |
| 52 | + "execution_id": payload.get("execution_id", ""), |
| 53 | + "image": payload.get("docker_image", ""), |
| 54 | + # base_args is the `oz agent run ...` argv the runtime should exec. |
| 55 | + "command": payload.get("base_args", []), |
| 56 | + "env": payload.get("env", {}), |
| 57 | + "mounts": [ |
| 58 | + { |
| 59 | + "image": sidecar.get("image", ""), |
| 60 | + "path": sidecar.get("mount_path", ""), |
| 61 | + "read_write": sidecar.get("read_write", False), |
| 62 | + } |
| 63 | + for sidecar in (payload.get("sidecars") or []) |
| 64 | + ], |
| 65 | + "callback_url": payload.get("server_root_url", ""), |
| 66 | + "metadata": { |
| 67 | + "worker_id": payload.get("worker_id", ""), |
| 68 | + "payload_version": payload.get("version"), |
| 69 | + "title": task.get("title", ""), |
| 70 | + "prompt": definition.get("prompt", ""), |
| 71 | + }, |
| 72 | + } |
| 73 | + } |
| 74 | + |
| 75 | + |
| 76 | +def main(): |
| 77 | + url = os.environ.get("OZ_DISPATCH_URL") |
| 78 | + if not url: |
| 79 | + sys.stderr.write("OZ_DISPATCH_URL must be set\n") |
| 80 | + return 2 |
| 81 | + timeout = float(os.environ.get("OZ_DISPATCH_TIMEOUT_SECS", "30")) |
| 82 | + |
| 83 | + try: |
| 84 | + payload = json.load(sys.stdin) |
| 85 | + except json.JSONDecodeError as exc: |
| 86 | + sys.stderr.write(f"invalid dispatch payload on stdin: {exc}\n") |
| 87 | + return 1 |
| 88 | + |
| 89 | + body = json.dumps(transform(payload)).encode("utf-8") |
| 90 | + |
| 91 | + request = urllib.request.Request(url, data=body, method="POST") |
| 92 | + request.add_header("Content-Type", "application/json") |
| 93 | + request.add_header("X-Oz-Task-Id", os.environ.get("OZ_TASK_ID", "")) |
| 94 | + auth = os.environ.get("OZ_DISPATCH_AUTH_HEADER") |
| 95 | + if auth: |
| 96 | + request.add_header("Authorization", auth) |
| 97 | + |
| 98 | + try: |
| 99 | + # Any 2xx is success; urllib raises HTTPError for status >= 400. |
| 100 | + with urllib.request.urlopen(request, timeout=timeout) as response: |
| 101 | + response.read() |
| 102 | + except urllib.error.HTTPError as exc: |
| 103 | + detail = exc.read().decode("utf-8", errors="replace") |
| 104 | + sys.stderr.write(f"dispatch endpoint returned HTTP {exc.code}: {detail}\n") |
| 105 | + return 1 |
| 106 | + except urllib.error.URLError as exc: |
| 107 | + sys.stderr.write(f"failed to reach dispatch endpoint: {exc}\n") |
| 108 | + return 1 |
| 109 | + return 0 |
| 110 | + |
| 111 | + |
| 112 | +if __name__ == "__main__": |
| 113 | + sys.exit(main()) |
0 commit comments