Merge pull request #98 from fal-ai/strict-cache

feat: strict cache when external verification isn't possible

src/isolate/backends/settings.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import shutil
 import tempfile
 from contextlib import contextmanager
@@ -16,13 +17,15 @@
     from isolate.backends import BaseEnvironment
 
 _SYSTEM_TEMP_DIR = Path(tempfile.gettempdir())
+_STRICT_CACHE = os.getenv("ISOLATE_STRICT_CACHE", "0") == "1"
 
 
 @dataclass(frozen=True)
 class IsolateSettings:
     cache_dir: Path = Path(user_cache_dir("isolate", "isolate"))
     serialization_method: str = "pickle"
     log_hook: Callable[[Log], None] = print
+    strict_cache: bool = _STRICT_CACHE
 
     def log(self, log: Log) -> None:
         self.log_hook(log)
@@ -82,6 +85,9 @@ def cache_dir_for(self, backend: BaseEnvironment) -> Path:
         environment_base_path.mkdir(exist_ok=True, parents=True)
         return environment_base_path / backend.key
 
+    def completion_marker_for(self, path: Path) -> Path:
+        return path / ".isolate.completed"
+
     replace = replace
 
 

src/isolate/backends/virtualenv.py

@@ -127,9 +127,15 @@ def create(self, *, force: bool = False) -> Path:
         virtualenv = optional_import("virtualenv")
 
         venv_path = self.settings.cache_dir_for(self)
+        completion_marker = self.settings.completion_marker_for(venv_path)
         with self.settings.cache_lock_for(venv_path):
-            if venv_path.exists() and not force:
-                return venv_path
+            if not force:
+                is_cached = venv_path.exists()
+                if self.settings.strict_cache:
+                    is_cached &= completion_marker.exists()
+
+                if is_cached:
+                    return venv_path
 
             self.log(f"Creating the environment at '{venv_path}'")
 
@@ -145,6 +151,7 @@ def create(self, *, force: bool = False) -> Path:
                 )
 
             self.install_requirements(venv_path)
+            completion_marker.touch()
 
         self.log(f"New environment cached at '{venv_path}'")
         return venv_path

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, ...]: