Skip to content

feat(runtime): FlowMesh-backed text EmbeddingOp with per-row artifact output#48

Merged
timzsu merged 3 commits into
mainfrom
feat/embedding-op-vllm
Jul 5, 2026
Merged

feat(runtime): FlowMesh-backed text EmbeddingOp with per-row artifact output#48
timzsu merged 3 commits into
mainfrom
feat/embedding-op-vllm

Conversation

@timzsu

@timzsu timzsu commented Jul 4, 2026

Copy link
Copy Markdown
Collaborator

Purpose

Make the text-embedding path work end-to-end: an EmbeddingOp job now dispatches a FlowMesh vLLM embedding task, returns one artifact ref per input doc, and the downloaded embeddings.safetensors contains real, per-doc vectors. This unblocks routing the skill-integration system's embeddings through Lumilake→FlowMesh instead of a local sentence-transformers model. Also brings the CI workflows onto GitHub-hosted runners with current tool/action versions.

Changes

Embedding op

  • src/lumilake_server/runtime/runtime_manager/flowmesh.py — fix the artifact fetch: the download path was f"artifacts/{path}", but the FlowMesh results route already prepends artifacts/, so it requested artifacts/artifacts/embeddings.safetensors → 404. Pass the bare path. Per-row archive now dedups by path (one download per shared safetensors, replicated per row) and surfaces per-row errors explicitly.
  • src/lumilake_server/ops/embedding_ops.py — EmbeddingOp → FlowMesh embedding task (model.vllm.convert: embed), per-row vectors output.
  • src/lumilake_server/routes/jobs.pyjob artifact now JSON-decodes per-row embedding output refs before the artifact-path membership check (was a raw substring check, so --path <uri> always failed for embedding outputs).
  • src/lumilake_server/runtime/server.py_relocate_artifacts_for_request now actually relocates the nested output uri inside a JSON-encoded row output (was a blind text exec-…/artifacts/req-…/artifacts/ string-replace that rewrote the URI without copying the bytes, leaving the file only under the exec--scoped key).
  • tests/runtime/test_embedding_op.py — fetch-locally proof (loads the archived artifact, asserts [N, dim] real floats).
  • tests/runtime/server/test_run_batch.py — downstream-consumption proof (3 docs → 3 per-row outputs reach a consumer via the real aggregate path).

Optimizer

  • src/lumilake_server/runtime/optimizer/schedule/model_size.py — register BAAI/bge-small-en-v1.5 (0.033B). It has no size suffix, so Halo's _model_suffix_pattern can't infer it and refuses to schedule; suffixed models (e.g. Qwen3-Embedding-0.6B) still infer automatically and need no entry.

Docs

  • docs/OPS.md — EmbeddingOp contract: the vectors output is a per-row list (one entry per input doc, each with its row index + a ref to the shared embeddings.safetensors, tensor key embeddings, shape [count, dim] float32). The concrete workflow YAML is intentionally not committed — the contract is here.

CI (GitHub-hosted runners)

  • .github/workflows/*.yml (10 files) — self-hostedubuntu-latest; bump uv 0.11.8→0.11.26, zizmor 1.24.1→1.26.1, pip-audit 2.9.0→2.10.1, docker actions + gh-action-pypi-publish to latest (SHA-pinned); enable-cache: false on the 3 release/publish setup-uv steps (zizmor 1.26.1 cache-poisoning hardening; non-release jobs keep caching).

Design

The failure only appeared under a live job — unit tests used in-memory fakes and never exercised the real FlowMesh artifact fetch. Root cause was a doubled artifacts/ path segment; confirmed by reading the worker container's manifest.json (the file sits at a single-nested artifacts/embeddings.safetensors). A native FlowMesh EmbeddingTask for the same model already succeeds on the box, so the fix aligns our fetch with the server's existing route contract rather than changing the executor. Two red herrings (destination type: http vs local; a hardcoded timeoutSec deep-merge bug) were investigated and reverted — the double-prefix was the sole, sufficient cause of the fetch 404. Verifying the fetch through Lumilake's product surface (lumilake job artifact) rather than the FlowMesh CLI then surfaced two further bugs the CLI shortcut had masked (the JSON-encoded per-row ref was neither collected nor relocated correctly) — both fixed here, so the artifact is fetchable the way a real consumer fetches it.

Test Plan

End-to-end command chain (local stack), submitting a real 3-doc job and verifying the fetched vectors:

# 1. Build + restart the stack with the change
lumilake deploy build
lumilake deploy restart

# 2. SUBMIT: a text-embedding workflow (contract in docs/OPS.md §EmbeddingOp),
#    3 docs = 3 slices, model BAAI/bge-small-en-v1.5, default halo optimizer
lumilake job submit <text-embedding-workflow>.yaml -f yaml \
  --output-type s3 --output-prefix skill-integration/e2e/ \
  --input 'Docs=alpha text,beta text,gamma text'
#   -> req-XXXX

# 3. Poll to completion
lumilake job info req-XXXX          # status: completed

# 4. VERIFY the result: 3 per-row vector refs, NO 404
lumilake job result req-XXXX
#   outputs["text-embedding"]["vectors"] = [
#     {"output": ".../embeddings.safetensors", "model": "BAAI/bge-small-en-v1.5", "row": 0},
#     {... "row": 1}, {... "row": 2}]

# 5. VERIFY the vectors: fetch the artifact through Lumilake's own surface
#    (the s3 output written under --output-prefix), then load it.
lumilake job artifact req-XXXX --path embeddings.safetensors -o /tmp/emb
#    (equivalently, read the same s3 object via lumid-data-app: DataRetrievalOp
#     type: lumid, mode: s3 — the path is the `output` uri in the vectors ref)
python -c "from safetensors.torch import load_file; \
  t=load_file('/tmp/emb/embeddings.safetensors')['embeddings']; \
  print(t.shape, t.dtype)"

The fetch goes through Lumilake (job artifact / lumid-data-app), never by probing FlowMesh directly — the artifact is surfaced from the job's s3 output.

What is verified:

  • The job schedules under halo (bge-small registered → no "Halo model size is unknown" crash) and completes.
  • vectors is a per-row list of 3 refs (one per doc, each with its row), and none carry a 404 (artifact not found) — the fetch bug is gone.
  • embeddings.safetensors loads as shape [3, 384] float32, all finite, unit-norm, and the 3 rows are distinct (pairwise cosine < 0.999) → one real embedding per doc.
  • Downstream consumption: test_run_batch.py asserts 3 docs → 3 per-row outputs reach a consumer through the real _aggregate_output_node path.
  • Regression: full suite + pre-commit + zizmor.

Test Result

# job result (step 4)
vectors: 3 entries, rows 0/1/2, model BAAI/bge-small-en-v1.5, no 404

# artifact fetched via Lumilake surface (step 5: lumilake job artifact)
shape: (3, 384)   dtype: float32   finite: True   norms: ~1.0
pairwise cosine: 0.841 / 0.806 / 0.786   (all < 0.999 → distinct per doc)

# full
pytest tests/                                                    -> 740 passed, 1 skipped
uvx --from zizmor==1.26.1 zizmor --persona pedantic .github/workflows -> No findings
uv run pre-commit run --all-files                               -> all hooks green

Pre-submission Checklist
  • I have read CONTRIBUTING.md.
  • I have run uv run pre-commit run --all-files and fixed any issues.
  • I have added or updated tests covering my changes (fetch-locally + downstream-consumption).
  • I have verified that uv run pytest tests/ passes locally (739 passed, 1 skipped).
  • If I changed the SDK or CLI, I have verified the affected interface works locally (live e2e job above).
  • If this is a breaking change, I have prefixed the PR title with [BREAKING] — n/a.
  • I have updated documentation (docs/OPS.md EmbeddingOp contract) since user-facing behavior changed.

timzsu added 2 commits July 4, 2026 21:40
Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
@timzsu
timzsu force-pushed the feat/embedding-op-vllm branch from d25afa9 to 3c683b0 Compare July 4, 2026 13:40
Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
@timzsu timzsu changed the title feat: Add embedding op feat(runtime): FlowMesh-backed text EmbeddingOp with per-row artifact output Jul 4, 2026
@timzsu
timzsu merged commit e386262 into main Jul 5, 2026
13 of 14 checks passed
@timzsu
timzsu deleted the feat/embedding-op-vllm branch July 5, 2026 03:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant