@@ -119,16 +119,8 @@ def check_folder(folder: str, file: str) -> bool:
119119 logger .info ("All required submodules are present." )
120120
121121
122- def main (args ):
123- if not python_is_compatible ():
124- sys .exit (1 )
125-
122+ def build_args_parser () -> argparse .ArgumentParser :
126123 # Parse options.
127-
128- EXECUTORCH_BUILD_PYBIND = ""
129- CMAKE_ARGS = os .getenv ("CMAKE_ARGS" , "" )
130- use_pytorch_nightly = True
131-
132124 parser = argparse .ArgumentParser ()
133125 parser .add_argument (
134126 "--pybind" ,
@@ -146,30 +138,48 @@ def main(args):
146138 action = "store_true" ,
147139 help = "build from the pinned PyTorch commit instead of nightly" ,
148140 )
149- args = parser .parse_args (args )
150- if args .pybind :
151- # Flatten list of lists.
152- args .pybind = list (itertools .chain (* args .pybind ))
153- if "off" in args .pybind :
154- if len (args .pybind ) != 1 :
141+ return parser
142+
143+
144+ def handle_pybind (args , cmake_args , executorch_build_pybind ):
145+ # Flatten list of lists.
146+ args .pybind = list (itertools .chain (* args .pybind ))
147+ if "off" in args .pybind :
148+ if len (args .pybind ) != 1 :
149+ raise Exception (f"Cannot combine `off` with other pybinds: { args .pybind } " )
150+ executorch_build_pybind = "OFF"
151+ else :
152+ for pybind_arg in args .pybind :
153+ if pybind_arg not in VALID_PYBINDS :
155154 raise Exception (
156- f"Cannot combine `off` with other pybinds : { args . pybind } "
155+ f"Unrecognized pybind argument { pybind_arg } ; valid options are : { ', ' . join ( VALID_PYBINDS ) } "
157156 )
158- EXECUTORCH_BUILD_PYBIND = "OFF"
159- else :
160- for pybind_arg in args .pybind :
161- if pybind_arg not in VALID_PYBINDS :
162- raise Exception (
163- f"Unrecognized pybind argument { pybind_arg } ; valid options are: { ', ' .join (VALID_PYBINDS )} "
164- )
165- if pybind_arg == "training" :
166- CMAKE_ARGS += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON"
167- os .environ ["EXECUTORCH_BUILD_TRAINING" ] = "ON"
168- elif pybind_arg == "mps" :
169- CMAKE_ARGS += " -DEXECUTORCH_BUILD_MPS=ON"
170- else :
171- CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{ pybind_arg .upper ()} =ON"
172- EXECUTORCH_BUILD_PYBIND = "ON"
157+ if pybind_arg == "training" :
158+ cmake_args += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON"
159+ os .environ ["EXECUTORCH_BUILD_TRAINING" ] = "ON"
160+ elif pybind_arg == "mps" :
161+ cmake_args += " -DEXECUTORCH_BUILD_MPS=ON"
162+ else :
163+ cmake_args += f" -DEXECUTORCH_BUILD_{ pybind_arg .upper ()} =ON"
164+ executorch_build_pybind = "ON"
165+ return executorch_build_pybind , cmake_args
166+
167+
168+ def main (args ):
169+ if not python_is_compatible ():
170+ sys .exit (1 )
171+
172+ parser = build_args_parser ()
173+ args = parser .parse_args ()
174+
175+ EXECUTORCH_BUILD_PYBIND = ""
176+ CMAKE_ARGS = os .getenv ("CMAKE_ARGS" , "" )
177+ use_pytorch_nightly = True
178+
179+ if args .pybind :
180+ EXECUTORCH_BUILD_PYBIND , CMAKE_ARGS = handle_pybind (
181+ args , CMAKE_ARGS , EXECUTORCH_BUILD_PYBIND
182+ )
173183
174184 if args .clean :
175185 clean ()
0 commit comments