Configure and invoke algorithms

The two-stage contract

DRTML params_defaults / user overlay
  → merged flat compute.* params
  → parse_compute_plan()
  → ComputePlan
  → processor or framework mode invokes algorithm registry

DRTML configures methods and parameters. It does not replace the processor that prepares MaterialRow, chooses K, or turns results into logical dataset rows.

Compute configuration

params_defaults:
  compute.embedding_method: sgns_v1
  compute.embed_dim: 128
  compute.embed_epochs: 3
  compute.embed_mode: train_infer
  compute.embed_lr: 0.025
  compute.embed_seed: 42
  compute.embed_negative_samples: 5
  compute.embed_warm_checkpoint: ""
  compute.embed_save_checkpoint: vectors-v1

  compute.cluster_method: minibatch_kmeans_v1
  compute.cluster_batch_size: 512
  compute.cluster_max_iter: 30
  compute.cluster_seed: 42

  compute.parallelism.backend: local
  compute.parallelism.threads: 8
  compute.parallelism.processes: 1

  compute.materialize.row_id_field: descriptor_id
  compute.materialize.token_ids_field: token_ids

ComputePlan contains embed, cluster, materialize, and algorithm parallelism. runtime.parallelism.* is a separate worker-dispatch concern.

Direct embedding from a processor

from drtoller.framework.algorithms import embed_rows
from drtoller.framework.algorithms.types import MaterialRow
from drtoller.framework.drtml import (
    parse_compute_plan,
    embed_algorithm_kwargs,
)

plan = parse_compute_plan(params)
material = [
    MaterialRow(
        row_id=str(row[plan.materialize.row_id_field]),
        token_ids=tuple(row[plan.materialize.token_ids_field]),
    )
    for row in input_rows
]

result = embed_rows(
    material,
    method_id=plan.embed.method_id,
    params=embed_algorithm_kwargs(plan),
)

For sgns_v1, required values must be present: dimension, epochs, learning rate, negative samples and seed. embed_train is preferable when W is run-global/stateful and inputs must stream rather than fit in one processor call.

Production embed_train currently hard-requires compute.embedding_method: sgns_v1. Single-table transports use SgnsVocabTrainer; Ray with workers≥2 uses sgns_ps_step against sharded actors. Registering another embed method is not enough to make it available there.

Stateful embed_train invocation

A DRTML v4 step declares run.type: embed_train, plus compute.*, train-pair/count inputs, an output vector dataset, and runtime streaming/parallelism params. Builtin run types do not require a custom run.entry. The runner:

1. loads per-group token counts;
2. constructs the vocabulary;
3. streams pair batches for each epoch;
4. updates W through inprocess / processes / Ray (single actor or PS shards);
5. writes vector snapshots through StorageSession;
6. emits observations/geometry metrics;
7. returns train_stats via EmbedTrainStats.as_public_dict().

The shared run_train_over_groups owns the outer group lifecycle; transports do not duplicate it. Details: embed_train.

Clustering with domain-owned K

from drtoller.framework.algorithms import cluster_with_k
from drtoller.framework.drtml import (
    parse_compute_plan,
    cluster_algorithm_kwargs,
)

plan = parse_compute_plan(params)
k = choose_k_from_domain_policy(probe_results, params)
result = cluster_with_k(
    vectors,
    row_ids,
    k,
    method_id=plan.cluster.method_id,
    params=cluster_algorithm_kwargs(plan),
)

Do not add compute.cluster_k as hidden generic policy when the step is responsible for probing/selection. The method implements clustering for a supplied K. At corpus scale prefer streaming fit/assign over materializing the full matrix into cluster_with_k. Candidate search belongs in the probe loop.

Writing results

Algorithms return arrays/DTOs. Domain code converts them to logical rows and returns ProcessResult; framework writes the declared output dataset.

rows = [
    {"row_id": row_id, "cluster_id": int(label)}
    for row_id, label in zip(result.row_ids, result.labels)
]
return [ProcessResult(ok=True, rows={"assignments": rows})]

Add a method

Full recipe: New algorithm method.

Embedding/clustering checklist:

1. implement a data-in/result-out function in the proper package;
2. add a *_from_params adapter;
3. register with register_embed_method or register_cluster_method;
4. add compute.* keys and defaults where required;
5. test direct API, plan kwargs, invalid method, and runtime integration;
6. document inputs, result type and resource behavior.

Dataset metric methods use register_method(MetricMethod(...)); see New metric method.