diff --git a/module/power_distributor/src/mod_power_distributor.c b/module/power_distributor/src/mod_power_distributor.c
index 0e19f64b203f9bf133ab4782ae86015ec7e21d6a..494349b07992bf6043b4a03df8a90794af9df26b 100644
--- a/module/power_distributor/src/mod_power_distributor.c
+++ b/module/power_distributor/src/mod_power_distributor.c
@@ -377,9 +377,63 @@ static int set_budgets()
     return FWK_SUCCESS;
 }
 
+static inline void clear_domains_budget(void)
+{
+    for (size_t i = 0; i < power_distributor_ctx.domain_count; ++i) {
+        struct mod_power_distributor_domain_ctx *domain_ctx = get_domain_ctx(i);
+        domain_ctx->node.data.power_budget = 0;
+    }
+}
+
+static inline void set_root_budget(void)
+{
+    uint32_t root_idx = power_distributor_ctx.tree_traverse_order_table[0];
+    struct mod_power_distributor_domain_ctx *root_ctx =
+        get_domain_ctx(root_idx);
+
+    root_ctx->node.data.power_budget = root_ctx->node.data.power_limit;
+}
+
+static inline int domains_power_distribute(void)
+{
+    int status = FWK_SUCCESS;
+
+    for (size_t i = 0; i < power_distributor_ctx.domain_count; ++i) {
+        uint32_t domain_idx =
+            power_distributor_ctx.tree_traverse_order_table[i];
+        struct mod_power_distributor_domain_ctx *domain_ctx =
+            get_domain_ctx(domain_idx);
+
+        fwk_assert(domain_ctx);
+
+        status = domain_power_distribute(domain_ctx);
+        if (status != FWK_SUCCESS) {
+            fwk_id_t domain_id =
+                FWK_ID_ELEMENT(FWK_MODULE_IDX_POWER_DISTRIBUTOR, domain_idx);
+            FWK_LOG_ERR(
+                MOD_NAME "Failed to distribute domain %s (index: %u)",
+                fwk_module_get_element_name(domain_id),
+                domain_idx);
+            break;
+        }
+    }
+    return status;
+}
+
 static int system_power_distribute(void)
 {
+    int status = FWK_SUCCESS;
+
     calculate_power_attributes();
+    clear_domains_budget();
+    set_root_budget();
+
+    status = domains_power_distribute();
+    if (status != FWK_SUCCESS) {
+        clear_domains_budget();
+        return status;
+    }
+
     return set_budgets();
 }
 
diff --git a/module/power_distributor/test/mod_power_distributor_unit_test.c b/module/power_distributor/test/mod_power_distributor_unit_test.c
index ed0fdc4ed57a6300b9d6633f88cc246287498775..97e96c88d3f9ae6572663304998e67a4529fc83d 100644
--- a/module/power_distributor/test/mod_power_distributor_unit_test.c
+++ b/module/power_distributor/test/mod_power_distributor_unit_test.c
@@ -759,6 +759,44 @@ void utest_mod_distributor_domain_power_distribute(void)
     }
 }
 #undef INIT_LEAF_NODE
+void utest_clear_domains_budget(void)
+{
+    struct mod_power_distributor_domain_ctx test_domains[TEST_DOMAIN_COUNT] = {
+        0
+    };
+
+    power_distributor_ctx.domain_count = TEST_DOMAIN_COUNT;
+    power_distributor_ctx.domain = test_domains;
+
+    for (size_t i = 0; i < TEST_DOMAIN_COUNT; i++) {
+        test_domains[i].node.data.power_budget = 100; /* Assign some value */
+    }
+
+    clear_domains_budget();
+
+    for (size_t i = 0; i < TEST_DOMAIN_COUNT; i++) {
+        TEST_ASSERT_EQUAL(0, test_domains[i].node.data.power_budget);
+    }
+}
+
+void utest_set_root_budget(void)
+{
+    struct mod_power_distributor_domain_ctx test_domains[TEST_DOMAIN_COUNT] = {
+        0
+    };
+    uint32_t tree_traverse_order_table[TEST_DOMAIN_COUNT];
+    tree_traverse_order_table[0] = TEST_DOMAIN_SOC;
+
+    power_distributor_ctx.domain_count = TEST_DOMAIN_COUNT;
+    power_distributor_ctx.domain = test_domains;
+    power_distributor_ctx.tree_traverse_order_table = tree_traverse_order_table;
+
+    test_domains[TEST_DOMAIN_SOC].node.data.power_limit = 200;
+
+    set_root_budget();
+    TEST_ASSERT_EQUAL(
+        200, test_domains[TEST_DOMAIN_SOC].node.data.power_budget);
+}
 
 void tearDown(void)
 {
@@ -789,6 +827,8 @@ int power_distributor_test_main(void)
     RUN_TEST(utest_mod_distributor_calculate_power_attributes);
     RUN_TEST(utest_mod_distributor_calculate_power_attributes_max_power);
     RUN_TEST(utest_mod_distributor_domain_power_distribute);
+    RUN_TEST(utest_clear_domains_budget);
+    RUN_TEST(utest_set_root_budget);
 
     return UNITY_END();
 }