(config)
| 51 | |
| 52 | |
| 53 | def main(config): |
| 54 | pathwaysutils.initialize() |
| 55 | |
| 56 | # No devices for local cpu test. A None for prefill and a None for generate. |
| 57 | devices = server_lib.get_devices() |
| 58 | server_config = maxengine_config.get_server_config(config.inference_server, config) |
| 59 | |
| 60 | metrics_server_config: config_lib.MetricsServerConfig | None = None |
| 61 | if config.prometheus_port != 0: |
| 62 | metrics_server_config = config_lib.MetricsServerConfig(port=config.prometheus_port) |
| 63 | |
| 64 | # We separate credential from run so that we can unit test it with |
| 65 | # local credentials. |
| 66 | # TODO: Add grpc credentials for OSS. |
| 67 | # pylint: disable=unexpected-keyword-arg |
| 68 | jetstream_server = server_lib.run( |
| 69 | threads=256, |
| 70 | port=9000, |
| 71 | config=server_config, |
| 72 | devices=devices, |
| 73 | metrics_server_config=metrics_server_config, |
| 74 | enable_jax_profiler=config.enable_jax_profiler if config.enable_jax_profiler else False, |
| 75 | jax_profiler_port=config.jax_profiler_port if config.jax_profiler_port else 9999, |
| 76 | enable_model_warmup=config.enable_model_warmup if config.enable_model_warmup else False, |
| 77 | lora_input_adapters_path=config.lora_input_adapters_path, |
| 78 | multi_sampling=config.multi_sampling if config.multi_sampling else False, |
| 79 | prefix_caching_config=_create_prefix_caching_config(config), |
| 80 | ) |
| 81 | jetstream_server.wait_for_termination() |
| 82 | |
| 83 | |
| 84 | if __name__ == "__main__": |
no test coverage detected