embed_train

How it differs

Not a partition_loop map. Shared mutable embedding table W is trained over multi-pass pair streams:

load vocab counts → filter min_token_count
  → stream positive pairs (bounded pending)
  → SGNS update W (inprocess | processes SHM | Ray)
  → snapshot vectors → StorageSession (parquet or qdrant)
  → observation / geometry metrics

SGNS math lives in framework/algorithms/embed/ (sgns.py, parameter-server step in sgns_ps.py).
Orchestration lives in processing/embed_train/ + engine under processing/embed_train/engine/ (former embed_vocab/ package).

DRTML mode contract

execution:
  entry: drtoller.framework.processing.embed_train.run:run_embed_train

embed_train:
  train_pairs:
    dataset: train_pairs
    group_column: kind
    group_exclude_values: [literal]
    center_column: center_token_id
    context_column: context_token_id
    count_column: count
    pair_source_column: pair_source
    policy_version_column: policy_version
    policy_version: dep_weight_v1
  vocab_counts:
    dataset: vocabulary_counts
    token_column: token_id
    count_column: count
  observation_dataset: train_observations

The compiled EmbedTrainModeSpec requires:

train_pairs and vocab_counts input contracts with matching declared columns, group_by, and sum_columns;
exactly one output dataset with kind: vector_vocab and uniq;
observation_dataset, when present, to be kind: observation with backend: postgres;
transport and queue capacities from runtime.parallelism.* and runtime.streaming.*.

Production currently accepts only compute.embedding_method: sgns_v1 and one included embed kind per run.

Shared train loop

processing/embed_train/engine/training/train_loop.py:

material_groups — sorted kinds from compute.embed_kinds (empty → error);
skip_numpy_geometry_for_output — skip host geometry when the vector sink is Qdrant;
run_train_over_groups(backend, groups, train_one_group=…) — plan log, per-group train, aggregate stats.

Backends (inprocess, processes, ray_actor) implement only train_one_group + transport.
Do not duplicate the outer loop / plan-done logging.

Ray: single W vs parameter-server shards

`runtime.parallelism.workers`	Path	Backend label
`1`	one `EmbedVocabActor` owns full `W`; `update_pairs.remote`	`ray`
`≥ 2`	`W` sharded across `workers` actors (`row_id % n_shards`)	`ray_ps`

Sharded path (backends/ray_w_shards.py + algorithms/embed/sgns_ps.py):

gather needed rows from owning shards;
run SGNS locally on the gathered slice (sgns_ps_step);
push row deltas back to shard actors;
merge exports for snapshot write and optional geometry.

Inflight is bounded by BoundedPending / compute.embed_ray_max_inflight.
Step03 defaults are backend: ray / workers: 8 — multi-shard PS is the production train path.

See also Ray integration and Parallelism.

Streaming contract

Capacities come from flattened runtime.streaming.* (DRTML nested params.runtime.streaming):

Param	Role (typical default)
`runtime.streaming.pending_per_worker`	in-flight pair batches (2)
`runtime.streaming.observation_q_max`	live observation channel (32)
`runtime.streaming.worker_metrics_q_max`	worker metric queue (64)

Modules:

streaming/pending.py — bounded in-flight;
streaming/async_sink.py — observation channel (drop-oldest for telemetry);
streaming/backpressure.py — put_with_backpressure.

Live observations use a separate StorageSession / PG connection.
The primary session writes the terminal vector snapshot. The auxiliary channel is bounded and drop-oldest for live telemetry, then writes one terminal observation before its dedicated session closes.

Metrics and public results

Live / final gauge names live in contracts/embed_train_metric_catalog.py and are shared by:

Prometheus emit (telemetry/prometheus/emit/embed_train.py);
observation writers (processing/embed_train/observation_write.py).

run_embed_train returns EmbedTrainStats.as_public_dict() (JSON-safe; heavy fields such as full token_counts omitted) for Airflow XCom / UI consumers.

Geometry helpers (L2, hubness) run after snapshot when the sink allows host numpy; Qdrant skips host geometry and relies on sink-side / offline evaluation.