Upload blender_session.py with huggingface_hub

blender_session.py

@@ -215,6 +215,7 @@ class BlenderSession:
         self.last_validation: dict | None = None
         self.checkpoint_path = os.path.join(self.workdir, "checkpoint.blend")
         self._last_checkpoint = 0.0
+        self._checkpoint_busy = threading.Lock()
         self._log_path = os.path.join(self.workdir, "blender.log")
         self._spawn()
 
@@ -239,16 +240,18 @@ class BlenderSession:
         checkpoint must never fail the tool call that triggered it. Large
         checkpoints are time-throttled to keep heavy scenes responsive.
         """
+        if not self._checkpoint_busy.acquire(blocking=False):
+            return  # one already in flight; the newer state gets the next save
         try:
-            size = os.path.getsize(self.checkpoint_path)
-        except OSError:
-            size = 0
-        if (
-            size > _CHECKPOINT_LARGE_BYTES
-            and time.time() - self._last_checkpoint < _CHECKPOINT_MIN_INTERVAL
-        ):
-            return
-        try:
+            try:
+                size = os.path.getsize(self.checkpoint_path)
+            except OSError:
+                size = 0
+            if (
+                size > _CHECKPOINT_LARGE_BYTES
+                and time.time() - self._last_checkpoint < _CHECKPOINT_MIN_INTERVAL
+            ):
+                return
             self.send_code(
                 "import bpy\n"
                 "bpy.ops.wm.save_as_mainfile("
@@ -260,6 +263,8 @@ class BlenderSession:
             self._last_checkpoint = time.time()
         except Exception:  # noqa: BLE001
             log.warning("Checkpoint save failed for session %s", self.session_id, exc_info=True)
+        finally:
+            self._checkpoint_busy.release()
 
     def restart(self) -> bool:
         """
@@ -424,6 +429,26 @@ _sessions: dict[str, BlenderSession] = {}
 _registry_lock = threading.Lock()
 _reaper_started = False
 
+# One spare session is kept warm in the background so a visitor's first message
+# doesn't pay Blender's ~10s cold start.
+_prewarmed: BlenderSession | None = None
+_prewarm_lock = threading.Lock()
+
+
+def _claim_prewarmed(session_id: str) -> BlenderSession | None:
+    global _prewarmed
+    with _prewarm_lock:
+        session = _prewarmed
+        _prewarmed = None
+    if session is None:
+        return None
+    if not session.alive:
+        session.close()
+        return None
+    session.session_id = session_id
+    session.touch()
+    return session
+
 
 def get_session(session_id: str) -> BlenderSession:
     """Return the live session for this id, creating one if needed."""
@@ -448,13 +473,15 @@ def get_session(session_id: str) -> BlenderSession:
                 f"All {MAX_SESSIONS} Blender slots are busy right now - please try again "
                 "in a few minutes."
             )
-        limit, used = container_memory()
-        if limit and (limit - used) < limit * _MIN_FREE_MEM_FRACTION:
-            raise RuntimeError(
-                "The server is low on memory right now (existing scenes are very "
-                "heavy) - please try again in a few minutes."
-            )
-        session = BlenderSession(session_id)
+        session = _claim_prewarmed(session_id)
+        if session is None:
+            limit, used = container_memory()
+            if limit and (limit - used) < limit * _MIN_FREE_MEM_FRACTION:
+                raise RuntimeError(
+                    "The server is low on memory right now (existing scenes are very "
+                    "heavy) - please try again in a few minutes."
+                )
+            session = BlenderSession(session_id)
         _sessions[session_id] = session
         return session
 
@@ -490,8 +517,34 @@ def _reaper_loop() -> None:
             session.close()
 
 
+def _prewarm_loop() -> None:
+    global _prewarmed
+    while True:
+        try:
+            with _prewarm_lock:
+                if _prewarmed is not None and not _prewarmed.alive:
+                    _prewarmed.close()
+                    _prewarmed = None
+                have_spare = _prewarmed is not None
+            with _registry_lock:
+                live = sum(1 for s in _sessions.values() if s.alive)
+            limit, used = container_memory()
+            memory_ok = not limit or (limit - used) >= limit * _MIN_FREE_MEM_FRACTION
+            # Leave one slot of headroom so the spare never blocks a real visitor.
+            if not have_spare and live < MAX_SESSIONS - 1 and memory_ok:
+                spare = BlenderSession("prewarm")
+                with _prewarm_lock:
+                    _prewarmed = spare
+                log.info("Prewarmed spare Blender session (port %d)", spare.port)
+        except Exception:  # noqa: BLE001
+            log.exception("Prewarm attempt failed")
+            time.sleep(30)
+        time.sleep(5)
+
+
 def start_reaper() -> None:
     global _reaper_started
     if not _reaper_started:
         threading.Thread(target=_reaper_loop, daemon=True, name="blender-reaper").start()
+        threading.Thread(target=_prewarm_loop, daemon=True, name="blender-prewarm").start()
         _reaper_started = True