fix: critical bugs — unsafe torch.load (RCE), UnboundLocalError crash, thread leak, broken tests

air_llm/airllm/airllm_base.py

@@ -744,8 +744,31 @@ def _post_hook(self, module, args, output):
         clean_memory()
         return output
 
+    # ---- lifecycle management ----------------------------------------------------------------
+
+    def close(self):
+        """Shut down the prefetch executor and release background resources.
+
+        Call this when you are done with the model to avoid leaking threads.
+        Safe to call multiple times.
+        """
+        if self._prefetch_future is not None:
+            self._prefetch_future.cancel()
+            self._prefetch_future = None
+        if self._executor is not None:
+            self._executor.shutdown(wait=False)
+            self._executor = None
+
+    def __del__(self):
+        # Safety net: clean up if the user forgot to call close().
+        try:
+            self.close()
+        except Exception:
+            pass
+
     # ---- delegation to the underlying transformers model ------------------------------------
 
+
     def generate(self, *args, **kwargs):
         return self.model.generate(*args, **kwargs)
 

air_llm/airllm/utils.py

@@ -80,7 +80,8 @@ def clean_memory():
     except Exception as ex:
         # maybe platform
         pass
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
 
 def uncompress_layer_state_dict(layer_state_dict):
@@ -196,12 +197,13 @@ def compress_layer_state_dict(layer_state_dict, compression=None):
     return compressed_layer_state_dict if compressed_layer_state_dict is not None else layer_state_dict
 
 def remove_real_and_linked_file(to_delete):
-    if (os.path.realpath(to_delete) != to_delete):
+    targetpath = None
+    if os.path.realpath(to_delete) != to_delete:
         targetpath = os.path.realpath(to_delete)
 
     os.remove(to_delete)
-    if (targetpath):
-         os.remove(targetpath)
+    if targetpath:
+        os.remove(targetpath)
 
 
 
@@ -270,7 +272,7 @@ def split_and_save_layers(checkpoint_path, layer_shards_saving_path=None, splitt
     elif os.path.exists(checkpoint_path / 'pytorch_model.bin'):
         # single-file torch checkpoint: map every tensor to that one file
         safetensors_format = False
-        single_sd = torch.load(checkpoint_path / 'pytorch_model.bin', map_location='cpu')
+        single_sd = torch.load(checkpoint_path / 'pytorch_model.bin', map_location='cpu', weights_only=True)
         index = {k: 'pytorch_model.bin' for k in single_sd.keys()}
         del single_sd
     else:
@@ -431,7 +433,7 @@ def _last_shard_of(layer):
                                                     token=hf_token)
 
                 if not safetensors_format:
-                    state_dict.update(torch.load(to_load, map_location='cpu'))
+                    state_dict.update(torch.load(to_load, map_location='cpu', weights_only=True))
                 else:
                     state_dict.update(load_file(to_load, device='cpu'))
 
@@ -445,7 +447,7 @@ def _last_shard_of(layer):
                 huggingface_hub.snapshot_download(repo_id, allow_patterns=os.path.basename(to_load),
                                                 token=hf_token)
             if not safetensors_format:
-                state_dict.update(torch.load(to_load, map_location='cpu'))
+                state_dict.update(torch.load(to_load, map_location='cpu', weights_only=True))
             else:
                 state_dict.update(load_file(to_load, device='cpu'))
 

air_llm/tests/test_automodel.py

@@ -15,17 +15,20 @@ def tearDown(self):
 
     def test_auto_model_should_return_correct_model(self):
         mapping_dict = {
-            'garage-bAInd/Platypus2-7B': 'AirLLMLlama2',
+            # These architectures have dedicated subclasses in ARCH_OVERRIDES
             'Qwen/Qwen-7B': 'AirLLMQWen',
-            'internlm/internlm-chat-7b': 'AirLLMInternLM',
             'THUDM/chatglm3-6b-base': 'AirLLMChatGLM',
             'baichuan-inc/Baichuan2-7B-Base': 'AirLLMBaichuan',
-            'mistralai/Mistral-7B-Instruct-v0.1': 'AirLLMMistral',
-            'mistralai/Mixtral-8x7B-v0.1': 'AirLLMMixtral'
+            'internlm/internlm-chat-7b': 'AirLLMInternLM',
+            # Standard architectures now use the generic streaming model
+            'garage-bAInd/Platypus2-7B': 'AirLLMBaseModel',
+            'mistralai/Mistral-7B-Instruct-v0.1': 'AirLLMBaseModel',
+            'mistralai/Mixtral-8x7B-v0.1': 'AirLLMBaseModel',
         }
 
 
         for k,v in mapping_dict.items():
             module, cls = AutoModel.get_module_class(k)
-            self.assertEqual(cls, v, f"expecting {v}")
+            self.assertEqual(cls, v, f"expecting {v} for {k}, got {cls}")
+