Spaces:
Paused
Paused
| import os | |
| import sys | |
| import argparse | |
| from traceback import print_exc | |
| import pickle | |
| import tqdm | |
| import pandas as pd | |
| from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor | |
| import torch | |
| import bittensor | |
| from meta_utils import load_metagraphs | |
| #TODO: make line charts and other cool stuff for each metagraph snapshot | |
| def process(block, netuid=1, lite=True, difficulty=False, prune_weights=False, return_graph=False, half=True, subtensor=None): | |
| if subtensor is None: | |
| subtensor = bittensor.subtensor(network='finney') | |
| try: | |
| metagraph = subtensor.metagraph(block=block, netuid=netuid, lite=lite) | |
| if difficulty: | |
| metagraph.difficulty = subtensor.difficulty(block=block, netuid=netuid) | |
| if not lite: | |
| if half: | |
| metagraph.weights = torch.nn.Parameter(metagraph.weights.half(), requires_grad=False) | |
| if prune_weights: | |
| metagraph.weights = metagraph.weights[metagraph.weights.sum(axis=1) > 0] | |
| with open(f'data/metagraph/{netuid}/{block}.pkl', 'wb') as f: | |
| pickle.dump(metagraph, f) | |
| return metagraph if return_graph else True | |
| except Exception as e: | |
| print(f'Error processing block {block}: {e}') | |
| def parse_arguments(): | |
| parser = argparse.ArgumentParser(description='Process metagraphs for a given network.') | |
| parser.add_argument('--netuid', type=int, default=1, help='Network UID to use.') | |
| parser.add_argument('--difficulty', action='store_true', help='Include difficulty in metagraph.') | |
| parser.add_argument('--prune_weights', action='store_true', help='Prune weights in metagraph.') | |
| parser.add_argument('--return_graph', action='store_true', help='Return metagraph instead of True.') | |
| parser.add_argument('--no_dataframe', action='store_true', help='Do not create dataframe.') | |
| parser.add_argument('--max_workers', type=int, default=32, help='Max workers to use.') | |
| parser.add_argument('--start_block', type=int, default=1_500_000, help='Start block.') | |
| parser.add_argument('--end_block', type=int, default=600_000, help='End block.') | |
| parser.add_argument('--step_size', type=int, default=100, help='Step size.') | |
| return parser.parse_args() | |
| if __name__ == '__main__': | |
| subtensor = bittensor.subtensor(network='finney') | |
| print(f'Current block: {subtensor.block}') | |
| args = parse_arguments() | |
| netuid=args.netuid | |
| difficulty=args.difficulty | |
| overwrite=False | |
| return_graph=args.return_graph | |
| step_size = args.step_size | |
| start_block = args.start_block | |
| start_block = (min(subtensor.block, start_block)//step_size)*step_size # round to nearest step_size | |
| end_block = args.end_block | |
| blocks = range(start_block, end_block, -step_size) | |
| # only get weights for multiple of 500 blocks | |
| lite=lambda x: x%500!=0 | |
| max_workers = min(args.max_workers, len(blocks)) | |
| datadir = f'data/metagraph/{netuid}' | |
| os.makedirs(datadir, exist_ok=True) | |
| if not overwrite: | |
| blocks = [block for block in blocks if not os.path.exists(f'data/metagraph/{netuid}/{block}.pkl')] | |
| metagraphs = [] | |
| if len(blocks)>0: | |
| print(f'Processing {len(blocks)} blocks from {blocks[0]}-{blocks[-1]} using {max_workers} workers.') | |
| with ProcessPoolExecutor(max_workers=max_workers) as executor: | |
| futures = [ | |
| executor.submit(process, block, lite=lite(block), netuid=netuid, difficulty=difficulty) | |
| for block in blocks | |
| ] | |
| success = 0 | |
| with tqdm.tqdm(total=len(futures)) as pbar: | |
| for block, future in zip(blocks,futures): | |
| try: | |
| metagraphs.append(future.result()) | |
| success += 1 | |
| except Exception as e: | |
| print(f'generated an exception: {print_exc(e)}') | |
| pbar.update(1) | |
| pbar.set_description(f'Processed {success} blocks. Current block: {block}') | |
| if not success: | |
| raise ValueError('No blocks were successfully processed.') | |
| print(f'Processed {success} blocks.') | |
| if return_graph: | |
| for metagraph in metagraphs: | |
| print(f'{metagraph.block}: {metagraph.n.item()} nodes, difficulty={getattr(metagraph, "difficulty", None)}, weights={metagraph.weights.shape if hasattr(metagraph, "weights") else None}') | |
| print(metagraphs[-1]) | |
| else: | |
| print(f'No blocks to process. Current block: {subtensor.block}') | |
| if not args.no_dataframe: | |
| save_path = f'data/metagraph/{netuid}/df.parquet' | |
| blocks = range(start_block, end_block, step_size) | |
| df_loaded = None | |
| if os.path.exists(save_path): | |
| df_loaded = pd.read_parquet(save_path) | |
| blocks = [block for block in blocks if block not in df_loaded.block.unique()] | |
| print(f'Loaded dataframe from {save_path!r}. {len(df_loaded)} rows. {len(blocks)} blocks to process.') | |
| if len(blocks)==0: | |
| print('No blocks to process.') | |
| sys.exit(0) | |
| df = load_metagraphs(blocks[0], blocks[-1], block_step=step_size, datadir=datadir) | |
| if df_loaded is not None: | |
| df = pd.concat([df, df_loaded], ignore_index=True) | |
| df.to_parquet(save_path) | |
| print(f'Saved dataframe to {save_path!r}') |