Skip to content

Commit

Permalink
feat: use rust shuffle (#7120)
Browse files Browse the repository at this point in the history
* feat: add temp-deps to test on feat group

* feat: add build_temp_deps.sh process

* feat: use async for shuffling in epoch transition

* feat: use v0.0.1 instead of temp-deps

* chore: lint and check-types

* fix: log error context

* refactor: use toHex

* fix: address computeEpochShuffling refactor comments

* test: remove perf tests that were moved to swp-or-not-shuffle package

* test: add strict equal check for sync/async computeEpochShuffling impls

* Revert "refactor: use toHex"

This reverts commit 9d64b67.

* fix: EpochShuffling ssz type issue

* feat: upgrade swap-or-not to 0.0.2

* refactor: buildCommitteesFromShuffling

* docs: add TODO about removing state from shuffling computation

* docs: add TODO about removing state from shuffling computation
  • Loading branch information
matthewkeil authored Oct 11, 2024
1 parent 105a388 commit 911a3f5
Show file tree
Hide file tree
Showing 12 changed files with 161 additions and 399 deletions.
21 changes: 13 additions & 8 deletions packages/beacon-node/src/chain/shufflingCache.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,11 @@ import {
IShufflingCache,
ShufflingBuildProps,
computeEpochShuffling,
computeEpochShufflingAsync,
} from "@lodestar/state-transition";
import {Epoch, RootHex} from "@lodestar/types";
import {LodestarError, Logger, MapDef, pruneSetToMax} from "@lodestar/utils";
import {Metrics} from "../metrics/metrics.js";
import {callInNextEventLoop} from "../util/eventLoop.js";

/**
* Same value to CheckpointBalancesCache, with the assumption that we don't have to use it for old epochs. In the worse case:
Expand Down Expand Up @@ -178,14 +178,19 @@ export class ShufflingCache implements IShufflingCache {
this.insertPromise(epoch, decisionRoot);
/**
* TODO: (@matthewkeil) This will get replaced by a proper build queue and a worker to do calculations
* on a NICE thread with a rust implementation
* on a NICE thread
*/
callInNextEventLoop(() => {
const timer = this.metrics?.shufflingCache.shufflingCalculationTime.startTimer({source: "build"});
const shuffling = computeEpochShuffling(state, activeIndices, epoch);
timer?.();
this.set(shuffling, decisionRoot);
});
const timer = this.metrics?.shufflingCache.shufflingCalculationTime.startTimer({source: "build"});
computeEpochShufflingAsync(state, activeIndices, epoch)
.then((shuffling) => {
this.set(shuffling, decisionRoot);
})
.catch((err) =>
this.logger?.error(`error building shuffling for epoch ${epoch} at decisionRoot ${decisionRoot}`, {}, err)
)
.finally(() => {
timer?.();
});
}

/**
Expand Down
17 changes: 10 additions & 7 deletions packages/beacon-node/test/spec/presets/shuffling.test.ts
Original file line number Diff line number Diff line change
@@ -1,24 +1,27 @@
import path from "node:path";
import {unshuffleList} from "@lodestar/state-transition";
import {unshuffleList} from "@chainsafe/swap-or-not-shuffle";
import {InputType} from "@lodestar/spec-test-util";
import {bnToNum, fromHex} from "@lodestar/utils";
import {ACTIVE_PRESET} from "@lodestar/params";
import {ACTIVE_PRESET, SHUFFLE_ROUND_COUNT} from "@lodestar/params";
import {RunnerType, TestRunnerFn} from "../utils/types.js";
import {ethereumConsensusSpecsTests} from "../specTestVersioning.js";
import {specTestIterator} from "../utils/specTestIterator.js";

const shuffling: TestRunnerFn<ShufflingTestCase, number[]> = () => {
const shuffling: TestRunnerFn<ShufflingTestCase, string> = () => {
return {
testFunction: (testcase) => {
const seed = fromHex(testcase.mapping.seed);
const output = Array.from({length: bnToNum(testcase.mapping.count)}, (_, i) => i);
unshuffleList(output, seed);
return output;
const output = unshuffleList(
Uint32Array.from(Array.from({length: bnToNum(testcase.mapping.count)}, (_, i) => i)),
seed,
SHUFFLE_ROUND_COUNT
);
return Buffer.from(output).toString("hex");
},
options: {
inputTypes: {mapping: InputType.YAML},
timeout: 10000,
getExpected: (testCase) => testCase.mapping.mapping.map((value) => bnToNum(value)),
getExpected: (testCase) => Buffer.from(testCase.mapping.mapping.map((value) => bnToNum(value))).toString("hex"),
// Do not manually skip tests here, do it in packages/beacon-node/test/spec/presets/index.test.ts
},
};
Expand Down
1 change: 1 addition & 0 deletions packages/state-transition/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
"@chainsafe/persistent-merkle-tree": "^0.8.0",
"@chainsafe/persistent-ts": "^0.19.1",
"@chainsafe/ssz": "^0.17.1",
"@chainsafe/swap-or-not-shuffle": "^0.0.2",
"@lodestar/config": "^1.22.0",
"@lodestar/params": "^1.22.0",
"@lodestar/types": "^1.22.0",
Expand Down
59 changes: 41 additions & 18 deletions packages/state-transition/src/util/epochShuffling.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import {asyncUnshuffleList, unshuffleList} from "@chainsafe/swap-or-not-shuffle";
import {Epoch, RootHex, ssz, ValidatorIndex} from "@lodestar/types";
import {GaugeExtra, intDiv, Logger, NoLabels, toRootHex} from "@lodestar/utils";
import {
Expand All @@ -6,11 +7,11 @@ import {
MAX_COMMITTEES_PER_SLOT,
SLOTS_PER_EPOCH,
TARGET_COMMITTEE_SIZE,
SHUFFLE_ROUND_COUNT,
} from "@lodestar/params";
import {BeaconConfig} from "@lodestar/config";
import {BeaconStateAllForks} from "../types.js";
import {getSeed} from "./seed.js";
import {unshuffleList} from "./shuffle.js";
import {computeStartSlotAtEpoch} from "./epoch.js";
import {getBlockRootAtSlot} from "./blockRoot.js";
import {computeAnchorCheckpoint} from "./computeAnchorCheckpoint.js";
Expand Down Expand Up @@ -102,42 +103,64 @@ export function computeCommitteeCount(activeValidatorCount: number): number {
return Math.max(1, Math.min(MAX_COMMITTEES_PER_SLOT, committeesPerSlot));
}

export function computeEpochShuffling(
state: BeaconStateAllForks,
activeIndices: Uint32Array,
epoch: Epoch
): EpochShuffling {
const activeValidatorCount = activeIndices.length;

const shuffling = activeIndices.slice();
const seed = getSeed(state, epoch, DOMAIN_BEACON_ATTESTER);
unshuffleList(shuffling, seed);

function buildCommitteesFromShuffling(shuffling: Uint32Array): Uint32Array[][] {
const activeValidatorCount = shuffling.length;
const committeesPerSlot = computeCommitteeCount(activeValidatorCount);

const committeeCount = committeesPerSlot * SLOTS_PER_EPOCH;

const committees: Uint32Array[][] = [];
const committees = new Array<Uint32Array[]>(SLOTS_PER_EPOCH);
for (let slot = 0; slot < SLOTS_PER_EPOCH; slot++) {
const slotCommittees: Uint32Array[] = [];
const slotCommittees = new Array<Uint32Array>(committeesPerSlot);

for (let committeeIndex = 0; committeeIndex < committeesPerSlot; committeeIndex++) {
const index = slot * committeesPerSlot + committeeIndex;
const startOffset = Math.floor((activeValidatorCount * index) / committeeCount);
const endOffset = Math.floor((activeValidatorCount * (index + 1)) / committeeCount);
if (!(startOffset <= endOffset)) {
throw new Error(`Invalid offsets: start ${startOffset} must be less than or equal end ${endOffset}`);
}
slotCommittees.push(shuffling.subarray(startOffset, endOffset));
slotCommittees[committeeIndex] = shuffling.subarray(startOffset, endOffset);
}
committees.push(slotCommittees);

committees[slot] = slotCommittees;
}

return committees;
}

export function computeEpochShuffling(
// TODO: (@matthewkeil) remove state/epoch and pass in seed to clean this up
state: BeaconStateAllForks,
activeIndices: Uint32Array,
epoch: Epoch
): EpochShuffling {
const seed = getSeed(state, epoch, DOMAIN_BEACON_ATTESTER);
const shuffling = unshuffleList(activeIndices, seed, SHUFFLE_ROUND_COUNT);
const committees = buildCommitteesFromShuffling(shuffling);
return {
epoch,
activeIndices,
shuffling,
committees,
committeesPerSlot: committees[0].length,
};
}

export async function computeEpochShufflingAsync(
// TODO: (@matthewkeil) remove state/epoch and pass in seed to clean this up
state: BeaconStateAllForks,
activeIndices: Uint32Array,
epoch: Epoch
): Promise<EpochShuffling> {
const seed = getSeed(state, epoch, DOMAIN_BEACON_ATTESTER);
const shuffling = await asyncUnshuffleList(activeIndices, seed, SHUFFLE_ROUND_COUNT);
const committees = buildCommitteesFromShuffling(shuffling);
return {
epoch,
activeIndices,
shuffling,
committees,
committeesPerSlot,
committeesPerSlot: committees[0].length,
};
}

Expand Down
1 change: 0 additions & 1 deletion packages/state-transition/src/util/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ export * from "./genesis.js";
export * from "./interop.js";
export * from "./rootCache.js";
export * from "./seed.js";
export * from "./shuffle.js";
export * from "./shufflingDecisionRoot.js";
export * from "./signatureSets.js";
export * from "./signingRoot.js";
Expand Down
Loading

0 comments on commit 911a3f5

Please sign in to comment.