Add isolate stuff

fal-ai · Nov 7, 2023 · 27342bd · 27342bd
1 parent c0a0531
commit 27342bd
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/isolate/server/server.py b/src/isolate/server/server.py
@@ -32,6 +32,8 @@
 from isolate.server.health_server import HealthServicer
 from isolate.server.interface import from_grpc, to_grpc
 
+MAX_GRPC_WAIT_TIMEOUT = float(os.getenv("ISOLATE_MAX_GRPC_WAIT_TIMEOUT", 10.0))
+
 # Whether to inherit all the packages from the current environment or not.
 INHERIT_FROM_LOCAL = os.getenv("ISOLATE_INHERIT_FROM_LOCAL") == "1"
 
@@ -132,7 +134,7 @@ def _allocate_new_agent(
                     agent.terminate()
 
         bound_context = ExitStack()
-        stub = bound_context.enter_context(connection._establish_bridge())
+        stub = bound_context.enter_context(connection._establish_bridge(max_wait_timeout=MAX_GRPC_WAIT_TIMEOUT))
         return RunnerAgent(stub, queue, bound_context)
 
     def _identify(self, connection: LocalPythonGRPC) -> Tuple[Any, ...]: