diff --git a/python/treelite/contrib/__init__.py b/python/treelite/contrib/__init__.py index 85247e22..531e6328 100644 --- a/python/treelite/contrib/__init__.py +++ b/python/treelite/contrib/__init__.py @@ -179,7 +179,8 @@ def generate_cmakelists(dirpath, options=None): ''', file=f) -def create_shared(toolchain, dirpath, nthread=None, verbose=False, options=None): +def create_shared(toolchain, dirpath, *, nthread=None, verbose=False, options=None, + long_build_time_warning=True): """Create shared library. Parameters @@ -199,6 +200,8 @@ def create_shared(toolchain, dirpath, nthread=None, verbose=False, options=None) options : :py:class:`list ` of :py:class:`str `, \ optional Additional options to pass to toolchain + long_build_time_warning : :py:class:`bool `, optional + If set to False, suppress the warning about potentially long build time Returns ------- @@ -253,18 +256,19 @@ def create_shared(toolchain, dirpath, nthread=None, verbose=False, options=None) else: options = [] - # write warning for potentially long compile time - long_time_warning = False - for source in recipe['sources']: - if int(source['length']) > 10000: - long_time_warning = True - break - if long_time_warning: - log_info(__file__, lineno(), - '\033[1;31mWARNING: some of the source files are long. ' + \ - 'Expect long compilation time.\u001B[0m ' + \ - 'You may want to adjust the parameter ' + \ - '\x1B[33mparallel_comp\u001B[0m.\n') + # Write warning for potentially long compile time + if long_build_time_warning: + warn = False + for source in recipe['sources']: + if int(source['length']) > 10000: + warn = True + break + if warn: + log_info(__file__, lineno(), + '\033[1;31mWARNING: some of the source files are long. ' + \ + 'Expect long build time.\u001B[0m ' + \ + 'You may want to adjust the parameter ' + \ + '\x1B[33mparallel_comp\u001B[0m.\n') tstart = time.time() _toolchain_exist_check(toolchain) diff --git a/python/treelite/frontend.py b/python/treelite/frontend.py index 75426b6b..ccf62942 100644 --- a/python/treelite/frontend.py +++ b/python/treelite/frontend.py @@ -173,15 +173,20 @@ def export_lib(self, toolchain, libpath, params=None, compiler='ast_native', model.compile(dirpath='/temporary/directory', params={}, verbose=True) treelite.create_shared(toolchain='msvc', dirpath='/temporary/directory', - verbose=True) + verbose=True) # move the library out of the temporary directory shutil.move('/temporary/directory/mymodel.dll', './mymodel.dll') """ _toolchain_exist_check(toolchain) + + _params = dict(params) if isinstance(params, list) else params + long_build_time_warning = not (_params and 'parallel_comp' in _params) + with TemporaryDirectory(dir=os.path.dirname(libpath)) as temp_dir: self.compile(temp_dir, params, compiler, verbose) - temp_libpath = create_shared(toolchain, temp_dir, nthread, - verbose, options) + temp_libpath = create_shared( + toolchain, temp_dir, nthread=nthread, verbose=verbose, options=options, + long_build_time_warning=long_build_time_warning) if os.path.exists(libpath) and os.path.isfile(libpath): os.remove(libpath) shutil.move(temp_libpath, libpath)