cp: fix: Raise exception in nemo.export instead of allowing pickle.loads (#15266)

nemo/export/tensorrt_llm.py

@@ -1667,14 +1667,7 @@ def _prep_ptuning_table(self):
             self.p_table = None
 
     def _load_prompt_tables(self):
-        if self.model_dir is not None:
-            pt_path = Path(os.path.join(self.model_dir, 'prompt_tables.pkl'))
-            if pt_path.exists():
-                with open(pt_path, 'rb') as f:
-                    self.ptuning_tables = pickle.load(f)
-                self._prep_ptuning_table()
-            else:
-                self.ptuning_tables = []
+        raise Exception("nemo.export is deprecated. Please use the repo https://github.com/NVIDIA-NeMo/Export-Deploy.")
 
     def _get_prompt_embedding_table_ckpt(self, prompt_embeddings_checkpoint_path):
         with TarPath(prompt_embeddings_checkpoint_path) as checkpoint_archive:

nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py

@@ -17,7 +17,6 @@
 import json
 import logging
 import os
-import pickle
 import shutil
 from io import BytesIO
 from pathlib import Path
@@ -52,19 +51,7 @@ def load_extra_state_from_bytes(val: Optional[Union[torch.Tensor, BytesIO]]) ->
     Returns:
         Optional[dict]: Deserialized extra_state, or None if the bytes storage is empty.
     """
-    if val is None:
-        return None
-
-    # TransformerEngine shifted from storing extra_states bytes storage from _io.BytesIO to torch.Tensor
-    if isinstance(val, torch.Tensor):
-        if val.numel() == 0:
-            return None
-
-        val = val.detach().numpy(force=True).tobytes()
-        return pickle.loads(val)
-
-    val.seek(0)
-    return torch.load(val, weights_only=True)
+    raise Exception("nemo.export is deprecated. Please use the repo https://github.com/NVIDIA-NeMo/Export-Deploy.")
 
 
 def preprocess_scaling_factors_for_local_export(state_dict: Dict[str, Any]) -> Dict[str, Any]: