Move main script of http_server into async main.

runhouse/servers/http/http_server.py

@@ -852,11 +852,9 @@ def _log_cluster_data(data: dict, labels: dict):
             )
 
 
-if __name__ == "__main__":
+async def main():
     import uvicorn
 
-    ssl_keyfile, ssl_certfile = None, None
-
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--host",
@@ -940,18 +938,19 @@ def _log_cluster_data(data: dict, labels: dict):
     # We only want to forcibly start a Ray cluster if asked.
     # We connect this to the "base" env, which we'll initialize later,
     # so writes to the obj_store within the server get proxied to the "base" env.
-    obj_store.initialize(
+    await obj_store.ainitialize(
         "base",
         setup_cluster_servlet=ClusterServletSetupOption.FORCE_CREATE,
     )
 
-    cluster_config = obj_store.get_cluster_config()
+    cluster_config = await obj_store.aget_cluster_config()
     if not cluster_config:
         logger.warning(
             "Cluster config is not set. Using default values where possible."
         )
     else:
         logger.info("Loaded cluster config from Ray.")
+        global suspend_autostop
         suspend_autostop = cluster_config.get("autostop_mins", -1) > 0
 
     ########################################
@@ -1114,18 +1113,18 @@ def _log_cluster_data(data: dict, labels: dict):
         # a local `runhouse start`
         cluster_config["server_connection_type"] = "tls" if use_https else "none"
 
-    obj_store.set_cluster_config(cluster_config)
+    await obj_store.aset_cluster_config(cluster_config)
     logger.info("Updated cluster config with parsed argument values.")
 
-    HTTPServer.initialize(
+    await HTTPServer.ainitialize(
         conda_env=conda_name,
         enable_local_span_collection=use_local_telemetry
         or configs.data_collection_enabled(),
     )
 
     if den_auth:
         # Update den auth if enabled - keep as a class attribute to be referenced by the validator decorator
-        HTTPServer.enable_den_auth()
+        await HTTPServer.aenable_den_auth()
 
     if use_https and not domain:
         # If using https (whether or not Caddy is being used) and no domain is specified, need to provide both
@@ -1187,10 +1186,16 @@ def _log_cluster_data(data: dict, labels: dict):
     uvicorn_cert = parsed_ssl_certfile if not use_caddy and use_https else None
     uvicorn_key = parsed_ssl_keyfile if not use_caddy and use_https else None
 
-    uvicorn.run(
+    config = uvicorn.Config(
         app,
         host=host,
         port=daemon_port,
         ssl_certfile=uvicorn_cert,
         ssl_keyfile=uvicorn_key,
     )
+    server = uvicorn.Server(config)
+    await server.serve()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

runhouse/servers/obj_store.py

@@ -367,9 +367,6 @@ async def aset_cluster_config(self, config: Dict[str, Any]):
             self.cluster_servlet, "set_cluster_config", config
         )
 
-    def set_cluster_config(self, config: Dict[str, Any]):
-        return sync_function(self.aset_cluster_config)(config)
-
     async def aset_cluster_config_value(self, key: str, value: Any):
         return await self.acall_actor_method(
             self.cluster_servlet, "set_cluster_config_value", key, value