diff --git a/ventis/controller/cloud_provider_logic/EC2/_runtime.py b/ventis/controller/cloud_provider_logic/EC2/_runtime.py index 4c74005..4d5f766 100644 --- a/ventis/controller/cloud_provider_logic/EC2/_runtime.py +++ b/ventis/controller/cloud_provider_logic/EC2/_runtime.py @@ -239,12 +239,15 @@ def _bootstrap_instance(host, spec, replica_index, cfg, redis_host, redis_port, result = subprocess.run( "set -o pipefail; " f"docker save {shlex.quote(image)} | ssh -o StrictHostKeyChecking=no " - f"-o IdentitiesOnly=yes -i {shlex.quote(key)} " + f"-o IdentitiesOnly=yes -o ConnectTimeout=10 " + f"-o ServerAliveInterval=10 -o ServerAliveCountMax=3 " + f"-i {shlex.quote(key)} " f"{shlex.quote(f'{ssh_user}@{host}')} 'sudo docker load'", shell=True, capture_output=True, text=True, executable="/bin/bash", + timeout=180, ) logger.info( "docker save|load returncode=%s stdout=%s stderr=%s", diff --git a/ventis/controller/global_controller.py b/ventis/controller/global_controller.py index b00cef6..a2cff0a 100644 --- a/ventis/controller/global_controller.py +++ b/ventis/controller/global_controller.py @@ -24,6 +24,7 @@ send_agent_information, ) from ventis.utils.redis_client import RedisClient +from ventis.utils.grpc_options import GRPC_CHANNEL_OPTIONS # Add generated grpc_stubs from the local project to the path sys.path.insert(0, os.path.abspath("grpc_stubs")) @@ -558,7 +559,7 @@ def _on_routing_table_updated(self, table): def _get_lc_stub(self, endpoint): """Get or create a cached gRPC stub for a local controller endpoint.""" if endpoint not in self._lc_stubs: - channel = grpc.insecure_channel(endpoint) + channel = grpc.insecure_channel(endpoint, options=GRPC_CHANNEL_OPTIONS) self._lc_stubs[endpoint] = local_controler_pb2_grpc.LocalControllerStub( channel ) @@ -637,7 +638,9 @@ def _run_cmd(self, cmd, host, user=None): """ is_local = _is_local_host(host) if is_local: - return subprocess.run(cmd, capture_output=True, text=True) + return subprocess.run( + cmd, capture_output=True, text=True, timeout=180 + ) else: ssh_key_path = os.path.expanduser( self.config.get("ec2", {}).get( @@ -657,6 +660,10 @@ def _run_cmd(self, cmd, host, user=None): "IdentitiesOnly=yes", "-o", "ConnectTimeout=10", + "-o", + "ServerAliveInterval=10", + "-o", + "ServerAliveCountMax=3", "-i", ssh_key_path, ssh_target, @@ -664,6 +671,7 @@ def _run_cmd(self, cmd, host, user=None): ], capture_output=True, text=True, + timeout=180, ) def launch_docker_agents(self): diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index cc89c11..8b5942e 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -19,10 +19,12 @@ from ventis.controller.local_controller_frontend import start_server from ventis.controller.utils.gpu_metrics import read_gpu_percent from ventis.utils.redis_client import RedisClient + from ventis.utils.grpc_options import GRPC_CHANNEL_OPTIONS except ImportError: from gpu_metrics import read_gpu_percent from local_controller_frontend import start_server from redis_client import RedisClient + from grpc_options import GRPC_CHANNEL_OPTIONS # Add local generated grpc_stubs to path (Docker context copies them directly to /app) sys.path.insert(0, ".") @@ -658,7 +660,9 @@ def _execute_locally( def _get_remote_stub(self, endpoint): """Get or create a cached gRPC stub for a remote controller.""" if endpoint not in self._remote_stubs: - self._remote_channels[endpoint] = grpc.insecure_channel(endpoint) + self._remote_channels[endpoint] = grpc.insecure_channel( + endpoint, options=GRPC_CHANNEL_OPTIONS + ) self._remote_stubs[endpoint] = local_controler_pb2_grpc.LocalControllerStub( self._remote_channels[endpoint] ) diff --git a/ventis/controller/local_controller_frontend.py b/ventis/controller/local_controller_frontend.py index bf440c8..a5fcc25 100644 --- a/ventis/controller/local_controller_frontend.py +++ b/ventis/controller/local_controller_frontend.py @@ -151,9 +151,16 @@ def _cleanup_request(self, request_id): def start_server(port=50051, my_endpoint="unknown"): """Start the gRPC server.""" + try: + from ventis.utils.grpc_options import GRPC_SERVER_OPTIONS + except ImportError: + from grpc_options import GRPC_SERVER_OPTIONS + servicer = LocalControllerServicer(my_endpoint=my_endpoint) - server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) + server = grpc.server( + futures.ThreadPoolExecutor(max_workers=1), options=GRPC_SERVER_OPTIONS + ) local_controler_pb2_grpc.add_LocalControllerServicer_to_server(servicer, server) server.add_insecure_port(f"[::]:{port}") server.start() diff --git a/ventis/deploy.py b/ventis/deploy.py index 2c2531b..2775932 100644 --- a/ventis/deploy.py +++ b/ventis/deploy.py @@ -29,6 +29,7 @@ def my_workflow(query: str): import uuid from flask import Flask, request, jsonify +from werkzeug.serving import WSGIRequestHandler # Try to import from absolute package (local install) or fallback to flat file (Docker container) try: @@ -288,4 +289,4 @@ def get_status(request_id): ) logger.info("Status endpoint: GET http://%s:%d/status/", host, port) - app.run(host=host, port=port) + app.run(host=host, port=port, threaded=True, request_handler=type("_TimeoutWSGIRequestHandler", (WSGIRequestHandler,), {"timeout": 30})) diff --git a/ventis/future.py b/ventis/future.py index ff7ac7c..68615f1 100644 --- a/ventis/future.py +++ b/ventis/future.py @@ -12,6 +12,11 @@ except ImportError: import ventis_context +try: + from ventis.utils.grpc_options import GRPC_CHANNEL_OPTIONS +except ImportError: + from grpc_options import GRPC_CHANNEL_OPTIONS + # Add generated grpc_stubs to path (Docker context copies them directly to /app, and local relies on project dir) sys.path.insert(0, ".") sys.path.insert(0, "/app") @@ -45,7 +50,7 @@ def _get_stub(cls): """Get or create the cached gRPC stub for the local controller.""" if cls._stub is None: endpoint = f"{cls._lc_host}:{cls._lc_port}" - cls._channel = grpc.insecure_channel(endpoint) + cls._channel = grpc.insecure_channel(endpoint, options=GRPC_CHANNEL_OPTIONS) cls._stub = local_controler_pb2_grpc.LocalControllerStub(cls._channel) logger.info("Connected to local controller at %s", endpoint) return cls._stub @@ -188,6 +193,39 @@ def is_available(self): """ return self.result is not None + def _get_consumers(self): + """Return the list of consumers from Redis.""" + return self.redis.smembers(self._consumers_key()) + + def _notify_consumers(self): + """Push this future's result to all registered consumer endpoints via gRPC WriteResult.""" + consumers = self._get_consumers() + if not consumers: + return + for endpoint in consumers: + try: + if not self.result: + logger.warning( + "Future %s is notifying consumer %s with an empty/None result", + self.id, + endpoint, + ) + channel = grpc.insecure_channel(endpoint, options=GRPC_CHANNEL_OPTIONS) + stub = local_controler_pb2_grpc.LocalControllerStub(channel) + payload = json.dumps({"future_id": self.id, "result": self.result}) + request = local_controler_pb2.JsonResponse(resonse=payload) + stub.WriteResult(request) + logger.info( + "Notified consumer %s with result for future %s", endpoint, self.id + ) + except Exception as e: + logger.error( + "Failed to notify consumer %s for future %s: %s", + endpoint, + self.id, + e, + ) + def _add_consumer(self, consumer): """Add a consumer.""" self.consumers.append(consumer) diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index cabc032..e408553 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -328,6 +328,7 @@ def generate_docker( "local_controller_frontend.py", ), (os.path.join(script_dir, "utils", "redis_client.py"), "redis_client.py"), + (os.path.join(script_dir, "utils", "grpc_options.py"), "grpc_options.py"), ( os.path.join(script_dir, "controller", "utils", "gpu_metrics.py"), "gpu_metrics.py", @@ -445,6 +446,7 @@ def generate_workflow_docker( "local_controller_frontend.py", ), (os.path.join(script_dir, "utils", "redis_client.py"), "redis_client.py"), + (os.path.join(script_dir, "utils", "grpc_options.py"), "grpc_options.py"), *[ (os.path.join(script_dir, "controller", "utils", name), name) for name in ("gpu_metrics.py", "session_logging.py") diff --git a/ventis/utils/grpc_options.py b/ventis/utils/grpc_options.py new file mode 100644 index 0000000..cedc1d8 --- /dev/null +++ b/ventis/utils/grpc_options.py @@ -0,0 +1,23 @@ +# Keepalive options for gRPC channels/servers so idle connections get detected instead of stalling, see GRPC_STALLING_FIX.md + +""" +grpc.keepalive_time_ms (int) - time period sender pings the server. +grpc.keepalive_timeout_ms (int) - time server waits for the ping before erroring +grpc.keepalive_permit_without_calls (bool) - allows for sender to send requests even when no active stream is open +grpc.http2.max_pings_without_data (int) - allows sender to send # of pings in a row without sending real data, putting it at 0 removes the limit +grpc.http2.min_ping_interval_without_data_ms (int) - the period of time the server would accept pings without blocking the sender (default: 5 minutes) +""" +GRPC_CHANNEL_OPTIONS = [ + ("grpc.keepalive_time_ms", 30000), + ("grpc.keepalive_timeout_ms", 10000), + ("grpc.keepalive_permit_without_calls", 1), + ("grpc.http2.max_pings_without_data", 0), +] + +GRPC_SERVER_OPTIONS = [ + ("grpc.keepalive_time_ms", 30000), + ("grpc.keepalive_timeout_ms", 10000), + ("grpc.keepalive_permit_without_calls", 1), + ("grpc.http2.max_pings_without_data", 0), + ("grpc.http2.min_ping_interval_without_data_ms", 10000), +]