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)

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

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

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

Parameters:

  • rid (str)

  • node_name (str)

  • graph_walk (str)

  • matched_len (int)

Return type:

None

build_cuda_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

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

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

Parameters:
Return type:

None

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

open state on all resources; overrides are per-resource

Parameters:
Return type:

None

property order: tuple[str, ...]#
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]

publish(request_ids, node_name=None, graph_walk=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:

  • request_ids (list[str])

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

dict[str, dict[str, PublishedInfo]]

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

Publish resources configured for the end of a dynamic loop.

Parameters:

  • request_ids (list[str])

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

dict[str, dict[str, PublishedInfo]]

remove_request(rid)[source]#
Parameters:

rid (str)

Return type:

None

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

The longest prefix 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:

int

property resources: Mapping[str, Resource]#
property staged: bool#

A pre-plan is staged and not yet promoted.

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

resources (Mapping[str, Resource])

Return type:

tuple[str, …]