gRPC timeout fix

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",

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,13 +660,18 @@ 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,
                     remote_cmd,
                 ],
                 capture_output=True,
                 text=True,
+                timeout=180,
             )
 
     def launch_docker_agents(self):

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]
             )

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()

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/<request_id>", host, port)
 
-    app.run(host=host, port=port)
+    app.run(host=host, port=port, threaded=True, request_handler=type("_TimeoutWSGIRequestHandler", (WSGIRequestHandler,), {"timeout": 30}))

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)

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")

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),
+]