disssid commited on
Commit
b77f959
·
1 Parent(s): f8e9264

Reactive mesh · Phase M3: mesh Modal endpoint + app wiring (flag-gated)

Browse files

modal_mesh_endpoint.py (new): LLaMA-Mesh on Modal (trikona-mesh app, scale-to-
zero), returns raw mesh text; geometry runs in the backend, not here. DEPLOY
FLAGGED — never run by the loop. app.py: _mesh_emit POSTs to MESH_ENDPOINT_URL
(server-side, hidden) and _codegen_generate injects mesh_fn into pipeline.build
only when TRIKONA_MESH + URL are set. Offline tests via canned _model_emit/
_mesh_emit: routes to a valid mesh contract when on; unchanged (guided, mesh
never called) when off; endpoint module parses. Suite green.

Files changed (3) hide show
  1. app.py +39 -2
  2. modal_mesh_endpoint.py +73 -0
  3. tests/test_mesh_integration.py +71 -0
app.py CHANGED
@@ -40,6 +40,12 @@ MODAL_TIMEOUT = int(os.environ.get("MODAL_TIMEOUT", "150"))
40
  CODEGEN = bool(os.environ.get("TRIKONA_CODEGEN"))
41
  _SPEC_CACHE: dict[str, dict] = {} # prompt -> last build_spec (for model-free regen)
42
 
 
 
 
 
 
 
43
  EXAMPLE_PROMPTS = [
44
  "dodecahedron",
45
  "octahedron",
@@ -99,6 +105,36 @@ def _model_emit(prompt: str, params: dict | None, feedback: str | None = None) -
99
  return obj if isinstance(obj, dict) else {"error": "model_invalid_json"}
100
 
101
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
102
  def _regen_modelfree(spec: dict, params: dict | None) -> str | None:
103
  """Re-run a cached generator with new params — no model call. None on failure."""
104
  from geometry import assembler, metrics, primitives
@@ -123,9 +159,10 @@ def _regen_modelfree(spec: dict, params: dict | None) -> str | None:
123
 
124
  def _codegen_generate(prompt: str, params: dict | None = None,
125
  cache_key: str | None = None) -> str:
126
- """Full codegen path via the pipeline (gate -> retry -> primitive fallback)."""
127
  from geometry import pipeline
128
- result, meta = pipeline.build(_model_emit, prompt, params=params)
 
129
  if cache_key is not None and meta.get("spec") is not None:
130
  _SPEC_CACHE[cache_key] = meta["spec"]
131
  return json.dumps(result)
 
40
  CODEGEN = bool(os.environ.get("TRIKONA_CODEGEN"))
41
  _SPEC_CACHE: dict[str, dict] = {} # prompt -> last build_spec (for model-free regen)
42
 
43
+ # Reactive mesh fallback (SPEC_REACTIVE_MESH): when codegen can't construct a
44
+ # shape, route to LLaMA-Mesh. Opt-in via TRIKONA_MESH + MESH_ENDPOINT_URL; the
45
+ # mesh endpoint URL stays server-side, never exposed to the browser.
46
+ MESH_URL = os.environ.get("MESH_ENDPOINT_URL")
47
+ MESH_ON = bool(os.environ.get("TRIKONA_MESH"))
48
+
49
  EXAMPLE_PROMPTS = [
50
  "dodecahedron",
51
  "octahedron",
 
105
  return obj if isinstance(obj, dict) else {"error": "model_invalid_json"}
106
 
107
 
108
+ def _mesh_emit(prompt: str) -> str | None:
109
+ """Reactive mesh tier: POST to the LLaMA-Mesh endpoint, return raw mesh text.
110
+
111
+ Returns None on any failure (the pipeline then falls through to guided). Keeps
112
+ the mesh endpoint URL server-side, like the codegen endpoint.
113
+ """
114
+ if not MESH_URL or requests is None:
115
+ return None
116
+ deadline = time.monotonic() + MODAL_TIMEOUT
117
+ try:
118
+ while True:
119
+ resp = requests.post(MESH_URL, json={"prompt": prompt},
120
+ timeout=MODAL_TIMEOUT, allow_redirects=False)
121
+ if resp.is_redirect or resp.status_code in (301, 302, 303, 307, 308):
122
+ if time.monotonic() >= deadline:
123
+ return None
124
+ time.sleep(2)
125
+ continue
126
+ data = resp.content.decode("utf-8", "replace")
127
+ try:
128
+ obj = json.loads(data)
129
+ if isinstance(obj, dict) and isinstance(obj.get("text"), str):
130
+ return obj["text"]
131
+ except (json.JSONDecodeError, TypeError):
132
+ pass
133
+ return data # endpoint returned raw text rather than {"text": ...}
134
+ except Exception:
135
+ return None
136
+
137
+
138
  def _regen_modelfree(spec: dict, params: dict | None) -> str | None:
139
  """Re-run a cached generator with new params — no model call. None on failure."""
140
  from geometry import assembler, metrics, primitives
 
159
 
160
  def _codegen_generate(prompt: str, params: dict | None = None,
161
  cache_key: str | None = None) -> str:
162
+ """Full codegen path via the pipeline (gate -> retry -> primitive -> mesh)."""
163
  from geometry import pipeline
164
+ mesh_fn = _mesh_emit if (MESH_ON and MESH_URL) else None
165
+ result, meta = pipeline.build(_model_emit, prompt, params=params, mesh_fn=mesh_fn)
166
  if cache_key is not None and meta.get("spec") is not None:
167
  _SPEC_CACHE[cache_key] = meta["spec"]
168
  return json.dumps(result)
modal_mesh_endpoint.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Trikona mesh endpoint — LLaMA-Mesh on Modal (SPEC_REACTIVE_MESH §6).
2
+
3
+ The reactive fallback for ORGANIC / freeform shapes that the codegen path can't
4
+ express. Returns the model's RAW mesh text ({"text": ...}); the Gradio backend
5
+ (geometry.mesh_adapter) parses it into the frozen contract — no geometry runs
6
+ here. Structured like modal_endpoint.py: on-demand (scale-to-zero), weights on a
7
+ Volume, model loaded inside @modal.enter on the GPU container.
8
+
9
+ FLAGGED — the build loop writes this file but MUST NOT deploy it. To enable:
10
+ modal deploy modal_mesh_endpoint.py
11
+ then copy the URL into the Space secret MESH_ENDPOINT_URL and set TRIKONA_MESH=1.
12
+ Confirm the exact LLaMA-Mesh HF repo id below before deploying.
13
+ """
14
+ import json
15
+ from pathlib import Path
16
+
17
+ import modal
18
+ from pydantic import BaseModel
19
+
20
+ # LLaMA-Mesh: ~8B (LLaMA-3.1-8B fine-tune), within the 32B cap, fits an A100-40GB.
21
+ # VERIFY this repo id before deploy (e.g. the original "Zhengyi/LLaMA-Mesh" or an
22
+ # NVIDIA mirror).
23
+ MODEL_NAME = "Zhengyi/LLaMA-Mesh"
24
+ MAX_MODEL_LEN = 8192 # meshes are long token sequences
25
+ GPU_MEMORY_UTILIZATION = 0.90
26
+ TEMPERATURE = 0.6 # some diversity helps organic mesh generation
27
+ MAX_TOKENS = 7000 # room for a sizeable mesh; adapter enforces hard limits
28
+
29
+ app = modal.App("trikona-mesh")
30
+
31
+ hf_cache = modal.Volume.from_name("trikona-mesh-hf-cache", create_if_missing=True)
32
+ HF_CACHE_DIR = "/root/.cache/huggingface"
33
+
34
+ image = (
35
+ modal.Image.debian_slim(python_version="3.11")
36
+ .pip_install("vllm==0.6.3", "transformers==4.46.3", "fastapi[standard]", "pydantic>=2")
37
+ .env({"HF_HOME": HF_CACHE_DIR})
38
+ )
39
+
40
+
41
+ class MeshRequest(BaseModel):
42
+ prompt: str
43
+
44
+
45
+ @app.cls(
46
+ gpu="A100",
47
+ image=image,
48
+ volumes={HF_CACHE_DIR: hf_cache},
49
+ scaledown_window=300,
50
+ timeout=180,
51
+ max_containers=1,
52
+ startup_timeout=600,
53
+ )
54
+ class MeshModel:
55
+ @modal.enter()
56
+ def load(self):
57
+ from vllm import LLM, SamplingParams
58
+
59
+ self.llm = LLM(
60
+ model=MODEL_NAME,
61
+ max_model_len=MAX_MODEL_LEN,
62
+ gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
63
+ )
64
+ self.sampling_params = SamplingParams(temperature=TEMPERATURE, max_tokens=MAX_TOKENS)
65
+
66
+ @modal.fastapi_endpoint(method="POST")
67
+ def generate(self, request: MeshRequest) -> dict:
68
+ # LLaMA-Mesh emits OBJ-like text when asked to create a 3D model. The
69
+ # backend (mesh_adapter) extracts v/f lines, so we return raw text.
70
+ messages = [{"role": "user",
71
+ "content": f"Create a 3D obj model of {request.prompt}."}]
72
+ outputs = self.llm.chat(messages, self.sampling_params)
73
+ return {"text": outputs[0].outputs[0].text}
tests/test_mesh_integration.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Phase M3 — mesh endpoint + app wiring (SPEC_REACTIVE_MESH §6, §7). Offline only.
2
+
3
+ Uses canned _model_emit / _mesh_emit (no live Modal). Verifies the codegen path
4
+ routes to mesh when the flag is on, stays unchanged when off, and that the mesh
5
+ endpoint module parses. Live deploy + live run are FLAGGED (not tested).
6
+ """
7
+ import ast
8
+ import importlib.util
9
+ import json
10
+ import sys
11
+ from pathlib import Path
12
+
13
+ ROOT = Path(__file__).resolve().parent.parent
14
+ if str(ROOT) not in sys.path:
15
+ sys.path.insert(0, str(ROOT))
16
+
17
+ import app # noqa: E402
18
+
19
+ _spec = importlib.util.spec_from_file_location("schema_mod", ROOT / "tests" / "test_schema.py")
20
+ _schema = importlib.util.module_from_spec(_spec)
21
+ _spec.loader.exec_module(_schema)
22
+ validate_shape = _schema.validate_shape
23
+
24
+ CUBE_OBJ = (ROOT / "tests" / "fixtures" / "mesh" / "cube_quads.obj").read_text()
25
+
26
+ # An organic build_spec whose generator yields a degenerate mesh (fails the gate)
27
+ # with no matching primitive -> codegen exhausts -> mesh tier.
28
+ DEGEN_ORGANIC = {
29
+ "type": "build_spec", "shape_id": "blob_creature", "regularity": "freeform",
30
+ "generator": {"lang": "python", "entry": "build",
31
+ "code": "def build(p):\n return [[0,0,0],[1,0,0],[0.5,1e-9,0]], [[0,1,2,'b']]\n"},
32
+ "components": [{"id": "b", "label": "b", "param_dependencies": []}],
33
+ "params": [], "constraints": [], "construction_steps": [],
34
+ }
35
+
36
+
37
+ def test_mesh_flag_defaults_off():
38
+ assert app.MESH_ON is False
39
+
40
+
41
+ def test_mesh_emit_returns_none_without_url(monkeypatch):
42
+ monkeypatch.setattr(app, "MESH_URL", None)
43
+ assert app._mesh_emit("anything") is None
44
+
45
+
46
+ def test_codegen_routes_to_mesh_when_enabled(monkeypatch):
47
+ monkeypatch.setattr(app, "MESH_ON", True)
48
+ monkeypatch.setattr(app, "MESH_URL", "http://stub")
49
+ monkeypatch.setattr(app, "_model_emit", lambda prompt, params, feedback=None: DEGEN_ORGANIC)
50
+ monkeypatch.setattr(app, "_mesh_emit", lambda prompt: CUBE_OBJ)
51
+ contract = json.loads(app._codegen_generate("a blob creature", cache_key="k"))
52
+ assert contract.get("type") != "guided_response"
53
+ assert validate_shape(contract) == []
54
+ assert contract["params"] == []
55
+ assert contract["shape_id"].startswith("mesh_")
56
+
57
+
58
+ def test_codegen_stays_guided_when_mesh_off(monkeypatch):
59
+ monkeypatch.setattr(app, "MESH_ON", False) # flag off
60
+ monkeypatch.setattr(app, "_model_emit", lambda prompt, params, feedback=None: DEGEN_ORGANIC)
61
+ called = {"mesh": False}
62
+ monkeypatch.setattr(app, "_mesh_emit", lambda prompt: called.__setitem__("mesh", True) or CUBE_OBJ)
63
+ contract = json.loads(app._codegen_generate("a blob creature"))
64
+ assert contract.get("type") == "guided_response" # no mesh tier
65
+ assert called["mesh"] is False # _mesh_emit never used
66
+
67
+
68
+ def test_mesh_endpoint_module_parses():
69
+ src = (ROOT / "modal_mesh_endpoint.py").read_text()
70
+ ast.parse(src) # syntax-valid; deploy is FLAGGED, not exercised here
71
+ assert "trikona-mesh" in src and "MeshRequest" in src