mstar.engine.resources.runner

mstar.engine.resources.runner#

sequencing for resource call cycle

runner only moves .plan outputs by putting under StepContext.plan_results and handing to next. (in fact, maybe plan should do this and runner only moves StepContext?)

Functions

topo_sort(resources)

Keys in dependency order; built resources at load, their specs when the conductor validates a deployment.

Classes

StepRunner(resources[, node_resources, ...])

drives resources through per-step cycle
class mstar.engine.resources.runner.StepRunner(resources, node_resources=None, enable_nvtx=False)[source]#

Bases: object

drives resources through per-step cycle

Parameters:

  • resources (Mapping[str, Resource])

  • node_resources (Mapping[str, Collection[str]] | None)

  • enable_nvtx (bool)

admit(step)[source]#

Reserve capacity for step. A refused admit gives back everything the step reserved.

Parameters:

step (SubmoduleStep)

Return type:

FullAdmitOutcome

admit_retrieve(rid, node_name, graph_walk, published=None)[source]#

bring published state in

gives ready=False when still has inflights; shortcircuit on failure

Swept over node_name’s own resources. Asking a node about resources it does not own is not just wasted work: KVManager.get_labels answers for an unrecognised node with the default [“main”], so the node ends up gated on — and able to allocate against — another node’s cache.

Parameters:
Return type:

FullAdmitOutcome

adopt_cached_prefix(rid, node_name, graph_walk, leases, caught_up)[source]#
Parameters:
Return type:

bool

apply_cached_prefix(rid, node_name, graph_walk, inputs, prefix)[source]#

Give every resource the agreed prefix, with the untrimmed inputs.

Parameters:
Return type:

None

build_accelerator_graph_buffers(slots, max_bs, max_seq_len, node_name=None)[source]#

resources preallocate the static buffers captured replays will read

Swept over node_name’s own resources: another node’s slot count and batch size say nothing about a resource this node’s graphs never touch.

Parameters:
Return type:

None

clear_preplan()[source]#

Drop the staged pre-plan on every resource, and the record of it.

Return type:

None

commit(step)[source]#

record step consumption

Parameters:

step (SubmoduleStep)

Return type:

None

complete_cached_walk(rid, node_name, graph_walk)[source]#

Tell the node’s resources that the cache served this walk whole.

Parameters:
Return type:

None

extend_prefix_chains(rid, node_name, graph_walk, outputs, keys=None)[source]#

Offer this step’s sampled tokens to the node’s own resources.

Parameters:
Return type:

None

extend_prefix_chains_batch(request_ids, node_name, walk_of, host_rows, outputs, keys)[source]#

extend_prefix_chains for a step: a resource with a batch form takes the host rows; the others are offered each request’s host copy.

walk_of rather than one walk for the step, so a combined walk comes through here too instead of falling back to a call per row; a step of one walk passes that walk itself.

Parameters:
Return type:

None

ingest_capture_row(rid)[source]#

open a capture padding row on all resources

Parameters:

rid (str)

Return type:

None

ingest_request(rid, overrides=None)[source]#

open state on all resources; overrides are per-resource. All or nothing: a resource that raises leaves the request on none of them.

Parameters:
Return type:

None

lease_cached_prefix(rid, node_name, graph_walk)[source]#
Parameters:
Return type:

dict[str, list[int]]

property order: tuple[str, ...]#
page_trails(node_name)[source]#
Parameters:

node_name (str)

Return type:

dict[str, tuple[int, int]]

plan(step)[source]#

plan in dependency order

place plan in step.ctx.plan_results before next plan runs again could possibly move that into plan itself

Parameters:

step (SubmoduleStep)

Return type:

dict[str, Any]

pre_admit(step)[source]#

admit over the pre-planning subset, a step ahead

the later full admit covers the rest; these resources see their own state as already reserved and no-op

Unwound on refusal exactly as admit is.

Parameters:

step (SubmoduleStep)

Return type:

FullAdmitOutcome

pre_plan(step)[source]#

plan the pre-planning subset, a step ahead

ctx.is_preplan sends each one’s output to its pending slot rather than the live one; the later full plan promotes it. Results are left in ctx.plan_results for the dependents in this same subset.

Parameters:

step (SubmoduleStep)

Return type:

dict[str, Any]

prefix_chain_keys(node_name)[source]#

The node’s resources with a prefix cache open right now. Empty in the common case, which lets the caller skip a per-request sweep that was six no-op calls per request per step.

Parameters:

node_name (str | None)

Return type:

list[str]

publish(request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

durable state outward publish

non-publish resources noop; pairs with admit_retrieve

Scoped to the node whose step just ran; no other node’s state moved, and merge_publish_info keeps the entry it already published.

Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

publish_after_stop(request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

Publish resources configured for the end of a dynamic loop.

Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

publish_from_snapshots(snapshots, request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

The published info of request_ids from their snapshots, each under its own walk (see publish_snapshot).

Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

publish_snapshot(request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

publish, deferred: per rid the snapshot each publisher took.

Finished by publish_from_snapshots for the rids that need it. A rid with nothing to publish has no entry. Under a combined walk each rid is snapshotted under its own walk (request_walks), as publish does: a publisher may key what it publishes on the walk.

Parameters:
Return type:

dict[str, dict[str, Any]]

release_cached_prefix(rid, node_name, graph_walk)[source]#
Parameters:
Return type:

None

remove_request(rid)[source]#
Parameters:

rid (str)

Return type:

None

resolve_cached_prefix(rid, node_name, graph_walk, inputs)[source]#

The longest prefix of this walk every resource on node_name can serve.

The smallest answer wins, because a step is one span: a resource that holds less than another has to be caught up by the tokens that run. A resource with no opinion is not an answer of zero.

Parameters:
Return type:

CachedPrefix | None

property resources: Mapping[str, Resource]#
rollback_admit(step)[source]#

Give back everything an admitted step reserved, for a step that will not run: the unwind a refused admit does, over every resource the step declares.

Parameters:

step (SubmoduleStep)

Return type:

None

property staged: bool#

A pre-plan is staged and not yet promoted.

mstar.engine.resources.runner.topo_sort(resources)[source]#

Keys in dependency order; built resources at load, their specs when the conductor validates a deployment.

Parameters:

resources (Mapping[str, Resource | NodeResourceSpec])

Return type:

tuple[str, …]