Skip to main content

netget/server/yarn/
actions.rs

1//! Hadoop YARN ResourceManager REST API protocol actions and event types.
2//!
3//! The LLM roleplays a YARN cluster's ResourceManager: it invents the applications,
4//! nodes and metrics the cluster reports. Every response an action produces is wrapped
5//! in the exact JSON envelope a real YARN RM emits (see the Hadoop
6//! `ResourceManagerRest` docs), so `curl` / a real YARN client parses it unchanged.
7//!
8//! No storage: the model supplies all cluster state per request. There is one event
9//! type (`yarn_request`) and it is emitted for every LLM-handled request; the purely
10//! mechanical `GET /ws/v1/cluster/info` version banner is answered statically in
11//! `mod.rs` and never reaches the model.
12
13use crate::llm::actions::protocol_trait::{ActionResult, Protocol, Server};
14use crate::llm::actions::{ActionDefinition, Parameter};
15use crate::protocol::log_template::LogTemplate;
16use crate::protocol::EventType;
17use crate::state::app_state::AppState;
18use anyhow::Result;
19use serde_json::{json, Value};
20use std::sync::LazyLock;
21
22/// YARN ResourceManager protocol handler.
23pub struct YarnProtocol {}
24
25impl YarnProtocol {
26    pub fn new() -> Self {
27        Self {}
28    }
29}
30
31impl Default for YarnProtocol {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37/// Fired for every YARN RM REST request that requires an invented answer (everything
38/// except the static `/ws/v1/cluster/info` banner).
39pub static YARN_REQUEST_EVENT: LazyLock<EventType> = LazyLock::new(|| {
40    EventType::new(
41        "yarn_request",
42        "Hadoop YARN ResourceManager REST request received",
43        json!({
44            "type": "send_yarn_apps",
45            "apps": [{
46                "id": "application_1476912658570_0002",
47                "user": "dr.who",
48                "name": "word count",
49                "queue": "default",
50                "state": "FINISHED",
51                "finalStatus": "SUCCEEDED",
52                "progress": 100.0,
53                "applicationType": "MAPREDUCE"
54            }]
55        }),
56    )
57    .with_parameters(vec![
58        Parameter {
59            name: "method".to_string(),
60            type_hint: "string".to_string(),
61            description: "HTTP method (GET, POST)".to_string(),
62            required: true,
63        },
64        Parameter {
65            name: "path".to_string(),
66            type_hint: "string".to_string(),
67            description: "Request path, e.g. /ws/v1/cluster/apps".to_string(),
68            required: true,
69        },
70        Parameter {
71            name: "operation".to_string(),
72            type_hint: "string".to_string(),
73            description: "Detected operation: metrics, scheduler, apps, submit, \
74                          new_application, app, nodes, node"
75                .to_string(),
76            required: true,
77        },
78        Parameter {
79            name: "app_id".to_string(),
80            type_hint: "string".to_string(),
81            description: "Application id from the path (for the 'app' operation)".to_string(),
82            required: false,
83        },
84        Parameter {
85            name: "node_id".to_string(),
86            type_hint: "string".to_string(),
87            description: "Node id from the path (for the 'node' operation)".to_string(),
88            required: false,
89        },
90        Parameter {
91            name: "request_body".to_string(),
92            type_hint: "string".to_string(),
93            description: "JSON request body (for application submission)".to_string(),
94            required: false,
95        },
96    ])
97    .with_actions(yarn_actions())
98    .with_log_template(
99        LogTemplate::new()
100            .with_info("YARN {client_ip} {method} {path}")
101            .with_debug("YARN {method} {path} op={operation} from {client_ip}:{client_port}")
102            .with_trace("YARN: {json_pretty(.)}"),
103    )
104});
105
106/// The full action set the model may emit for a `yarn_request`.
107fn yarn_actions() -> Vec<ActionDefinition> {
108    vec![
109        send_yarn_metrics_action(),
110        send_yarn_scheduler_action(),
111        send_yarn_apps_action(),
112        send_yarn_app_action(),
113        send_yarn_nodes_action(),
114        send_yarn_new_application_action(),
115        send_yarn_submit_response_action(),
116        send_yarn_error_action(),
117    ]
118}
119
120fn send_yarn_metrics_action() -> ActionDefinition {
121    ActionDefinition {
122        name: "send_yarn_metrics".to_string(),
123        description: "Answer GET /ws/v1/cluster/metrics with the cluster's aggregate metrics. \
124                      All fields are numbers describing the (invented) cluster; keep them \
125                      self-consistent (allocatedMB <= totalMB, activeNodes <= totalNodes). \
126                      Wrapped in the {\"clusterMetrics\": {...}} envelope automatically."
127            .to_string(),
128        parameters: vec![Parameter {
129            name: "metrics".to_string(),
130            type_hint: "object".to_string(),
131            description: "Object of clusterMetrics fields: appsSubmitted, appsRunning, \
132                              appsPending, appsCompleted, appsFailed, appsKilled, totalMB, \
133                              availableMB, allocatedMB, reservedMB, totalVirtualCores, \
134                              availableVirtualCores, allocatedVirtualCores, containersAllocated, \
135                              totalNodes, activeNodes, lostNodes, unhealthyNodes, \
136                              decommissionedNodes. Any field you omit defaults to 0."
137                .to_string(),
138            required: true,
139        }],
140        example: json!({
141            "type": "send_yarn_metrics",
142            "metrics": {
143                "appsSubmitted": 4, "appsRunning": 1, "appsCompleted": 3,
144                "totalMB": 32768, "availableMB": 24576, "allocatedMB": 8192,
145                "totalVirtualCores": 16, "availableVirtualCores": 12, "allocatedVirtualCores": 4,
146                "containersAllocated": 4, "totalNodes": 2, "activeNodes": 2
147            }
148        }),
149        log_template: Some(
150            LogTemplate::new().with_info("-> YARN metrics {metrics.totalNodes} nodes"),
151        ),
152    }
153}
154
155fn send_yarn_scheduler_action() -> ActionDefinition {
156    ActionDefinition {
157        name: "send_yarn_scheduler".to_string(),
158        description: "Answer GET /ws/v1/cluster/scheduler with the scheduler and its queue tree. \
159                      Rendered in the CapacityScheduler's own shape, {\"scheduler\": \
160                      {\"schedulerInfo\": {\"type\": \"capacityScheduler\", \"queueName\": \"root\", \
161                      ..., \"queues\": {\"queue\": [...]}}}}, with each queue's children nested the \
162                      same way; or, with scheduler_type fifoScheduler, the FifoScheduler's flat \
163                      summary. Keep it consistent with the apps you report (a queue's \
164                      numApplications)."
165            .to_string(),
166        parameters: vec![
167            Parameter {
168                name: "scheduler_type".to_string(),
169                type_hint: "string".to_string(),
170                description: "capacityScheduler (default) or fifoScheduler".to_string(),
171                required: false,
172            },
173            Parameter {
174                name: "root".to_string(),
175                type_hint: "object".to_string(),
176                description: "The root queue's (or the FIFO scheduler's) own numbers: capacity \
177                              (default 100), usedCapacity (default 0), maxCapacity (default \
178                              100); for fifoScheduler also qstate, numNodes, numContainers, \
179                              usedNodeCapacity, availNodeCapacity, totalNodeCapacity, \
180                              minQueueMemoryCapacity, maxQueueMemoryCapacity."
181                    .to_string(),
182                required: false,
183            },
184            Parameter {
185                name: "queues".to_string(),
186                type_hint: "array".to_string(),
187                description: format!(
188                    "capacityScheduler only, required: the queues under root. Each is an \
189                     object with queueName (required) and capacity, usedCapacity, maxCapacity, \
190                     absoluteCapacity, absoluteUsedCapacity, absoluteMaxCapacity, state \
191                     (RUNNING/STOPPED), numApplications, numActiveApplications, \
192                     numPendingApplications, numContainers, maxApplications, userLimit, \
193                     userLimitFactor; a parent queue carries its children as a `queues` array \
194                     of the same objects, at most {MAX_QUEUE_DEPTH} levels below root. \
195                     queuePath is filled in from the names."
196                ),
197                required: false,
198            },
199        ],
200        example: json!({
201            "type": "send_yarn_scheduler",
202            "root": {"capacity": 100.0, "usedCapacity": 25.0, "maxCapacity": 100.0},
203            "queues": [
204                {"queueName": "default", "capacity": 60.0, "usedCapacity": 0.0,
205                 "maxCapacity": 100.0, "numApplications": 0},
206                {"queueName": "etl", "capacity": 40.0, "usedCapacity": 62.5,
207                 "maxCapacity": 100.0, "numApplications": 1,
208                 "queues": [{"queueName": "nightly", "capacity": 100.0,
209                             "usedCapacity": 62.5, "numApplications": 1}]}
210            ]
211        }),
212        log_template: Some(LogTemplate::new().with_info("-> YARN scheduler {queues_len} queues")),
213    }
214}
215
216fn send_yarn_apps_action() -> ActionDefinition {
217    ActionDefinition {
218        name: "send_yarn_apps".to_string(),
219        description: "Answer GET /ws/v1/cluster/apps with the list of applications. Supply the \
220                      `apps` array (empty array = no applications, rendered as the YARN idiom \
221                      \"apps\": null). Each element is an application object; wrapped as \
222                      {\"apps\": {\"app\": [...]}} automatically."
223            .to_string(),
224        parameters: vec![Parameter {
225            name: "apps".to_string(),
226            type_hint: "array".to_string(),
227            description: "Array of application objects. Common fields per app: id \
228                          (e.g. application_1476912658570_0002), user, name, queue, state \
229                          (NEW/SUBMITTED/ACCEPTED/RUNNING/FINISHED/FAILED/KILLED), finalStatus \
230                          (UNDEFINED/SUCCEEDED/FAILED/KILLED), progress (0-100), \
231                          applicationType (MAPREDUCE/SPARK/...), startedTime, finishedTime, \
232                          elapsedTime, allocatedMB, allocatedVCores, runningContainers."
233                .to_string(),
234            required: true,
235        }],
236        example: json!({
237            "type": "send_yarn_apps",
238            "apps": [{
239                "id": "application_1476912658570_0002", "user": "dr.who", "name": "word count",
240                "queue": "default", "state": "RUNNING", "finalStatus": "UNDEFINED",
241                "progress": 62.5, "applicationType": "MAPREDUCE", "allocatedMB": 4096,
242                "allocatedVCores": 2, "runningContainers": 2
243            }]
244        }),
245        log_template: Some(LogTemplate::new().with_info("-> YARN {apps_len} apps")),
246    }
247}
248
249fn send_yarn_app_action() -> ActionDefinition {
250    ActionDefinition {
251        name: "send_yarn_app".to_string(),
252        description: "Answer GET /ws/v1/cluster/apps/{appid} with a single application. Supply \
253                      the `app` object; wrapped as {\"app\": {...}} automatically. Emit \
254                      send_yarn_error with status 404 if the application does not exist."
255            .to_string(),
256        parameters: vec![Parameter {
257            name: "app".to_string(),
258            type_hint: "object".to_string(),
259            description: "A single application object (same fields as an element of \
260                          send_yarn_apps.apps)."
261                .to_string(),
262            required: true,
263        }],
264        example: json!({
265            "type": "send_yarn_app",
266            "app": {
267                "id": "application_1476912658570_0002", "user": "dr.who", "name": "word count",
268                "queue": "default", "state": "FINISHED", "finalStatus": "SUCCEEDED",
269                "progress": 100.0, "applicationType": "MAPREDUCE"
270            }
271        }),
272        log_template: Some(LogTemplate::new().with_info("-> YARN app {app.id}")),
273    }
274}
275
276fn send_yarn_nodes_action() -> ActionDefinition {
277    ActionDefinition {
278        name: "send_yarn_nodes".to_string(),
279        description: "Answer GET /ws/v1/cluster/nodes with the NodeManager list. Supply the \
280                      `nodes` array; wrapped as {\"nodes\": {\"node\": [...]}} automatically."
281            .to_string(),
282        parameters: vec![Parameter {
283            name: "nodes".to_string(),
284            type_hint: "array".to_string(),
285            description: "Array of node objects. Common fields per node: id (host:port, e.g. \
286                          host.domain.com:8041), nodeHostName, nodeHTTPAddress, rack \
287                          (/default-rack), state (NEW/RUNNING/UNHEALTHY/DECOMMISSIONED/LOST), \
288                          healthReport, numContainers, usedMemoryMB, availMemoryMB, \
289                          usedVirtualCores, availableVirtualCores, version, lastHealthUpdate."
290                .to_string(),
291            required: true,
292        }],
293        example: json!({
294            "type": "send_yarn_nodes",
295            "nodes": [{
296                "id": "host1.example.com:8041", "nodeHostName": "host1.example.com",
297                "nodeHTTPAddress": "host1.example.com:8042", "rack": "/default-rack",
298                "state": "RUNNING", "healthReport": "", "numContainers": 2,
299                "usedMemoryMB": 4096, "availMemoryMB": 12288, "usedVirtualCores": 2,
300                "availableVirtualCores": 6, "version": "3.3.6"
301            }]
302        }),
303        log_template: Some(LogTemplate::new().with_info("-> YARN {nodes_len} nodes")),
304    }
305}
306
307fn send_yarn_new_application_action() -> ActionDefinition {
308    ActionDefinition {
309        name: "send_yarn_new_application".to_string(),
310        description: "Answer POST /ws/v1/cluster/apps/new-application (the first step of the \
311                      submission flow) with a freshly allocated application id and the cluster's \
312                      maximum resource capability."
313            .to_string(),
314        parameters: vec![
315            Parameter {
316                name: "application_id".to_string(),
317                type_hint: "string".to_string(),
318                description: "Newly allocated application id, e.g. \
319                              application_1476912658570_0005."
320                    .to_string(),
321                required: true,
322            },
323            Parameter {
324                name: "max_memory_mb".to_string(),
325                type_hint: "number".to_string(),
326                description: "maximum-resource-capability memory in MB (default 8192)".to_string(),
327                required: false,
328            },
329            Parameter {
330                name: "max_vcores".to_string(),
331                type_hint: "number".to_string(),
332                description: "maximum-resource-capability vCores (default 4)".to_string(),
333                required: false,
334            },
335        ],
336        example: json!({
337            "type": "send_yarn_new_application",
338            "application_id": "application_1476912658570_0005",
339            "max_memory_mb": 8192,
340            "max_vcores": 4
341        }),
342        log_template: Some(
343            LogTemplate::new().with_info("-> YARN new-application {application_id}"),
344        ),
345    }
346}
347
348fn send_yarn_submit_response_action() -> ActionDefinition {
349    ActionDefinition {
350        name: "send_yarn_submit_response".to_string(),
351        description: "Answer POST /ws/v1/cluster/apps (submit an application). On acceptance YARN \
352                      replies 202 Accepted with an empty body and a Location header pointing at \
353                      the new application; set accepted=true. To reject the submission set \
354                      accepted=false and supply a message (rendered as a 400 RemoteException)."
355            .to_string(),
356        parameters: vec![
357            Parameter {
358                name: "accepted".to_string(),
359                type_hint: "boolean".to_string(),
360                description: "true = 202 Accepted, false = 400 rejection".to_string(),
361                required: true,
362            },
363            Parameter {
364                name: "application_id".to_string(),
365                type_hint: "string".to_string(),
366                description: "The submitted application id (used to build the Location header \
367                              when accepted)."
368                    .to_string(),
369                required: false,
370            },
371            Parameter {
372                name: "message".to_string(),
373                type_hint: "string".to_string(),
374                description: "Rejection reason (when accepted=false)".to_string(),
375                required: false,
376            },
377        ],
378        example: json!({
379            "type": "send_yarn_submit_response",
380            "accepted": true,
381            "application_id": "application_1476912658570_0005"
382        }),
383        log_template: Some(LogTemplate::new().with_info("-> YARN submit accepted={accepted}")),
384    }
385}
386
387fn send_yarn_error_action() -> ActionDefinition {
388    ActionDefinition {
389        name: "send_yarn_error".to_string(),
390        description: "Return a YARN RemoteException error envelope with an explicit HTTP status. \
391                      Use for a missing application/node (404), a bad request (400) or any \
392                      condition without a dedicated action. status must be 100-599."
393            .to_string(),
394        parameters: vec![
395            Parameter {
396                name: "status".to_string(),
397                type_hint: "number".to_string(),
398                description: "HTTP status code (400, 404, 500, ...)".to_string(),
399                required: true,
400            },
401            Parameter {
402                name: "exception".to_string(),
403                type_hint: "string".to_string(),
404                description: "Exception class name, e.g. NotFoundException, BadRequestException"
405                    .to_string(),
406                required: false,
407            },
408            Parameter {
409                name: "message".to_string(),
410                type_hint: "string".to_string(),
411                description: "Human-readable error message".to_string(),
412                required: true,
413            },
414        ],
415        example: json!({
416            "type": "send_yarn_error",
417            "status": 404,
418            "exception": "NotFoundException",
419            "message": "app with id 'application_1_0009' not found"
420        }),
421        log_template: Some(LogTemplate::new().with_info("-> YARN error {status} {exception}")),
422    }
423}
424
425pub fn get_yarn_event_types() -> Vec<EventType> {
426    vec![YARN_REQUEST_EVENT.clone()]
427}
428
429/// How many levels of queues below root `send_yarn_scheduler` renders. A CapacityScheduler
430/// hierarchy is rarely more than three or four deep; the bound keeps the conversion of the
431/// model-supplied tree (which recurses once per level) from following an arbitrarily deep one.
432pub const MAX_QUEUE_DEPTH: usize = 8;
433
434/// One model-supplied queue in the CapacityScheduler's shape: `queuePath` from the names, the
435/// children (`queues`, an array) wrapped as `{"queue": [...]}`, `queueType` and, for a leaf,
436/// the leaf queue `type` the real scheduler reports. `depth` is this queue's level below root.
437fn capacity_queue(queue: &Value, parent_path: &str, depth: usize) -> Result<Value> {
438    if depth > MAX_QUEUE_DEPTH {
439        return Err(anyhow::anyhow!(
440            "the queue tree is deeper than {MAX_QUEUE_DEPTH} levels below root"
441        ));
442    }
443    let mut out = queue
444        .as_object()
445        .cloned()
446        .ok_or_else(|| anyhow::anyhow!("each queue is an object"))?;
447    let name = out
448        .get("queueName")
449        .and_then(|v| v.as_str())
450        .filter(|n| !n.is_empty())
451        .ok_or_else(|| anyhow::anyhow!("each queue needs a queueName"))?
452        .to_string();
453    let path = format!("{parent_path}.{name}");
454    out.entry("queuePath").or_insert_with(|| json!(path));
455    out.entry("state").or_insert_with(|| json!("RUNNING"));
456    match out.remove("queues") {
457        Some(Value::Array(children)) if !children.is_empty() => {
458            let children = children
459                .iter()
460                .map(|c| capacity_queue(c, &path, depth + 1))
461                .collect::<Result<Vec<_>>>()?;
462            out.insert("queueType".into(), json!("parent"));
463            out.insert("queues".into(), json!({ "queue": children }));
464        }
465        None | Some(Value::Null) | Some(Value::Array(_)) => {
466            out.insert("queueType".into(), json!("leaf"));
467            out.entry("type")
468                .or_insert_with(|| json!("capacitySchedulerLeafQueueInfo"));
469        }
470        Some(_) => return Err(anyhow::anyhow!("a queue's `queues` is an array of queues")),
471    }
472    Ok(Value::Object(out))
473}
474
475/// Serialize a response body without a panic path (serializing a Value cannot fail).
476fn body_of(value: &Value) -> String {
477    serde_json::to_string(value).unwrap_or_else(|_| value.to_string())
478}
479
480/// Validate a model-supplied HTTP status, rejecting out-of-range values where the
481/// message still reaches the model rather than panicking later.
482fn parse_status(action: &Value, key: &str) -> Result<u16> {
483    let raw = action
484        .get(key)
485        .and_then(|v| v.as_u64())
486        .ok_or_else(|| anyhow::anyhow!("Missing or invalid '{key}': expected a number"))?;
487    if !(100..=599).contains(&raw) {
488        return Err(anyhow::anyhow!(
489            "Invalid {key} {raw}: must be an HTTP status between 100 and 599"
490        ));
491    }
492    Ok(raw as u16)
493}
494
495/// Build the YARN RemoteException envelope every client parses on error.
496fn remote_exception(exception: &str, message: &str) -> Value {
497    json!({
498        "RemoteException": {
499            "exception": exception,
500            "message": message,
501            "javaClassName": format!("org.apache.hadoop.yarn.webapp.{exception}"),
502        }
503    })
504}
505
506impl Protocol for YarnProtocol {
507    fn get_startup_parameters(&self) -> Vec<crate::llm::actions::ParameterDefinition> {
508        use crate::llm::actions::ParameterDefinition;
509        vec![
510            ParameterDefinition {
511                name: "resource_manager_version".to_string(),
512                type_hint: "string".to_string(),
513                description: "Version string reported by the static GET /ws/v1/cluster/info \
514                              banner (resourceManagerVersion / hadoopVersion). Default 3.3.6."
515                    .to_string(),
516                required: false,
517                example: json!("3.3.6"),
518                default: None,
519            },
520            ParameterDefinition {
521                name: "cluster_id".to_string(),
522                type_hint: "string".to_string(),
523                description: "Numeric cluster id reported in the /ws/v1/cluster/info banner \
524                              (the RM start epoch-ms). Default 1476912658570."
525                    .to_string(),
526                required: false,
527                example: json!("1476912658570"),
528                default: None,
529            },
530        ]
531    }
532    fn get_async_actions(&self, _state: &AppState) -> Vec<ActionDefinition> {
533        vec![]
534    }
535    fn get_sync_actions(&self) -> Vec<ActionDefinition> {
536        yarn_actions()
537    }
538    fn protocol_name(&self) -> &'static str {
539        "YARN"
540    }
541    fn get_event_types(&self) -> Vec<EventType> {
542        get_yarn_event_types()
543    }
544    fn stack_name(&self) -> &'static str {
545        "ETH>IP>TCP>HTTP>YARN"
546    }
547    fn keywords(&self) -> Vec<&'static str> {
548        vec!["yarn", "hadoop yarn", "yarn resourcemanager"]
549    }
550    fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
551        use crate::protocol::metadata::{DevelopmentState, ProtocolMetadataV2};
552        ProtocolMetadataV2::builder()
553            .state(DevelopmentState::Experimental)
554            .well_known_port(8088)
555            .implementation("hyper v1 HTTP/1.1 server, manual YARN RM REST envelopes")
556            .llm_control("Applications, nodes and cluster metrics invented by the LLM")
557            .e2e_testing(
558                "tests/server/yarn/datadog_check_test.rs runs Datadog's YARN integration \
559                 (datadog_checks.yarn, integrations-core tag yarn-8.4.2) once against the server \
560                 and asserts the cluster, application, node and queue metrics, tags, service \
561                 checks and version it submitted from the model's answers to /metrics, /apps, \
562                 /nodes and /scheduler. e2e_test.rs checks the documented envelopes with \
563                 reqwest. NetGet's own YARN client is tested against real Hadoop, not against \
564                 this server",
565            )
566            .notes(
567                "Virtual cluster, no storage. GET /ws/v1/cluster/info is answered \
568                    statically; metrics/scheduler/apps/nodes are LLM-driven, the scheduler in \
569                    the CapacityScheduler's nested queues.queue shape (or the FifoScheduler's). \
570                    Fail-closed: LLM failure returns 503/500 RemoteException, never an \
571                    empty-but-200 cluster. Experimental: one independent consumer (Datadog's \
572                    check), which reads the four list endpoints only; nothing third-party has \
573                    driven an application or node by id, new-application or submission.",
574            )
575            .max_inbound_bytes(crate::server::yarn::MAX_REQUEST_BODY_BYTES)
576            .request_only("The YARN REST API is HTTP request/response; a server cannot send a peer anything unprompted")
577            .build()
578    }
579    fn description(&self) -> &'static str {
580        "Hadoop YARN ResourceManager REST API"
581    }
582    fn example_prompt(&self) -> &'static str {
583        "Start a Hadoop YARN ResourceManager on port 8088 with a couple of running Spark apps"
584    }
585    fn group_name(&self) -> &'static str {
586        "Application"
587    }
588
589    fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
590        use crate::llm::actions::StartupExamples;
591
592        // Deterministic: report fixed cluster metrics for every request, no LLM
593        // call.
594        let script = r#"import json, sys
595data = json.load(sys.stdin)
596event = data["event"]
597if data["event_type_id"] == "yarn_request":
598    actions = [{"type": "send_yarn_metrics",
599                "metrics": {"appsRunning": 0, "availableMB": 8192}}]
600else:
601    actions = []
602print(json.dumps({"actions": actions}))"#;
603
604        StartupExamples::new(
605            json!({
606                "type": "open_server",
607                "port": 8088,
608                "base_stack": "yarn",
609                "instruction": "Hadoop YARN ResourceManager reporting a small MapReduce cluster"
610            }),
611            json!({
612                "type": "open_server",
613                "port": 8088,
614                "base_stack": "yarn",
615                "event_handlers": [{
616                    "event_pattern": "yarn_request",
617                    "handler": { "type": "script", "language": "python", "code": script }
618                }]
619            }),
620            json!({
621                "type": "open_server",
622                "port": 8088,
623                "base_stack": "yarn",
624                "event_handlers": [{
625                    "event_pattern": "yarn_request",
626                    "handler": {
627                        "type": "static",
628                        "actions": [{
629                            "type": "send_yarn_metrics",
630                            "metrics": { "totalNodes": 1, "activeNodes": 1, "totalMB": 8192,
631                                         "availableMB": 8192 }
632                        }]
633                    }
634                }]
635            }),
636        )
637    }
638}
639
640impl Server for YarnProtocol {
641    fn spawn(
642        &self,
643        ctx: crate::protocol::SpawnContext,
644    ) -> std::pin::Pin<
645        Box<dyn std::future::Future<Output = anyhow::Result<std::net::SocketAddr>> + Send>,
646    > {
647        Box::pin(async move {
648            // Read the optional static-banner params before moving ctx fields.
649            let rm_version = ctx
650                .startup_params
651                .as_ref()
652                .and_then(|p| {
653                    p.get_optional_string("resource_manager_version")
654                        .ok()
655                        .flatten()
656                })
657                .unwrap_or_else(|| "3.3.6".to_string());
658            let cluster_id = ctx
659                .startup_params
660                .as_ref()
661                .and_then(|p| p.get_optional_string("cluster_id").ok().flatten())
662                .unwrap_or_else(|| "1476912658570".to_string());
663
664            crate::server::yarn::YarnServer::spawn_with_llm_actions(
665                ctx.legacy_listen_addr(),
666                ctx.llm_client,
667                ctx.state,
668                ctx.status_tx,
669                ctx.server_id,
670                rm_version,
671                cluster_id,
672            )
673            .await
674        })
675    }
676
677    fn execute_action(&self, action: Value) -> Result<ActionResult> {
678        let action_type = action
679            .get("type")
680            .and_then(|v| v.as_str())
681            .ok_or_else(|| anyhow::anyhow!("Missing action type"))?;
682
683        let custom = |status: u16, body: Value| {
684            Ok(ActionResult::Custom {
685                name: "yarn_response".to_string(),
686                data: json!({ "status": status, "body": body_of(&body), "location": Value::Null }),
687            })
688        };
689
690        match action_type {
691            "send_yarn_metrics" => {
692                let metrics = action
693                    .get("metrics")
694                    .cloned()
695                    .filter(|v| v.is_object())
696                    .ok_or_else(|| anyhow::anyhow!("Missing 'metrics' object"))?;
697                // Fill the mandatory keys clients read so a partial object still parses.
698                let mut m = metrics;
699                for key in [
700                    "appsSubmitted",
701                    "appsCompleted",
702                    "appsPending",
703                    "appsRunning",
704                    "appsFailed",
705                    "appsKilled",
706                    "reservedMB",
707                    "availableMB",
708                    "allocatedMB",
709                    "totalMB",
710                    "reservedVirtualCores",
711                    "availableVirtualCores",
712                    "allocatedVirtualCores",
713                    "totalVirtualCores",
714                    "containersAllocated",
715                    "containersReserved",
716                    "containersPending",
717                    "totalNodes",
718                    "activeNodes",
719                    "lostNodes",
720                    "unhealthyNodes",
721                    "decommissionedNodes",
722                    "decommissioningNodes",
723                    "rebootedNodes",
724                    "shutdownNodes",
725                ] {
726                    if m.get(key).is_none() {
727                        m[key] = json!(0);
728                    }
729                }
730                custom(200, json!({ "clusterMetrics": m }))
731            }
732            "send_yarn_scheduler" => {
733                let kind = action
734                    .get("scheduler_type")
735                    .and_then(|v| v.as_str())
736                    .unwrap_or("capacityScheduler");
737                let mut info = action
738                    .get("root")
739                    .and_then(|v| v.as_object())
740                    .cloned()
741                    .unwrap_or_default();
742                info.insert("type".into(), json!(kind));
743                for (key, default) in [
744                    ("capacity", 100.0),
745                    ("usedCapacity", 0.0),
746                    ("maxCapacity", 100.0),
747                ] {
748                    info.entry(key).or_insert_with(|| json!(default));
749                }
750                match kind {
751                    "capacityScheduler" => {
752                        let queues = action
753                            .get("queues")
754                            .and_then(|v| v.as_array())
755                            .filter(|q| !q.is_empty())
756                            .ok_or_else(|| {
757                                anyhow::anyhow!(
758                                    "a capacityScheduler needs a non-empty 'queues' array"
759                                )
760                            })?
761                            .iter()
762                            .map(|q| capacity_queue(q, "root", 1))
763                            .collect::<Result<Vec<_>>>()?;
764                        info.insert("queueName".into(), json!("root"));
765                        info.insert("queuePath".into(), json!("root"));
766                        info.insert("queueType".into(), json!("parent"));
767                        info.insert("queues".into(), json!({ "queue": queues }));
768                    }
769                    "fifoScheduler" => {
770                        info.entry("qstate").or_insert_with(|| json!("RUNNING"));
771                    }
772                    other => {
773                        return Err(anyhow::anyhow!(
774                            "scheduler_type {other:?}: capacityScheduler or fifoScheduler"
775                        ))
776                    }
777                }
778                custom(200, json!({ "scheduler": { "schedulerInfo": info } }))
779            }
780            "send_yarn_apps" => {
781                let apps = action
782                    .get("apps")
783                    .and_then(|v| v.as_array())
784                    .ok_or_else(|| anyhow::anyhow!("Missing 'apps' array"))?;
785                // YARN renders an empty application list as "apps": null, not {"app": []}.
786                let body = if apps.is_empty() {
787                    json!({ "apps": Value::Null })
788                } else {
789                    json!({ "apps": { "app": apps } })
790                };
791                custom(200, body)
792            }
793            "send_yarn_app" => {
794                let app = action
795                    .get("app")
796                    .cloned()
797                    .filter(|v| v.is_object())
798                    .ok_or_else(|| anyhow::anyhow!("Missing 'app' object"))?;
799                custom(200, json!({ "app": app }))
800            }
801            "send_yarn_nodes" => {
802                let nodes = action
803                    .get("nodes")
804                    .and_then(|v| v.as_array())
805                    .ok_or_else(|| anyhow::anyhow!("Missing 'nodes' array"))?;
806                let body = if nodes.is_empty() {
807                    json!({ "nodes": Value::Null })
808                } else {
809                    json!({ "nodes": { "node": nodes } })
810                };
811                custom(200, body)
812            }
813            "send_yarn_new_application" => {
814                let app_id = action
815                    .get("application_id")
816                    .and_then(|v| v.as_str())
817                    .ok_or_else(|| anyhow::anyhow!("Missing 'application_id'"))?;
818                let mem = action
819                    .get("max_memory_mb")
820                    .and_then(|v| v.as_u64())
821                    .unwrap_or(8192);
822                let vcores = action
823                    .get("max_vcores")
824                    .and_then(|v| v.as_u64())
825                    .unwrap_or(4);
826                custom(
827                    200,
828                    json!({
829                        "application-id": app_id,
830                        "maximum-resource-capability": { "memory": mem, "vCores": vcores }
831                    }),
832                )
833            }
834            "send_yarn_submit_response" => {
835                let accepted = action
836                    .get("accepted")
837                    .and_then(|v| v.as_bool())
838                    .ok_or_else(|| anyhow::anyhow!("Missing 'accepted' boolean"))?;
839                if accepted {
840                    let app_id = action
841                        .get("application_id")
842                        .and_then(|v| v.as_str())
843                        .unwrap_or("");
844                    // 202 Accepted with empty body + Location header, as YARN does.
845                    Ok(ActionResult::Custom {
846                        name: "yarn_response".to_string(),
847                        data: json!({
848                            "status": 202,
849                            "body": "",
850                            "location": format!("/ws/v1/cluster/apps/{app_id}"),
851                        }),
852                    })
853                } else {
854                    let message = action
855                        .get("message")
856                        .and_then(|v| v.as_str())
857                        .unwrap_or("application submission rejected");
858                    custom(400, remote_exception("BadRequestException", message))
859                }
860            }
861            "send_yarn_error" => {
862                let status = parse_status(&action, "status")?;
863                let message = action
864                    .get("message")
865                    .and_then(|v| v.as_str())
866                    .ok_or_else(|| anyhow::anyhow!("Missing 'message'"))?;
867                let exception = action
868                    .get("exception")
869                    .and_then(|v| v.as_str())
870                    .unwrap_or("WebApplicationException");
871                custom(status, remote_exception(exception, message))
872            }
873            _ => Err(anyhow::anyhow!("Unknown action type: {}", action_type)),
874        }
875    }
876}