MILAB-6480: optimize embedding resource requests with fluent formulas API

[blackcat]

What & why

Migrates the embedding execs to the fluent exec.formula / .resources() API (workflow-tengo 6.7.x) and sizes GPU jobs to the smallest node that safely fits, so embedding no longer over-provisions onto the scarce, expensive g6e.12xlarge (the node type that sat Pending under capacity pressure).

Changes

Workflow

• compute-embeddings now declares .resources({ onCPU | onGPU }) instead of a flat 16 CPU / 32 GiB / 16 GiB-VRAM request:
  • GPU path: host 2 CPU / 8 GiB (single streaming process — per-chunk CPU work is interleaved with the forward pass, and the heavy PColumn→TSV shaping is already upstream CPU execs), VRAM model-tiered: 3 GiB (dim ≤ 768) → gpu-3g (g6f.xlarge), 6 GiB (dim ≥ 1024) → gpu-6g (g6f.2xlarge). Never forces g6e.12xlarge.
  • CPU path: cores from size("batch") clamped [4,16], RAM 2 GiB/core [8,32], .staticFallback = the old fixed 16 / 32 GiB (for backends that can't evaluate formulas).
• compute-stats sizes RAM from the source-TSV size (it full-loads the TSV via read_csv), replacing the flat 4 GiB that could OOM on large inputs.

Software (main.py)

• On CUDA the per-forward token budget is auto-sized from allocated VRAM (PLATFORMA_GPU_MEMORY), mirroring how --max-memory-gb sizes the host-RAM path: a larger VRAM request → larger batches (higher throughput), a smaller one stays safe. --token-budget still overrides; the halve-on-OOM retry remains the backstop.

Model / UI

• The mem / cpu Advanced-Settings fields are now opt-in (init no longer hardcodes 32/16), so the workflow's automatic sizing applies unless the user sets a value (which then overrides per dimension). Existing projects keep any value they already had.

Pins @platforma-sdk/workflow-tengo 6.6.5 → 6.7.2.

Verification

Built + checked clean; validated live on a GPU backend (ESM-2 150M / CDR3):

	before	after (observed live)
Request	16 CPU / 32 GiB / 16 GiB VRAM	2 CPU / 8 GiB / 3 GiB VRAM
Landed on	g6e.12xlarge (4×L40S, ~$10.59/hr)	g6f.xlarge (fractional L4, ~$0.24/hr)

The batch pods provisioned g6f.xlarge (gpu-3g) nodes, scheduled, ran to completion, and the stats pass succeeded.

Greptile Summary

This PR replaces the flat 16 CPU / 32 GiB / 16 GiB-VRAM embedding resource request with the fluent exec.formula / .resources() API (workflow-tengo 6.7.x), right-sizing GPU jobs to cheap fractional-L4 nodes and scaling CPU jobs from actual batch volume. It also wires the Python runtime to auto-size the per-forward token budget from the allocated VRAM (PLATFORMA_GPU_MEMORY), and makes mem/cpu in the UI opt-in overrides instead of hard defaults.

• Workflow \u2014 compute-embeddings declares onGPU: { cpu: 2, ram: 8GiB, vram: 3|6GiB } and onCPU: { cpu: f.size(\"batch\")/2MiB clamped [4,16], ram: cpu\u00d72GiB clamped [8,32] } with .staticFallback; compute-stats sizes its exec RAM from f.size(\"source\")\u00d73+2GiB replacing the flat 4 GiB.
• Software \u2014 resolve_gpu_token_budget() reads PLATFORMA_GPU_MEMORY, subtracts estimated weight footprint and a 1 GB reserve, then scales by GPU_TOKENS_PER_GB=4096; explicit --token-budget still wins and halve-on-OOM is the backstop.
• Model / UI \u2014 dataModel.ts no longer seeds mem: 32, cpu: 16; existing projects retain any saved value, new projects get automatic sizing.

Key terms introduced or changed:

• exec.formula / f \u2014 Fluent formula API for runtime-computed resource requests. New; replaces flat static values.
• onGPU / onCPU \u2014 Keys in .resources({}) for per-path resource bundles. New.
• gpuVramGiB \u2014 New function returning VRAM tier: 3 GiB (dim \u2264 768) or 6 GiB (dim \u2265 1024).
• PLATFORMA_GPU_MEMORY \u2014 Env var now consumed by main.py to auto-size token budget.
• GPU_TOKENS_PER_GB \u2014 New constant (4096) in main.py.
• resolve_gpu_token_budget \u2014 New function computing per-forward token budget from allocated VRAM.
• parse_gpu_memory_env \u2014 New parser for PLATFORMA_GPU_MEMORY (bytes or SI quantities).
• staticFallback \u2014 Formula method providing a static fallback for older backends.
• BlockDataV*.mem / .cpu \u2014 Changed: "Undefined \u2192 32 GiB / 16 cores" \u2192 "Undefined \u2192 auto-sized".
• dataModel.ts initial state \u2014 Changed: { embedding: {}, mem: 32, cpu: 16 } \u2192 { embedding: {} }.

Confidence Score: 4/5

The resource-sizing changes are well-reasoned with safe static fallbacks on every dimension and a halve-on-OOM retry backstop; GPU VRAM tiers are conservatively chosen and validated live.

The core sizing logic is safe and tested end-to-end. The token-budget sentinel comparison against DEFAULT_TOKEN_BUDGET is a latent API fragility, and the resources()/addFile() ordering inconsistency in compute-stats is worth a second look before the next SDK bump, but neither is a current runtime failure.

software/src_python/main.py (token-budget sentinel logic) and workflow/src/compute-stats.tpl.tengo (builder chain ordering) deserve a second look.

Important Files Changed

Filename	Overview
workflow/src/compute-embeddings.tpl.tengo	Migrates resource requests to fluent `.resources({ onGPU
workflow/src/compute-stats.tpl.tengo	Adds formula-based resources for the Python counting exec (sized from source-TSV file size); seqTb builder retains hardcoded 1-core/4-GiB via metaInputs. Minor: .resources() is declared before the tagged .addFile() in the builder chain.
software/src_python/main.py	Adds parse_gpu_memory_env, resolve_gpu_token_budget, and related constants to auto-size the per-forward token budget from PLATFORMA_GPU_MEMORY on CUDA. Uses != DEFAULT_TOKEN_BUDGET as the explicit-override sentinel, which is fragile.
workflow/src/models.lib.tengo	Adds gpuVramGiB returning the VRAM scheduling tier (3 GiB for dim ≤ 768, 6 GiB for dim ≥ 1024) with safe fallback to 6 for unknown tags.
model/src/dataModel.ts	Removes hardcoded mem: 32, cpu: 16 from initial block state, making resource overrides opt-in. Existing projects retain saved values; new projects get automatic sizing.
workflow/src/main.tpl.tengo	compute-stats render still passes hardcoded metaInputs: { cpu: 1, mem: 4 } (correct — these go to the seqTb TSV builder, not the Python exec); embedMeta passes cpu/mem overrides as metaExtra to compute-embeddings.
model/src/types.ts	Updates JSDoc on mem/cpu fields across BlockDataV1/V2/V3 and BlockArgs to reflect auto-sizing instead of fixed defaults.
ui/src/pages/MainPage.vue	Tooltip copy updated for mem/cpu fields to say 'Leave empty to size automatically'; no functional changes.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[wf.body in main.tpl.tengo] --> B{exec.hasGpu?}
    B -- gpu --> C[compute-embeddings: onGPU]
    B -- cpu --> D[compute-embeddings: onCPU]
    C --> C1["cpu: cpuOverride ?? 2 / ram: ramOverride ?? 8GiB / vram: gpuVramGiB"]
    C1 --> C2["gpu-3g g6f.xlarge or gpu-6g g6f.2xlarge"]
    D --> D1["cores = f.size(batch)/2MiB clamped 4-16 / ram = cores x 2GiB"]
    D1 --> D2["staticFallback: 16 cores / 32GiB"]
    A --> E[compute-stats: onCPU]
    E --> E1["cpu = f.size(source)/1GiB clamped 1-4 / ram = size x3 + 2GiB clamped 4-32GiB"]
    C1 --> F[PLATFORMA_GPU_MEMORY env]
    F --> G[parse_gpu_memory_env]
    G --> H["resolve_gpu_token_budget: avail = VRAM - weights - 1GB / budget = avail x 4096"]
    H --> I[clamp 4096 to 131072 tokens]
    I --> J[Embedder forward pass with halve-on-OOM backstop]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A[wf.body in main.tpl.tengo] --> B{exec.hasGpu?}
    B -- gpu --> C[compute-embeddings: onGPU]
    B -- cpu --> D[compute-embeddings: onCPU]
    C --> C1["cpu: cpuOverride ?? 2 / ram: ramOverride ?? 8GiB / vram: gpuVramGiB"]
    C1 --> C2["gpu-3g g6f.xlarge or gpu-6g g6f.2xlarge"]
    D --> D1["cores = f.size(batch)/2MiB clamped 4-16 / ram = cores x 2GiB"]
    D1 --> D2["staticFallback: 16 cores / 32GiB"]
    A --> E[compute-stats: onCPU]
    E --> E1["cpu = f.size(source)/1GiB clamped 1-4 / ram = size x3 + 2GiB clamped 4-32GiB"]
    C1 --> F[PLATFORMA_GPU_MEMORY env]
    F --> G[parse_gpu_memory_env]
    G --> H["resolve_gpu_token_budget: avail = VRAM - weights - 1GB / budget = avail x 4096"]
    H --> I[clamp 4096 to 131072 tokens]
    I --> J[Embedder forward pass with halve-on-OOM backstop]

Loading

Prompt To Fix All With AI

Fix the following 3 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 3
software/src_python/main.py:257-260
The explicit-override sentinel compares against `DEFAULT_TOKEN_BUDGET` (16384), so passing `--token-budget 16384` explicitly is silently treated as "not set" and VRAM auto-sizing takes over. If `DEFAULT_TOKEN_BUDGET` is ever changed, existing invocations that explicitly pin the old value will silently change behaviour. The `--max-memory-gb` argument uses `0` as the "not set" sentinel — the same pattern would be cleaner here (default 0, any positive value wins).

```suggestion
    if args.token_budget and args.token_budget > 0:
        return args.token_budget            # explicit operator override
    if device.type != "cuda":               # budget only bounds VRAM; CPU/MPS keep the default
        return DEFAULT_TOKEN_BUDGET
```

### Issue 2 of 3
workflow/src/compute-stats.tpl.tengo:92-100
**`resources()` declared before the tagged `addFile()`** — `f.size("source")` references the `"source"` tag, but `.addFile("source.tsv", sourceTsv, { tag: "source" })` appears later in the same chain. In `compute-embeddings.tpl.tengo` the tagged file is always registered first. While formula evaluation is presumably lazy (resolved at job-submission time, not at builder-chain time), the inconsistent ordering is fragile: if any SDK version evaluates the tag lookup eagerly at `.resources()` call time, the reference will be unresolved. Moving `.addFile(...)` before `.resources(...)` costs nothing and removes the ambiguity.

### Issue 3 of 3
workflow/src/compute-embeddings.tpl.tengo:116-122
**GPU path exposes no `vram` override** — `cpuOverride` and `ramOverride` are respected on the GPU path so a user can tune cores/RAM, but `vram` is always the auto-computed tier from `gpuVramGiB`. If a sequence batch OOMs even after the halve-retry (e.g., an unusually long single sequence), the user has no escape hatch. Adding an optional `vramOverride` wired from a UI Advanced-Settings field (similar to `mem`/`cpu`) would let power users move from `gpu-3g` to `gpu-6g` without a code change.

_{Reviews (1): Last reviewed commit: "MILAB-6480: honor explicit --token-budge..." | Re-trigger Greptile}

Greptile also left 3 inline comments on this PR.

Context used:

• Context used - Terms is a types in codebase. Provide the list of ... (source)

[gemini-code-assist]

Code Review

This pull request optimizes compute-resource requests for sequence embeddings by transitioning from static defaults to dynamic, data-driven resource allocation. CPU and RAM requirements are now automatically sized based on input volume, and GPU VRAM is tiered by model size, allowing jobs to run on cheaper fractional-L4 instances. Additionally, on CUDA, the token budget is auto-sized from the allocated VRAM. Feedback on the changes highlights a flaw in the GPU token budget resolution logic where explicitly setting the budget to the default value would incorrectly trigger auto-sizing; updating the default argument to None is recommended to properly handle user overrides.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.