diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 673f7f8..39464d5 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -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, diff --git a/tests/unit/system/test_prompt_transform.py b/tests/unit/system/test_prompt_transform.py index d6158bd..1d1148e 100644 --- a/tests/unit/system/test_prompt_transform.py +++ b/tests/unit/system/test_prompt_transform.py @@ -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):