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:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user