feat(fork-network): override configs at SetValidatorsCmd (#12575)

I've made two observations:
* `num_seats` is not propagated to forknet in `set-validators` command,
like it was done in `amend-genesis`
* the actual node dir preparation happens at `set-validators`, not
`init_configs`

So I suggest to move epoch config overrides to `set-validators` as well.
It looks very natural. I think we can always assume that
`genesis_protocol_version` is the first version, and then pass epoch
config overrides through `new-test` to `set-validators`. I used this for
forknet resharding and it worked well.

If epoch config overrides happen at init_configs, then, after init and
before new-test, forknet nodes will be in weird state where epoch
configs are already overridden but genesis isn't.

chain/indexer/src/lib.rs

@@ -172,6 +172,5 @@ pub fn indexer_init_configs(
         params.download_config_url.as_deref(),
         params.boot_nodes.as_deref(),
         params.max_gas_burnt_view,
-        None,
     )
 }

nearcore/src/config.rs

@@ -37,8 +37,6 @@ use near_jsonrpc::RpcConfig;
 use near_network::config::NetworkConfig;
 use near_network::tcp;
 use near_o11y::log_config::LogConfig;
-use near_primitives::chains::MAINNET;
-use near_primitives::epoch_manager::EpochConfigStore;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::ShardLayout;
 use near_primitives::test_utils::create_test_signer;
@@ -48,7 +46,7 @@ use near_primitives::types::{
 };
 use near_primitives::utils::{from_timestamp, get_num_seats_per_shard};
 use near_primitives::validator_signer::{InMemoryValidatorSigner, ValidatorSigner};
-use near_primitives::version::{ProtocolVersion, PROTOCOL_VERSION};
+use near_primitives::version::PROTOCOL_VERSION;
 #[cfg(feature = "rosetta_rpc")]
 use near_rosetta_rpc::RosettaRpcConfig;
 use near_store::config::{
@@ -792,7 +790,6 @@ pub fn init_configs(
     download_config_url: Option<&str>,
     boot_nodes: Option<&str>,
     max_gas_burnt_view: Option<Gas>,
-    dump_epoch_config: Option<ProtocolVersion>,
 ) -> anyhow::Result<()> {
     fs::create_dir_all(dir).with_context(|| anyhow!("Failed to create directory {:?}", dir))?;
 
@@ -987,20 +984,6 @@ pub fn init_configs(
         }
     }
 
-    if let Some(first_version) = dump_epoch_config {
-        let epoch_config_dir = dir.join("epoch_configs");
-        fs::create_dir_all(epoch_config_dir.clone())
-            .with_context(|| anyhow!("Failed to create directory {:?}", epoch_config_dir))?;
-        EpochConfigStore::for_chain_id(MAINNET, None)
-            .expect("Could not load the EpochConfigStore for mainnet.")
-            .dump_epoch_configs_between(
-                &first_version,
-                &PROTOCOL_VERSION,
-                epoch_config_dir.to_str().unwrap(),
-            );
-        info!(target: "near", "Generated epoch configs files in {}", epoch_config_dir.display());
-    }
-
     Ok(())
 }
 
@@ -1526,7 +1509,6 @@ mod tests {
             None,
             None,
             None,
-            None,
         )
         .unwrap();
         let genesis = Genesis::from_file(
@@ -1584,7 +1566,6 @@ mod tests {
             None,
             None,
             None,
-            None,
         )
         .unwrap();
 
@@ -1618,7 +1599,6 @@ mod tests {
             None,
             None,
             None,
-            None,
         )
         .unwrap();
 

neard/src/cli.rs

@@ -255,6 +255,7 @@ pub(super) enum NeardSubCommand {
     ReplayArchive(ReplayArchiveCommand),
 }
 
+#[allow(unused)]
 #[derive(Debug, Clone)]
 enum FirstProtocolVersion {
     Since(ProtocolVersion),
@@ -320,10 +321,6 @@ pub(super) struct InitCmd {
     /// from genesis configuration will be taken.
     #[clap(long)]
     max_gas_burnt_view: Option<Gas>,
-    /// Dump epoch config from the given protocol version onwards.
-    /// Can be a number or the word "latest".
-    #[clap(long)]
-    dump_epoch_config: Option<FirstProtocolVersion>,
 }
 
 /// Warns if unsupported build of the executable is used on mainnet or testnet.
@@ -381,11 +378,6 @@ impl InitCmd {
             None
         };
 
-        let dump_epoch_config = self.dump_epoch_config.map(|first| match first {
-            FirstProtocolVersion::Since(version) => version,
-            FirstProtocolVersion::Latest => near_primitives::version::PROTOCOL_VERSION,
-        });
-
         nearcore::init_configs(
             home_dir,
             self.chain_id,
@@ -401,7 +393,6 @@ impl InitCmd {
             self.download_config_url.as_deref(),
             self.boot_nodes.as_deref(),
             self.max_gas_burnt_view,
-            dump_epoch_config,
         )
         .context("Failed to initialize configs")
     }

pytest/tests/mocknet/helpers/neard_runner.py

@@ -1080,12 +1080,19 @@ def network_init(self):
             self.set_state(TestState.AMEND_GENESIS)
         else:
             cmd = [
-                self.data['binaries'][0]['system_path'], '--home',
-                self.target_near_home_path(), 'fork-network', 'set-validators',
+                self.data['binaries'][0]['system_path'],
+                '--home',
+                self.target_near_home_path(),
+                'fork-network',
+                'set-validators',
                 '--validators',
-                self.home_path('validators.json'), '--epoch-length',
-                str(n['epoch_length']), '--genesis-time',
-                str(n['genesis_time'])
+                self.home_path('validators.json'),
+                '--epoch-length',
+                str(n['epoch_length']),
+                '--genesis-time',
+                str(n['genesis_time']),
+                '--num-seats',
+                str(n['num_seats']),
             ]
             if new_chain_id is not None:
                 cmd.append('--chain-id')

runtime/runtime-params-estimator/src/main.rs

@@ -181,7 +181,6 @@ fn run_estimation(cli_args: CliArgs) -> anyhow::Result<Option<CostTable>> {
             None,
             None,
             None,
-            None,
         )
         .expect("failed to init config");
 

tools/fork-network/src/cli.rs

@@ -14,7 +14,7 @@ use near_parameters::{RuntimeConfig, RuntimeConfigStore};
 use near_primitives::account::id::AccountType;
 use near_primitives::account::{AccessKey, AccessKeyPermission, Account};
 use near_primitives::borsh;
-use near_primitives::epoch_manager::EpochConfigStore;
+use near_primitives::epoch_manager::{EpochConfig, EpochConfigStore};
 use near_primitives::hash::CryptoHash;
 use near_primitives::serialize::dec_format;
 use near_primitives::shard_layout::ShardUId;
@@ -23,7 +23,7 @@ use near_primitives::state_record::StateRecord;
 use near_primitives::trie_key::col;
 use near_primitives::trie_key::trie_key_parsers::parse_account_id_from_account_key;
 use near_primitives::types::{
-    AccountId, AccountInfo, Balance, BlockHeight, EpochId, NumBlocks, ShardId, StateRoot,
+    AccountId, AccountInfo, Balance, BlockHeight, EpochId, NumBlocks, NumSeats, ShardId, StateRoot,
 };
 use near_primitives::version::{ProtocolVersion, PROTOCOL_VERSION};
 use near_store::adapter::StoreAdapter;
@@ -36,7 +36,7 @@ use near_store::{
 use nearcore::{load_config, open_storage, NearConfig, NightshadeRuntime, NightshadeRuntimeExt};
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use serde::Deserialize;
-use std::collections::HashSet;
+use std::collections::{BTreeMap, HashSet};
 use std::fs::File;
 use std::io::BufReader;
 use std::path::{Path, PathBuf};
@@ -122,6 +122,9 @@ struct SetValidatorsCmd {
     /// will be used.
     #[arg(long)]
     pub protocol_version: Option<ProtocolVersion>,
+    /// Number of validator seats.
+    #[clap(long)]
+    pub num_seats: Option<NumSeats>,
 }
 
 const FORKED_ROOTS_KEY_PREFIX: &str = "FORK_TOOL_SHARD_ID:";
@@ -211,12 +214,14 @@ impl ForkNetworkCommand {
                 epoch_length,
                 chain_id_suffix,
                 chain_id,
+                num_seats,
             }) => {
                 self.set_validators(
                     genesis_time.unwrap_or_else(chrono::Utc::now),
                     *protocol_version,
                     validators,
                     *epoch_length,
+                    num_seats,
                     chain_id_suffix,
                     chain_id,
                     near_config,
@@ -395,6 +400,7 @@ impl ForkNetworkCommand {
         protocol_version: Option<ProtocolVersion>,
         validators: &Path,
         epoch_length: u64,
+        num_seats: &Option<NumSeats>,
         chain_id_suffix: &str,
         chain_id: &Option<String>,
         near_config: &mut NearConfig,
@@ -421,7 +427,7 @@ impl ForkNetworkCommand {
         let runtime_config = runtime_config_store.get_config(PROTOCOL_VERSION);
 
         let storage_mutator =
-            StorageMutator::new(epoch_manager.clone(), &runtime, epoch_id, prev_state_roots)?;
+            StorageMutator::new(epoch_manager, &runtime, epoch_id, prev_state_roots)?;
         let (new_state_roots, new_validator_accounts) =
             self.add_validator_accounts(validators, runtime_config, home_dir, storage_mutator)?;
 
@@ -431,13 +437,12 @@ impl ForkNetworkCommand {
             genesis_time,
             protocol_version,
             epoch_length,
+            num_seats,
             block_height,
             chain_id_suffix,
             chain_id,
-            &epoch_id,
             new_state_roots.clone(),
             new_validator_accounts.clone(),
-            epoch_manager,
             home_dir,
             near_config,
         )?;
@@ -527,6 +532,47 @@ impl ForkNetworkCommand {
         return Ok(());
     }
 
+    /// Creates epoch config overrides since `first_version` and places them
+    /// in `home_dir`.
+    fn override_epoch_configs(
+        &self,
+        first_version: ProtocolVersion,
+        num_seats: &Option<NumSeats>,
+        home_dir: &Path,
+    ) -> anyhow::Result<EpochConfig> {
+        let epoch_config_dir = home_dir.join("epoch_configs");
+        if epoch_config_dir.exists() {
+            std::fs::remove_dir_all(epoch_config_dir.clone())?;
+        }
+        std::fs::create_dir_all(epoch_config_dir.clone()).with_context(|| {
+            anyhow::anyhow!("Failed to create directory {:?}", epoch_config_dir)
+        })?;
+
+        let base_epoch_config_store =
+            EpochConfigStore::for_chain_id(near_primitives::chains::MAINNET, None)
+                .expect("Could not load the EpochConfigStore for mainnet.");
+        let mut new_epoch_configs = BTreeMap::new();
+        for version in first_version..=PROTOCOL_VERSION {
+            let mut config = base_epoch_config_store.get_config(version).as_ref().clone();
+            if let Some(num_seats) = num_seats {
+                config.num_block_producer_seats = *num_seats;
+                config.validator_selection_config.num_chunk_producer_seats = *num_seats;
+                config.validator_selection_config.num_chunk_validator_seats = *num_seats;
+            }
+            new_epoch_configs.insert(version, Arc::new(config));
+        }
+        let first_config = new_epoch_configs.get(&first_version).unwrap().as_ref().clone();
+        let epoch_config_store = EpochConfigStore::test(new_epoch_configs);
+
+        epoch_config_store.dump_epoch_configs_between(
+            &first_version,
+            &PROTOCOL_VERSION,
+            epoch_config_dir.to_str().unwrap(),
+        );
+        tracing::info!(target: "near", "Generated epoch configs files in {}", epoch_config_dir.display());
+        Ok(first_config)
+    }
+
     fn prepare_shard_state(
         &self,
         batch_size: u64,
@@ -824,13 +870,12 @@ impl ForkNetworkCommand {
         genesis_time: DateTime<Utc>,
         protocol_version: Option<ProtocolVersion>,
         epoch_length: u64,
+        num_seats: &Option<NumSeats>,
         height: BlockHeight,
         chain_id_suffix: &str,
         chain_id: &Option<String>,
-        epoch_id: &EpochId,
         new_state_roots: Vec<StateRoot>,
         new_validator_accounts: Vec<AccountInfo>,
-        epoch_manager: Arc<EpochManagerHandle>,
         home_dir: &Path,
         near_config: &mut NearConfig,
     ) -> anyhow::Result<()> {
@@ -850,18 +895,9 @@ impl ForkNetworkCommand {
         // This is based on the assumption that epoch length is part of genesis config and not epoch config.
         near_config.genesis.config.epoch_length = epoch_length;
 
-        let epoch_config_dir = home_dir.join("epoch_configs");
-        let epoch_config = if epoch_config_dir.exists() {
-            tracing::info!(target: "fork-network", "Loading epoch config from {:?}", epoch_config_dir);
-            EpochConfigStore::for_chain_id(&new_chain_id, Some(epoch_config_dir))
-                .unwrap()
-                .get_config(genesis_protocol_version)
-                .as_ref()
-                .clone()
-        } else {
-            tracing::info!(target: "fork-network", "Loading epoch config from epoch manager");
-            epoch_manager.get_epoch_config(epoch_id)?
-        };
+        let epoch_config =
+            self.override_epoch_configs(genesis_protocol_version, num_seats, home_dir)?;
+
         let original_config = near_config.genesis.config.clone();
 
         let new_config = GenesisConfig {