def _ray_start_cluster(**kwargs): init_kwargs = get_default_fixture_ray_kwargs() num_nodes = 0 do_init = False # num_nodes & do_init are not arguments for ray.init, so delete them. if "num_nodes" in kwargs: num_nodes = kwargs["num_nodes"] del kwargs["num_nodes"] if "do_init" in kwargs: do_init = kwargs["do_init"] del kwargs["do_init"] elif num_nodes > 0: do_init = True init_kwargs.update(kwargs) cluster = Cluster() remote_nodes = [] for i in range(num_nodes): if i > 0 and "_system_config" in init_kwargs: del init_kwargs["_system_config"] remote_nodes.append(cluster.add_node(**init_kwargs)) # We assume driver will connect to the head (first node), # so ray init will be invoked if do_init is true if len(remote_nodes) == 1 and do_init: if client_test_enabled(): ray_client.ray.init(address=cluster.address) else: ray.init(address=cluster.address) yield cluster # The code after the yield will run as teardown code. if client_test_enabled(): ray_client.ray.disconnect() ray_client._stop_test_server(1) ray_client.reset_api() ray.shutdown() cluster.shutdown()
def ray_start_client_server(): server = ray_client_server.serve("localhost:50051", test_mode=True) ray.connect("localhost:50051") yield ray ray.disconnect() server.stop(0) reset_api()
def shutdown_only(): yield None # The code after the yield will run as teardown code. if client_test_enabled(): ray_client.ray.disconnect() ray_client._stop_test_server(1) ray_client.reset_api() ray.shutdown()
def _ray_start(**kwargs): init_kwargs = get_default_fixture_ray_kwargs() init_kwargs.update(kwargs) # Start the Ray processes. if client_test_enabled(): address_info = ray_client.ray.init(**init_kwargs) else: address_info = ray.init(**init_kwargs) yield address_info # The code after the yield will run as teardown code. if client_test_enabled(): ray_client.ray.disconnect() ray_client._stop_test_server(1) ray_client.reset_api() ray.shutdown()