Fix DistributedSeed and DistributedValue not receiving enabled_worker_ids

The orchestration layer never injected enabled_worker_ids into seed and
value nodes, so workers could not resolve their index and fell back to
returning the original seed unchanged.
This commit is contained in:
Robert Wojciechowski
2026-03-07 05:33:49 +00:00
parent dc0945e0f2
commit ce3b2c63f3
2 changed files with 16 additions and 4 deletions
+6 -4
View File
@@ -447,7 +447,7 @@ def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str
return job_map
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id):
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json):
"""Configure DistributedSeed nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
node = prompt_copy.get(node_id)
@@ -455,6 +455,7 @@ def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id):
continue
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
inputs["enabled_worker_ids"] = enabled_json
if is_master:
inputs["worker_id"] = ""
else:
@@ -521,7 +522,7 @@ def _override_upscale_nodes(
inputs["worker_id"] = participant_id
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id):
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json):
"""Configure DistributedValue nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedValue"):
node = prompt_copy.get(node_id)
@@ -529,6 +530,7 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id):
continue
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
inputs["enabled_worker_ids"] = enabled_json
if is_master:
inputs["worker_id"] = ""
else:
@@ -672,8 +674,8 @@ def apply_participant_overrides(
is_master = participant_id == "master"
enabled_json = json.dumps(enabled_worker_ids)
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id)
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id)
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json)
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json)
_override_collector_nodes(
prompt_copy,
prompt_index,
@@ -486,6 +486,11 @@ class ApplyOverridesSeedTests(unittest.TestCase):
result = _apply(self._seed_prompt(), "master")
self.assertEqual(result["1"]["inputs"]["worker_id"], "")
def test_enabled_worker_ids_is_set(self):
enabled = ["worker-a", "worker-b"]
result = _apply(self._seed_prompt(), "worker-a", enabled_worker_ids=enabled)
self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
# ---------------------------------------------------------------------------
# apply_participant_overrides – UltimateSDUpscaleDistributed
@@ -544,6 +549,11 @@ class ApplyOverridesValueTests(unittest.TestCase):
result = _apply(self._value_prompt(), "master")
self.assertEqual(result["1"]["inputs"]["worker_id"], "")
def test_enabled_worker_ids_is_set(self):
enabled = ["worker-a", "worker-b"]
result = _apply(self._value_prompt(), "worker-a", enabled_worker_ids=enabled)
self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
class ApplyOverridesBranchTests(unittest.TestCase):
def test_master_gets_branch_zero_worker_gets_branch_one(self):