beaupreda commited on
Commit
86778a6
·
verified ·
1 Parent(s): df371e2

Upload sensAI-Generic-Object-Detection with upload_repo.py

Browse files
Files changed (7) hide show
  1. Dockerfile +2 -0
  2. README.md +11 -6
  3. app.py +57 -95
  4. install_eve.py +25 -2
  5. mod_models.py +24 -15
  6. shared/eve_app_tabs.py +0 -132
  7. shared/eve_worker_pool.py +1119 -1121
Dockerfile CHANGED
@@ -13,6 +13,8 @@ RUN --mount=type=secret,id=MODEL_ACCESS_TOKEN \
13
  RUN useradd -m -u 1000 user
14
  USER user
15
  ENV PATH="/home/user/.local/bin:$PATH"
 
 
16
 
17
  WORKDIR /app
18
 
 
13
  RUN useradd -m -u 1000 user
14
  USER user
15
  ENV PATH="/home/user/.local/bin:$PATH"
16
+ # install_eve.py (above, as root) fetched the gated MOD weights here.
17
+ ENV MOD_MODELS_DIR="/opt/eve-models"
18
 
19
  WORKDIR /app
20
 
README.md CHANGED
@@ -1,7 +1,7 @@
1
  ---
2
  title: Lattice sensAI Generic Object Detection
3
  slug: sensAI-Generic-Object-Detection
4
- short_description: Object detection — GMOD-80 / AMOD-8 / OMOD on the EVE SDK
5
  emoji: 📦
6
  colorFrom: blue
7
  colorTo: indigo
@@ -15,11 +15,14 @@ tags:
15
  - embedded
16
  - low-power
17
  - object-detection
18
- - gmod
19
- - amod
20
- - omod
 
21
  - npu
22
  - soc
 
 
23
  ---
24
 
25
  The Lattice sensAI Generic Object Detection demo runs the EVE SDK's MOD
@@ -30,11 +33,11 @@ an uploaded video, or a live webcam feed. Pick exactly one of three models:
30
  | --- | --- |
31
  | **GMOD-80** | Generic 80-class object detector |
32
  | **AMOD-8** | 8-class automotive object detector — person, bicycle, car, motorcycle, truck, bus, traffic light, stop sign |
33
- | **OMOD** | 8-class office-objects detector — bottle, cup, potted plant, laptop, mouse, keyboard, cell phone, book |
34
 
35
  Each model is a self-contained `.tflite` loaded through the EVE SDK's custom
36
  object-detection model path. GMOD-80 uses the SDK's bundled weights; AMOD-8 and
37
- OMOD load from `models/`. Selecting a model in either tab loads it on the worker.
38
 
39
  To preview the demo, use the tabs:
40
 
@@ -45,6 +48,8 @@ To preview the demo, use the tabs:
45
  > Note: For demo purposes, the AI pipeline and image-draw operations run on
46
  > a Hugging Face CPU server. Performance varies with concurrent users.
47
 
 
 
48
  For SDK access, fill the form on
49
  [this page](https://huggingface.co/LatticeSemi/sensAI-Edge-Vision-Engine-SDK-Packages)
50
  and follow the download instructions. Support: evehelp@latticesemi.com.
 
1
  ---
2
  title: Lattice sensAI Generic Object Detection
3
  slug: sensAI-Generic-Object-Detection
4
+ short_description: Generic object detection (80 classes)
5
  emoji: 📦
6
  colorFrom: blue
7
  colorTo: indigo
 
15
  - embedded
16
  - low-power
17
  - object-detection
18
+ - object-detector
19
+ - detector
20
+ - automotive
21
+ - office
22
  - npu
23
  - soc
24
+ - fpga
25
+ - sensai
26
  ---
27
 
28
  The Lattice sensAI Generic Object Detection demo runs the EVE SDK's MOD
 
33
  | --- | --- |
34
  | **GMOD-80** | Generic 80-class object detector |
35
  | **AMOD-8** | 8-class automotive object detector — person, bicycle, car, motorcycle, truck, bus, traffic light, stop sign |
36
+ | **OMOD-8** | 8-class office-objects detector — bottle, cup, potted plant, laptop, mouse, keyboard, cell phone, book |
37
 
38
  Each model is a self-contained `.tflite` loaded through the EVE SDK's custom
39
  object-detection model path. GMOD-80 uses the SDK's bundled weights; AMOD-8 and
40
+ OMOD-8 load from `models/`. Selecting a model in either tab loads it on the worker.
41
 
42
  To preview the demo, use the tabs:
43
 
 
48
  > Note: For demo purposes, the AI pipeline and image-draw operations run on
49
  > a Hugging Face CPU server. Performance varies with concurrent users.
50
 
51
+ For model access, visit the **Files and versions** tab at the top of this page.
52
+
53
  For SDK access, fill the form on
54
  [this page](https://huggingface.co/LatticeSemi/sensAI-Edge-Vision-Engine-SDK-Packages)
55
  and follow the download instructions. Support: evehelp@latticesemi.com.
app.py CHANGED
@@ -10,7 +10,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent / "shared"))
10
 
11
  from env_utils import load_dotenv_if_present, require_secrets
12
  from eula_tab import build_eula_tab
13
- from eve_app_tabs import build_image_or_video_offline_tab, build_live_inference_tab
14
  from eve_inference_handlers import EveAppHandlers, patch_video_for_external_urls
15
  from eve_worker_pool import EveWorkerPool
16
  from live_inference import (
@@ -28,7 +28,6 @@ from session_tracker import SessionTracker
28
  from usage_analytics import UsageTracker
29
  from video_processing import VideoLimits
30
 
31
- _IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
32
  _VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv", ".webm")
33
 
34
 
@@ -41,8 +40,8 @@ def _build_feature_radio(hint: str = "") -> gr.Radio:
41
  choices=MOD_MODELS,
42
  value=DEFAULT_MOD_MODEL,
43
  label=label,
44
- info="Pick one detector: GMOD-80 (generic 80-class), AMOD-8 "
45
- "(automotive), or OMOD (office objects).",
46
  interactive=True,
47
  )
48
 
@@ -111,13 +110,12 @@ if __name__ == "__main__":
111
  atexit.register(_shutdown)
112
 
113
  examples_dir = Path(__file__).resolve().parent / "examples"
114
- image_examples = _scan_examples(examples_dir, _IMAGE_EXTS)
115
  video_examples = _scan_examples(examples_dir, _VIDEO_EXTS)
116
  video_limits = VideoLimits()
117
  rtc_config_provider = RtcConfigProvider()
118
 
119
  with gr.Blocks(
120
- title="GMOD / AMOD / OMOD Object Detection Demo",
121
  theme=gr.themes.Default(
122
  text_size=gr.themes.sizes.text_lg,
123
  primary_hue=gr.themes.colors.yellow,
@@ -155,17 +153,28 @@ if __name__ == "__main__":
155
  ) as demo:
156
  gr.Markdown("# Lattice sensAI — Generic Object Detection")
157
  gr.Markdown(
158
- "Run object detection on images, videos, or a live camera feed using "
159
  "one of three EVE SDK models:\n\n"
160
- "- **GMOD-80** — generic 80-class detector\n"
161
- "- **AMOD-8** 8-class automotive object detector "
162
- "(person, bicycle, car, motorcycle, truck, bus, traffic light, "
163
- "stop sign)\n"
164
- "- **OMOD** 8-class office-objects detector "
165
- "(bottle, cup, potted plant, laptop, mouse, keyboard, cell phone, "
166
- "book)\n\n"
167
- "Only one model is active at a time. Pick the model with the radio "
168
- "in each tab."
 
 
 
 
 
 
 
 
 
 
 
169
  )
170
 
171
  # Required by EveAppHandlers signatures (Face ID gallery state). Empty
@@ -185,108 +194,61 @@ if __name__ == "__main__":
185
  height=camera_height,
186
  )
187
 
188
- offline = build_image_or_video_offline_tab(
189
  feature_checkbox_builder=_build_feature_radio,
190
- image_examples=image_examples,
191
- video_examples=video_examples,
192
  video_limits=video_limits,
193
  )
194
- patch_video_for_external_urls(offline.video_output)
195
- offline_radio = offline.checkboxes # gr.Radio in this demo
196
 
197
  build_eula_tab()
198
 
199
- # --- Mutually exclusive accordions on the offline tab ---
200
- offline.image_accordion.expand(
201
- fn=lambda: gr.update(open=False),
202
- outputs=[offline.video_accordion],
203
- )
204
- offline.video_accordion.expand(
205
- fn=lambda: gr.update(open=False),
206
- outputs=[offline.image_accordion],
207
- )
208
-
209
  # --- Examples wire-up: clicking a sample loads it into the matching input ---
210
- if offline.image_example_dataset is not None:
211
- offline.image_example_dataset.click(
212
- fn=lambda sample: sample[0],
213
- inputs=[offline.image_example_dataset],
214
- outputs=[offline.image_input],
215
- )
216
- if offline.video_example_dataset is not None:
217
- offline.video_example_dataset.click(
218
- fn=lambda sample: sample[0],
219
- inputs=[offline.video_example_dataset],
220
- outputs=[offline.video_input],
221
- )
222
 
223
  # --- Process button gating ---
224
- def _on_input_change(image_path: str | None, video_path: str | None) -> dict:
225
- return gr.update(interactive=bool(image_path) or bool(video_path))
226
-
227
- for src in (offline.image_input, offline.video_input):
228
- src.change(
229
- fn=_on_input_change,
230
- inputs=[offline.image_input, offline.video_input],
231
- outputs=[offline.process_btn],
232
- )
233
 
234
- # --- Process dispatcher: image vs video ---
235
  def _process_dispatch(
236
- image_path: str | None,
237
  video_path: str | None,
238
  mod_model: str,
239
  registry: dict,
240
  request: gr.Request,
241
  progress: gr.Progress = gr.Progress(),
242
  ):
243
- if image_path:
244
- annotated, registry = handlers.run_eve_image_inference(
245
- image_path,
246
- False,
247
- False,
248
- False,
249
- False,
250
- registry,
251
- mod_model,
252
- request,
253
- progress,
254
- )
255
- return (
256
- gr.update(value=annotated, visible=True),
257
- gr.update(value=None, visible=False),
258
- registry,
259
- )
260
- if video_path:
261
- output_path, registry = handlers.run_eve_inference(
262
- video_path,
263
- False,
264
- False,
265
- False,
266
- False,
267
- registry,
268
- mod_model,
269
- request,
270
- progress,
271
- )
272
- return (
273
- gr.update(value=None, visible=False),
274
- gr.update(value=output_path, visible=True),
275
- registry,
276
- )
277
- raise gr.Error("Please upload an image or video.")
278
-
279
- offline.process_btn.click(
280
  fn=_process_dispatch,
281
  inputs=[
282
- offline.image_input,
283
- offline.video_input,
284
  offline_radio,
285
  session_registry,
286
  ],
287
  outputs=[
288
- offline.image_output,
289
- offline.video_output,
290
  session_registry,
291
  ],
292
  concurrency_limit=pool.worker_count,
 
10
 
11
  from env_utils import load_dotenv_if_present, require_secrets
12
  from eula_tab import build_eula_tab
13
+ from eve_app_tabs import build_offline_inference_tab, build_live_inference_tab
14
  from eve_inference_handlers import EveAppHandlers, patch_video_for_external_urls
15
  from eve_worker_pool import EveWorkerPool
16
  from live_inference import (
 
28
  from usage_analytics import UsageTracker
29
  from video_processing import VideoLimits
30
 
 
31
  _VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv", ".webm")
32
 
33
 
 
40
  choices=MOD_MODELS,
41
  value=DEFAULT_MOD_MODEL,
42
  label=label,
43
+ info="Pick one object detector: GMOD-80 (generic 80-class), Automotive finetune (8-class)"
44
+ ", or Office finetune (8-class).",
45
  interactive=True,
46
  )
47
 
 
110
  atexit.register(_shutdown)
111
 
112
  examples_dir = Path(__file__).resolve().parent / "examples"
 
113
  video_examples = _scan_examples(examples_dir, _VIDEO_EXTS)
114
  video_limits = VideoLimits()
115
  rtc_config_provider = RtcConfigProvider()
116
 
117
  with gr.Blocks(
118
+ title="Object Detection Demo",
119
  theme=gr.themes.Default(
120
  text_size=gr.themes.sizes.text_lg,
121
  primary_hue=gr.themes.colors.yellow,
 
153
  ) as demo:
154
  gr.Markdown("# Lattice sensAI — Generic Object Detection")
155
  gr.Markdown(
156
+ "Run object detection on videos or a live camera feed using "
157
  "one of three EVE SDK models:\n\n"
158
+ "- **GMOD-80** — 80-class generic multi-object detector.\n\n"
159
+ " - Classes: person, bicycle, car, motorcycle, airplane, bus, train, truck, boat, "
160
+ "traffic light, fire hydrant, stop sign, parking meter, bench, bird, cat, dog, horse, "
161
+ "sheep, cow, elephant, bear, zebra, giraffe, backpack, umbrella, handbag, tie, suitcase, "
162
+ "frisbee, skis, snowboard, sports ball, kite, baseball bat, baseball glove, skateboard, "
163
+ "surfboard, tennis racket, bottle, wine glass, cup, fork, knife, spoon, bowl, banana, "
164
+ "apple, sandwich, orange, broccoli, carrot, hot dog, pizza, donut, cake, chair, couch, "
165
+ "potted plant, bed, dining table, toilet, tv, laptop, mouse, remote, keyboard, cell phone, "
166
+ "microwave, oven, toaster, sink, refrigerator, book, clock, vase, scissors, teddy bear, "
167
+ "hair drier and toothbrush.\n\n"
168
+ "- **Automotive finetune** — 8-class automotive object detector finetuned from GMOD-80.\n\n"
169
+ " - Classes: person, bicycle, car, motorcycle, bus, truck, traffic light and stop sign.\n"
170
+ "- **Office finetune** — 8-class office object detector finetuned from GMOD-80.\n\n"
171
+ " - Classes: bottle, cup, potted plant, laptop, mouse, keyboard, cell phone and book.\n\n"
172
+ "Only one model is active at a time. Pick the model with the radio button below."
173
+ )
174
+ gr.Markdown(
175
+ "A guide on how to finetune the GMOD-80 model with your own data will be available soon, along with the release 7.3 of the EVE SDK.\n\n"
176
+ "For any questions or support, please reach out to us at evehelp@latticesemi.com.\n\n"
177
+ "> Note: For demo purposes, execution of the AI pipeline and image draw operations are all performed on a Hugging Face CPU server. Performance may vary based on the number of concurrent users."
178
  )
179
 
180
  # Required by EveAppHandlers signatures (Face ID gallery state). Empty
 
194
  height=camera_height,
195
  )
196
 
197
+ _, video_input, video_output, offline_radio, process_btn, video_example_dataset = build_offline_inference_tab(
198
  feature_checkbox_builder=_build_feature_radio,
199
+ example_videos=video_examples,
 
200
  video_limits=video_limits,
201
  )
202
+ patch_video_for_external_urls(video_output)
 
203
 
204
  build_eula_tab()
205
 
 
 
 
 
 
 
 
 
 
 
206
  # --- Examples wire-up: clicking a sample loads it into the matching input ---
207
+ video_example_dataset.click(
208
+ fn=lambda sample: sample[0],
209
+ inputs=[video_example_dataset],
210
+ outputs=[video_input],
211
+ )
 
 
 
 
 
 
 
212
 
213
  # --- Process button gating ---
214
+ video_input.change(
215
+ fn=lambda video_path: gr.update(interactive=bool(video_path)),
216
+ inputs=[video_input],
217
+ outputs=[process_btn],
218
+ )
 
 
 
 
219
 
220
+ # --- Process dispatcher ---
221
  def _process_dispatch(
 
222
  video_path: str | None,
223
  mod_model: str,
224
  registry: dict,
225
  request: gr.Request,
226
  progress: gr.Progress = gr.Progress(),
227
  ):
228
+ if not video_path:
229
+ raise gr.Error("Please upload a video.")
230
+ output_path, registry = handlers.run_eve_inference(
231
+ video_path,
232
+ False,
233
+ False,
234
+ False,
235
+ False,
236
+ registry,
237
+ mod_model,
238
+ request,
239
+ progress,
240
+ )
241
+ return output_path, registry
242
+
243
+ process_btn.click(
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
244
  fn=_process_dispatch,
245
  inputs=[
246
+ video_input,
 
247
  offline_radio,
248
  session_registry,
249
  ],
250
  outputs=[
251
+ video_output,
 
252
  session_registry,
253
  ],
254
  concurrency_limit=pool.worker_count,
install_eve.py CHANGED
@@ -21,11 +21,19 @@ import sys
21
  from huggingface_hub import hf_hub_download
22
 
23
  EVE_REPO = "LatticeSemi/EdgeVisionEngine-Files"
24
- EVE_DEB = "LINUX_X86-10-7.2-eve-sensai_7.2.10~git20260515.8139998_amd64.deb"
25
  EVE_LICENSE = "libEveDevLicense.so"
26
  SECRET_PATH = "/run/secrets/MODEL_ACCESS_TOKEN"
27
  DOWNLOAD_DIR = "/tmp/eve"
28
 
 
 
 
 
 
 
 
 
29
 
30
  def get_token():
31
  """Return an HF token from the first available source, or None."""
@@ -61,6 +69,19 @@ def get_license_destination_path() -> str:
61
  return f"/opt/EVE-{version}-Source/lib"
62
 
63
 
 
 
 
 
 
 
 
 
 
 
 
 
 
64
  def main():
65
  token, auth_source = get_token()
66
 
@@ -95,8 +116,10 @@ def main():
95
  print(f"Installing {license_path}...")
96
  subprocess.run(["mv", license_path, destination_path], check=True)
97
 
 
 
98
  shutil.rmtree(DOWNLOAD_DIR, ignore_errors=True)
99
- print("Eve SDK installed successfully.")
100
 
101
 
102
  if __name__ == "__main__":
 
21
  from huggingface_hub import hf_hub_download
22
 
23
  EVE_REPO = "LatticeSemi/EdgeVisionEngine-Files"
24
+ EVE_DEB = "LINUX_X86-4-7.3-eve-sensai_7.3.4~git20260623.fd5012d_amd64.deb"
25
  EVE_LICENSE = "libEveDevLicense.so"
26
  SECRET_PATH = "/run/secrets/MODEL_ACCESS_TOKEN"
27
  DOWNLOAD_DIR = "/tmp/eve"
28
 
29
+ # Finetuned MOD weights live in the same gated repo as the SDK so they never
30
+ # ship inside the (public) Space. mod_models.py reads them at runtime from
31
+ # MOD_MODELS_DIR. Kept outside /app: this script runs as root before the
32
+ # user-owned WORKDIR /app is created, and writing into /app here would leave it
33
+ # root-owned and break the later user-stage download_examples.py write.
34
+ MOD_MODELS = ("automotive-640x640.tflite", "office-640x640.tflite")
35
+ MOD_MODELS_DIR = "/opt/eve-models"
36
+
37
 
38
  def get_token():
39
  """Return an HF token from the first available source, or None."""
 
69
  return f"/opt/EVE-{version}-Source/lib"
70
 
71
 
72
+ def download_mod_models(token: str | None) -> None:
73
+ """Download the finetuned MOD ``.tflite`` weights into ``MOD_MODELS_DIR``."""
74
+ os.makedirs(MOD_MODELS_DIR, exist_ok=True)
75
+ for model in MOD_MODELS:
76
+ print(f"Downloading {model} from {EVE_REPO}...")
77
+ hf_hub_download(
78
+ repo_id=EVE_REPO,
79
+ filename=model,
80
+ local_dir=MOD_MODELS_DIR,
81
+ token=token,
82
+ )
83
+
84
+
85
  def main():
86
  token, auth_source = get_token()
87
 
 
116
  print(f"Installing {license_path}...")
117
  subprocess.run(["mv", license_path, destination_path], check=True)
118
 
119
+ download_mod_models(token)
120
+
121
  shutil.rmtree(DOWNLOAD_DIR, ignore_errors=True)
122
+ print("Eve SDK and MOD models installed successfully.")
123
 
124
 
125
  if __name__ == "__main__":
mod_models.py CHANGED
@@ -6,30 +6,35 @@ one at new weights, is a one-line change here — nothing else needs to move.
6
 
7
  - **GMOD-80** uses the EVE SDK's bundled default model, so it carries no path
8
  and no class table.
9
- - **AMOD-8** and **OMOD** load their ``.tflite`` from ``models/`` and ship
10
- their own class labels (in output-index order).
 
 
11
  """
12
 
13
  from __future__ import annotations
14
 
 
15
  from pathlib import Path
16
 
17
  from eve_messages import ModelConfig
18
 
19
- _MODELS_DIR = Path(__file__).resolve().parent / "models"
 
 
20
 
21
  # Class labels in the order the model emits them (output index -> label).
22
- AMOD_CLASS_NAMES: tuple[str, ...] = (
23
  "person",
24
  "bicycle",
25
  "car",
26
  "motorcycle",
27
- "truck",
28
  "bus",
 
29
  "traffic light",
30
  "stop sign",
31
  )
32
- OMOD_CLASS_NAMES: tuple[str, ...] = (
33
  "bottle",
34
  "cup",
35
  "potted plant",
@@ -41,16 +46,20 @@ OMOD_CLASS_NAMES: tuple[str, ...] = (
41
  )
42
 
43
  MOD_MODEL_REGISTRY: dict[str, ModelConfig] = {
44
- "GMOD-80": ModelConfig(name="GMOD-80"),
45
- "AMOD-8": ModelConfig(
46
- name="AMOD-8",
47
- model_path=str(_MODELS_DIR / "amod.tflite"),
48
- class_names=AMOD_CLASS_NAMES,
 
 
49
  ),
50
- "OMOD": ModelConfig(
51
- name="OMOD",
52
- model_path=str(_MODELS_DIR / "omod.tflite"),
53
- class_names=OMOD_CLASS_NAMES,
 
 
54
  ),
55
  }
56
 
 
6
 
7
  - **GMOD-80** uses the EVE SDK's bundled default model, so it carries no path
8
  and no class table.
9
+ - **Automotive** and **Office** load their ``.tflite`` from ``MOD_MODELS_DIR``
10
+ (the gated weights fetched at build time, see ``install_eve.py``; falls back to
11
+ the demo-relative ``models/`` locally) and ship their own class labels (in
12
+ output-index order).
13
  """
14
 
15
  from __future__ import annotations
16
 
17
+ import os
18
  from pathlib import Path
19
 
20
  from eve_messages import ModelConfig
21
 
22
+ # In the Space the gated weights are fetched to MOD_MODELS_DIR by install_eve.py
23
+ # at Docker build time; local runs fall back to the demo-relative ``models/``.
24
+ _MODELS_DIR = Path(os.environ.get("MOD_MODELS_DIR") or Path(__file__).resolve().parent / "models")
25
 
26
  # Class labels in the order the model emits them (output index -> label).
27
+ AUTOMOTIVE_CLASS_NAMES: tuple[str, ...] = (
28
  "person",
29
  "bicycle",
30
  "car",
31
  "motorcycle",
 
32
  "bus",
33
+ "truck",
34
  "traffic light",
35
  "stop sign",
36
  )
37
+ OFFICE_CLASS_NAMES: tuple[str, ...] = (
38
  "bottle",
39
  "cup",
40
  "potted plant",
 
46
  )
47
 
48
  MOD_MODEL_REGISTRY: dict[str, ModelConfig] = {
49
+ "GMOD-80": ModelConfig(name="GMOD-80", nms_threshold=0.2),
50
+ "Automotive": ModelConfig(
51
+ name="Automotive",
52
+ model_path=str(_MODELS_DIR / "automotive-640x640.tflite"),
53
+ class_names=AUTOMOTIVE_CLASS_NAMES,
54
+ nms_threshold=0.4,
55
+ iou_threshold=0.4
56
  ),
57
+ "Office": ModelConfig(
58
+ name="Office",
59
+ model_path=str(_MODELS_DIR / "office-640x640.tflite"),
60
+ class_names=OFFICE_CLASS_NAMES,
61
+ nms_threshold=0.4,
62
+ iou_threshold=0.4
63
  ),
64
  }
65
 
shared/eve_app_tabs.py CHANGED
@@ -15,7 +15,6 @@ events (``process_btn.click``, ``webrtc_stream.stream``, etc.) outside.
15
 
16
  from __future__ import annotations
17
 
18
- from dataclasses import dataclass
19
  from typing import Callable, TypeVar
20
 
21
  import gradio as gr
@@ -99,137 +98,6 @@ def build_offline_inference_tab(
99
  return video_tab, input_video, output_video, checkboxes, process_btn, example_dataset
100
 
101
 
102
- @dataclass
103
- class ImageOrVideoOfflineTab:
104
- """Components returned by :func:`build_image_or_video_offline_tab`.
105
-
106
- The image/video accordions are exposed so the caller can wire mutual
107
- exclusion (expanding one collapses the other) — same pattern as
108
- :class:`face_id_tab.FaceIdTab`.
109
- """
110
-
111
- tab: gr.TabItem
112
- image_input: gr.Image
113
- video_input: gr.Video
114
- image_output: gr.Image
115
- video_output: gr.Video
116
- checkboxes: object # demo-specific (radio, checkbox tuple, etc.)
117
- process_btn: gr.Button
118
- image_example_dataset: gr.Dataset | None
119
- video_example_dataset: gr.Dataset | None
120
- image_accordion: gr.Accordion
121
- video_accordion: gr.Accordion
122
-
123
-
124
- def build_image_or_video_offline_tab(
125
- *,
126
- feature_checkbox_builder: Callable[..., TCheckboxes],
127
- image_examples: list | None,
128
- video_examples: list | None,
129
- video_limits: VideoLimits,
130
- tab_label: str = "Offline Inference",
131
- feature_hint: str = "applied when processing starts",
132
- ) -> ImageOrVideoOfflineTab:
133
- """Build an Offline Inference tab that accepts image OR video.
134
-
135
- Pattern: two mutually-exclusive accordions on the input side ("Input
136
- from an Image" / "Input from a Video"), one Process button, two
137
- output components (image + video) shown side-by-side. Caller is
138
- responsible for:
139
-
140
- - Wiring ``image_accordion`` / ``video_accordion`` mutual exclusion
141
- (one-liner per side, see ``face_id_tab.FaceIdTab.wire``).
142
- - Routing ``process_btn.click`` to a handler that dispatches by
143
- which input is populated.
144
- - Toggling output visibility based on which branch ran.
145
-
146
- Args:
147
- feature_checkbox_builder: Callable that builds the demo's feature
148
- checkboxes/radio inside the tab. Called with keyword
149
- ``hint=feature_hint``.
150
- image_examples: Pre-loaded image examples (list of ``[path]``
151
- rows) or ``None`` to skip the examples accordion.
152
- video_examples: Same for videos.
153
- video_limits: Upload constraints rendered inside the video
154
- accordion.
155
- tab_label: Tab label text.
156
- feature_hint: Short hint shown next to the Features heading.
157
- """
158
- with gr.TabItem(tab_label) as tab:
159
- with gr.Accordion("Instructions", open=False):
160
- gr.Markdown(
161
- "1. Select the model that will be used for object detection\n"
162
- "2. Choose between processing an **Image** or a **Video**:\n"
163
- " - For an image: select an example or upload your own, then press "
164
- "**Process**.\n"
165
- " - For a video: expand the video section, select an example or "
166
- "upload your own, then press **Process**.\n\n"
167
- "Once processing is complete, the annotated result appears on the "
168
- "right (image or video, depending on the input)."
169
- )
170
-
171
- checkboxes = feature_checkbox_builder(hint=feature_hint)
172
-
173
- with gr.Row():
174
- # --- Left column: input ---
175
- with gr.Column(scale=5):
176
- image_example_dataset: gr.Dataset | None = None
177
- video_example_dataset: gr.Dataset | None = None
178
-
179
- with gr.Accordion(
180
- "Input from an Image", open=True
181
- ) as image_accordion:
182
- if image_examples:
183
- with gr.Accordion("Image Examples", open=True):
184
- image_example_dataset = gr.Dataset(
185
- components=[gr.Image(visible=False)],
186
- samples=image_examples,
187
- show_label=False,
188
- )
189
- image_input = gr.Image(
190
- label="Input Image",
191
- sources=["upload", "webcam"],
192
- type="filepath",
193
- )
194
-
195
- with gr.Accordion(
196
- "Input from a Video", open=False
197
- ) as video_accordion:
198
- if video_examples:
199
- with gr.Accordion("Video Examples", open=True):
200
- video_example_dataset = gr.Dataset(
201
- components=[gr.Video(visible=False)],
202
- samples=video_examples,
203
- show_label=False,
204
- )
205
- build_video_constraints_accordion(video_limits)
206
- video_input = gr.Video(
207
- label="Input Video", sources=["upload", "webcam"]
208
- )
209
-
210
- process_btn = gr.Button(
211
- "Process", variant="primary", interactive=False
212
- )
213
-
214
- # --- Right column: output (image OR video, toggled by handler) ---
215
- with gr.Column(scale=5):
216
- image_output = gr.Image(label="Output Image", visible=True)
217
- video_output = gr.Video(label="Output Video", visible=False)
218
-
219
- return ImageOrVideoOfflineTab(
220
- tab=tab,
221
- image_input=image_input,
222
- video_input=video_input,
223
- image_output=image_output,
224
- video_output=video_output,
225
- checkboxes=checkboxes,
226
- process_btn=process_btn,
227
- image_example_dataset=image_example_dataset,
228
- video_example_dataset=video_example_dataset,
229
- image_accordion=image_accordion,
230
- video_accordion=video_accordion,
231
- )
232
-
233
 
234
  def build_live_inference_tab(
235
  *,
 
15
 
16
  from __future__ import annotations
17
 
 
18
  from typing import Callable, TypeVar
19
 
20
  import gradio as gr
 
98
  return video_tab, input_video, output_video, checkboxes, process_btn, example_dataset
99
 
100
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
101
 
102
  def build_live_inference_tab(
103
  *,
shared/eve_worker_pool.py CHANGED
@@ -1,1121 +1,1119 @@
1
- """Multi-process Eve SDK worker pool.
2
-
3
- Spawns N child processes, each hosting an independent ``EveWrapper`` instance
4
- with its own ``ctypes.CDLL`` handle. The main Gradio process communicates with
5
- workers via ``multiprocessing.Pipe`` (one bidirectional pipe per worker).
6
-
7
- Frame data is serialised as raw bytes (``ndarray.tobytes()``) rather than
8
- pickled ``ndarray`` objects for speed (~1-2 ms for a 1280x720x3 frame).
9
-
10
- Usage::
11
-
12
- pool = EveWorkerPool(max_workers=4, ram_headroom_gb=2.0)
13
- worker = pool.acquire(session_hash="abc123")
14
- try:
15
- result = worker.send_inference(frame, features)
16
- finally:
17
- pool.release(worker)
18
- """
19
-
20
- import atexit
21
- import logging
22
- import multiprocessing as mp
23
- import os
24
- import threading
25
- import time
26
- from collections.abc import Callable
27
- from dataclasses import dataclass
28
- from multiprocessing.connection import Connection
29
-
30
- import numpy as np
31
- from eve_messages import (
32
- CalibrateNewUserCmd,
33
- CalibrateOkResponse,
34
- CalibrationResultMsg,
35
- EnableFaceIdCmd,
36
- ErrorResponse,
37
- FeatureFlags,
38
- GalleryRestoredResponse,
39
- GetProfileStatsCmd,
40
- GetTimingStatsCmd,
41
- HeartbeatResponse,
42
- InferenceCmd,
43
- InferenceResponse,
44
- OkResponse,
45
- ProcessVideoCmd,
46
- ProfileStatsResponse,
47
- ProgressResponse,
48
- ReadyResponse,
49
- RemoveAllUsersCmd,
50
- RemoveUsersOkResponse,
51
- RestoreGalleryCmd,
52
- RestoreGalleryOkResponse,
53
- SerializedFrame,
54
- ShutdownCmd,
55
- StartProfilingCmd,
56
- StopProfilingCmd,
57
- TimingStatsResponse,
58
- VideoProcessingDoneResponse,
59
- WorkerCmd,
60
- WorkerResponse,
61
- )
62
- from log_utils import setup_logger
63
-
64
- logger = setup_logger("EveWorkerPool")
65
-
66
- # Cached per-process so the probe runs once per worker, not per video.
67
- _h264_encoder_cache: tuple[str, dict[str, str]] | None = None
68
-
69
-
70
- def _get_h264_encoder() -> tuple[str, dict[str, str]]:
71
- """Return (codec_name, codec_options) for H.264 encoding.
72
-
73
- Tries NVENC (GPU) first, falls back to libx264 (CPU).
74
- """
75
- global _h264_encoder_cache
76
- if _h264_encoder_cache is not None:
77
- return _h264_encoder_cache
78
-
79
- import logging
80
-
81
- import av
82
-
83
- log = logging.getLogger("EveWorkerPool")
84
- try:
85
- ctx = av.CodecContext.create("h264_nvenc", "w")
86
- ctx.width = 64
87
- ctx.height = 64
88
- ctx.pix_fmt = "yuv420p"
89
- ctx.open()
90
- ctx.close()
91
- log.info("Using GPU encoder: h264_nvenc")
92
- _h264_encoder_cache = (
93
- "h264_nvenc",
94
- {"profile": "baseline", "preset": "p1", "delay": "0"},
95
- )
96
- except Exception:
97
- log.info("GPU encoder unavailable, falling back to libx264")
98
- _h264_encoder_cache = (
99
- "libx264",
100
- {"profile": "baseline", "level": "3.1", "preset": "ultrafast", "threads": "1"},
101
- )
102
- return _h264_encoder_cache
103
-
104
-
105
- def log_worker_activity(
106
- log: logging.Logger,
107
- action: str,
108
- feature: str,
109
- pool: "EveWorkerPool",
110
- worker_id: int | None = None,
111
- note: str = "",
112
- ) -> None:
113
- """Single-line worker activity log for tracking acquire/release/queue events."""
114
- busy = pool.worker_count - pool.idle_count
115
- w = f" w{worker_id}" if worker_id is not None else ""
116
- n = f" ({note})" if note else ""
117
- log.info(
118
- f"[worker] {action}{w} for {feature}{n} "
119
- f"[busy={busy} idle={pool.idle_count} wait={pool.waiting_count} "
120
- f"sessions={pool.session_count}]"
121
- )
122
-
123
-
124
- def compute_fps_sleep(
125
- elapsed: float,
126
- max_fps: float | None,
127
- min_fps: float | None,
128
- load: float,
129
- ) -> float:
130
- """Compute how long to sleep after one frame to stay within the FPS cap.
131
-
132
- Linearly interpolates the target cycle time between ``1/max_fps``
133
- (at *load* 0) and ``1/min_fps`` (at *load* 1), then subtracts the
134
- time already spent.
135
-
136
- Used by both the live-inference bridge and the video-processing loop
137
- so that the throttle logic lives in one place.
138
- """
139
- if max_fps is None:
140
- return 0.0
141
- effective_min = min_fps if min_fps is not None else max_fps
142
- max_period = 1.0 / max_fps
143
- min_period = 1.0 / effective_min
144
- target_period = max_period + load * (min_period - max_period)
145
- return max(0.0, target_period - elapsed)
146
-
147
-
148
- # Use 'spawn' on all platforms so each child gets a fresh interpreter
149
- # (required on Windows; avoids CDLL sharing on Linux/fork).
150
- _mp_ctx = mp.get_context("spawn")
151
-
152
- DO_PROBE_WORKER = (
153
- False # Set to False to skip the RAM probe and just use an estimate of the RAM used
154
- )
155
-
156
-
157
- # ---------------------------------------------------------------------------
158
- # Worker process entry point (runs in the child)
159
- # ---------------------------------------------------------------------------
160
-
161
-
162
- def _eve_worker_main(
163
- conn: Connection,
164
- eve_bin_path: str,
165
- eve_lib_path: str,
166
- worker_id: int,
167
- max_jobs_per_worker: int | None = None,
168
- shared_load: "mp.Value | None" = None,
169
- ) -> None:
170
- """Top-level function executed inside each worker ``Process``.
171
-
172
- Communicates with the main process exclusively through *conn*.
173
- *shared_load* is a cross-process float (0.0–1.0) updated by the pool
174
- on acquire/release, used by the video processing loop to throttle.
175
- """
176
- import ctypes
177
- import gc
178
- import platform
179
- import signal
180
- import sys
181
- from pathlib import Path
182
-
183
- signal.signal(signal.SIGINT, signal.SIG_IGN)
184
-
185
- # CPU affinity (for stress testing)
186
- affinity_env = os.environ.get(f"WORKER{worker_id}_AFFINITY")
187
- if affinity_env:
188
- cpu_id = int(affinity_env.removeprefix("CPU"))
189
- import psutil
190
-
191
- psutil.Process().cpu_affinity([cpu_id])
192
- logger.info(f"Worker {worker_id} pinned to CPU {cpu_id}")
193
-
194
- # Ensure shared package is importable inside the child
195
- shared_dir = str(Path(__file__).resolve().parent)
196
- if shared_dir not in sys.path:
197
- sys.path.insert(0, shared_dir)
198
-
199
- from eve_wrapper import EveWrapper # noqa: E402 (deferred import)
200
- from frame_utils import load_media_frames_raw # noqa: E402
201
-
202
- eve: EveWrapper | None = None
203
- try:
204
- eve = EveWrapper(eve_bin_path, eve_lib_path)
205
- conn.send(ReadyResponse(pid=os.getpid()))
206
- except Exception as exc:
207
- conn.send(ErrorResponse(error=str(exc)))
208
- return
209
-
210
- # Optional cProfile (enabled via ENABLE_PROFILER env var)
211
- profiler: "cProfile.Profile | None" = None
212
- profiling_enabled = os.environ.get("ENABLE_PROFILER", "").strip() not in ("", "0", "false")
213
- if profiling_enabled:
214
- import cProfile
215
-
216
- # Use process_time so I/O wait (pipe polling) doesn't dominate the profile
217
- profiler = cProfile.Profile(timer=time.process_time)
218
-
219
- job_count = 0
220
- while True:
221
- try:
222
- if not conn.poll(timeout=2.0):
223
- # No command received — send heartbeat
224
- try:
225
- conn.send(HeartbeatResponse())
226
- except (BrokenPipeError, OSError):
227
- break
228
- continue
229
-
230
- cmd: WorkerCmd = conn.recv()
231
- except (EOFError, OSError):
232
- break
233
-
234
- if isinstance(cmd, ShutdownCmd):
235
- conn.send(OkResponse())
236
- break
237
-
238
- try:
239
- if isinstance(cmd, InferenceCmd):
240
- frame = np.frombuffer(cmd.frame_bytes, dtype=cmd.dtype).reshape(cmd.shape)
241
- features = cmd.features
242
- del cmd # free recv'd bytes early; numpy holds its own ref
243
-
244
- eve.enable_face_and_person_detection(
245
- faceEnabled=features.face_detection,
246
- personEnabled=features.person_detection,
247
- )
248
- eve.enable_face_id(enabled=features.face_id)
249
- eve.enable_hand_gesture(enabled=features.hand_gesture)
250
- eve.enable_object_detection(config=features.mod_model_config)
251
- eve.enable_mirror(True)
252
- result = eve.inference(frame)
253
- del frame
254
-
255
- conn.send(
256
- InferenceResponse(
257
- frame_bytes=result.tobytes(),
258
- shape=result.shape,
259
- dtype=str(result.dtype),
260
- )
261
- )
262
- del result
263
-
264
- elif isinstance(cmd, CalibrateNewUserCmd):
265
- frames = _deserialize_frames(cmd.frames_data)
266
- result = eve.calibrate_new_user(frames)
267
- conn.send(
268
- CalibrateOkResponse(
269
- result=CalibrationResultMsg(result.success, result.user_id, result.message),
270
- )
271
- )
272
-
273
- elif isinstance(cmd, RemoveAllUsersCmd):
274
- ok = eve.remove_all_users()
275
- conn.send(RemoveUsersOkResponse(result=ok))
276
-
277
- elif isinstance(cmd, RestoreGalleryCmd):
278
- frames_per_user = [_deserialize_frames(fd) for fd in cmd.frames_per_user_data]
279
- eve.remove_all_users()
280
- results = eve.restore_gallery(frames_per_user)
281
- conn.send(
282
- RestoreGalleryOkResponse(
283
- results=[
284
- CalibrationResultMsg(r.success, r.user_id, r.message) for r in results
285
- ],
286
- )
287
- )
288
-
289
- elif isinstance(cmd, EnableFaceIdCmd):
290
- eve.enable_face_id(enabled=cmd.enabled)
291
- conn.send(OkResponse())
292
-
293
- elif isinstance(cmd, ProcessVideoCmd):
294
- import av
295
- import cv2
296
- from fractions import Fraction
297
-
298
- features = cmd.features
299
-
300
- # Workers are reused across jobs, so reset the EVE pipeline
301
- # before each video otherwise tracking / ideal-user state
302
- # leaks from the previous run and the second video shows no
303
- # ideal user.
304
- eve.reset_pipeline()
305
-
306
- # Gallery restore (if requested) — must happen before
307
- # applying user feature flags because restore_gallery()
308
- # internally enables face/person detection for calibration.
309
- gallery_results: list[CalibrationResultMsg] = []
310
- if cmd.remove_all_users:
311
- eve.remove_all_users()
312
- if cmd.gallery_paths:
313
- frames_per_user = [load_media_frames_raw(p) for p in cmd.gallery_paths]
314
- raw_results = eve.restore_gallery(frames_per_user)
315
- del frames_per_user
316
- gallery_results = [
317
- CalibrationResultMsg(r.success, r.user_id, r.message) for r in raw_results
318
- ]
319
- del raw_results
320
- conn.send(GalleryRestoredResponse(results=gallery_results))
321
-
322
- # Apply user's feature flags after gallery restore so they
323
- # aren't overridden by restore_gallery's internal SDK calls.
324
- eve.enable_face_and_person_detection(
325
- faceEnabled=features.face_detection,
326
- personEnabled=features.person_detection,
327
- )
328
- eve.enable_face_id(enabled=features.face_id)
329
- eve.enable_hand_gesture(enabled=features.hand_gesture)
330
- eve.enable_object_detection(config=features.mod_model_config)
331
- eve.enable_mirror(False)
332
-
333
- # Video processing loop
334
- cap = cv2.VideoCapture(cmd.input_path)
335
- codec_name, codec_opts = _get_h264_encoder()
336
- container = av.open(cmd.output_path, mode="w", options={"movflags": "faststart"})
337
- stream = container.add_stream(codec_name, rate=round(cmd.fps))
338
- stream.time_base = Fraction(1, round(cmd.fps))
339
- stream.width = cmd.width
340
- stream.height = cmd.height
341
- stream.pix_fmt = "yuv420p"
342
- stream.codec_context.options = codec_opts
343
-
344
- frame_count = 0
345
- v_max_fps = cmd.max_fps
346
- v_min_fps = cmd.min_fps if cmd.min_fps is not None else v_max_fps
347
- try:
348
- while True:
349
- ret, frame = cap.read()
350
- if not ret:
351
- break
352
- t0 = time.monotonic()
353
- result = eve.inference(frame)
354
- vf = av.VideoFrame.from_ndarray(result, format="bgr24")
355
- vf.pts = frame_count
356
- for pkt in stream.encode(vf):
357
- container.mux(pkt)
358
- del pkt
359
- del result, vf, frame
360
- frame_count += 1
361
- if frame_count % 10 == 0:
362
- conn.send(
363
- ProgressResponse(
364
- current=frame_count,
365
- total=cmd.total_frames,
366
- )
367
- )
368
- load = shared_load.value if shared_load is not None else 0.0
369
- # FPS throttle (only if there's some load already)
370
- if load > 0.0:
371
- sleep = compute_fps_sleep(
372
- time.monotonic() - t0, v_max_fps, v_min_fps, load
373
- )
374
- if sleep > 0:
375
- time.sleep(sleep)
376
- # Flush encoder
377
- for pkt in stream.encode():
378
- container.mux(pkt)
379
- del pkt
380
- finally:
381
- cap.release()
382
- container.close()
383
- del (
384
- cap,
385
- container,
386
- stream,
387
- )
388
-
389
- job_count += 1
390
- conn.send(
391
- VideoProcessingDoneResponse(
392
- frames_processed=frame_count,
393
- gallery_results=gallery_results,
394
- recycle=max_jobs_per_worker is not None
395
- and job_count >= max_jobs_per_worker,
396
- )
397
- )
398
-
399
- del cmd
400
- gc.collect()
401
- if platform.system() == "Linux":
402
- try:
403
- ctypes.CDLL("libc.so.6").malloc_trim(0)
404
- except Exception:
405
- pass
406
-
407
- if max_jobs_per_worker is not None and job_count >= max_jobs_per_worker:
408
- logger.info(f"Worker {worker_id} recycling after {job_count} jobs")
409
- break
410
-
411
- elif isinstance(cmd, StartProfilingCmd):
412
- if profiler is not None:
413
- profiler.enable()
414
- conn.send(OkResponse())
415
-
416
- elif isinstance(cmd, StopProfilingCmd):
417
- if profiler is not None:
418
- profiler.disable()
419
- conn.send(OkResponse())
420
-
421
- elif isinstance(cmd, GetProfileStatsCmd):
422
- if profiler is not None:
423
- import io
424
- import marshal
425
- import pstats
426
-
427
- profiler.disable()
428
- stream = io.StringIO()
429
- ps = pstats.Stats(profiler, stream=stream)
430
- conn.send(ProfileStatsResponse(stats_data=marshal.dumps(ps.stats)))
431
- else:
432
- conn.send(ProfileStatsResponse(stats_data=b""))
433
-
434
- elif isinstance(cmd, GetTimingStatsCmd):
435
- conn.send(TimingStatsResponse(stats=eve.get_timing_stats(reset=cmd.reset)))
436
-
437
- else:
438
- conn.send(ErrorResponse(error=f"Unknown command: {cmd}"))
439
-
440
- except Exception as exc:
441
- logger.error(f"Worker {worker_id} error handling '{type(cmd).__name__}': {exc}")
442
- try:
443
- conn.send(ErrorResponse(error=str(exc)))
444
- except (BrokenPipeError, OSError):
445
- break
446
-
447
- # Clean shutdown
448
- if eve is not None:
449
- eve.shutdown()
450
-
451
-
452
- def _serialize_frames(frames: list[np.ndarray]) -> list[SerializedFrame]:
453
- """Convert a list of ndarrays to a picklable representation."""
454
- return [SerializedFrame(data=f.tobytes(), shape=f.shape, dtype=str(f.dtype)) for f in frames]
455
-
456
-
457
- def _deserialize_frames(data: list[SerializedFrame]) -> list[np.ndarray]:
458
- """Reconstruct ndarrays from their serialized form."""
459
- return [np.frombuffer(d.data, dtype=d.dtype).reshape(d.shape) for d in data]
460
-
461
-
462
- # ---------------------------------------------------------------------------
463
- # EveWorker — main-process handle for one child process
464
- # ---------------------------------------------------------------------------
465
-
466
-
467
- class EveWorker:
468
- """Main-process proxy for a single worker process.
469
-
470
- Not thread-safe by itself — the ``EveWorkerPool`` ensures only one thread
471
- accesses a worker at a time.
472
- """
473
-
474
- def __init__(
475
- self,
476
- process: mp.Process,
477
- conn: Connection,
478
- worker_id: int,
479
- ):
480
- self.process = process
481
- self._conn = conn
482
- self.worker_id = worker_id
483
- self.status: str = "idle" # idle | busy | dead
484
- self.session_hash: str | None = None
485
- self.busy_since: float | None = None
486
- self.last_heartbeat: float = time.monotonic()
487
- self.is_live_stream: bool = False
488
- self.pending_recycle: bool = False
489
-
490
- # -- command helpers ---------------------------------------------------
491
-
492
- def _recv_response(self) -> WorkerResponse:
493
- """Receive the next non-heartbeat response from the worker."""
494
- while True:
495
- try:
496
- resp: WorkerResponse = self._conn.recv()
497
- except (EOFError, OSError):
498
- raise RuntimeError(f"Worker {self.worker_id} pipe closed")
499
- except Exception as exc:
500
- self.pending_recycle = True
501
- raise RuntimeError(
502
- f"Worker {self.worker_id} pipe corrupt: {exc}"
503
- ) from exc
504
- if isinstance(resp, HeartbeatResponse):
505
- self.last_heartbeat = time.monotonic()
506
- continue
507
- return resp
508
-
509
- def _expect_ok(self, context: str) -> WorkerResponse:
510
- """Receive a response and raise on error."""
511
- resp = self._recv_response()
512
- if isinstance(resp, ErrorResponse):
513
- raise RuntimeError(f"Worker {self.worker_id} {context} error: {resp.error}")
514
- return resp
515
-
516
- def send_inference(self, frame: np.ndarray, features: FeatureFlags) -> np.ndarray:
517
- """Send a frame for inference and return the annotated result."""
518
- self._conn.send(
519
- InferenceCmd(
520
- frame_bytes=frame.tobytes(),
521
- shape=frame.shape,
522
- dtype=str(frame.dtype),
523
- features=features,
524
- )
525
- )
526
- resp = self._expect_ok("inference")
527
- if not isinstance(resp, InferenceResponse):
528
- raise RuntimeError(f"Worker {self.worker_id} inference: unexpected {resp}")
529
- return np.frombuffer(resp.frame_bytes, dtype=resp.dtype).reshape(resp.shape)
530
-
531
-
532
- def send_calibrate_new_user(self, frames: list[np.ndarray]) -> CalibrationResultMsg:
533
- self._conn.send(CalibrateNewUserCmd(frames_data=_serialize_frames(frames)))
534
- resp = self._expect_ok("calibrate")
535
- if not isinstance(resp, CalibrateOkResponse):
536
- raise RuntimeError(f"Worker {self.worker_id} calibrate: unexpected {resp}")
537
- return resp.result
538
-
539
- def send_remove_all_users(self) -> bool:
540
- self._conn.send(RemoveAllUsersCmd())
541
- resp = self._expect_ok("remove")
542
- if not isinstance(resp, RemoveUsersOkResponse):
543
- raise RuntimeError(f"Worker {self.worker_id} remove: unexpected {resp}")
544
- return resp.result
545
-
546
- def send_restore_gallery(
547
- self, frames_per_user: list[list[np.ndarray]]
548
- ) -> list[CalibrationResultMsg]:
549
- self._conn.send(
550
- RestoreGalleryCmd(
551
- frames_per_user_data=[_serialize_frames(fu) for fu in frames_per_user],
552
- )
553
- )
554
- resp = self._expect_ok("restore")
555
- if not isinstance(resp, RestoreGalleryOkResponse):
556
- raise RuntimeError(f"Worker {self.worker_id} restore: unexpected {resp}")
557
- return resp.results
558
-
559
- def send_enable_face_id(self, enabled: bool) -> None:
560
- self._conn.send(EnableFaceIdCmd(enabled=enabled))
561
- self._expect_ok("enable_face_id")
562
-
563
- def send_process_video(
564
- self,
565
- input_path: str,
566
- output_path: str,
567
- features: FeatureFlags,
568
- gallery_paths: list[str],
569
- remove_all_users: bool,
570
- fps: float,
571
- width: int,
572
- height: int,
573
- total_frames: int,
574
- progress: Callable[..., None] | None = None,
575
- max_fps: float | None = None,
576
- min_fps: float | None = None,
577
- ) -> tuple[int, list[CalibrationResultMsg]]:
578
- """Run the full video processing loop inside the worker process.
579
-
580
- The worker reads the input video, runs Eve inference on each frame,
581
- encodes the result with PyAV, and writes the output video. Only
582
- small progress messages travel over the pipe — no frame data.
583
-
584
- Args:
585
- input_path: Path to the input video file.
586
- output_path: Path where the output video will be written.
587
- features: Feature toggles for the Eve SDK.
588
- gallery_paths: Media paths for Face ID gallery restore (may be empty).
589
- remove_all_users: Whether to clear the Face ID gallery first.
590
- fps: Output video frame rate.
591
- width: Output video width in pixels.
592
- height: Output video height in pixels.
593
- total_frames: Total frame count (for progress reporting).
594
- progress: Optional ``gr.Progress``-compatible callback.
595
- max_fps: FPS cap when idle (``None`` = unlimited).
596
- min_fps: FPS cap under full load (defaults to ``max_fps``).
597
-
598
- Returns:
599
- Tuple of (frames_processed, gallery_results).
600
-
601
- Raises:
602
- RuntimeError: If the worker reports an error.
603
- """
604
- self._conn.send(
605
- ProcessVideoCmd(
606
- input_path=input_path,
607
- output_path=output_path,
608
- features=features,
609
- gallery_paths=gallery_paths,
610
- remove_all_users=remove_all_users,
611
- fps=fps,
612
- width=width,
613
- height=height,
614
- total_frames=total_frames,
615
- max_fps=max_fps,
616
- min_fps=min_fps,
617
- )
618
- )
619
-
620
- gallery_results: list[CalibrationResultMsg] = []
621
-
622
- while True:
623
- resp = self._recv_response()
624
-
625
- if isinstance(resp, GalleryRestoredResponse):
626
- gallery_results = resp.results
627
-
628
- elif isinstance(resp, ProgressResponse):
629
- if progress is not None:
630
- progress(
631
- (resp.current, resp.total),
632
- desc=f"Processing frame {resp.current}/{resp.total}",
633
- )
634
-
635
- elif isinstance(resp, VideoProcessingDoneResponse):
636
- if resp.recycle:
637
- self.pending_recycle = True
638
- return resp.frames_processed, gallery_results
639
-
640
- elif isinstance(resp, ErrorResponse):
641
- raise RuntimeError(f"Worker {self.worker_id} process_video error: {resp.error}")
642
-
643
- else:
644
- raise RuntimeError(f"Worker {self.worker_id} unexpected response: {resp}")
645
-
646
- def send_start_profiling(self) -> None:
647
- self._conn.send(StartProfilingCmd())
648
- self._expect_ok("start_profiling")
649
-
650
- def send_stop_profiling(self) -> None:
651
- self._conn.send(StopProfilingCmd())
652
- self._expect_ok("stop_profiling")
653
-
654
- def send_get_profile_stats(self) -> bytes:
655
- """Request profiling stats from the worker. Returns marshalled pstats data."""
656
- self._conn.send(GetProfileStatsCmd())
657
- resp = self._recv_response()
658
- if isinstance(resp, ErrorResponse):
659
- raise RuntimeError(f"Worker {self.worker_id} get_profile_stats error: {resp.error}")
660
- if not isinstance(resp, ProfileStatsResponse):
661
- raise RuntimeError(f"Worker {self.worker_id} get_profile_stats: unexpected {resp}")
662
- return resp.stats_data
663
-
664
- def send_get_timing_stats(self, reset: bool = True) -> dict[str, tuple[int, float]]:
665
- """Request per-SDK-call timing stats from the worker."""
666
- self._conn.send(GetTimingStatsCmd(reset=reset))
667
- resp = self._recv_response()
668
- if isinstance(resp, ErrorResponse):
669
- raise RuntimeError(f"Worker {self.worker_id} get_timing_stats error: {resp.error}")
670
- if not isinstance(resp, TimingStatsResponse):
671
- raise RuntimeError(f"Worker {self.worker_id} get_timing_stats: unexpected {resp}")
672
- return resp.stats
673
-
674
- def send_shutdown(self) -> None:
675
- try:
676
- self._conn.send(ShutdownCmd())
677
- self._conn.recv()
678
- except (EOFError, OSError, BrokenPipeError):
679
- pass
680
-
681
- def drain_heartbeats(self) -> None:
682
- """Consume any pending heartbeat messages on the pipe."""
683
- while self._conn.poll(timeout=0):
684
- try:
685
- msg = self._conn.recv()
686
- if isinstance(msg, HeartbeatResponse):
687
- self.last_heartbeat = time.monotonic()
688
- except (EOFError, OSError):
689
- break
690
- except Exception:
691
- # Pipe has non-pickle data (e.g. native library output) — the
692
- # framing is now unreliable so schedule this worker for recycling.
693
- logger.warning(
694
- "Worker %d: corrupt data on pipe, scheduling recycle",
695
- self.worker_id,
696
- )
697
- self.pending_recycle = True
698
- break
699
-
700
-
701
- # ---------------------------------------------------------------------------
702
- # EveWorkerPool — manages all workers
703
- # ---------------------------------------------------------------------------
704
-
705
-
706
- @dataclass
707
- class _PoolConfig:
708
- max_workers: int = 8
709
- max_ram_gb: float = 32.0
710
- ram_headroom_gb: float = 2.0
711
- rss_safety_factor: float = 2.5
712
- stuck_timeout_s: float = 300.0 # 5 min for video processing
713
- health_check_interval_s: float = 5.0
714
- max_jobs_per_worker: int | None = 50
715
-
716
-
717
- class EveWorkerPool:
718
- """Pool of Eve SDK worker processes.
719
-
720
- Measures available RAM at startup, spawns as many workers as fit
721
- (with safety margins), and provides acquire / release semantics.
722
-
723
- Args:
724
- max_workers: Upper bound on the number of workers.
725
- max_ram_gb: Cap on available RAM (GB) considered for worker sizing.
726
- Useful on shared machines where reported available RAM is
727
- much larger than what the demo should consume.
728
- ram_headroom_gb: Free RAM (GB) to reserve for main process / OS.
729
- rss_safety_factor: Multiplier applied to the measured init RSS to
730
- estimate peak runtime memory per worker.
731
- eve_bin_path: Forwarded to ``EveWrapper``.
732
- eve_lib_path: Forwarded to ``EveWrapper``.
733
- """
734
-
735
- def __init__(
736
- self,
737
- max_workers: int = 8,
738
- max_ram_gb: float = 32.0,
739
- ram_headroom_gb: float = 2.0,
740
- rss_safety_factor: float = 2.5,
741
- eve_bin_path: str = "",
742
- eve_lib_path: str = "",
743
- max_jobs_per_worker: int | None = 50,
744
- ):
745
- self._cfg = _PoolConfig(
746
- max_workers=max_workers,
747
- max_ram_gb=max_ram_gb,
748
- ram_headroom_gb=ram_headroom_gb,
749
- rss_safety_factor=rss_safety_factor,
750
- max_jobs_per_worker=max_jobs_per_worker,
751
- )
752
- self._eve_bin_path = eve_bin_path
753
- self._eve_lib_path = eve_lib_path
754
- self._lock = threading.Condition(threading.Lock())
755
- self._workers: list[EveWorker] = []
756
- self._shutting_down = False
757
- self._next_id = 0
758
- self._waiting_count = 0
759
- self.session_count = 0
760
- # Shared load signal readable by worker processes (0.0–1.0).
761
- # lock=False is safe: single writer (main process), readers get
762
- # a slightly stale value at worst.
763
- self._shared_load: mp.Value = _mp_ctx.Value("d", 0.0, lock=False)
764
-
765
- # Spawn workers based on RAM capacity
766
- count = self._compute_worker_count()
767
- self._spawn_workers_parallel(count)
768
-
769
- logger.info(f"EveWorkerPool started with {len(self._workers)} workers")
770
-
771
- # Start health-check daemon
772
- self._health_thread = threading.Thread(target=self._health_monitor, daemon=True)
773
- self._health_thread.start()
774
-
775
- # Hourly memory report
776
- from memory_monitor import start_memory_reporter
777
-
778
- self._mem_thread = start_memory_reporter(
779
- get_workers=lambda: list(self._workers),
780
- lock=self._lock,
781
- shutdown_flag=lambda: self._shutting_down,
782
- max_ram_gb=self._cfg.max_ram_gb,
783
- )
784
-
785
- atexit.register(self.shutdown_all)
786
-
787
- # -- public API --------------------------------------------------------
788
-
789
- @property
790
- def worker_count(self) -> int:
791
- """Number of live workers (for setting ``concurrency_limit``)."""
792
- with self._lock:
793
- return sum(1 for w in self._workers if w.status != "dead")
794
-
795
- @property
796
- def idle_count(self) -> int:
797
- """Number of idle workers."""
798
- with self._lock:
799
- return sum(1 for w in self._workers if w.status == "idle")
800
-
801
- @property
802
- def waiting_count(self) -> int:
803
- """Number of callers blocked in ``acquire()`` waiting for a worker."""
804
- with self._lock:
805
- return self._waiting_count
806
-
807
- def try_acquire(self, session_hash: str) -> EveWorker | None:
808
- """Non-blocking acquire returns an idle worker or ``None``.
809
-
810
- Unlike :meth:`acquire`, this never blocks or increments the
811
- waiting count. Used by :class:`LiveStreamManager` so that
812
- WebRTC frame handlers can return immediately when no worker
813
- is available.
814
- """
815
- with self._lock:
816
- return self._try_acquire(session_hash)
817
-
818
- def _try_acquire(self, session_hash: str) -> EveWorker | None:
819
- """Try to grab an idle worker (must hold ``self._lock``).
820
-
821
- Returns the worker if successful, ``None`` otherwise.
822
- """
823
- # Prefer session-affiliated idle worker (Face ID gallery affinity)
824
- for w in self._workers:
825
- if w.status == "idle" and not w.pending_recycle and w.session_hash == session_hash:
826
- w.status = "busy"
827
- w.busy_since = time.monotonic()
828
- w.drain_heartbeats()
829
- self._update_shared_load()
830
- return w
831
- # Any idle worker
832
- for w in self._workers:
833
- if w.status == "idle" and not w.pending_recycle:
834
- w.status = "busy"
835
- w.session_hash = session_hash
836
- w.busy_since = time.monotonic()
837
- w.drain_heartbeats()
838
- self._update_shared_load()
839
- return w
840
- return None
841
-
842
- def acquire(
843
- self,
844
- session_hash: str,
845
- timeout: float = 300.0,
846
- progress: Callable[..., None] | None = None,
847
- eta_fn: Callable[[int], float | None] | None = None,
848
- ) -> EveWorker:
849
- """Reserve a worker for *session_hash*, blocking until one is free.
850
-
851
- Prefers a worker already affiliated with the session (Face ID
852
- gallery affinity). Falls back to any idle worker. Blocks up to
853
- *timeout* seconds if all workers are busy.
854
-
855
- Args:
856
- session_hash: Gradio session hash for worker affinity.
857
- timeout: Max seconds to wait for a free worker.
858
- progress: Optional ``gr.Progress``-compatible callback invoked
859
- while waiting so the UI can show queue position.
860
- eta_fn: Optional callable accepting a 1-based queue position
861
- and returning estimated seconds until a worker frees up.
862
- The return value is shown in the progress description.
863
-
864
- Raises:
865
- RuntimeError: If no worker becomes available within *timeout*.
866
- """
867
- deadline = time.monotonic() + timeout
868
-
869
- # Fast path — try without incrementing waiting count
870
- with self._lock:
871
- worker = self._try_acquire(session_hash)
872
- if worker is not None:
873
- return worker
874
- self._waiting_count += 1
875
-
876
- # Slow path — block until a worker is released
877
- try:
878
- while True:
879
- # Send progress update OUTSIDE the lock to avoid blocking release()
880
- if progress is not None:
881
- with self._lock:
882
- pos = self._waiting_count
883
- eta = eta_fn(pos) if eta_fn is not None else None
884
- if eta is not None:
885
- eta_int = int(eta)
886
- eta_mins = max(0.5, round(eta / 30) * 0.5)
887
- progress(
888
- (0, eta_int),
889
- desc=f"⏳ In queue ~{eta_mins:g} minutes",
890
- )
891
- else:
892
- progress((0, 1), desc=f"⏳ In queue")
893
-
894
- with self._lock:
895
- worker = self._try_acquire(session_hash)
896
- if worker is not None:
897
- return worker
898
-
899
- remaining = deadline - time.monotonic()
900
- if remaining <= 0:
901
- raise RuntimeError(
902
- "No idle Eve workers available — timed out after "
903
- f"{timeout:.0f}s. Try again shortly."
904
- )
905
- self._lock.wait(timeout=min(remaining, 1.0))
906
- finally:
907
- with self._lock:
908
- self._waiting_count -= 1
909
-
910
- def release(self, worker: EveWorker) -> None:
911
- """Return a worker to the idle pool and wake any waiting acquirers."""
912
- with self._lock:
913
- worker.status = "idle"
914
- worker.busy_since = None
915
- worker.is_live_stream = False
916
- self._update_shared_load()
917
- self._lock.notify_all()
918
-
919
- def get_live_workers(self) -> list[EveWorker]:
920
- """Return a snapshot of non-dead workers (thread-safe)."""
921
- with self._lock:
922
- return [w for w in self._workers if w.status != "dead"]
923
-
924
- def shutdown_all(self) -> None:
925
- """Gracefully shut down all workers."""
926
- self._shutting_down = True
927
- with self._lock:
928
- for w in self._workers:
929
- if w.status != "dead":
930
- try:
931
- w.send_shutdown()
932
- except Exception:
933
- pass
934
- w.process.join(timeout=5)
935
- if w.process.is_alive():
936
- w.process.kill()
937
- w.process.join(timeout=2)
938
- w.status = "dead"
939
-
940
- def _update_shared_load(self) -> None:
941
- """Refresh the shared load signal (must hold ``self._lock``)."""
942
- total = sum(1 for w in self._workers if w.status != "dead")
943
- if total <= 1:
944
- self._shared_load.value = 0.0
945
- else:
946
- idle = sum(1 for w in self._workers if w.status == "idle")
947
- self._shared_load.value = 1.0 - idle / total
948
-
949
- # -- internals ---------------------------------------------------------
950
-
951
- def _compute_worker_count(self) -> int:
952
- """Determine how many workers fit in available RAM."""
953
- try:
954
- import psutil
955
- except ImportError:
956
- logger.warning("psutil not installed — defaulting to 1 worker")
957
- return 1
958
-
959
- raw_available_mb = psutil.virtual_memory().available / (1024 * 1024)
960
- cap_mb = self._cfg.max_ram_gb * 1024
961
- available_mb = min(raw_available_mb, cap_mb)
962
- headroom_mb = self._cfg.ram_headroom_gb * 1024
963
-
964
- # Spawn a probe worker to measure init RSS
965
-
966
- if DO_PROBE_WORKER:
967
- probe = self._spawn_worker(probe=True)
968
- try:
969
- import psutil as _ps
970
-
971
- probe_rss_mb = _ps.Process(probe.process.pid).memory_info().rss / (1024 * 1024)
972
- except Exception:
973
- logger.warning("Could not measure worker RSS — defaulting to 1 worker")
974
- return 1
975
- else:
976
- # Estimated 500MB RSS for an idle worker without running inference (based on observed values)
977
- # It's around 200MB idle, and there's the video being loaded in memory, depends on the size of the video.
978
- probe_rss_mb = 250
979
-
980
- estimated_peak_mb = probe_rss_mb * self._cfg.rss_safety_factor
981
- usable_mb = available_mb - headroom_mb
982
- # +1 because the probe worker already consumed memory
983
- max_by_ram = max(1, int(usable_mb / estimated_peak_mb) + 1)
984
- if DO_PROBE_WORKER:
985
- actual = min(self._cfg.max_workers - 1, max_by_ram)
986
- else:
987
- actual = min(self._cfg.max_workers, max_by_ram)
988
-
989
- logger.info(
990
- f"RAM estimation: init_rss={probe_rss_mb:.0f}MB, "
991
- f"estimated_peak={estimated_peak_mb:.0f}MB, "
992
- f"raw_available={raw_available_mb:.0f}MB, "
993
- f"available={available_mb:.0f}MB (cap={cap_mb:.0f}MB), "
994
- f"headroom={headroom_mb:.0f}MB, "
995
- f"max_by_ram={max_by_ram}, max_workers={self._cfg.max_workers}, "
996
- f"actual={actual}"
997
- )
998
- return actual
999
-
1000
- def _launch_worker(self) -> tuple[int, mp.Process, Connection]:
1001
- """Start a worker process without waiting for readiness.
1002
-
1003
- Returns:
1004
- Tuple of (worker_id, process, parent_conn).
1005
- """
1006
- wid = self._next_id
1007
- self._next_id += 1
1008
-
1009
- parent_conn, child_conn = _mp_ctx.Pipe()
1010
- proc = _mp_ctx.Process(
1011
- target=_eve_worker_main,
1012
- args=(
1013
- child_conn,
1014
- self._eve_bin_path,
1015
- self._eve_lib_path,
1016
- wid,
1017
- self._cfg.max_jobs_per_worker,
1018
- self._shared_load,
1019
- ),
1020
- daemon=True,
1021
- )
1022
- proc.start()
1023
- return wid, proc, parent_conn
1024
-
1025
- def _wait_for_ready(self, wid: int, proc: mp.Process, parent_conn: Connection) -> EveWorker:
1026
- """Block until a launched worker sends its "ready" message.
1027
-
1028
- Raises:
1029
- RuntimeError: If the worker doesn't become ready within 120 s.
1030
- """
1031
- if not parent_conn.poll(timeout=120):
1032
- proc.kill()
1033
- proc.join(timeout=5)
1034
- raise RuntimeError(f"Worker {wid} did not start within 120 s")
1035
-
1036
- msg: WorkerResponse = parent_conn.recv()
1037
- if not isinstance(msg, ReadyResponse):
1038
- proc.kill()
1039
- proc.join(timeout=5)
1040
- error = msg.error if isinstance(msg, ErrorResponse) else str(msg)
1041
- raise RuntimeError(f"Worker {wid} failed to initialise: {error}")
1042
-
1043
- worker = EveWorker(proc, parent_conn, wid)
1044
- with self._lock:
1045
- self._workers.append(worker)
1046
- self._lock.notify_all()
1047
- logger.info(f"Worker {wid} (pid={proc.pid}) ready")
1048
- return worker
1049
-
1050
- def _spawn_workers_parallel(self, count: int) -> None:
1051
- """Launch *count* workers in parallel and wait for all to be ready."""
1052
- pending = [self._launch_worker() for _ in range(count)]
1053
- for wid, proc, conn in pending:
1054
- try:
1055
- self._wait_for_ready(wid, proc, conn)
1056
- except RuntimeError:
1057
- logger.error(f"Worker {wid} failed to start skipping")
1058
-
1059
- def _spawn_worker(self, probe: bool = False) -> EveWorker:
1060
- """Spawn a single worker and wait for it to become ready.
1061
-
1062
- Used for the RAM probe and for respawning crashed workers.
1063
- """
1064
- wid, proc, conn = self._launch_worker()
1065
- return self._wait_for_ready(wid, proc, conn)
1066
-
1067
- def _respawn_worker(self, dead_worker: EveWorker) -> None:
1068
- """Replace a dead worker with a fresh one."""
1069
- if self._shutting_down:
1070
- return
1071
- logger.warning(
1072
- f"Respawning worker {dead_worker.worker_id} " f"(was pid={dead_worker.process.pid})"
1073
- )
1074
- dead_worker.process.join(timeout=0)
1075
- with self._lock:
1076
- self._workers.remove(dead_worker)
1077
- try:
1078
- self._spawn_worker()
1079
- except RuntimeError as exc:
1080
- logger.error(f"Failed to respawn worker: {exc}")
1081
-
1082
- def _health_monitor(self) -> None:
1083
- """Background daemon that detects dead / stuck workers."""
1084
- while not self._shutting_down:
1085
- time.sleep(self._cfg.health_check_interval_s)
1086
- if self._shutting_down:
1087
- break
1088
-
1089
- with self._lock:
1090
- workers_snapshot = list(self._workers)
1091
-
1092
- for w in workers_snapshot:
1093
- if w.status == "dead":
1094
- continue
1095
-
1096
- # Detect crashed / zombie processes
1097
- if not w.process.is_alive():
1098
- if w.pending_recycle:
1099
- logger.info(f"Worker {w.worker_id} (pid={w.process.pid}) recycled cleanly")
1100
- else:
1101
- logger.error(
1102
- f"Worker {w.worker_id} (pid={w.process.pid}) died unexpectedly"
1103
- )
1104
- w.status = "dead"
1105
- self._respawn_worker(w)
1106
- continue
1107
-
1108
- # Detect stuck workers (skip live streams — they're long by design)
1109
- if w.status == "busy" and not w.is_live_stream and w.busy_since is not None:
1110
- elapsed = time.monotonic() - w.busy_since
1111
- if elapsed > self._cfg.stuck_timeout_s * 2:
1112
- logger.error(f"Worker {w.worker_id} stuck for {elapsed:.0f}s — killing")
1113
- w.process.kill()
1114
- w.process.join(timeout=5)
1115
- w.status = "dead"
1116
- self._respawn_worker(w)
1117
- elif elapsed > self._cfg.stuck_timeout_s:
1118
- logger.warning(
1119
- f"Worker {w.worker_id} busy for {elapsed:.0f}s "
1120
- f"(threshold={self._cfg.stuck_timeout_s}s)"
1121
- )
 
1
+ """Multi-process Eve SDK worker pool.
2
+
3
+ Spawns N child processes, each hosting an independent ``EveWrapper`` instance
4
+ with its own ``ctypes.CDLL`` handle. The main Gradio process communicates with
5
+ workers via ``multiprocessing.Pipe`` (one bidirectional pipe per worker).
6
+
7
+ Frame data is serialised as raw bytes (``ndarray.tobytes()``) rather than
8
+ pickled ``ndarray`` objects for speed (~1-2 ms for a 1280x720x3 frame).
9
+
10
+ Usage::
11
+
12
+ pool = EveWorkerPool(max_workers=4, ram_headroom_gb=2.0)
13
+ worker = pool.acquire(session_hash="abc123")
14
+ try:
15
+ result = worker.send_inference(frame, features)
16
+ finally:
17
+ pool.release(worker)
18
+ """
19
+
20
+ import atexit
21
+ import logging
22
+ import multiprocessing as mp
23
+ import os
24
+ import threading
25
+ import time
26
+ from collections.abc import Callable
27
+ from dataclasses import dataclass
28
+ from multiprocessing.connection import Connection
29
+
30
+ import numpy as np
31
+ from eve_messages import (
32
+ CalibrateNewUserCmd,
33
+ CalibrateOkResponse,
34
+ CalibrationResultMsg,
35
+ EnableFaceIdCmd,
36
+ ErrorResponse,
37
+ FeatureFlags,
38
+ GalleryRestoredResponse,
39
+ GetProfileStatsCmd,
40
+ GetTimingStatsCmd,
41
+ HeartbeatResponse,
42
+ InferenceCmd,
43
+ InferenceResponse,
44
+ OkResponse,
45
+ ProcessVideoCmd,
46
+ ProfileStatsResponse,
47
+ ProgressResponse,
48
+ ReadyResponse,
49
+ RemoveAllUsersCmd,
50
+ RemoveUsersOkResponse,
51
+ RestoreGalleryCmd,
52
+ RestoreGalleryOkResponse,
53
+ SerializedFrame,
54
+ ShutdownCmd,
55
+ StartProfilingCmd,
56
+ StopProfilingCmd,
57
+ TimingStatsResponse,
58
+ VideoProcessingDoneResponse,
59
+ WorkerCmd,
60
+ WorkerResponse,
61
+ )
62
+ from log_utils import setup_logger
63
+
64
+ logger = setup_logger("EveWorkerPool")
65
+
66
+ # Cached per-process so the probe runs once per worker, not per video.
67
+ _h264_encoder_cache: tuple[str, dict[str, str]] | None = None
68
+
69
+
70
+ def _get_h264_encoder() -> tuple[str, dict[str, str]]:
71
+ """Return (codec_name, codec_options) for H.264 encoding.
72
+
73
+ Tries NVENC (GPU) first, falls back to libx264 (CPU).
74
+ """
75
+ global _h264_encoder_cache
76
+ if _h264_encoder_cache is not None:
77
+ return _h264_encoder_cache
78
+
79
+ import logging
80
+
81
+ import av
82
+
83
+ log = logging.getLogger("EveWorkerPool")
84
+ try:
85
+ ctx = av.CodecContext.create("h264_nvenc", "w")
86
+ ctx.width = 64
87
+ ctx.height = 64
88
+ ctx.pix_fmt = "yuv420p"
89
+ ctx.open()
90
+ ctx.close()
91
+ log.info("Using GPU encoder: h264_nvenc")
92
+ _h264_encoder_cache = (
93
+ "h264_nvenc",
94
+ {"profile": "baseline", "preset": "p1", "delay": "0"},
95
+ )
96
+ except Exception:
97
+ log.info("GPU encoder unavailable, falling back to libx264")
98
+ _h264_encoder_cache = (
99
+ "libx264",
100
+ {"profile": "baseline", "level": "3.1", "preset": "ultrafast", "threads": "1"},
101
+ )
102
+ return _h264_encoder_cache
103
+
104
+
105
+ def log_worker_activity(
106
+ log: logging.Logger,
107
+ action: str,
108
+ feature: str,
109
+ pool: "EveWorkerPool",
110
+ worker_id: int | None = None,
111
+ note: str = "",
112
+ ) -> None:
113
+ """Single-line worker activity log for tracking acquire/release/queue events."""
114
+ busy = pool.worker_count - pool.idle_count
115
+ w = f" w{worker_id}" if worker_id is not None else ""
116
+ n = f" ({note})" if note else ""
117
+ log.info(
118
+ f"[worker] {action}{w} for {feature}{n} "
119
+ f"[busy={busy} idle={pool.idle_count} wait={pool.waiting_count} "
120
+ f"sessions={pool.session_count}]"
121
+ )
122
+
123
+
124
+ def compute_fps_sleep(
125
+ elapsed: float,
126
+ max_fps: float | None,
127
+ min_fps: float | None,
128
+ load: float,
129
+ ) -> float:
130
+ """Compute how long to sleep after one frame to stay within the FPS cap.
131
+
132
+ Linearly interpolates the target cycle time between ``1/max_fps``
133
+ (at *load* 0) and ``1/min_fps`` (at *load* 1), then subtracts the
134
+ time already spent.
135
+
136
+ Used by both the live-inference bridge and the video-processing loop
137
+ so that the throttle logic lives in one place.
138
+ """
139
+ if max_fps is None:
140
+ return 0.0
141
+ effective_min = min_fps if min_fps is not None else max_fps
142
+ max_period = 1.0 / max_fps
143
+ min_period = 1.0 / effective_min
144
+ target_period = max_period + load * (min_period - max_period)
145
+ return max(0.0, target_period - elapsed)
146
+
147
+
148
+ # Use 'spawn' on all platforms so each child gets a fresh interpreter
149
+ # (required on Windows; avoids CDLL sharing on Linux/fork).
150
+ _mp_ctx = mp.get_context("spawn")
151
+
152
+ DO_PROBE_WORKER = (
153
+ False # Set to False to skip the RAM probe and just use an estimate of the RAM used
154
+ )
155
+
156
+
157
+ # ---------------------------------------------------------------------------
158
+ # Worker process entry point (runs in the child)
159
+ # ---------------------------------------------------------------------------
160
+
161
+
162
+ def _eve_worker_main(
163
+ conn: Connection,
164
+ eve_bin_path: str,
165
+ eve_lib_path: str,
166
+ worker_id: int,
167
+ max_jobs_per_worker: int | None = None,
168
+ shared_load: "mp.Value | None" = None,
169
+ ) -> None:
170
+ """Top-level function executed inside each worker ``Process``.
171
+
172
+ Communicates with the main process exclusively through *conn*.
173
+ *shared_load* is a cross-process float (0.0–1.0) updated by the pool
174
+ on acquire/release, used by the video processing loop to throttle.
175
+ """
176
+ import ctypes
177
+ import gc
178
+ import platform
179
+ import signal
180
+ import sys
181
+ from pathlib import Path
182
+
183
+ signal.signal(signal.SIGINT, signal.SIG_IGN)
184
+
185
+ # CPU affinity (for stress testing)
186
+ affinity_env = os.environ.get(f"WORKER{worker_id}_AFFINITY")
187
+ if affinity_env:
188
+ cpu_id = int(affinity_env.removeprefix("CPU"))
189
+ import psutil
190
+
191
+ psutil.Process().cpu_affinity([cpu_id])
192
+ logger.info(f"Worker {worker_id} pinned to CPU {cpu_id}")
193
+
194
+ # Ensure shared package is importable inside the child
195
+ shared_dir = str(Path(__file__).resolve().parent)
196
+ if shared_dir not in sys.path:
197
+ sys.path.insert(0, shared_dir)
198
+
199
+ from eve_wrapper import EveWrapper # noqa: E402 (deferred import)
200
+ from frame_utils import load_media_frames_raw # noqa: E402
201
+
202
+ eve: EveWrapper | None = None
203
+ try:
204
+ eve = EveWrapper(eve_bin_path, eve_lib_path)
205
+ conn.send(ReadyResponse(pid=os.getpid()))
206
+ except Exception as exc:
207
+ conn.send(ErrorResponse(error=str(exc)))
208
+ return
209
+
210
+ # Optional cProfile (enabled via ENABLE_PROFILER env var)
211
+ profiler: "cProfile.Profile | None" = None
212
+ profiling_enabled = os.environ.get("ENABLE_PROFILER", "").strip() not in ("", "0", "false")
213
+ if profiling_enabled:
214
+ import cProfile
215
+
216
+ # Use process_time so I/O wait (pipe polling) doesn't dominate the profile
217
+ profiler = cProfile.Profile(timer=time.process_time)
218
+
219
+ job_count = 0
220
+ while True:
221
+ try:
222
+ if not conn.poll(timeout=2.0):
223
+ # No command received — send heartbeat
224
+ try:
225
+ conn.send(HeartbeatResponse())
226
+ except (BrokenPipeError, OSError):
227
+ break
228
+ continue
229
+
230
+ cmd: WorkerCmd = conn.recv()
231
+ except (EOFError, OSError):
232
+ break
233
+
234
+ if isinstance(cmd, ShutdownCmd):
235
+ conn.send(OkResponse())
236
+ break
237
+
238
+ try:
239
+ if isinstance(cmd, InferenceCmd):
240
+ frame = np.frombuffer(cmd.frame_bytes, dtype=cmd.dtype).reshape(cmd.shape)
241
+ features = cmd.features
242
+ del cmd # free recv'd bytes early; numpy holds its own ref
243
+
244
+ eve.enable_face_and_person_detection(
245
+ faceEnabled=features.face_detection,
246
+ personEnabled=features.person_detection,
247
+ )
248
+ eve.enable_face_id(enabled=features.face_id)
249
+ eve.enable_hand_gesture(enabled=features.hand_gesture)
250
+ eve.enable_object_detection(config=features.mod_model_config)
251
+ eve.enable_mirror(True)
252
+ result = eve.inference(frame)
253
+ del frame
254
+
255
+ conn.send(
256
+ InferenceResponse(
257
+ frame_bytes=result.tobytes(),
258
+ shape=result.shape,
259
+ dtype=str(result.dtype),
260
+ )
261
+ )
262
+ del result
263
+
264
+ elif isinstance(cmd, CalibrateNewUserCmd):
265
+ frames = _deserialize_frames(cmd.frames_data)
266
+ result = eve.calibrate_new_user(frames)
267
+ conn.send(
268
+ CalibrateOkResponse(
269
+ result=CalibrationResultMsg(result.success, result.user_id, result.message),
270
+ )
271
+ )
272
+
273
+ elif isinstance(cmd, RemoveAllUsersCmd):
274
+ ok = eve.remove_all_users()
275
+ conn.send(RemoveUsersOkResponse(result=ok))
276
+
277
+ elif isinstance(cmd, RestoreGalleryCmd):
278
+ frames_per_user = [_deserialize_frames(fd) for fd in cmd.frames_per_user_data]
279
+ eve.remove_all_users()
280
+ results = eve.restore_gallery(frames_per_user)
281
+ conn.send(
282
+ RestoreGalleryOkResponse(
283
+ results=[
284
+ CalibrationResultMsg(r.success, r.user_id, r.message) for r in results
285
+ ],
286
+ )
287
+ )
288
+
289
+ elif isinstance(cmd, EnableFaceIdCmd):
290
+ eve.enable_face_id(enabled=cmd.enabled)
291
+ conn.send(OkResponse())
292
+
293
+ elif isinstance(cmd, ProcessVideoCmd):
294
+ import av
295
+ import cv2
296
+ from fractions import Fraction
297
+
298
+ features = cmd.features
299
+
300
+ # Gallery restore (if requested) must happen before
301
+ # applying user feature flags because restore_gallery()
302
+ # internally enables face/person detection for calibration.
303
+ gallery_results: list[CalibrationResultMsg] = []
304
+ if cmd.remove_all_users:
305
+ eve.remove_all_users()
306
+ if cmd.gallery_paths:
307
+ frames_per_user = [load_media_frames_raw(p) for p in cmd.gallery_paths]
308
+ raw_results = eve.restore_gallery(frames_per_user)
309
+ del frames_per_user
310
+ gallery_results = [
311
+ CalibrationResultMsg(r.success, r.user_id, r.message) for r in raw_results
312
+ ]
313
+ del raw_results
314
+ conn.send(GalleryRestoredResponse(results=gallery_results))
315
+
316
+ # Reset needs to be after the restore_gallery() function since we activate
317
+ # features of EVE inside, which can mess up the ideal user.
318
+ eve.reset_pipeline()
319
+
320
+ # Apply user's feature flags after gallery restore so they
321
+ # aren't overridden by restore_gallery's internal SDK calls.
322
+ eve.enable_face_and_person_detection(
323
+ faceEnabled=features.face_detection,
324
+ personEnabled=features.person_detection,
325
+ )
326
+ eve.enable_face_id(enabled=features.face_id)
327
+ eve.enable_hand_gesture(enabled=features.hand_gesture)
328
+ eve.enable_object_detection(config=features.mod_model_config)
329
+ eve.enable_mirror(False)
330
+
331
+ # Video processing loop
332
+ cap = cv2.VideoCapture(cmd.input_path)
333
+ codec_name, codec_opts = _get_h264_encoder()
334
+ container = av.open(cmd.output_path, mode="w", options={"movflags": "faststart"})
335
+ stream = container.add_stream(codec_name, rate=round(cmd.fps))
336
+ stream.time_base = Fraction(1, round(cmd.fps))
337
+ stream.width = cmd.width
338
+ stream.height = cmd.height
339
+ stream.pix_fmt = "yuv420p"
340
+ stream.codec_context.options = codec_opts
341
+
342
+ frame_count = 0
343
+ v_max_fps = cmd.max_fps
344
+ v_min_fps = cmd.min_fps if cmd.min_fps is not None else v_max_fps
345
+ try:
346
+ while True:
347
+ ret, frame = cap.read()
348
+ if not ret:
349
+ break
350
+ t0 = time.monotonic()
351
+ result = eve.inference(frame)
352
+ vf = av.VideoFrame.from_ndarray(result, format="bgr24")
353
+ vf.pts = frame_count
354
+ for pkt in stream.encode(vf):
355
+ container.mux(pkt)
356
+ del pkt
357
+ del result, vf, frame
358
+ frame_count += 1
359
+ if frame_count % 10 == 0:
360
+ conn.send(
361
+ ProgressResponse(
362
+ current=frame_count,
363
+ total=cmd.total_frames,
364
+ )
365
+ )
366
+ load = shared_load.value if shared_load is not None else 0.0
367
+ # FPS throttle (only if there's some load already)
368
+ if load > 0.0:
369
+ sleep = compute_fps_sleep(
370
+ time.monotonic() - t0, v_max_fps, v_min_fps, load
371
+ )
372
+ if sleep > 0:
373
+ time.sleep(sleep)
374
+ # Flush encoder
375
+ for pkt in stream.encode():
376
+ container.mux(pkt)
377
+ del pkt
378
+ finally:
379
+ cap.release()
380
+ container.close()
381
+ del (
382
+ cap,
383
+ container,
384
+ stream,
385
+ )
386
+
387
+ job_count += 1
388
+ conn.send(
389
+ VideoProcessingDoneResponse(
390
+ frames_processed=frame_count,
391
+ gallery_results=gallery_results,
392
+ recycle=max_jobs_per_worker is not None
393
+ and job_count >= max_jobs_per_worker,
394
+ )
395
+ )
396
+
397
+ del cmd
398
+ gc.collect()
399
+ if platform.system() == "Linux":
400
+ try:
401
+ ctypes.CDLL("libc.so.6").malloc_trim(0)
402
+ except Exception:
403
+ pass
404
+
405
+ if max_jobs_per_worker is not None and job_count >= max_jobs_per_worker:
406
+ logger.info(f"Worker {worker_id} recycling after {job_count} jobs")
407
+ break
408
+
409
+ elif isinstance(cmd, StartProfilingCmd):
410
+ if profiler is not None:
411
+ profiler.enable()
412
+ conn.send(OkResponse())
413
+
414
+ elif isinstance(cmd, StopProfilingCmd):
415
+ if profiler is not None:
416
+ profiler.disable()
417
+ conn.send(OkResponse())
418
+
419
+ elif isinstance(cmd, GetProfileStatsCmd):
420
+ if profiler is not None:
421
+ import io
422
+ import marshal
423
+ import pstats
424
+
425
+ profiler.disable()
426
+ stream = io.StringIO()
427
+ ps = pstats.Stats(profiler, stream=stream)
428
+ conn.send(ProfileStatsResponse(stats_data=marshal.dumps(ps.stats)))
429
+ else:
430
+ conn.send(ProfileStatsResponse(stats_data=b""))
431
+
432
+ elif isinstance(cmd, GetTimingStatsCmd):
433
+ conn.send(TimingStatsResponse(stats=eve.get_timing_stats(reset=cmd.reset)))
434
+
435
+ else:
436
+ conn.send(ErrorResponse(error=f"Unknown command: {cmd}"))
437
+
438
+ except Exception as exc:
439
+ logger.error(f"Worker {worker_id} error handling '{type(cmd).__name__}': {exc}")
440
+ try:
441
+ conn.send(ErrorResponse(error=str(exc)))
442
+ except (BrokenPipeError, OSError):
443
+ break
444
+
445
+ # Clean shutdown
446
+ if eve is not None:
447
+ eve.shutdown()
448
+
449
+
450
+ def _serialize_frames(frames: list[np.ndarray]) -> list[SerializedFrame]:
451
+ """Convert a list of ndarrays to a picklable representation."""
452
+ return [SerializedFrame(data=f.tobytes(), shape=f.shape, dtype=str(f.dtype)) for f in frames]
453
+
454
+
455
+ def _deserialize_frames(data: list[SerializedFrame]) -> list[np.ndarray]:
456
+ """Reconstruct ndarrays from their serialized form."""
457
+ return [np.frombuffer(d.data, dtype=d.dtype).reshape(d.shape) for d in data]
458
+
459
+
460
+ # ---------------------------------------------------------------------------
461
+ # EveWorker — main-process handle for one child process
462
+ # ---------------------------------------------------------------------------
463
+
464
+
465
+ class EveWorker:
466
+ """Main-process proxy for a single worker process.
467
+
468
+ Not thread-safe by itself the ``EveWorkerPool`` ensures only one thread
469
+ accesses a worker at a time.
470
+ """
471
+
472
+ def __init__(
473
+ self,
474
+ process: mp.Process,
475
+ conn: Connection,
476
+ worker_id: int,
477
+ ):
478
+ self.process = process
479
+ self._conn = conn
480
+ self.worker_id = worker_id
481
+ self.status: str = "idle" # idle | busy | dead
482
+ self.session_hash: str | None = None
483
+ self.busy_since: float | None = None
484
+ self.last_heartbeat: float = time.monotonic()
485
+ self.is_live_stream: bool = False
486
+ self.pending_recycle: bool = False
487
+
488
+ # -- command helpers ---------------------------------------------------
489
+
490
+ def _recv_response(self) -> WorkerResponse:
491
+ """Receive the next non-heartbeat response from the worker."""
492
+ while True:
493
+ try:
494
+ resp: WorkerResponse = self._conn.recv()
495
+ except (EOFError, OSError):
496
+ raise RuntimeError(f"Worker {self.worker_id} pipe closed")
497
+ except Exception as exc:
498
+ self.pending_recycle = True
499
+ raise RuntimeError(
500
+ f"Worker {self.worker_id} pipe corrupt: {exc}"
501
+ ) from exc
502
+ if isinstance(resp, HeartbeatResponse):
503
+ self.last_heartbeat = time.monotonic()
504
+ continue
505
+ return resp
506
+
507
+ def _expect_ok(self, context: str) -> WorkerResponse:
508
+ """Receive a response and raise on error."""
509
+ resp = self._recv_response()
510
+ if isinstance(resp, ErrorResponse):
511
+ raise RuntimeError(f"Worker {self.worker_id} {context} error: {resp.error}")
512
+ return resp
513
+
514
+ def send_inference(self, frame: np.ndarray, features: FeatureFlags) -> np.ndarray:
515
+ """Send a frame for inference and return the annotated result."""
516
+ self._conn.send(
517
+ InferenceCmd(
518
+ frame_bytes=frame.tobytes(),
519
+ shape=frame.shape,
520
+ dtype=str(frame.dtype),
521
+ features=features,
522
+ )
523
+ )
524
+ resp = self._expect_ok("inference")
525
+ if not isinstance(resp, InferenceResponse):
526
+ raise RuntimeError(f"Worker {self.worker_id} inference: unexpected {resp}")
527
+ return np.frombuffer(resp.frame_bytes, dtype=resp.dtype).reshape(resp.shape)
528
+
529
+
530
+ def send_calibrate_new_user(self, frames: list[np.ndarray]) -> CalibrationResultMsg:
531
+ self._conn.send(CalibrateNewUserCmd(frames_data=_serialize_frames(frames)))
532
+ resp = self._expect_ok("calibrate")
533
+ if not isinstance(resp, CalibrateOkResponse):
534
+ raise RuntimeError(f"Worker {self.worker_id} calibrate: unexpected {resp}")
535
+ return resp.result
536
+
537
+ def send_remove_all_users(self) -> bool:
538
+ self._conn.send(RemoveAllUsersCmd())
539
+ resp = self._expect_ok("remove")
540
+ if not isinstance(resp, RemoveUsersOkResponse):
541
+ raise RuntimeError(f"Worker {self.worker_id} remove: unexpected {resp}")
542
+ return resp.result
543
+
544
+ def send_restore_gallery(
545
+ self, frames_per_user: list[list[np.ndarray]]
546
+ ) -> list[CalibrationResultMsg]:
547
+ self._conn.send(
548
+ RestoreGalleryCmd(
549
+ frames_per_user_data=[_serialize_frames(fu) for fu in frames_per_user],
550
+ )
551
+ )
552
+ resp = self._expect_ok("restore")
553
+ if not isinstance(resp, RestoreGalleryOkResponse):
554
+ raise RuntimeError(f"Worker {self.worker_id} restore: unexpected {resp}")
555
+ return resp.results
556
+
557
+ def send_enable_face_id(self, enabled: bool) -> None:
558
+ self._conn.send(EnableFaceIdCmd(enabled=enabled))
559
+ self._expect_ok("enable_face_id")
560
+
561
+ def send_process_video(
562
+ self,
563
+ input_path: str,
564
+ output_path: str,
565
+ features: FeatureFlags,
566
+ gallery_paths: list[str],
567
+ remove_all_users: bool,
568
+ fps: float,
569
+ width: int,
570
+ height: int,
571
+ total_frames: int,
572
+ progress: Callable[..., None] | None = None,
573
+ max_fps: float | None = None,
574
+ min_fps: float | None = None,
575
+ ) -> tuple[int, list[CalibrationResultMsg]]:
576
+ """Run the full video processing loop inside the worker process.
577
+
578
+ The worker reads the input video, runs Eve inference on each frame,
579
+ encodes the result with PyAV, and writes the output video. Only
580
+ small progress messages travel over the pipe no frame data.
581
+
582
+ Args:
583
+ input_path: Path to the input video file.
584
+ output_path: Path where the output video will be written.
585
+ features: Feature toggles for the Eve SDK.
586
+ gallery_paths: Media paths for Face ID gallery restore (may be empty).
587
+ remove_all_users: Whether to clear the Face ID gallery first.
588
+ fps: Output video frame rate.
589
+ width: Output video width in pixels.
590
+ height: Output video height in pixels.
591
+ total_frames: Total frame count (for progress reporting).
592
+ progress: Optional ``gr.Progress``-compatible callback.
593
+ max_fps: FPS cap when idle (``None`` = unlimited).
594
+ min_fps: FPS cap under full load (defaults to ``max_fps``).
595
+
596
+ Returns:
597
+ Tuple of (frames_processed, gallery_results).
598
+
599
+ Raises:
600
+ RuntimeError: If the worker reports an error.
601
+ """
602
+ self._conn.send(
603
+ ProcessVideoCmd(
604
+ input_path=input_path,
605
+ output_path=output_path,
606
+ features=features,
607
+ gallery_paths=gallery_paths,
608
+ remove_all_users=remove_all_users,
609
+ fps=fps,
610
+ width=width,
611
+ height=height,
612
+ total_frames=total_frames,
613
+ max_fps=max_fps,
614
+ min_fps=min_fps,
615
+ )
616
+ )
617
+
618
+ gallery_results: list[CalibrationResultMsg] = []
619
+
620
+ while True:
621
+ resp = self._recv_response()
622
+
623
+ if isinstance(resp, GalleryRestoredResponse):
624
+ gallery_results = resp.results
625
+
626
+ elif isinstance(resp, ProgressResponse):
627
+ if progress is not None:
628
+ progress(
629
+ (resp.current, resp.total),
630
+ desc=f"Processing frame {resp.current}/{resp.total}",
631
+ )
632
+
633
+ elif isinstance(resp, VideoProcessingDoneResponse):
634
+ if resp.recycle:
635
+ self.pending_recycle = True
636
+ return resp.frames_processed, gallery_results
637
+
638
+ elif isinstance(resp, ErrorResponse):
639
+ raise RuntimeError(f"Worker {self.worker_id} process_video error: {resp.error}")
640
+
641
+ else:
642
+ raise RuntimeError(f"Worker {self.worker_id} unexpected response: {resp}")
643
+
644
+ def send_start_profiling(self) -> None:
645
+ self._conn.send(StartProfilingCmd())
646
+ self._expect_ok("start_profiling")
647
+
648
+ def send_stop_profiling(self) -> None:
649
+ self._conn.send(StopProfilingCmd())
650
+ self._expect_ok("stop_profiling")
651
+
652
+ def send_get_profile_stats(self) -> bytes:
653
+ """Request profiling stats from the worker. Returns marshalled pstats data."""
654
+ self._conn.send(GetProfileStatsCmd())
655
+ resp = self._recv_response()
656
+ if isinstance(resp, ErrorResponse):
657
+ raise RuntimeError(f"Worker {self.worker_id} get_profile_stats error: {resp.error}")
658
+ if not isinstance(resp, ProfileStatsResponse):
659
+ raise RuntimeError(f"Worker {self.worker_id} get_profile_stats: unexpected {resp}")
660
+ return resp.stats_data
661
+
662
+ def send_get_timing_stats(self, reset: bool = True) -> dict[str, tuple[int, float]]:
663
+ """Request per-SDK-call timing stats from the worker."""
664
+ self._conn.send(GetTimingStatsCmd(reset=reset))
665
+ resp = self._recv_response()
666
+ if isinstance(resp, ErrorResponse):
667
+ raise RuntimeError(f"Worker {self.worker_id} get_timing_stats error: {resp.error}")
668
+ if not isinstance(resp, TimingStatsResponse):
669
+ raise RuntimeError(f"Worker {self.worker_id} get_timing_stats: unexpected {resp}")
670
+ return resp.stats
671
+
672
+ def send_shutdown(self) -> None:
673
+ try:
674
+ self._conn.send(ShutdownCmd())
675
+ self._conn.recv()
676
+ except (EOFError, OSError, BrokenPipeError):
677
+ pass
678
+
679
+ def drain_heartbeats(self) -> None:
680
+ """Consume any pending heartbeat messages on the pipe."""
681
+ while self._conn.poll(timeout=0):
682
+ try:
683
+ msg = self._conn.recv()
684
+ if isinstance(msg, HeartbeatResponse):
685
+ self.last_heartbeat = time.monotonic()
686
+ except (EOFError, OSError):
687
+ break
688
+ except Exception:
689
+ # Pipe has non-pickle data (e.g. native library output) — the
690
+ # framing is now unreliable so schedule this worker for recycling.
691
+ logger.warning(
692
+ "Worker %d: corrupt data on pipe, scheduling recycle",
693
+ self.worker_id,
694
+ )
695
+ self.pending_recycle = True
696
+ break
697
+
698
+
699
+ # ---------------------------------------------------------------------------
700
+ # EveWorkerPool — manages all workers
701
+ # ---------------------------------------------------------------------------
702
+
703
+
704
+ @dataclass
705
+ class _PoolConfig:
706
+ max_workers: int = 8
707
+ max_ram_gb: float = 32.0
708
+ ram_headroom_gb: float = 2.0
709
+ rss_safety_factor: float = 2.5
710
+ stuck_timeout_s: float = 300.0 # 5 min for video processing
711
+ health_check_interval_s: float = 5.0
712
+ max_jobs_per_worker: int | None = 50
713
+
714
+
715
+ class EveWorkerPool:
716
+ """Pool of Eve SDK worker processes.
717
+
718
+ Measures available RAM at startup, spawns as many workers as fit
719
+ (with safety margins), and provides acquire / release semantics.
720
+
721
+ Args:
722
+ max_workers: Upper bound on the number of workers.
723
+ max_ram_gb: Cap on available RAM (GB) considered for worker sizing.
724
+ Useful on shared machines where reported available RAM is
725
+ much larger than what the demo should consume.
726
+ ram_headroom_gb: Free RAM (GB) to reserve for main process / OS.
727
+ rss_safety_factor: Multiplier applied to the measured init RSS to
728
+ estimate peak runtime memory per worker.
729
+ eve_bin_path: Forwarded to ``EveWrapper``.
730
+ eve_lib_path: Forwarded to ``EveWrapper``.
731
+ """
732
+
733
+ def __init__(
734
+ self,
735
+ max_workers: int = 8,
736
+ max_ram_gb: float = 32.0,
737
+ ram_headroom_gb: float = 2.0,
738
+ rss_safety_factor: float = 2.5,
739
+ eve_bin_path: str = "",
740
+ eve_lib_path: str = "",
741
+ max_jobs_per_worker: int | None = 50,
742
+ ):
743
+ self._cfg = _PoolConfig(
744
+ max_workers=max_workers,
745
+ max_ram_gb=max_ram_gb,
746
+ ram_headroom_gb=ram_headroom_gb,
747
+ rss_safety_factor=rss_safety_factor,
748
+ max_jobs_per_worker=max_jobs_per_worker,
749
+ )
750
+ self._eve_bin_path = eve_bin_path
751
+ self._eve_lib_path = eve_lib_path
752
+ self._lock = threading.Condition(threading.Lock())
753
+ self._workers: list[EveWorker] = []
754
+ self._shutting_down = False
755
+ self._next_id = 0
756
+ self._waiting_count = 0
757
+ self.session_count = 0
758
+ # Shared load signal readable by worker processes (0.0–1.0).
759
+ # lock=False is safe: single writer (main process), readers get
760
+ # a slightly stale value at worst.
761
+ self._shared_load: mp.Value = _mp_ctx.Value("d", 0.0, lock=False)
762
+
763
+ # Spawn workers based on RAM capacity
764
+ count = self._compute_worker_count()
765
+ self._spawn_workers_parallel(count)
766
+
767
+ logger.info(f"EveWorkerPool started with {len(self._workers)} workers")
768
+
769
+ # Start health-check daemon
770
+ self._health_thread = threading.Thread(target=self._health_monitor, daemon=True)
771
+ self._health_thread.start()
772
+
773
+ # Hourly memory report
774
+ from memory_monitor import start_memory_reporter
775
+
776
+ self._mem_thread = start_memory_reporter(
777
+ get_workers=lambda: list(self._workers),
778
+ lock=self._lock,
779
+ shutdown_flag=lambda: self._shutting_down,
780
+ max_ram_gb=self._cfg.max_ram_gb,
781
+ )
782
+
783
+ atexit.register(self.shutdown_all)
784
+
785
+ # -- public API --------------------------------------------------------
786
+
787
+ @property
788
+ def worker_count(self) -> int:
789
+ """Number of live workers (for setting ``concurrency_limit``)."""
790
+ with self._lock:
791
+ return sum(1 for w in self._workers if w.status != "dead")
792
+
793
+ @property
794
+ def idle_count(self) -> int:
795
+ """Number of idle workers."""
796
+ with self._lock:
797
+ return sum(1 for w in self._workers if w.status == "idle")
798
+
799
+ @property
800
+ def waiting_count(self) -> int:
801
+ """Number of callers blocked in ``acquire()`` waiting for a worker."""
802
+ with self._lock:
803
+ return self._waiting_count
804
+
805
+ def try_acquire(self, session_hash: str) -> EveWorker | None:
806
+ """Non-blocking acquire — returns an idle worker or ``None``.
807
+
808
+ Unlike :meth:`acquire`, this never blocks or increments the
809
+ waiting count. Used by :class:`LiveStreamManager` so that
810
+ WebRTC frame handlers can return immediately when no worker
811
+ is available.
812
+ """
813
+ with self._lock:
814
+ return self._try_acquire(session_hash)
815
+
816
+ def _try_acquire(self, session_hash: str) -> EveWorker | None:
817
+ """Try to grab an idle worker (must hold ``self._lock``).
818
+
819
+ Returns the worker if successful, ``None`` otherwise.
820
+ """
821
+ # Prefer session-affiliated idle worker (Face ID gallery affinity)
822
+ for w in self._workers:
823
+ if w.status == "idle" and not w.pending_recycle and w.session_hash == session_hash:
824
+ w.status = "busy"
825
+ w.busy_since = time.monotonic()
826
+ w.drain_heartbeats()
827
+ self._update_shared_load()
828
+ return w
829
+ # Any idle worker
830
+ for w in self._workers:
831
+ if w.status == "idle" and not w.pending_recycle:
832
+ w.status = "busy"
833
+ w.session_hash = session_hash
834
+ w.busy_since = time.monotonic()
835
+ w.drain_heartbeats()
836
+ self._update_shared_load()
837
+ return w
838
+ return None
839
+
840
+ def acquire(
841
+ self,
842
+ session_hash: str,
843
+ timeout: float = 300.0,
844
+ progress: Callable[..., None] | None = None,
845
+ eta_fn: Callable[[int], float | None] | None = None,
846
+ ) -> EveWorker:
847
+ """Reserve a worker for *session_hash*, blocking until one is free.
848
+
849
+ Prefers a worker already affiliated with the session (Face ID
850
+ gallery affinity). Falls back to any idle worker. Blocks up to
851
+ *timeout* seconds if all workers are busy.
852
+
853
+ Args:
854
+ session_hash: Gradio session hash for worker affinity.
855
+ timeout: Max seconds to wait for a free worker.
856
+ progress: Optional ``gr.Progress``-compatible callback invoked
857
+ while waiting so the UI can show queue position.
858
+ eta_fn: Optional callable accepting a 1-based queue position
859
+ and returning estimated seconds until a worker frees up.
860
+ The return value is shown in the progress description.
861
+
862
+ Raises:
863
+ RuntimeError: If no worker becomes available within *timeout*.
864
+ """
865
+ deadline = time.monotonic() + timeout
866
+
867
+ # Fast path try without incrementing waiting count
868
+ with self._lock:
869
+ worker = self._try_acquire(session_hash)
870
+ if worker is not None:
871
+ return worker
872
+ self._waiting_count += 1
873
+
874
+ # Slow path — block until a worker is released
875
+ try:
876
+ while True:
877
+ # Send progress update OUTSIDE the lock to avoid blocking release()
878
+ if progress is not None:
879
+ with self._lock:
880
+ pos = self._waiting_count
881
+ eta = eta_fn(pos) if eta_fn is not None else None
882
+ if eta is not None:
883
+ eta_int = int(eta)
884
+ eta_mins = max(0.5, round(eta / 30) * 0.5)
885
+ progress(
886
+ (0, eta_int),
887
+ desc=f"⏳ In queue ~{eta_mins:g} minutes",
888
+ )
889
+ else:
890
+ progress((0, 1), desc=f"⏳ In queue")
891
+
892
+ with self._lock:
893
+ worker = self._try_acquire(session_hash)
894
+ if worker is not None:
895
+ return worker
896
+
897
+ remaining = deadline - time.monotonic()
898
+ if remaining <= 0:
899
+ raise RuntimeError(
900
+ "No idle Eve workers available — timed out after "
901
+ f"{timeout:.0f}s. Try again shortly."
902
+ )
903
+ self._lock.wait(timeout=min(remaining, 1.0))
904
+ finally:
905
+ with self._lock:
906
+ self._waiting_count -= 1
907
+
908
+ def release(self, worker: EveWorker) -> None:
909
+ """Return a worker to the idle pool and wake any waiting acquirers."""
910
+ with self._lock:
911
+ worker.status = "idle"
912
+ worker.busy_since = None
913
+ worker.is_live_stream = False
914
+ self._update_shared_load()
915
+ self._lock.notify_all()
916
+
917
+ def get_live_workers(self) -> list[EveWorker]:
918
+ """Return a snapshot of non-dead workers (thread-safe)."""
919
+ with self._lock:
920
+ return [w for w in self._workers if w.status != "dead"]
921
+
922
+ def shutdown_all(self) -> None:
923
+ """Gracefully shut down all workers."""
924
+ self._shutting_down = True
925
+ with self._lock:
926
+ for w in self._workers:
927
+ if w.status != "dead":
928
+ try:
929
+ w.send_shutdown()
930
+ except Exception:
931
+ pass
932
+ w.process.join(timeout=5)
933
+ if w.process.is_alive():
934
+ w.process.kill()
935
+ w.process.join(timeout=2)
936
+ w.status = "dead"
937
+
938
+ def _update_shared_load(self) -> None:
939
+ """Refresh the shared load signal (must hold ``self._lock``)."""
940
+ total = sum(1 for w in self._workers if w.status != "dead")
941
+ if total <= 1:
942
+ self._shared_load.value = 0.0
943
+ else:
944
+ idle = sum(1 for w in self._workers if w.status == "idle")
945
+ self._shared_load.value = 1.0 - idle / total
946
+
947
+ # -- internals ---------------------------------------------------------
948
+
949
+ def _compute_worker_count(self) -> int:
950
+ """Determine how many workers fit in available RAM."""
951
+ try:
952
+ import psutil
953
+ except ImportError:
954
+ logger.warning("psutil not installed — defaulting to 1 worker")
955
+ return 1
956
+
957
+ raw_available_mb = psutil.virtual_memory().available / (1024 * 1024)
958
+ cap_mb = self._cfg.max_ram_gb * 1024
959
+ available_mb = min(raw_available_mb, cap_mb)
960
+ headroom_mb = self._cfg.ram_headroom_gb * 1024
961
+
962
+ # Spawn a probe worker to measure init RSS
963
+
964
+ if DO_PROBE_WORKER:
965
+ probe = self._spawn_worker(probe=True)
966
+ try:
967
+ import psutil as _ps
968
+
969
+ probe_rss_mb = _ps.Process(probe.process.pid).memory_info().rss / (1024 * 1024)
970
+ except Exception:
971
+ logger.warning("Could not measure worker RSS — defaulting to 1 worker")
972
+ return 1
973
+ else:
974
+ # Estimated 500MB RSS for an idle worker without running inference (based on observed values)
975
+ # It's around 200MB idle, and there's the video being loaded in memory, depends on the size of the video.
976
+ probe_rss_mb = 250
977
+
978
+ estimated_peak_mb = probe_rss_mb * self._cfg.rss_safety_factor
979
+ usable_mb = available_mb - headroom_mb
980
+ # +1 because the probe worker already consumed memory
981
+ max_by_ram = max(1, int(usable_mb / estimated_peak_mb) + 1)
982
+ if DO_PROBE_WORKER:
983
+ actual = min(self._cfg.max_workers - 1, max_by_ram)
984
+ else:
985
+ actual = min(self._cfg.max_workers, max_by_ram)
986
+
987
+ logger.info(
988
+ f"RAM estimation: init_rss={probe_rss_mb:.0f}MB, "
989
+ f"estimated_peak={estimated_peak_mb:.0f}MB, "
990
+ f"raw_available={raw_available_mb:.0f}MB, "
991
+ f"available={available_mb:.0f}MB (cap={cap_mb:.0f}MB), "
992
+ f"headroom={headroom_mb:.0f}MB, "
993
+ f"max_by_ram={max_by_ram}, max_workers={self._cfg.max_workers}, "
994
+ f"actual={actual}"
995
+ )
996
+ return actual
997
+
998
+ def _launch_worker(self) -> tuple[int, mp.Process, Connection]:
999
+ """Start a worker process without waiting for readiness.
1000
+
1001
+ Returns:
1002
+ Tuple of (worker_id, process, parent_conn).
1003
+ """
1004
+ wid = self._next_id
1005
+ self._next_id += 1
1006
+
1007
+ parent_conn, child_conn = _mp_ctx.Pipe()
1008
+ proc = _mp_ctx.Process(
1009
+ target=_eve_worker_main,
1010
+ args=(
1011
+ child_conn,
1012
+ self._eve_bin_path,
1013
+ self._eve_lib_path,
1014
+ wid,
1015
+ self._cfg.max_jobs_per_worker,
1016
+ self._shared_load,
1017
+ ),
1018
+ daemon=True,
1019
+ )
1020
+ proc.start()
1021
+ return wid, proc, parent_conn
1022
+
1023
+ def _wait_for_ready(self, wid: int, proc: mp.Process, parent_conn: Connection) -> EveWorker:
1024
+ """Block until a launched worker sends its "ready" message.
1025
+
1026
+ Raises:
1027
+ RuntimeError: If the worker doesn't become ready within 120 s.
1028
+ """
1029
+ if not parent_conn.poll(timeout=120):
1030
+ proc.kill()
1031
+ proc.join(timeout=5)
1032
+ raise RuntimeError(f"Worker {wid} did not start within 120 s")
1033
+
1034
+ msg: WorkerResponse = parent_conn.recv()
1035
+ if not isinstance(msg, ReadyResponse):
1036
+ proc.kill()
1037
+ proc.join(timeout=5)
1038
+ error = msg.error if isinstance(msg, ErrorResponse) else str(msg)
1039
+ raise RuntimeError(f"Worker {wid} failed to initialise: {error}")
1040
+
1041
+ worker = EveWorker(proc, parent_conn, wid)
1042
+ with self._lock:
1043
+ self._workers.append(worker)
1044
+ self._lock.notify_all()
1045
+ logger.info(f"Worker {wid} (pid={proc.pid}) ready")
1046
+ return worker
1047
+
1048
+ def _spawn_workers_parallel(self, count: int) -> None:
1049
+ """Launch *count* workers in parallel and wait for all to be ready."""
1050
+ pending = [self._launch_worker() for _ in range(count)]
1051
+ for wid, proc, conn in pending:
1052
+ try:
1053
+ self._wait_for_ready(wid, proc, conn)
1054
+ except RuntimeError:
1055
+ logger.error(f"Worker {wid} failed to start — skipping")
1056
+
1057
+ def _spawn_worker(self, probe: bool = False) -> EveWorker:
1058
+ """Spawn a single worker and wait for it to become ready.
1059
+
1060
+ Used for the RAM probe and for respawning crashed workers.
1061
+ """
1062
+ wid, proc, conn = self._launch_worker()
1063
+ return self._wait_for_ready(wid, proc, conn)
1064
+
1065
+ def _respawn_worker(self, dead_worker: EveWorker) -> None:
1066
+ """Replace a dead worker with a fresh one."""
1067
+ if self._shutting_down:
1068
+ return
1069
+ logger.warning(
1070
+ f"Respawning worker {dead_worker.worker_id} " f"(was pid={dead_worker.process.pid})"
1071
+ )
1072
+ dead_worker.process.join(timeout=0)
1073
+ with self._lock:
1074
+ self._workers.remove(dead_worker)
1075
+ try:
1076
+ self._spawn_worker()
1077
+ except RuntimeError as exc:
1078
+ logger.error(f"Failed to respawn worker: {exc}")
1079
+
1080
+ def _health_monitor(self) -> None:
1081
+ """Background daemon that detects dead / stuck workers."""
1082
+ while not self._shutting_down:
1083
+ time.sleep(self._cfg.health_check_interval_s)
1084
+ if self._shutting_down:
1085
+ break
1086
+
1087
+ with self._lock:
1088
+ workers_snapshot = list(self._workers)
1089
+
1090
+ for w in workers_snapshot:
1091
+ if w.status == "dead":
1092
+ continue
1093
+
1094
+ # Detect crashed / zombie processes
1095
+ if not w.process.is_alive():
1096
+ if w.pending_recycle:
1097
+ logger.info(f"Worker {w.worker_id} (pid={w.process.pid}) recycled cleanly")
1098
+ else:
1099
+ logger.error(
1100
+ f"Worker {w.worker_id} (pid={w.process.pid}) died unexpectedly"
1101
+ )
1102
+ w.status = "dead"
1103
+ self._respawn_worker(w)
1104
+ continue
1105
+
1106
+ # Detect stuck workers (skip live streams — they're long by design)
1107
+ if w.status == "busy" and not w.is_live_stream and w.busy_since is not None:
1108
+ elapsed = time.monotonic() - w.busy_since
1109
+ if elapsed > self._cfg.stuck_timeout_s * 2:
1110
+ logger.error(f"Worker {w.worker_id} stuck for {elapsed:.0f}s — killing")
1111
+ w.process.kill()
1112
+ w.process.join(timeout=5)
1113
+ w.status = "dead"
1114
+ self._respawn_worker(w)
1115
+ elif elapsed > self._cfg.stuck_timeout_s:
1116
+ logger.warning(
1117
+ f"Worker {w.worker_id} busy for {elapsed:.0f}s "
1118
+ f"(threshold={self._cfg.stuck_timeout_s}s)"
1119
+ )