(log_level: object = None)
| 46 | |
| 47 | |
| 48 | def _init(log_level: object = None) -> None: |
| 49 | global _inited |
| 50 | if _inited: |
| 51 | return |
| 52 | _inited = True |
| 53 | # Move to __init__ |
| 54 | if log_level is not None: |
| 55 | logger.set_level(log_level) |
| 56 | |
| 57 | if os.getenv("TRT_LLM_NO_LIB_INIT", "0") == "1": |
| 58 | logger.info("Skipping TensorRT LLM init.") |
| 59 | return |
| 60 | |
| 61 | logger.info("Starting TensorRT LLM init.") |
| 62 | |
| 63 | # load plugin lib |
| 64 | _load_plugin_lib() |
| 65 | |
| 66 | # load FT decoder layer and torch custom ops |
| 67 | project_dir = str(Path(__file__).parent.absolute()) |
| 68 | if platform.system() == "Windows": |
| 69 | ft_decoder_lib = project_dir + "/libs/th_common.dll" |
| 70 | else: |
| 71 | ft_decoder_lib = project_dir + "/libs/libth_common.so" |
| 72 | try: |
| 73 | torch.classes.load_library(ft_decoder_lib) |
| 74 | from ._torch.custom_ops import _register_fake |
| 75 | |
| 76 | _register_fake() |
| 77 | except Exception as e: |
| 78 | msg = ( |
| 79 | "\nFATAL: Decoding operators failed to load. This may be caused by an incompatibility " |
| 80 | "between PyTorch and TensorRT-LLM. Please rebuild and install TensorRT-LLM." |
| 81 | ) |
| 82 | raise ImportError(str(e) + msg) |
| 83 | |
| 84 | MpiComm.local_init() |
| 85 | |
| 86 | def _print_stacks(): |
| 87 | counter = 0 |
| 88 | while True: |
| 89 | time.sleep(print_stacks_period) |
| 90 | counter += 1 |
| 91 | logger.error(f"Printing stacks {counter} times") |
| 92 | print_all_stacks() |
| 93 | |
| 94 | print_stacks_period = int(os.getenv("TRTLLM_PRINT_STACKS_PERIOD", "-1")) |
| 95 | if print_stacks_period > 0: |
| 96 | print_stacks_thread = threading.Thread(target=_print_stacks, daemon=True) |
| 97 | print_stacks_thread.start() |
| 98 | |
| 99 | logger.info("TensorRT LLM inited.") |
| 100 | |
| 101 | |
| 102 | def default_net() -> Network: |
no test coverage detected