From 0a313bf6c5d82b57103b973f6c851e5186f91cb1 Mon Sep 17 00:00:00 2001
From: Alexander Borzunov <borzunov.alexander@gmail.com>
Date: Sun, 7 May 2023 14:57:05 +0400
Subject: [PATCH] Update hivemind to 1.1.8, enable efficient bfloat16 encoding
 (#311)

This PR:

1. Updates hivemind to 1.1.8 (includes https://github.com/learning-at-home/hivemind/pull/565)
2. Enables efficient bfloat16 serialization by default (`USE_LEGACY_BFLOAT16 = False`)
3. Removes logging code that was included to hivemind in https://github.com/learning-at-home/hivemind/pull/542
---
 setup.cfg                   |  2 +-
 src/petals/__init__.py      | 11 +++++++++++
 src/petals/utils/logging.py | 18 ------------------
 3 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index 786c8f5..8c237aa 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -37,7 +37,7 @@ install_requires =
     huggingface-hub>=0.11.1,<1.0.0
     transformers>=4.25.1,<5.0.0
     speedtest-cli==2.1.3
-    hivemind==1.1.7
+    hivemind==1.1.8
     tensor_parallel==1.0.23
     humanfriendly
     async-timeout>=4.0.2
diff --git a/src/petals/__init__.py b/src/petals/__init__.py
index 373c40a..7a39b49 100644
--- a/src/petals/__init__.py
+++ b/src/petals/__init__.py
@@ -1,6 +1,17 @@
+import os
+
+import hivemind
+
 from petals.client import *
 from petals.utils.logging import initialize_logs as _initialize_logs
 
 __version__ = "1.1.4"
 
+
+def _override_bfloat16_mode_default():
+    if os.getenv("USE_LEGACY_BFLOAT16") is None:
+        hivemind.compression.base.USE_LEGACY_BFLOAT16 = False
+
+
 _initialize_logs()
+_override_bfloat16_mode_default()
diff --git a/src/petals/utils/logging.py b/src/petals/utils/logging.py
index 6fe099f..0574fa0 100644
--- a/src/petals/utils/logging.py
+++ b/src/petals/utils/logging.py
@@ -4,16 +4,6 @@ import os
 from hivemind.utils import logging as hm_logging
 
 
-def in_jupyter() -> bool:
-    """Check if the code is run in Jupyter or Colab"""
-
-    try:
-        __IPYTHON__
-        return True
-    except NameError:
-        return False
-
-
 def initialize_logs():
     """Initialize Petals logging tweaks. This function is called when you import the `petals` module."""
 
@@ -21,14 +11,6 @@ def initialize_logs():
     if os.getenv("PETALS_LOGGING", "True").lower() in ("false", "0"):
         return
 
-    if in_jupyter():
-        os.environ["HIVEMIND_COLORS"] = "True"
-    importlib.reload(hm_logging)
-
-    # Remove log handlers from previous import of hivemind.utils.logging and extra handlers on Colab
-    hm_logging.get_logger().handlers.clear()
-    hm_logging.get_logger("hivemind").handlers.clear()
-
     hm_logging.use_hivemind_log_handler("in_root_logger")
 
     # We suppress asyncio error logs by default since they are mostly not relevant for the end user,