WIP checkin

Author: thewhaleking

bittensor_cli/cli.py

@@ -618,6 +618,7 @@ def __init__(self):
             "wallet_hotkey": None,
             "network": None,
             "use_cache": True,
+            "disk_cache": False,
             "rate_tolerance": None,
             "safe_staking": True,
             "allow_partial_stake": False,
@@ -1087,6 +1088,7 @@ def initialize_chain(
                 "Verify this is intended.",
             )
             if not self.subtensor:
+                use_disk_cache = self.config.get("disk_cache", False)
                 if network:
                     network_ = None
                     for item in network:
@@ -1103,15 +1105,15 @@ def initialize_chain(
                             f"[{COLORS.G.ARG}]{', '.join(not_selected_networks)}[/{COLORS.G.ARG}]"
                         )
 
-                    self.subtensor = SubtensorInterface(network_)
+                    self.subtensor = SubtensorInterface(network_, use_disk_cache=use_disk_cache)
                 elif self.config["network"]:
                     console.print(
                         f"Using the specified network [{COLORS.G.LINKS}]{self.config['network']}"
                         f"[/{COLORS.G.LINKS}] from config"
                     )
-                    self.subtensor = SubtensorInterface(self.config["network"])
+                    self.subtensor = SubtensorInterface(self.config["network"], use_disk_cache=use_disk_cache)
                 else:
-                    self.subtensor = SubtensorInterface(defaults.subtensor.network)
+                    self.subtensor = SubtensorInterface(defaults.subtensor.network, use_disk_cache=use_disk_cache)
         return self.subtensor
 
     def _run_command(self, cmd: Coroutine, exit_early: bool = True):
@@ -1268,6 +1270,13 @@ def set_config(
             help="Disable caching of some commands. This will disable the `--reuse-last` and `--html` flags on "
             "commands such as `subnets metagraph`, `stake show` and `subnets list`.",
         ),
+        disk_cache: Optional[bool] = typer.Option(
+            None,
+            "--disk-cache/--no-disk-cache",
+            " /--no-disk-cache",
+            help="Enables or disables the caching on disk. Enabling this can significantly speed up commands run "
+                 "sequentially"
+        ),
         rate_tolerance: Optional[float] = typer.Option(
             None,
             "--tolerance",
@@ -1314,12 +1323,13 @@ def set_config(
             "wallet_hotkey": wallet_hotkey,
             "network": network,
             "use_cache": use_cache,
+            "disk_cache": disk_cache,
             "rate_tolerance": rate_tolerance,
             "safe_staking": safe_staking,
             "allow_partial_stake": allow_partial_stake,
             "dashboard_path": dashboard_path,
         }
-        bools = ["use_cache", "safe_staking", "allow_partial_stake"]
+        bools = ["use_cache", "disk_cache", "safe_staking", "allow_partial_stake"]
         if all(v is None for v in args.values()):
             # Print existing configs
             self.get_config()

bittensor_cli/src/__init__.py

@@ -94,6 +94,7 @@ class config:
             "wallet_name": None,
             "wallet_hotkey": None,
             "use_cache": True,
+            "disk_cache": False,
             "metagraph_cols": {
                 "UID": True,
                 "GLOBAL_STAKE": True,

bittensor_cli/src/bittensor/subtensor_interface.py

@@ -42,12 +42,6 @@
     get_hotkey_pub_ss58,
 )
 
-SubstrateClass = (
-    DiskCachedAsyncSubstrateInterface
-    if os.getenv("DISK_CACHE", "0") == "1"
-    else AsyncSubstrateInterface
-)
-
 
 class ParamWithTypes(TypedDict):
     name: str  # Name of the parameter.
@@ -81,7 +75,7 @@ class SubtensorInterface:
     Thin layer for interacting with Substrate Interface. Mostly a collection of frequently-used calls.
     """
 
-    def __init__(self, network):
+    def __init__(self, network, use_disk_cache: bool = False):
         if network in Constants.network_map:
             self.chain_endpoint = Constants.network_map[network]
             self.network = network
@@ -111,8 +105,12 @@ def __init__(self, network):
                 )
                 self.chain_endpoint = Constants.network_map[defaults.subtensor.network]
                 self.network = defaults.subtensor.network
-
-        self.substrate = SubstrateClass(
+        substrate_class = (
+            DiskCachedAsyncSubstrateInterface
+            if (use_disk_cache or os.getenv("DISK_CACHE", "0") == "1")
+            else AsyncSubstrateInterface
+        )
+        self.substrate = substrate_class(
             url=self.chain_endpoint,
             ss58_format=SS58_FORMAT,
             type_registry=TYPE_REGISTRY,