def test_logging_file_directory(): container = BentoMLContainer config = BentoMLConfiguration().as_dict() container.config.set(config) assert container.logging_file_directory.get() == os.path.join( mock_bentoml_home(), "logs" ) override_config = tempfile.NamedTemporaryFile(delete=False) override_config.write( b""" logging: file: directory: /tmp/logs """ ) override_config.close() config = BentoMLConfiguration(override_config_file=override_config.name).as_dict() container.config.set(config) assert container.logging_file_directory.get() == "/tmp/logs" os.remove(override_config.name)
def test_yatai_database_url(): container = BentoMLContainer config = BentoMLConfiguration().as_dict() container.config.set(config) assert container.yatai_database_url.get() == "{}:///{}".format( "sqlite", os.path.join(mock_bentoml_home(), "storage.db") ) override_config = tempfile.NamedTemporaryFile(delete=False) override_config.write( b""" yatai: database: url: customized_url """ ) override_config.close() config = BentoMLConfiguration(override_config_file=override_config.name).as_dict() container.config.set(config) assert container.yatai_database_url.get() == "customized_url" os.remove(override_config.name)
def inject_dependencies(): """Inject dependencies and configuration to BentoML packages""" from timeit import default_timer as timer start = timer() logger.debug("Start dependency injection") from bentoml.configuration.containers import BentoMLContainer, BentoMLConfiguration config_file = get_local_config_file() if config_file.endswith(".yml"): configuration = BentoMLConfiguration(override_config_file=config_file) else: configuration = BentoMLConfiguration() container = BentoMLContainer() container.config.from_dict(configuration.as_dict()) from bentoml import marshal, server, tracing, cli container.wire(packages=[marshal, server, tracing, cli]) end = timer() logger.debug("Dependency injection completed in %.3f seconds", end - start)
def test_yatai_logging_path(): container = BentoMLContainer config = BentoMLConfiguration().as_dict() container.config.set(config) assert container.yatai_logging_path.get() == os.path.join( mock_bentoml_home(), "logs", "yatai_web_server.log" ) override_config = tempfile.NamedTemporaryFile(delete=False) override_config.write( b""" yatai: logging: path: /tmp/customized.log """ ) override_config.close() config = BentoMLConfiguration(override_config_file=override_config.name).as_dict() container.config.set(config) assert container.yatai_logging_path.get() == "/tmp/customized.log" os.remove(override_config.name)
def _start_prod_server( saved_bundle_path: str, config: BentoMLConfiguration, port: Optional[int] = None, prometheus_lock: Optional[multiprocessing.Lock] = None, ): logger.info("Starting BentoML API server in production mode..") container = BentoMLContainer() container.config.from_dict(config.as_dict()) container.wire(packages=[sys.modules[__name__]]) from bentoml.server.gunicorn_server import GunicornBentoServer if port is None: gunicorn_app = GunicornBentoServer( saved_bundle_path, prometheus_lock=prometheus_lock, ) else: gunicorn_app = GunicornBentoServer( saved_bundle_path, port=port, prometheus_lock=prometheus_lock, ) gunicorn_app.run()
def _start_prod_batching_server( saved_bundle_path: str, api_server_port: int, config: BentoMLConfiguration, prometheus_lock: Optional[multiprocessing.Lock] = None, ): logger.info("Starting BentoML Batching server in production mode..") container = BentoMLContainer() container.config.from_dict(config.as_dict()) from bentoml import marshal from bentoml.server.marshal_server import GunicornMarshalServer container.wire(packages=[sys.modules[__name__], marshal]) # avoid load model before gunicorn fork marshal_server = GunicornMarshalServer( bundle_path=saved_bundle_path, prometheus_lock=prometheus_lock, outbound_host="localhost", outbound_port=api_server_port, ) marshal_server.run()
def test_legacy_compatibility(): config = tempfile.NamedTemporaryFile(delete=False) config.write(b""" bento_server: port: 0 workers: 0 max_request_size: 0 metrics: namespace: Null adapters: image_input: extensions: [] """) config.close() config_dict = BentoMLConfiguration( default_config_file=config.name, legacy_compatibility=True, validate_schema=False, ).as_dict() os.remove(config.name) assert config_dict is not None assert config_dict["bento_server"]["port"] == 5000 assert config_dict["bento_server"]["workers"] == 1 assert config_dict["bento_server"]["max_request_size"] == 20971520 assert config_dict["bento_server"]["metrics"]["namespace"] == "BENTOML" assert config_dict["adapters"]["image_input"]["default_extensions"] == [ '.jpg', '.png', '.jpeg', '.tiff', '.webp', '.bmp', ]
def test_legacy_compatibiltiy(): config = tempfile.NamedTemporaryFile(delete=False) config.write(b""" api_server: port: 0 max_request_size: 0 marshal_server: request_header_flag: Null yatai: url: Null instrument: namespace: Null """) config.close() config_dict = BentoMLConfiguration( default_config_file=config.name, legacy_compatibility=True, validate_schema=False, ).as_dict() os.remove(config.name) assert config_dict is not None assert config_dict["api_server"]["port"] == 5000 assert config_dict["api_server"]["max_request_size"] == 20971520 assert config_dict["marshal_server"]["request_header_flag"] == ( "BentoML-Is-Merged-Request") assert config_dict["yatai"]["url"] == "" assert config_dict["instrument"]["namespace"] == "BENTOML"
def test_prometheus_multiproc_dir(): container = BentoMLContainer config = BentoMLConfiguration().as_dict() container.config.set(config) assert container.prometheus_multiproc_dir.get() == os.path.join( mock_bentoml_home(), "prometheus_multiproc_dir" )
def inject_dependencies(): """Inject dependencies and configuration to BentoML packages""" from timeit import default_timer as timer start = timer() logger.debug("Start dependency injection") from bentoml.configuration.containers import BentoMLContainer, BentoMLConfiguration config_file = get_local_config_file() if config_file and config_file.endswith(".yml"): configuration = BentoMLConfiguration(override_config_file=config_file) else: configuration = BentoMLConfiguration() container = BentoMLContainer() container.config.from_dict(configuration.as_dict()) from bentoml import ( marshal, server, tracing, cli, adapters, saved_bundle, service, ) from bentoml.yatai import yatai_service from bentoml.yatai import yatai_service_impl from bentoml.yatai.repository import s3_repository, gcs_repository container.wire( modules=[ yatai_service, s3_repository, gcs_repository, yatai_service_impl ], packages=[ marshal, server, tracing, cli, adapters, saved_bundle, service ], ) end = timer() logger.debug("Dependency injection completed in %.3f seconds", end - start)
def test_yatai_tls_root_ca_cert(): container = BentoMLContainer config = BentoMLConfiguration().as_dict() container.config.set(config) assert container.yatai_tls_root_ca_cert.get() is None override_config = tempfile.NamedTemporaryFile(delete=False) override_config.write( b""" yatai: remote: tls: client_certificate_file: value1 """ ) override_config.close() config = BentoMLConfiguration(override_config_file=override_config.name).as_dict() container.config.set(config) assert container.yatai_tls_root_ca_cert.get() == "value1" os.remove(override_config.name) override_config = tempfile.NamedTemporaryFile(delete=False) override_config.write( b""" yatai: remote: tls: root_ca_cert: value1 client_certificate_file: value2 """ ) override_config.close() config = BentoMLConfiguration(override_config_file=override_config.name).as_dict() container.config.set(config) assert container.yatai_tls_root_ca_cert.get() == "value1" os.remove(override_config.name)
def run(api_name, config, run_args, bento=None): container = BentoMLContainer() config = BentoMLConfiguration(override_config_file=config) container.config.from_dict(config.as_dict()) from bentoml import tracing container.wire(modules=[tracing]) parser = argparse.ArgumentParser() parser.add_argument('--yatai-url', type=str, default=None) parsed_args, _ = parser.parse_known_args(run_args) yatai_url = parsed_args.yatai_url saved_bundle_path = resolve_bundle_path(bento, pip_installed_bundle_path, yatai_url) api = load_bento_service_api(saved_bundle_path, api_name) exit_code = api.handle_cli(run_args) sys.exit(exit_code)
def test_api_server_workers(): container = BentoMLContainer config_auto_workers = tempfile.NamedTemporaryFile(delete=False) config_auto_workers.write( b""" bento_server: workers: Null """ ) config_auto_workers.close() container.config.set( BentoMLConfiguration( default_config_file=config_auto_workers.name, validate_schema=False, ).as_dict(), ) os.remove(config_auto_workers.name) workers = container.api_server_workers.get() assert workers is not None assert workers > 0 config_manual_workers = tempfile.NamedTemporaryFile(delete=False) config_manual_workers.write( b""" bento_server: workers: 42 """ ) config_manual_workers.close() container.config.set( BentoMLConfiguration( default_config_file=config_manual_workers.name, validate_schema=False, ).as_dict(), ) os.remove(config_manual_workers.name) workers = container.api_server_workers.get() assert workers is not None assert workers == 42
def test_customized_bento_bundle_deployment_version(): override_config = tempfile.NamedTemporaryFile(delete=False) override_config.write( b""" bento_bundle: deployment_version: 0.0.1 """ ) override_config.close() container = BentoMLContainer config = BentoMLConfiguration(override_config_file=override_config.name).as_dict() container.config.set(config) assert container.bento_bundle_deployment_version.get() == "0.0.1" os.remove(override_config.name)
def test_validate_schema(): config = tempfile.NamedTemporaryFile(delete=False) config.write(b""" invalid_key1: invalid_key2: Null """) config.close() with pytest.raises(BentoMLConfigException) as e: BentoMLConfiguration( default_config_file=config.name, validate_schema=True, legacy_compatibility=False, ) assert e is not None os.remove(config.name)
def _start_dev_server( saved_bundle_path: str, api_server_port: int, config: BentoMLConfiguration, ): logger.info("Starting BentoML API server in development mode..") from bentoml.saved_bundle import load_from_dir bento_service = load_from_dir(saved_bundle_path) from bentoml.server.api_server import BentoAPIServer container = BentoMLContainer() container.config.from_dict(config.as_dict()) container.wire(packages=[sys.modules[__name__]]) api_server = BentoAPIServer(bento_service) api_server.start(port=api_server_port)
def test_config_file_override(): default_config_file = tempfile.NamedTemporaryFile(delete=False) default_config_file.write( b""" key1: key2: key3: value3 key4: value4 key5: value5 """ ) default_config_file.close() override_config_file = tempfile.NamedTemporaryFile(delete=False) override_config_file.write( b""" key1: key2: key3: override3 key5: override5 """ ) override_config_file.close() config = BentoMLConfiguration( default_config_file=default_config_file.name, override_config_file=override_config_file.name, validate_schema=False, ).as_dict() os.remove(default_config_file.name) os.remove(override_config_file.name) print(default_config_file.name) assert config is not None assert config["key1"] is not None assert config["key1"]["key2"] is not None assert config["key1"]["key2"]["key3"] == "override3" assert config["key1"]["key2"]["key4"] == "value4" assert config["key1"]["key2"]["key5"] == "override5"
def _start_dev_proxy( saved_bundle_path: str, api_server_port: int, config: BentoMLConfiguration, ): logger.info("Starting BentoML API proxy in development mode..") from bentoml import marshal container = BentoMLContainer() container.config.from_dict(config.as_dict()) container.wire(packages=[marshal]) from bentoml.marshal.marshal import MarshalService marshal_server = MarshalService( saved_bundle_path, outbound_host="localhost", outbound_port=api_server_port, ) marshal_server.fork_start_app()
def test_override_none_value(): config = BentoMLConfiguration(legacy_compatibility=False) config.override(["bento_server", "port"], None) config_dict = config.as_dict() assert config_dict is not None assert config_dict["bento_server"]["port"] == 5000
def start_prod_server( saved_bundle_path: str, port: Optional[int] = None, workers: Optional[int] = None, timeout: Optional[int] = None, enable_microbatch: Optional[bool] = None, enable_swagger: Optional[bool] = None, mb_max_batch_size: Optional[int] = None, mb_max_latency: Optional[int] = None, microbatch_workers: Optional[int] = None, config_file: Optional[str] = None, ): import psutil assert ( psutil.POSIX ), "BentoML API Server production mode only supports POSIX platforms" config = BentoMLConfiguration(override_config_file=config_file) config.override(["api_server", "port"], port) config.override(["api_server", "workers"], workers) config.override(["api_server", "timeout"], timeout) config.override(["api_server", "enable_microbatch"], enable_microbatch) config.override(["api_server", "enable_swagger"], enable_swagger) config.override(["marshal_server", "max_batch_size"], mb_max_batch_size) config.override(["marshal_server", "max_latency"], mb_max_latency) config.override(["marshal_server", "workers"], microbatch_workers) if config.config['api_server'].get('enable_microbatch'): prometheus_lock = multiprocessing.Lock() with reserve_free_port() as api_server_port: pass model_server_job = multiprocessing.Process( target=_start_prod_server, kwargs=dict( saved_bundle_path=saved_bundle_path, port=api_server_port, config=config, prometheus_lock=prometheus_lock, ), daemon=True, ) model_server_job.start() try: _start_prod_batching_server( saved_bundle_path=saved_bundle_path, config=config, api_server_port=api_server_port, prometheus_lock=prometheus_lock, ) finally: model_server_job.terminate() else: _start_prod_server(saved_bundle_path=saved_bundle_path, config=config)
def test_default_bento_bundle_deployment_version(): container = BentoMLContainer config = BentoMLConfiguration().as_dict() container.config.set(config) assert container.bento_bundle_deployment_version.get() is not None
def serve_gunicorn( port, workers, timeout, bento, enable_microbatch, mb_max_batch_size, mb_max_latency, microbatch_workers, yatai_url, enable_swagger, config, ): if not psutil.POSIX: _echo( "The `bentoml serve-gunicorn` command is only supported on POSIX. " "On windows platform, use `bentoml serve` for local API testing and " "docker for running production API endpoint: " "https://docs.docker.com/docker-for-windows/ ") return saved_bundle_path = resolve_bundle_path(bento, pip_installed_bundle_path, yatai_url) container = BentoMLContainer() config = BentoMLConfiguration(override_config_file=config) config.override(["api_server", "port"], port) config.override(["api_server", "workers"], workers) config.override(["api_server", "timeout"], timeout) config.override(["api_server", "enable_microbatch"], enable_microbatch) config.override(["api_server", "enable_swagger"], enable_swagger) config.override(["marshal_server", "max_batch_size"], mb_max_batch_size) config.override(["marshal_server", "max_latency"], mb_max_latency) config.override(["marshal_server", "workers"], microbatch_workers) container.config.from_dict(config.as_dict()) from bentoml import marshal, server container.wire(packages=[marshal, server]) start_prod_server(saved_bundle_path)
def serve( port, bento, enable_microbatch, mb_max_batch_size, mb_max_latency, run_with_ngrok, yatai_url, enable_swagger, config, ): saved_bundle_path = resolve_bundle_path( bento, pip_installed_bundle_path, yatai_url ) container = BentoMLContainer() config = BentoMLConfiguration(override_config_file=config) config.override(["api_server", "port"], port) config.override(["api_server", "enable_microbatch"], enable_microbatch) config.override(["api_server", "run_with_ngrok"], run_with_ngrok) config.override(["api_server", "enable_swagger"], enable_swagger) config.override(["marshal_server", "max_batch_size"], mb_max_batch_size) config.override(["marshal_server", "max_latency"], mb_max_latency) container.config.from_dict(config.as_dict()) from bentoml import marshal, server container.wire(packages=[marshal, server]) start_dev_server(saved_bundle_path)
def test_override(): config = BentoMLConfiguration() config.override(["bento_server", "port"], 6000) config_dict = config.as_dict() assert config_dict is not None assert config_dict["bento_server"]["port"] == 6000
def test_override_empty_key(): config = BentoMLConfiguration() with pytest.raises(BentoMLConfigException) as e: config.override([], 6000) assert e is not None
def start_dev_server( bundle_path: str, port: Optional[int] = None, enable_microbatch: Optional[bool] = None, mb_max_batch_size: Optional[int] = None, mb_max_latency: Optional[int] = None, run_with_ngrok: Optional[bool] = None, enable_swagger: Optional[bool] = None, config_file: Optional[str] = None, ): config = BentoMLConfiguration(override_config_file=config_file) config.override(["api_server", "port"], port) config.override(["api_server", "enable_microbatch"], enable_microbatch) config.override(["api_server", "enable_swagger"], enable_swagger) config.override(["marshal_server", "max_batch_size"], mb_max_batch_size) config.override(["marshal_server", "max_latency"], mb_max_latency) if run_with_ngrok: from threading import Timer from bentoml.utils.flask_ngrok import start_ngrok thread = Timer(1, start_ngrok, args=(port, )) thread.setDaemon(True) thread.start() with reserve_free_port() as api_server_port: # start server right after port released # to reduce potential race model_server_proc = multiprocessing.Process( target=_start_dev_server, kwargs=dict( api_server_port=api_server_port, saved_bundle_path=bundle_path, config=config, ), daemon=True, ) model_server_proc.start() try: _start_dev_proxy( api_server_port=api_server_port, saved_bundle_path=bundle_path, config=config, ) finally: model_server_proc.terminate()
def test_override(): config = BentoMLConfiguration(legacy_compatibility=False) config.override(["api_server", "port"], 6000) config_dict = config.as_dict() assert config_dict is not None assert config_dict["api_server"]["port"] == 6000
def test_override_schema_violation(): config = BentoMLConfiguration() with pytest.raises(BentoMLConfigException) as e: config.override(["api_server", "port"], "non-integer") assert e is not None
def test_override_empty_key(): config = BentoMLConfiguration(legacy_compatibility=False) with pytest.raises(BentoMLConfigException) as e: config.override([], 6000) assert e is not None
def test_override_nonexist_key(): config = BentoMLConfiguration() with pytest.raises(BentoMLConfigException) as e: config.override(["non-existent", "non-existent"], 6000) assert e is not None