hyrax.context

Run-scoped context shared by verbs, models, vector databases, and exporters.

A Hyrax verb run produces a results directory and a handful of other facts about the run (which verb, which runtime configuration, which distributed rank). Several pieces of Hyrax need those facts but are not handed them directly - a model’s prepare_inputs is a staticmethod called before any instance exists, so the context is the only way it can reach the config.

Each verb run gets its own context object. The Verb base class installs one for the duration of run() and releases it when the run ends; everything else reads the current one with get_context().

Typical usage in a model:

from hyrax import get_context

@hyrax_model
class MyModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

    def train_post_epoch(self):
        context = get_context()
        rank = context["rank"]
        np.save(context["results_dir"] / f"acts_rank{rank}.npy", self.acts)

Because contexts are per-run objects that are never reused, an object that outlives the run that created it can simply hold on to the context it was given rather than copying values out of it. hyrax.vector_dbs.VectorDB does exactly that, so a database handed back by database_connection keeps pointing at its own directory no matter what runs afterwards.

Classes

ContextKeys	The keys Hyrax puts in the run context.
RunContext	A plain dict that explains itself when a key is missing.

Functions

get_context(→ RunContext)	Get the run context for the run currently underway.
run_context(→ collections.abc.Generator[RunContext])	Install a fresh context for one verb run, releasing it when the run ends.
use_context(→ collections.abc.Generator[RunContext])	Re-enter a context captured earlier, for work that outlives its verb run.
update_context(→ RunContext)	Merge values into the current run context, leaving existing keys in place.
clear_context(→ None)	Release the current run context.

Module Contents

class ContextKeys

Bases: TypedDict

The keys Hyrax puts in the run context.

This is a documentation and type-annotation aid; the context is an ordinary dictionary and is not validated against this definition.

Note: total=False means that all keys are optional so that type checkers don’t complain about the context being empty before a verb has started a run.

Initialize self. See help(type(self)) for accurate signature.

results_dir: pathlib.Path

The results directory for the current run.

verb: str

The name of the verb that started the run, e.g. "train", "infer"

config: dict

The runtime configuration the verb is running with.

The Verb base class puts this in alongside verb. Verbs that run against a configuration other than their own - to_onnx, which uses the config saved next to the trained weights - overwrite it with update_context().

ml_framework: str

The source framework for a model export. Only set by the to_onnx verb.

class RunContext

Bases: dict

A plain dict that explains itself when a key is missing.

Behaves exactly like a dict in every other respect, so it can be passed anywhere a context dictionary is expected.

Initialize self. See help(type(self)) for accurate signature.

__missing__(key)

get_context() → RunContext

Get the run context for the run currently underway.

Safe to call at any time. Outside a verb run the context is empty, and reading a key from it raises a KeyError explaining why.

Returns:

The current run context. Each verb run gets a fresh object, so code that needs a context to outlive its run can hold this rather than copying values out of it.

Return type:

RunContext

run_context(verb: str, **extra) → collections.abc.Generator[RunContext]

Install a fresh context for one verb run, releasing it when the run ends.

The Verb base class wraps every verb’s run() in this, so verbs do not call it themselves; they fill in their results directory with update_context() once they have created it.

Verb runs are not expected to nest. The context is released to empty rather than to whatever was current beforehand, so a verb that calls another verb is left without a context once the inner call returns.

Parameters:

• verb (str) – The name of the verb starting the run.

• **extra – Any additional keys to place in the context.

Yields:

RunContext – The newly installed context.

use_context(context: RunContext) → collections.abc.Generator[RunContext]

Re-enter a context captured earlier, for work that outlives its verb run.

An object that keeps running model code after the verb that built it has returned - hyrax.verbs.infer_stream.InferStreamSession, for example - captures get_context() at construction and re-installs it around that later work, so models see the run they belong to rather than an empty context.

Parameters:

context (RunContext) – A context captured earlier with get_context().

Yields:

RunContext – The re-installed context.

update_context(**kwargs) → RunContext

Merge values into the current run context, leaving existing keys in place.

Verbs use this to record their results directory once they have created it. Existing keys are preserved, so anything a model has stashed in the context survives.

If results_dir is given it is coerced to Path. If rank and world_size are not given they are filled in from ignite’s distributed configuration.

Returns:

The current run context, for convenience.

Return type:

RunContext

clear_context() → None

Release the current run context.

Verbs release their own context when they finish, so this is only needed to clean up after code that populated a context outside a verb run, such as a test.