| 141 | is_python: bool |
| 142 | |
| 143 | def run_cmd(self, cmd_idx: int, venv) -> str: |
| 144 | output = "" |
| 145 | mpi_cmd = self.mpi_cmd |
| 146 | build_cmd_str = self.get_cmd_str(len(self.data_cmds)) |
| 147 | current_cmd_str = self.get_cmd_str(cmd_idx) |
| 148 | if cmd_idx <= len(self.data_cmds) - 1: |
| 149 | print_info(f'Running prepare dataset command') |
| 150 | prepare_cmd = self.data_cmds[cmd_idx] |
| 151 | prepare_cmd_str = " ".join(prepare_cmd) |
| 152 | envs = copy.deepcopy(os.environ) |
| 153 | prepare_cmds = prepare_cmd_str.split(';') |
| 154 | for prepare_cmd in prepare_cmds: |
| 155 | print_info(f'Now running prepare data command: "{prepare_cmd}"') |
| 156 | if '>' in prepare_cmd: |
| 157 | cmd = prepare_cmd.split('>')[0] |
| 158 | dataset_file = prepare_cmd.split('>')[1].split()[0] |
| 159 | else: |
| 160 | cmd = prepare_cmd |
| 161 | dataset_file = None |
| 162 | output += subprocess.check_output(cmd.split(), |
| 163 | env=envs).decode() |
| 164 | if dataset_file: |
| 165 | with open(f"{dataset_file}", 'w+') as f: |
| 166 | f.write(output) |
| 167 | |
| 168 | elif cmd_idx == len(self.data_cmds): |
| 169 | #running build |
| 170 | if len(self.build_cmd) == 0: |
| 171 | pass |
| 172 | elif self.build_cmd[0].endswith('.py'): |
| 173 | print_info( |
| 174 | f'Running engine building command: "{build_cmd_str}"') |
| 175 | if len(mpi_cmd) > 0: |
| 176 | output += venv_mpi_check_output(venv, mpi_cmd, |
| 177 | self.build_cmd) |
| 178 | else: |
| 179 | output += venv.run_cmd(self.build_cmd, caller=check_output) |
| 180 | else: |
| 181 | envs = copy.deepcopy(os.environ) |
| 182 | print_info( |
| 183 | f'Running engine building command: "{build_cmd_str}"') |
| 184 | command = self.build_cmd |
| 185 | output += subprocess.check_output(command, env=envs).decode() |
| 186 | else: |
| 187 | #running throughput |
| 188 | print_info(f'Now running benchmarking command: "{current_cmd_str}"') |
| 189 | command = self.benchmark_cmds[cmd_idx - 1 - len(self.data_cmds)] |
| 190 | if self.is_python: |
| 191 | if len(mpi_cmd) > 0: |
| 192 | output += venv_mpi_check_output(venv, mpi_cmd, command) |
| 193 | else: |
| 194 | output += venv.run_cmd(command, caller=check_output) |
| 195 | else: |
| 196 | envs = copy.deepcopy(os.environ) |
| 197 | # Set LD_LIBRARY_PATH to the directory where the binary is located to find libtensorrt_llm.so and |
| 198 | # libnvinfer_plugin_tensorrt_llm.so.x. |
| 199 | envs[ |
| 200 | "LD_LIBRARY_PATH"] = f'{get_trt_llm_lib_dir(venv)}:{os.path.dirname(command[0])}:{envs.get("LD_LIBRARY_PATH", "")}' |