-
Notifications
You must be signed in to change notification settings - Fork 454
/
test_cli.py
151 lines (124 loc) · 7.55 KB
/
test_cli.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
# coding=utf-8
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import os
import shutil
import subprocess
import tempfile
import unittest
from pathlib import Path
import optimum.commands
from optimum.onnxruntime.configuration import AutoQuantizationConfig, ORTConfig
CLI_WIH_CUSTOM_COMMAND_PATH = Path(__file__).parent / "cli_with_custom_command.py"
OPTIMUM_COMMANDS_DIR = Path(inspect.getfile(optimum.commands)).parent
REGISTERED_CLI_WITH_CUSTOM_COMMAND_PATH = OPTIMUM_COMMANDS_DIR / "register" / "cli_with_custom_command.py"
class TestCLI(unittest.TestCase):
def tearDown(self):
super().tearDown()
REGISTERED_CLI_WITH_CUSTOM_COMMAND_PATH.unlink(missing_ok=True)
def test_helps_no_raise(self):
commands = [
"optimum-cli --help",
"optimum-cli export --help",
"optimum-cli export onnx --help",
"optimum-cli env --help",
"optimum-cli onnxruntime quantize --help",
"optimum-cli onnxruntime optimize --help",
]
for command in commands:
subprocess.run(command, shell=True, check=True)
def test_env_commands(self):
subprocess.run("optimum-cli env", shell=True, check=True)
def test_export_commands(self):
with tempfile.TemporaryDirectory() as tempdir:
commands = [
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-vision_perceiver_conv --task image-classification {tempdir}/onnx",
f"optimum-cli export tflite --model hf-internal-testing/tiny-random-bert --task text-classification --sequence_length 128 {tempdir}/tflite",
]
for command in commands:
subprocess.run(command, shell=True, check=True)
def test_optimize_commands(self):
with tempfile.TemporaryDirectory() as tempdir:
# First export a tiny encoder, decoder only and encoder-decoder
export_commands = [
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-BertModel {tempdir}/encoder",
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-gpt2 {tempdir}/decoder",
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-bart {tempdir}/encoder-decoder",
]
optimize_commands = [
f"optimum-cli onnxruntime optimize --onnx_model {tempdir}/encoder -O1 -o {tempdir}/optimized_encoder",
f"optimum-cli onnxruntime optimize --onnx_model {tempdir}/decoder -O1 -o {tempdir}/optimized_decoder",
f"optimum-cli onnxruntime optimize --onnx_model {tempdir}/encoder-decoder -O1 -o {tempdir}/optimized_encoder_decoder",
]
for export, optimize in zip(export_commands, optimize_commands):
subprocess.run(export, shell=True, check=True)
subprocess.run(optimize, shell=True, check=True)
def test_quantize_commands(self):
with tempfile.TemporaryDirectory() as tempdir:
ort_config = ORTConfig(quantization=AutoQuantizationConfig.avx2(is_static=False))
ort_config.save_pretrained(tempdir)
# First export a tiny encoder, decoder only and encoder-decoder
export_commands = [
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-bert {tempdir}/encoder",
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-gpt2 {tempdir}/decoder",
f"optimum-cli export onnx --model hf-internal-testing/tiny-random-t5 {tempdir}/encoder-decoder",
]
quantize_commands = [
f"optimum-cli onnxruntime quantize --onnx_model {tempdir}/encoder --avx2 -o {tempdir}/quantized_encoder",
f"optimum-cli onnxruntime quantize --onnx_model {tempdir}/decoder --avx2 -o {tempdir}/quantized_decoder",
f"optimum-cli onnxruntime quantize --onnx_model {tempdir}/encoder-decoder --avx2 -o {tempdir}/quantized_encoder_decoder",
]
quantize_with_config_commands = [
f"optimum-cli onnxruntime quantize --onnx_model hf-internal-testing/tiny-random-bert --c {tempdir}/ort_config.json -o {tempdir}/quantized_encoder_with_config",
f"optimum-cli onnxruntime quantize --onnx_model hf-internal-testing/tiny-random-gpt2 --c {tempdir}/ort_config.json -o {tempdir}/quantized_decoder_with_config",
f"optimum-cli onnxruntime quantize --onnx_model hf-internal-testing/tiny-random-t5 --c {tempdir}/ort_config.json -o {tempdir}/quantized_encoder_decoder_with_config",
]
for export, quantize, quantize_with_config in zip(
export_commands, quantize_commands, quantize_with_config_commands
):
subprocess.run(export, shell=True, check=True)
subprocess.run(quantize, shell=True, check=True)
subprocess.run(quantize_with_config, shell=True, check=True)
def _run_command_and_check_content(self, command: str, content: str) -> bool:
proc = subprocess.Popen(command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = proc.communicate()
stdout = stdout.decode("utf-8")
print(stdout)
print(stderr)
return content in stdout
def test_register_command(self):
# Nothing was registered, it should fail.
custom_command = "optimum-cli blablabla"
command_content = "If the CI can read this, it means it worked!"
succeeded = self._run_command_and_check_content(custom_command, command_content)
self.assertFalse(succeeded, "The command should fail here since it is not registered yet.")
# As a "base" command in `optimum-cli`.
shutil.copy(CLI_WIH_CUSTOM_COMMAND_PATH, REGISTERED_CLI_WITH_CUSTOM_COMMAND_PATH)
# We check that the print_help method prints the registered command.
succeeded = self._run_command_and_check_content("optimum-cli", "blablabla")
self.assertTrue(succeeded, "The command name should appear in the help.")
succeeded = self._run_command_and_check_content(custom_command, command_content)
self.assertTrue(succeeded, "The command should succeed here since it is registered.")
REGISTERED_CLI_WITH_CUSTOM_COMMAND_PATH.unlink()
# As a subcommand of an existing command, `optimum-cli export` here.
shutil.copy(CLI_WIH_CUSTOM_COMMAND_PATH, REGISTERED_CLI_WITH_CUSTOM_COMMAND_PATH)
os.environ["TEST_REGISTER_COMMAND_WITH_SUBCOMMAND"] = "true"
# We check that the print_help method prints the registered command.
succeeded = self._run_command_and_check_content("optimum-cli export", "blablabla")
self.assertTrue(succeeded, "The command name should appear in the help.")
custom_command = "optimum-cli export blablabla"
succeeded = self._run_command_and_check_content(custom_command, command_content)
self.assertTrue(succeeded, "The command should succeed here since it is registered.")
REGISTERED_CLI_WITH_CUSTOM_COMMAND_PATH.unlink()