Add version check caching, fix version comparison

bittensor/cli.py

@@ -246,8 +246,8 @@ def __init__(
         # If no_version_checking is not set or set as False in the config, version checking is done.
         if not self.config.get("no_version_checking", d=True):
             try:
-                bittensor.utils.version_checking()
-            except:
+                bittensor.utils.check_version()
+            except bittensor.utils.VersionCheckError:
                 # If version checking fails, inform user with an exception.
                 raise RuntimeError(
                     "To avoid internet-based version checking, pass --no_version_checking while running the CLI."

bittensor/utils/__init__.py

@@ -20,11 +20,11 @@
 
 import bittensor
 import hashlib
-import requests
 import torch
 import scalecodec
 
 from .wallet_utils import *  # noqa F401
+from .version import version_checking, check_version, VersionCheckError
 
 RAOPERTAO = 1e9
 U16_MAX = 65535
@@ -59,34 +59,6 @@ def unbiased_topk(values, k, dim=0, sorted=True, largest=True):
     return topk, permutation[indices]
 
 
-def version_checking(timeout: int = 15):
-    try:
-        bittensor.logging.debug(
-            f"Checking latest Bittensor version at: {bittensor.__pipaddress__}"
-        )
-        response = requests.get(bittensor.__pipaddress__, timeout=timeout)
-        latest_version = response.json()["info"]["version"]
-        version_split = latest_version.split(".")
-        latest_version_as_int = (
-            (100 * int(version_split[0]))
-            + (10 * int(version_split[1]))
-            + (1 * int(version_split[2]))
-        )
-
-        if latest_version_as_int > bittensor.__version_as_int__:
-            print(
-                "\u001b[33mBittensor Version: Current {}/Latest {}\nPlease update to the latest version at your earliest convenience. "
-                "Run the following command to upgrade:\n\n\u001b[0mpython -m pip install --upgrade bittensor".format(
-                    bittensor.__version__, latest_version
-                )
-            )
-
-    except requests.exceptions.Timeout:
-        bittensor.logging.error("Version check failed due to timeout")
-    except requests.exceptions.RequestException as e:
-        bittensor.logging.error(f"Version check failed due to request failure: {e}")
-
-
 def strtobool_with_default(
     default: bool,
 ) -> Callable[[str], Union[bool, Literal["==SUPRESS=="]]]:

bittensor/utils/version.py

@@ -0,0 +1,103 @@
+from typing import Optional
+from pathlib import Path
+import time
+from packaging.version import Version
+
+import bittensor
+import requests
+
+VERSION_CHECK_THRESHOLD = 86400
+
+
+class VersionCheckError(Exception):
+    pass
+
+
+def _get_version_file_path() -> Path:
+    return Path.home() / ".bittensor" / ".last_known_version"
+
+
+def _get_version_from_file(version_file: Path) -> Optional[str]:
+    try:
+        mtime = version_file.stat().st_mtime
+        bittensor.logging.debug(f"Found version file, last modified: {mtime}")
+        diff = time.time() - mtime
+
+        if diff >= VERSION_CHECK_THRESHOLD:
+            bittensor.logging.debug("Version file expired")
+            return None
+
+        return version_file.read_text()
+    except FileNotFoundError:
+        bittensor.logging.debug("No bitensor version file found")
+        return None
+    except OSError:
+        bittensor.logging.exception("Failed to read version file")
+        return None
+
+
+def _get_version_from_pypi(timeout: int = 15) -> str:
+    bittensor.logging.debug(
+        f"Checking latest Bittensor version at: {bittensor.__pipaddress__}"
+    )
+    try:
+        response = requests.get(bittensor.__pipaddress__, timeout=timeout)
+        latest_version = response.json()["info"]["version"]
+        return latest_version
+    except requests.exceptions.RequestException:
+        bittensor.logging.exception("Failed to get latest version from pypi")
+        raise
+
+
+def get_and_save_latest_version(timeout: int = 15) -> str:
+    version_file = _get_version_file_path()
+
+    if last_known_version := _get_version_from_file(version_file):
+        return last_known_version
+
+    latest_version = _get_version_from_pypi(timeout)
+
+    try:
+        version_file.write_text(latest_version)
+    except OSError:
+        bittensor.logging.exception("Failed to save latest version to file")
+
+    return latest_version
+
+
+def check_version(timeout: int = 15):
+    """
+    Check if the current version of Bittensor is up to date with the latest version on PyPi.
+    Raises a VersionCheckError if the version check fails.
+    """
+
+    try:
+        latest_version = get_and_save_latest_version(timeout)
+
+        if Version(latest_version) > Version(bittensor.__version__):
+            print(
+                "\u001b[33mBittensor Version: Current {}/Latest {}\nPlease update to the latest version at your earliest convenience. "
+                "Run the following command to upgrade:\n\n\u001b[0mpython -m pip install --upgrade bittensor".format(
+                    bittensor.__version__, latest_version
+                )
+            )
+    except Exception as e:
+        raise VersionCheckError("Version check failed") from e
+
+
+def version_checking(timeout: int = 15):
+    """
+    Deprecated, kept for backwards compatibility. Use check_version() instead.
+    """
+
+    from warnings import warn
+
+    warn(
+        "version_checking() is deprecated, please use check_version() instead",
+        DeprecationWarning,
+    )
+
+    try:
+        check_version(timeout)
+    except VersionCheckError:
+        bittensor.logging.exception("Version check failed")

requirements/dev.txt

@@ -10,3 +10,4 @@ hypothesis==6.81.1
 flake8==7.0.0
 mypy==1.8.0
 types-retry==0.9.9.4
+freezegun==1.5.0

requirements/prod.txt

@@ -14,6 +14,7 @@ netaddr
 numpy
 msgpack-numpy-opentensor==0.5.0
 nest_asyncio
+packaging
 pycryptodome>=3.18.0,<4.0.0
 pyyaml
 password_strength

tests/integration_tests/test_cli.py

@@ -2360,18 +2360,11 @@ def test_delegate(self, _):
         self.assertAlmostEqual(new_balance.tao, old_balance.tao - 10.0, delta=1e-6)
 
 
-# Test directory for creating mock wallets
-TEST_DIR = "/tmp/test_bittensor_wallets"
-
-
 @pytest.fixture(scope="function")
-def setup_wallets():
-    # Arrange: Create a temporary directory to simulate wallet paths
-    if not os.path.exists(TEST_DIR):
-        os.makedirs(TEST_DIR)
-    yield
-    # Teardown: Remove the temporary directory after tests
-    shutil.rmtree(TEST_DIR)
+def wallets_dir_path(tmp_path):
+    wallets_dir = tmp_path / "wallets"
+    wallets_dir.mkdir()
+    yield wallets_dir
 
 
 @pytest.mark.parametrize(
@@ -2387,15 +2380,14 @@ def setup_wallets():
     ],
 )
 def test_get_coldkey_wallets_for_path(
-    test_id, wallet_names, expected_wallet_count, setup_wallets
+    test_id, wallet_names, expected_wallet_count, wallets_dir_path
 ):
     # Arrange: Create mock wallet directories
     for name in wallet_names:
-        wallet_path = os.path.join(TEST_DIR, name)
-        os.makedirs(wallet_path)
+        (wallets_dir_path / name).mkdir()
 
     # Act: Call the function with the test directory
-    wallets = _get_coldkey_wallets_for_path(TEST_DIR)
+    wallets = _get_coldkey_wallets_for_path(str(wallets_dir_path))
 
     # Assert: Check if the correct number of wallet objects are returned
     assert len(wallets) == expected_wallet_count
@@ -2536,11 +2528,6 @@ def test_set_identity_command(
             assert mock_exit.call_count == 0
 
 
-TEST_DIR = "/tmp/test_bittensor_wallets"
-if not os.path.exists(TEST_DIR):
-    os.makedirs(TEST_DIR)
-
-
 @pytest.fixture
 def setup_files(tmp_path):
     def _setup_files(files):
@@ -2582,11 +2569,5 @@ def test_get_coldkey_ss58_addresses_for_path(
     ), f"Test ID: {test_id} failed. Expected {expected}, got {result}"
 
 
-# Cleanup after tests
-def teardown_module(module):
-    with contextlib.suppress(FileNotFoundError):
-        shutil.rmtree(TEST_DIR)
-
-
 if __name__ == "__main__":
     unittest.main()