feat(wasm): execute task-local losses through scheduler

Author: GBeurier

crates/dag-ml-core/src/plan.rs

@@ -287,6 +287,39 @@ impl ExecutionPlan {
         Ok(plan)
     }
 
+    /// Replace the plan's training-loss roles with a validated canonical set.
+    ///
+    /// Roles are grouped by their declared node and sorted by output/phase so
+    /// every binding produces the same `NodeTask` requirement order. Nodes not
+    /// present in `roles` are explicitly left without a configurable loss.
+    pub fn with_training_losses(mut self, roles: Vec<TrainingLossRoleReference>) -> Result<Self> {
+        self.validate()?;
+        let mut roles_by_node = BTreeMap::<NodeId, Vec<TrainingLossRoleReference>>::new();
+        for role in roles {
+            role.validate()?;
+            if !self.node_plans.contains_key(&role.node_id) {
+                return Err(DagMlError::Planning(format!(
+                    "training loss references unknown plan node `{}`",
+                    role.node_id
+                )));
+            }
+            roles_by_node
+                .entry(role.node_id.clone())
+                .or_default()
+                .push(role);
+        }
+        for node_roles in roles_by_node.values_mut() {
+            node_roles.sort_by(|left, right| {
+                (&left.output_id, &left.phases).cmp(&(&right.output_id, &right.phases))
+            });
+        }
+        for (node_id, node_plan) in &mut self.node_plans {
+            node_plan.training_losses = roles_by_node.remove(node_id).unwrap_or_default();
+        }
+        self.validate()?;
+        Ok(self)
+    }
+
     pub fn validate(&self) -> Result<()> {
         self.graph_plan.graph.validate()?;
         self.campaign.validate()?;
@@ -2228,6 +2261,87 @@ mod tests {
         }
     }
 
+    fn custom_loss_role(node_id: &str, output_id: &str) -> TrainingLossRoleReference {
+        let fixture: serde_json::Value = serde_json::from_str(include_str!(
+            "../../../examples/fixtures/criteria/javascript_local_implementations.v1.json"
+        ))
+        .unwrap();
+        let mut role: TrainingLossRoleReference =
+            serde_json::from_value(fixture["training_loss_role"].clone()).unwrap();
+        role.node_id = NodeId::new(node_id).unwrap();
+        role.output_id = Some(output_id.to_string());
+        role
+    }
+
+    #[test]
+    fn execution_plan_lowers_training_losses_in_canonical_order() {
+        let mut loss_registry = ControllerRegistry::new();
+        loss_registry
+            .register(manifest("controller:transform", NodeKind::Transform))
+            .unwrap();
+        let mut model_manifest = manifest("controller:model", NodeKind::Model);
+        model_manifest.capabilities.extend([
+            ControllerCapability::SupportsConfigurableLoss,
+            ControllerCapability::SupportsCustomLoss,
+            ControllerCapability::SupportsDifferentiableLoss,
+        ]);
+        loss_registry.register(model_manifest).unwrap();
+
+        let plan = build_execution_plan(
+            "plan:training-loss-lowering",
+            graph(),
+            campaign("campaign:training-loss-lowering"),
+            &loss_registry,
+        )
+        .unwrap();
+        let role_b = custom_loss_role("model:pls", "b");
+        let role_a = custom_loss_role("model:pls", "a");
+        let bound = plan
+            .clone()
+            .with_training_losses(vec![role_b, role_a.clone()])
+            .unwrap();
+        let model = bound
+            .node_plans
+            .get(&NodeId::new("model:pls").unwrap())
+            .unwrap();
+        assert_eq!(
+            model
+                .training_losses
+                .iter()
+                .map(|role| role.output_id.as_deref())
+                .collect::<Vec<_>>(),
+            vec![Some("a"), Some("b")]
+        );
+
+        let cleared = bound.with_training_losses(Vec::new()).unwrap();
+        assert!(cleared
+            .node_plans
+            .values()
+            .all(|node| node.training_losses.is_empty()));
+
+        let mut unknown = role_a.clone();
+        unknown.node_id = NodeId::new("model:unknown").unwrap();
+        assert!(plan
+            .clone()
+            .with_training_losses(vec![unknown])
+            .unwrap_err()
+            .to_string()
+            .contains("unknown plan node"));
+
+        let incapable = build_execution_plan(
+            "plan:training-loss-incapable",
+            graph(),
+            campaign("campaign:training-loss-incapable"),
+            &registry(),
+        )
+        .unwrap();
+        assert!(incapable
+            .with_training_losses(vec![role_a])
+            .unwrap_err()
+            .to_string()
+            .contains("does not support configurable loss"));
+    }
+
     #[test]
     fn build_execution_plan_consumes_controller_data_requirements_for_bindings() {
         let model_id = NodeId::new("model:pls").unwrap();

crates/dag-ml-wasm/README.md

@@ -76,6 +76,43 @@ requirements exactly match its active roles. A controller may copy the returned
 attestation into `NodeResult.lineage.loss_attestations` only after `loss(...)`
 returns successfully.
 
+For scheduler execution, lower the roles into the native plan and execute that
+plan rather than rebuilding a loss-free campaign plan:
+
+```js
+const planJson = build_execution_plan_with_training_losses_json(
+  planId,
+  JSON.stringify(graph),
+  JSON.stringify(campaign),
+  JSON.stringify(controllerManifests),
+  JSON.stringify(trainingLossRoles),
+);
+
+const resultsJson = execute_execution_plan_phase_json(
+  planJson,
+  JSON.stringify(controllerManifests),
+  runId,
+  rootSeed,
+  "FIT_CV",
+  (controllerId, taskJson, exactSeed) =>
+    controller.invoke(controllerId, taskJson, exactSeed),
+);
+```
+
+The native lowerer replaces all plan loss roles, groups them by node and sorts
+them canonically before validating controller capabilities. At execution, every
+manifest embedded in the plan must exactly match the independently supplied
+trusted controller registry before any callback runs. The scheduler then checks
+that each callback result contains exactly the task's required loss attestations
+in the same order.
+
+The callback's `exactSeed` argument is a decimal string (or `null`), avoiding
+precision loss for native `u64` seeds beyond JavaScript's safe-integer range.
+The callback may set `NodeResult.lineage.seed` to `null`; the WASM bridge then
+injects the authoritative native seed before scheduler validation. The returned
+`resultsJson` still contains native numeric `u64` values; use a lossless JSON
+integer parser, or preserve the raw JSON, when inspecting lineage seeds exactly.
+
 JavaScript-local descriptors use `binding:javascript` and a `host_local` or
 `portable_registered` lifecycle. A Web Worker must populate its own registry;
 functions are not cloned, posted, or embedded in replay artifacts. Resolution