From 8e263b8b7d3af73cd991614740111984d2c8038f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 1 Apr 2024 04:38:12 +1100
Subject: [PATCH] Nightly (#293)

Env checking
---
 unsloth/models/_utils.py  | 39 +++++++++++++++++++++++++++++++++++++++
 unsloth/models/llama.py   |  1 +
 unsloth/models/mistral.py |  1 +
 3 files changed, 41 insertions(+)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 1989313ee6..8cdb5e384c 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -27,6 +27,8 @@ from platform import system as platform_system
 platform_system = platform_system()
 import math
 import numpy as np
+import os
+import psutil
 
 __version__ = "2024.3"
 
@@ -67,6 +69,7 @@ __all__ = [
     "HAS_FLASH_ATTENTION",
     "platform_system",
     "patch_tokenizer",
+    "get_statistics",
 ]
 
 
@@ -169,6 +172,42 @@ except:
 pass
 
 
+def get_statistics():
+    # We log some basic stats about which environment is being used.
+    # We simply download a README.md file from HF - all data is made public.
+    # This is simply so we can check if some envs are broken or not.
+    try:
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled
+        n_cpus = psutil.cpu_count(logical = False)
+
+        keynames = "\n" + "\n".join(os.environ.keys())
+        statistics = None
+        if   "\nCOLAB_"  in keynames and n_cpus == 1: statistics = "colab"
+        elif "\nCOLAB_"  in keynames: statistics = "colabpro"
+        elif "\nKAGGLE_" in keynames: statistics = "kaggle"
+        elif "\nRUNPOD_" in keynames: statistics = "runpod"
+        elif "\nAWS_"    in keynames: statistics = "aws"
+        elif "\nAZURE_"  in keynames: statistics = "azure"
+        elif "\nK_" in keynames or "\nFUNCTION_" in keynames: statistics = "gcp"
+        elif "\nINVOCATION_ID" in keynames: statistics = "lambda"
+
+        if statistics is not None:
+            disabled = False
+            if not are_progress_bars_disabled():
+                disable_progress_bars()
+                disabled = True
+            pass
+            hf_hub_download(f"unslothai/statistics-{statistics}", "README.md", force_download = True)
+            if disabled:
+                enable_progress_bars()
+            pass
+        pass
+    except:
+        pass
+pass
+
+
 def _calculate_n_gradient_checkpoints(
     n_layers : int,
     method   : Optional[Union[str, int]] = "sqrt",
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index bc558c2ddb..bfbd10eb89 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -964,6 +964,7 @@ class FastLlamaModel:
            f' "-____-"     Free Apache license: http://github.com/unslothai/unsloth'
         print(statistics)
         model_patcher.pre_patch()
+        get_statistics()
 
         if dtype is None:
             dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 9c73266575..c609d2ecad 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -318,6 +318,7 @@ class FastMistralModel(FastLlamaModel):
            f' "-____-"     Free Apache license: http://github.com/unslothai/unsloth'
         print(statistics)
         model_patcher.pre_patch()
+        get_statistics()
 
         if dtype is None:
             dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16