DeployedGraph

The DeployedGraph class orchestrates distributed execution of Falcon graphs using Ray.

Overview

DeployedGraph wraps a Graph and handles:

• one or more sample actors per node (SampleActor, named after the node), which serve prior, proposal and posterior samples;
• one train actor per estimator node (TrainActor, named <node>/train), which trains the networks and publishes each new best state to the node's sample actors;
• the simulation loop that fills the buffer while the train actors run.

Pass train=False to deploy sample actors only, as falcon sample does.

Class Reference

DeployedGraph

DeployedGraph(graph, import_dirs=None, log_config=None, train=True)

Deploy a graph as Ray actors.

Every node gets num_actors sample actors; estimator nodes also get one train actor, unless train is False (sampling only).

Note: This class uses falcon.info(), falcon.warning() etc. for logging. These functions use the module-level logger set by cli.py via set_logger().

Source code in src/falcon/core/deployed_graph.py

def __init__(self, graph, import_dirs=None, log_config=None, train=True):
    """Deploy a graph as Ray actors.

    Every node gets ``num_actors`` sample actors; estimator nodes also get
    one train actor, unless ``train`` is False (sampling only).

    Note: This class uses falcon.info(), falcon.warning() etc. for logging.
    These functions use the module-level logger set by cli.py via set_logger().
    """
    self.graph = graph
    self.import_dirs = import_dirs or []
    self.log_config = log_config or {}
    self.train = train
    self.samplers: Dict[str, SamplePool] = {}
    self.trainers: Dict[str, object] = {}  # node name -> TrainActor handle
    self._dataset_manager_actor = None

    self.deploy_nodes()

deploy_nodes

deploy_nodes()

Deploy all nodes in the graph as Ray actors.

Source code in src/falcon/core/deployed_graph.py

def deploy_nodes(self):
    """Deploy all nodes in the graph as Ray actors."""
    info("Spinning up graph...")
    options = self._actor_options()
    self._check_resource_budget(options)
    threads = self._thread_counts(options)

    # Create all actors (non-blocking)
    ready = {}  # display name -> list of ready refs
    labels = {}  # display name -> thread note
    for node in self.graph.node_list:
        sample_opts, train_opts = options[node.name]
        sample_threads, train_threads = threads[node.name]
        actors = []
        for i in range(node.num_actors):
            opts = dict(sample_opts)
            log_name = node.name if i == 0 else f"{node.name}/{i}"
            if i > 0 and "name" in opts:
                opts["name"] = f"{opts['name']}/{i}"
            actors.append(SampleActor.options(**opts).remote(
                node, self.import_dirs, self.log_config, log_name, sample_threads
            ))
        self.samplers[node.name] = SamplePool(actors)
        ready[node.name] = [a.__ray_ready__.remote() for a in actors]
        if sample_threads is not None:
            labels[node.name] = f" ({sample_threads} CPU threads)"
        if train_opts is not None:
            trainer = TrainActor.options(**train_opts).remote(
                node, self.import_dirs, self.log_config, train_threads
            )
            self.trainers[node.name] = trainer
            ready[train_stream_name(node.name)] = [trainer.__ray_ready__.remote()]
            labels[train_stream_name(node.name)] = f" ({train_threads} CPU threads)"

    # Wait for all actors to initialize
    for name, refs in ready.items():
        try:
            done, _ = ray.wait(refs, num_returns=len(refs), timeout=60.0)
            if len(done) < len(refs):
                raise RuntimeError(
                    f"Actor '{name}' did not initialize within 60 s. "
                    "This usually means Ray cannot schedule the actor — "
                    "check that ray.num_train_gpus / num_sample_gpus / num_cpus in your "
                    "config do not exceed available cluster resources."
                )
            ray.get(done)  # re-raise any actor-side exception
            info(f"  ✓ {name}{labels.get(name, '')}")
        except ray.exceptions.RayActorError as e:
            raise RuntimeError(f"Failed to initialize actor '{name}': {e}") from e

sample

sample(num_samples, conditions=None)

Run forward sampling through the graph.

Parameters:

Name	Type	Description	Default
num_samples		Number of samples to generate	required
conditions		Optional dict of pre-set conditions (arrays/tensors)	None

Returns:

Type	Description
	List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values

Source code in src/falcon/core/deployed_graph.py

def sample(self, num_samples, conditions=None):
    """Run forward sampling through the graph.

    Args:
        num_samples: Number of samples to generate
        conditions: Optional dict of pre-set conditions (arrays/tensors)

    Returns:
        List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values
    """
    condition_refs = self._arrays_to_condition_refs(conditions, num_samples) if conditions else {}
    return self._execute_graph(
        num_samples, self.graph.forward_order, condition_refs, "sample",
    )

sample_posterior

sample_posterior(num_samples, conditions=None)

Run posterior sampling through the inference graph.

Parameters:

Name	Type	Description	Default
num_samples		Number of samples to generate	required
conditions		Optional dict of pre-set conditions (arrays/tensors)	None

Returns:

Type	Description
	List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values

Source code in src/falcon/core/deployed_graph.py

def sample_posterior(self, num_samples, conditions=None):
    """Run posterior sampling through the inference graph.

    Args:
        num_samples: Number of samples to generate
        conditions: Optional dict of pre-set conditions (arrays/tensors)

    Returns:
        List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values
    """
    condition_refs = self._arrays_to_condition_refs(conditions, num_samples) if conditions else {}
    return self._execute_graph(
        num_samples, self.graph.backward_order, condition_refs, "sample_posterior",
    )

sample_proposal

sample_proposal(num_samples, conditions=None)

Run proposal sampling through the inference graph.

Parameters:

Name	Type	Description	Default
num_samples		Number of samples to generate	required
conditions		Optional dict of pre-set conditions (arrays/tensors)	None

Returns:

Type	Description
	List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values

Source code in src/falcon/core/deployed_graph.py

def sample_proposal(self, num_samples, conditions=None):
    """Run proposal sampling through the inference graph.

    Args:
        num_samples: Number of samples to generate
        conditions: Optional dict of pre-set conditions (arrays/tensors)

    Returns:
        List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values
    """
    condition_refs = self._arrays_to_condition_refs(conditions, num_samples) if conditions else {}
    return self._execute_graph(
        num_samples, self.graph.backward_order, condition_refs, "sample_proposal",
    )

sample_ppd

sample_ppd(num_samples, conditions=None)

Run posterior predictive distribution (PPD) sampling.

Two-phase: sample latent variables from the posterior, then forward-simulate observables from those posterior samples.

Parameters:

Name	Type	Description	Default
num_samples		Number of samples to generate	required
conditions		Observations dict (same as passed to sample_posterior)	None

Returns:

Type	Description
	List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values (both posterior latents and forward-simulated observables)

Source code in src/falcon/core/deployed_graph.py

def sample_ppd(self, num_samples, conditions=None):
    """Run posterior predictive distribution (PPD) sampling.

    Two-phase: sample latent variables from the posterior, then forward-simulate
    observables from those posterior samples.

    Args:
        num_samples: Number of samples to generate
        conditions: Observations dict (same as passed to sample_posterior)

    Returns:
        List[Dict[str, ObjectRef]]: One dict per sample with refs to all node values
            (both posterior latents and forward-simulated observables)
    """
    observation_refs = self._arrays_to_condition_refs(conditions, num_samples) if conditions else {}

    # Phase 1: theta ~ p(theta | x_obs)
    posterior_refs = self._execute_graph(
        num_samples, self.graph.backward_order, observation_refs, "sample_posterior",
    )

    # Phase 2: x_ppd ~ p(x | theta)  — forward-simulate fresh observables
    # Condition on posterior theta; observed nodes are NOT pre-set here so they
    # get re-simulated rather than returning the original observations.
    posterior_condition_refs = self._extract_value_refs(posterior_refs)
    forward_refs = self._execute_graph(
        num_samples, self.graph.forward_order, posterior_condition_refs, "sample",
    )

    # Merge: posterior_refs holds theta.value; forward_refs holds x_ppd.value.
    # Conditioned (theta) nodes are absent from forward_refs due to the skip in
    # _execute_graph, so update is safe — no key collisions.
    merged = [dict(p) for p in posterior_refs]
    for i, fwd_dict in enumerate(forward_refs):
        merged[i].update(fwd_dict)
    return merged

get_status

get_status(timeout=2.0)

Status of all actors, keyed by actor name (z, z/train), and of the buffer.

Source code in src/falcon/core/deployed_graph.py

def get_status(self, timeout: float = 2.0) -> dict:
    """Status of all actors, keyed by actor name (``z``, ``z/train``), and of the buffer."""
    refs = {}
    for node in self.graph.node_list:
        refs[node.name] = self.samplers[node.name].status_ref()
        if node.name in self.trainers:
            refs[train_stream_name(node.name)] = self.trainers[node.name].get_status.remote()
    nodes = {}
    for name, ref in refs.items():
        try:
            nodes[name] = ray.get(ref, timeout=timeout)
        except Exception as e:
            nodes[name] = {"status": "error", "error": str(e)}
    buffer = {}
    if self._dataset_manager_actor is not None:
        try:
            buffer = ray.get(self._dataset_manager_actor.get_store_stats.remote(), timeout=timeout)
        except Exception:
            pass
    return {"nodes": nodes, "buffer": buffer}

shutdown

shutdown()

Shut down the deployed graph and release resources.

Source code in src/falcon/core/deployed_graph.py

def shutdown(self):
    """Shut down the deployed graph and release resources."""
    refs = [t.shutdown.remote() for t in self.trainers.values()]
    for pool in self.samplers.values():
        refs.extend(pool.shutdown_refs())
    ray.get(refs)

launch

launch(dataset_manager, observations, graph_path=None, stop_check=None)

Launch training.

Parameters:

Name	Type	Description	Default
dataset_manager		Dataset manager for samples	required
observations		Observation data	required
graph_path		Path to save/load graph	None
stop_check		Optional callable that returns True when graceful stop is requested	None

Source code in src/falcon/core/deployed_graph.py

def launch(self, dataset_manager, observations, graph_path=None, stop_check=None):
    """Launch training.

    Args:
        dataset_manager: Dataset manager for samples
        observations: Observation data
        graph_path: Path to save/load graph
        stop_check: Optional callable that returns True when graceful stop is requested
    """
    self._launch(dataset_manager, observations, graph_path=graph_path, stop_check=stop_check)

save

save(graph_dir)

Save the checkpoints of all trained nodes.

Source code in src/falcon/core/deployed_graph.py

def save(self, graph_dir):
    """Save the checkpoints of all trained nodes."""
    graph_dir = Path(graph_dir).expanduser().resolve()
    graph_dir.mkdir(parents=True, exist_ok=True)
    ray.get([trainer.save.remote(graph_dir / name) for name, trainer in self.trainers.items()])

load

load(graph_dir)

Load the checkpoints of all estimator nodes into their actors.

Source code in src/falcon/core/deployed_graph.py

def load(self, graph_dir):
    """Load the checkpoints of all estimator nodes into their actors."""
    info(f"Loading deployed graph from: {graph_dir}")
    refs = []
    for node in self.graph.node_list:
        if node.estimator_cls is None:
            continue
        node_dir = Path(graph_dir) / node.name
        refs.extend(a.load.remote(node_dir) for a in self.samplers[node.name].actors)
        if node.name in self.trainers:
            refs.append(self.trainers[node.name].load.remote(node_dir))
    ray.get(refs)

resolve_actor_options

resolve_actor_options(actor_config, has_estimator)

Split a node's ray: config into sample-actor and train-actor options.

num_sample_gpus / num_train_gpus set the GPUs of each actor, and num_sample_cpus / num_train_cpus their CPUs (default: num_cpus); every other Ray option applies to both. The deprecated num_gpus is split evenly between the two actors, so the node keeps its GPU budget.

Returns:

Type	Description
Tuple[dict, Optional[dict], List[str]]	(sample_options, train_options or None, warnings)

Source code in src/falcon/core/deployed_graph.py

def resolve_actor_options(actor_config: dict, has_estimator: bool) -> Tuple[dict, Optional[dict], List[str]]:
    """Split a node's ``ray:`` config into sample-actor and train-actor options.

    ``num_sample_gpus`` / ``num_train_gpus`` set the GPUs of each actor, and
    ``num_sample_cpus`` / ``num_train_cpus`` their CPUs (default: ``num_cpus``);
    every other Ray option applies to both. The deprecated ``num_gpus`` is
    split evenly between the two actors, so the node keeps its GPU budget.

    Returns:
        (sample_options, train_options or None, warnings)
    """
    cfg = dict(actor_config or {})
    sample_gpus = cfg.pop("num_sample_gpus", None)
    train_gpus = cfg.pop("num_train_gpus", None)
    legacy_gpus = cfg.pop("num_gpus", None)
    sample_cpus = cfg.pop("num_sample_cpus", None)
    train_cpus = cfg.pop("num_train_cpus", None)
    notes = []

    if legacy_gpus is not None:
        if sample_gpus is not None or train_gpus is not None:
            raise ValueError(
                "ray.num_gpus cannot be combined with ray.num_sample_gpus / ray.num_train_gpus; "
                "use only the latter"
            )
        if has_estimator:
            sample_gpus = train_gpus = legacy_gpus / 2
            notes.append(
                f"ray.num_gpus is deprecated; using num_train_gpus={train_gpus:g} and "
                f"num_sample_gpus={sample_gpus:g}"
            )
        else:
            sample_gpus = legacy_gpus
            notes.append(f"ray.num_gpus is deprecated; using num_sample_gpus={sample_gpus:g}")

    if not has_estimator:
        for key, value in (("num_train_gpus", train_gpus), ("num_train_cpus", train_cpus)):
            if value is not None:
                notes.append(f"ray.{key} is ignored: the node has no estimator")

    unknown = sorted(set(cfg) - _RAY_ACTOR_KEYS)
    if unknown:
        notes.append(f"ignoring unsupported ray options {unknown}")
    common = {k: v for k, v in cfg.items() if k in _RAY_ACTOR_KEYS}

    sample = dict(common)
    if sample_gpus is not None:
        sample["num_gpus"] = sample_gpus
    if sample_cpus is not None:
        sample["num_cpus"] = sample_cpus
    if not has_estimator:
        return sample, None, notes

    train = dict(common)
    if train_gpus is not None:
        train["num_gpus"] = train_gpus
    if train_cpus is not None:
        train["num_cpus"] = train_cpus
    if "name" in common:
        train["name"] = train_stream_name(common["name"])
    return sample, train, notes