pygx.algo.early_stopping

Early stopping policies for tuning loops.

early_stopping

Early stopping policies for tuning loops.

Early stopping policies are evaluated against measurements reported during a trial and decide whether the trial should be abandoned before completion. They are typically passed to pg.sample via the early_stopping_policy argument.

This module exposes:

• StepWise — abandon a trial based on step-keyed predicates over its reported measurements.
• early_stop_by_value / early_stop_by_rank — convenience constructors for the two most common StepWise variants.
• And / Or / Not — Boolean combinators for composing policies.

And

And(*children: EarlyStoppingPolicy, **kwargs)

Bases: Composite

Logical AND as a composite early stopping policy.

Source code in pygx/algo/early_stopping/_base.py

def __init__(self, *children: pg.tuning.EarlyStoppingPolicy, **kwargs):
    """Accept children as positional varargs: `And(p1, p2, p3)`."""
    if children:
        kwargs.setdefault('children', list(children))
    super().__init__(**kwargs)

Composite

Composite(*children: EarlyStoppingPolicy, **kwargs)

Bases: EarlyStoppingPolicyBase

Base class for composite early stopping policies.

Accept children as positional varargs: And(p1, p2, p3).

Source code in pygx/algo/early_stopping/_base.py

def __init__(self, *children: pg.tuning.EarlyStoppingPolicy, **kwargs):
    """Accept children as positional varargs: `And(p1, p2, p3)`."""
    if children:
        kwargs.setdefault('children', list(children))
    super().__init__(**kwargs)

EarlyStoppingPolicyBase

EarlyStoppingPolicyBase(
    *,
    allow_partial: bool = False,
    sealed: bool | None = None,
    root_path: KeyPath | None = None,
    explicit_init: bool = False,
    **kwargs: Any
)

Bases: EarlyStoppingPolicy

An early stopping policy base class that supports composition.

Source code in pygx/symbolic/_object.py

@_filter_typing_traceback_on_raise
def __init__(
    self,
    *,
    allow_partial: bool = False,
    sealed: bool | None = None,
    root_path: topology.KeyPath | None = None,
    explicit_init: bool = False,
    **kwargs: Any,
):
    """Create an Object instance.

    ``pg.Object`` synthesizes a keyword-only ``__init__``. Subclasses
    that need positional arguments must override ``__init__``
    explicitly and translate to keyword arguments before forwarding
    to ``super().__init__(**kwargs)``.

    Args:
      allow_partial: If True, the object can be partial.
      sealed: If True, seal the object from future modification (unless under
        a `pg.as_sealed(False)` context manager). If False, treat the object as
        unsealed. If None, it's determined by the class's ``frozen`` option.
      root_path: The symbolic path for current object. By default it's None,
        which indicates that newly constructed object does not have a parent.
      explicit_init: Should set to `True` when `__init__` is called via
        `pg.Object.__init__` instead of `super().__init__`.
      **kwargs: key/value arguments that align with the schema. All required
        keys in the schema must be specified, and values should be acceptable
        according to their value spec.

    Raises:
      TypeError: When a required field is missing, a value has an
        unacceptable type, or an unknown keyword argument is given.
      ValueError: When a value violates its value spec's constraints
        (e.g. out of range).
    """
    # Placeholder for Google-internal usage instrumentation.

    cls = self.__class__
    # The ever-instantiated latch `apply_schema`'s post-instance guard
    # reads (#432: `update_schema` is only legal before instances exist).
    # Own-`__dict__`, latched on the construction ATTEMPT — mirroring the
    # native core, whose `__native_ctor_meta__` (the equivalent marker) is
    # created on the first instance op. One dict read per construct.
    meta = cls.__sym_meta__
    if not meta.instantiated:
        meta.instantiated = True
    if sealed is None:
        sealed = cls.__sym_options__.frozen

    if not isinstance(allow_partial, bool):
        raise TypeError(
            f"Expect bool type for argument 'allow_partial' in "
            f'symbolic.Object.__init__ but encountered {allow_partial}.'
        )

    # Fast path: const-keyed schema.
    if cls._sym_init_fields is not None and self._init_fast_path(
        kwargs,
        allow_partial,
        sealed,
        root_path,
        explicit_init,
    ):
        return

    # Slow path: handles varargs and schemas with non-const keys.

    # Model before-hook (#516) — the fast path was not taken
    # (`_sym_init_fields is None`), so the hook has not run yet.
    if cls._sym_preinit_overridden:
        kwargs = _apply_preinit(cls, kwargs)

    # We delay the seal attempt until members are all set.
    super().__init__(
        allow_partial=allow_partial,
        accessor_writable=cls.__sym_options__.attr_write,
        sealed=sealed,
        root_path=root_path,
        init_super=not explicit_init,
    )

    # Fill field_args and init_args from **kwargs.
    _, unmatched_keys = cls.__schema__.resolve(list(kwargs.keys()))
    # Fields opted out of `__init__` via `enable_init=False` resolve
    # against the schema but must still be rejected as kwargs.
    get_field = cls.__schema__.get_field
    unmatched_keys = list(unmatched_keys) + [
        k
        for k in kwargs
        if (f := get_field(k)) is not None and not f.enable_init
    ]
    if unmatched_keys:
        _screen_unexpected_init_keys(cls, kwargs, unmatched_keys)

    # Declared aliases are RESERVED wire spellings (#517, owner ruling):
    # an open (extra) schema must not accept an extra kwarg equal to a
    # field's alias — the wire would become ambiguous. (Const-keyed
    # schemas never reach here with an alias-spelled kwarg; their
    # unknown-key screen already raised — or, under `extra='ignore'`,
    # raised for alias keys before dropping the rest.)
    if cls._sym_alias_maps is not None:
        alias_to_name = cls._sym_alias_maps[0]
        for k in kwargs:
            if k in alias_to_name:
                raise TypeError(
                    f'{cls.__name__}.__init__() got keyword argument '
                    f'{k!r}, which is reserved as the alias of field '
                    f'{alias_to_name[k]!r}.'
                )

    # All field values come from kwargs — `pg.Object.__init__` is
    # keyword-only.
    field_args = dict(kwargs)

    # Check missing arguments when partial binding is disallowed.
    if not base.accepts_partial(self):
        missing_args = []
        for field in self.__class__.__schema__.fields.values():
            if (
                not field.value.has_default
                and isinstance(field.key, pg_typing.ConstStrKey)
                and field.key not in field_args
            ):
                missing_args.append(str(field.key))
        if missing_args:
            arg_phrase = formatting.auto_plural(
                len(missing_args), 'argument'
            )
            keys_str = formatting.comma_delimited_str(missing_args)
            raise TypeError(
                f'{self.__class__.__name__}.__init__() missing {len(missing_args)} '
                f'required {arg_phrase}: {keys_str}.'
            )

    # NOTE(daiyip): Accessor writable is honored by `Object.__setattr__` so
    # we can always make `_sym_attributes` accessor writable. This prevents
    # a child object's attribute access from being changed when it's
    # attached to a parent whose symbolic attributes are not writable.
    # Per-field validation is gated by `field.enable_validation` inside
    # `Schema.apply` / `Dict._formalized_value` — so a class-wide
    # `validate=False` (resolved into each field) and per-field overrides
    # both flow through naturally here.
    vars(self)['_sym_attributes'] = pg_dict.Dict(
        field_args,  # pyright: ignore[reportArgumentType]
        value_spec=self.__class__.sym_fields,
        allow_partial=allow_partial,
        sealed=sealed,
        accessor_writable=True,
        root_path=root_path,
        as_object_attributes_container=True,
        under_topo=cls.__sym_options__.topo,
    )
    # The write-time `needs_topo` refusal (see `_init_fast_path` Phase 5):
    # the slow path formalizes inside the container, so scan its items.
    if not cls.__sym_options__.topo:
        for k, v in self._sym_attributes.sym_items():
            base.refuse_tree_walking_inferential(cls, k, v)
    self._sym_attributes.topo_setparent(self)
    # Seed `init=False` fields from their defaults before `on_sym_post_init`
    # runs so `on_sym_ready`/`on_sym_bound` overrides start from a known baseline.
    _seed_non_symbolic_fields(self, cls._non_symbolic_fields)
    # Model after-validator (#516): fires before any lifecycle hook can
    # observe the object, gated like `on_sym_ready` (concrete only).
    if cls._sym_validate_overridden:
        abstract = (
            self._allow_partial and self.sym_partial
        ) or self.sym_puresymbolic
        if not abstract:
            self.on_sym_validate()
    self.on_sym_post_init()
    self.sym_seal(sealed)

Not

Not(*children: EarlyStoppingPolicy, **kwargs)

Bases: Composite

Logical NOT as a composite early stopping policy.

Source code in pygx/algo/early_stopping/_base.py

def __init__(self, *children: pg.tuning.EarlyStoppingPolicy, **kwargs):
    """Accept children as positional varargs: `And(p1, p2, p3)`."""
    if children:
        kwargs.setdefault('children', list(children))
    super().__init__(**kwargs)

Or

Or(*children: EarlyStoppingPolicy, **kwargs)

Bases: Composite

Logical OR as a composite early stopping policy.

Source code in pygx/algo/early_stopping/_base.py

def __init__(self, *children: pg.tuning.EarlyStoppingPolicy, **kwargs):
    """Accept children as positional varargs: `And(p1, p2, p3)`."""
    if children:
        kwargs.setdefault('children', list(children))
    super().__init__(**kwargs)

StepWise

StepWise(plan: list[tuple[int, Any]], **kwargs)

Bases: EarlyStoppingPolicyBase

Early stopping policy driven by a step-keyed plan.

The plan is a list of (gating step, predicate) tuples. When a trial reports a measurement at or beyond a gating step, that gate's predicate is called with the trial's latest measurement and the list of measurements that earlier trials reported at the same gate; returning True abandons the trial. Each trial is decided at most once per gate.

See early_stop_by_value and early_stop_by_rank for ready-made plans.

Source code in pygx/algo/early_stopping/_step_wise.py

def __init__(self, plan: list[tuple[int, Any]], **kwargs):
    super().__init__(**dict(plan=plan), **kwargs)

should_stop_early

should_stop_early(trial: Trial) -> bool

Returns True if a trial should be stopped early.

Source code in pygx/algo/early_stopping/_step_wise.py

def should_stop_early(self, trial: pg.tuning.Trial) -> bool:
    """Returns True if a trial should be stopped early."""
    if not trial.measurements:
        return False

    should_stop = False
    gate_index = self._get_gate_index(trial)
    if gate_index >= 0:
        if trial.id in self._trial_gate_decision:
            decision_gate_index, decision = self._trial_gate_decision[
                trial.id
            ]
            if decision_gate_index == gate_index:
                return decision
        gate_predicate = self.plan[gate_index][1]
        gate_history = self._gate_history[gate_index]
        m = trial.measurements[-1]
        if gate_predicate(m, gate_history):
            should_stop = True
        gate_history.append(m)
        self._trial_gate_decision[trial.id] = (gate_index, should_stop)
    return should_stop

recover

recover(history: Iterable[Trial])

Recovers the policy state based on history.

Source code in pygx/algo/early_stopping/_step_wise.py

def recover(self, history: Iterable[pg.tuning.Trial]):
    """Recovers the policy state based on history."""
    for t in history:
        prev_m = None
        next_gate = 0
        for i, m in enumerate(t.measurements):
            if (
                next_gate != len(self.plan)
                and (
                    prev_m is None or prev_m.step < self.plan[next_gate][0]
                )
                and m.step >= self.plan[next_gate][0]
            ):
                stopping_decision = False
                gate_history = self._gate_history[next_gate]
                if i == len(t.measurements) - 1:
                    if t.status == 'COMPLETED':
                        stopping_decision = t.infeasible
                    elif t.status == 'PENDING':
                        # For the last measurement of a pending trial, we use gate
                        # predicate to determine the stopping decision.
                        gate_predicate = self.plan[next_gate][1]
                        stopping_decision = gate_predicate(m, gate_history)
                gate_history.append(m)
                self._trial_gate_decision[t.id] = (
                    next_gate,
                    stopping_decision,
                )
                next_gate += 1
            prev_m = m

early_stop_by_rank

early_stop_by_rank(
    step_ranks: list[tuple[int, float | int, int]],
    metric: str | Callable[[Measurement], float] = "reward",
    maximize: bool = True,
) -> StepWise

Step-wise early stopping policy based on the rank of reward/metric.

Example:

policy = early_stop_by_rank([
  # Stop at step 1 if accuracy is less than top 80% previous trials at
  # this step, enabled when there are at least 5 previous trials reported
  # at this step.
  (1, 0.8, 5),

  # Stop at step 2 if accuracy is less than top 20% previous trials at
  # this step, enabled when there are at least 10 previous trials reported
  # at this step.
  (2, 0.2, 10),

  # Stop at step 3 if accuracy is less than the 3rd best trial at this step,
  # enabled when there are at least 3 previous trials reported at this step.
  (3, 3, 3)
], metric='accuracy')()

Parameters:

Name	Type	Description	Default
step_ranks	list[tuple[int, float | int, int]]	A list of tuple (gating step, rank threshold, trigger histogram size). gating step - At which step this rule will be triggered. rank threshold - A float number in range [0.0, 1.0] indicating the rank percentage or an integer (> 0) indicating the absolute rank as the threshold for early stopping. trigger histogram size - The minimal number of historical trials reported at current step for this rule to trigger.	required
metric	str | Callable[[Measurement], float]	Based on which metric the rank will be computed. Use str for metric name or a callable object that takes a measurement object at a given step as input and returns a float value.	'reward'
maximize	bool	If True, trials whose value ranks below the threshold are stopped; otherwise trials whose value ranks above the threshold are stopped.	True

Returns:

Type	Description
StepWise	A StepWise early stopping policy.

Source code in pygx/algo/early_stopping/_step_wise.py

@pg.symbolize
def early_stop_by_rank(
    step_ranks: list[
        tuple[
            int,  # Gating step.
            (
                float  # Rank percentage in (0.0, 1.0).
                | int
            ),  # Absolute rank, below which will be stopped.
            int,
        ]
    ],  # Min histogram size at the step to trigger stopping.
    metric: str | Callable[[pg.tuning.Measurement], float] = 'reward',
    maximize: bool = True,
) -> StepWise:
    """Step-wise early stopping policy based on the rank of reward/metric.

    **Example:**

    ```python
    policy = early_stop_by_rank([
      # Stop at step 1 if accuracy is less than top 80% previous trials at
      # this step, enabled when there are at least 5 previous trials reported
      # at this step.
      (1, 0.8, 5),

      # Stop at step 2 if accuracy is less than top 20% previous trials at
      # this step, enabled when there are at least 10 previous trials reported
      # at this step.
      (2, 0.2, 10),

      # Stop at step 3 if accuracy is less than the 3rd best trial at this step,
      # enabled when there are at least 3 previous trials reported at this step.
      (3, 3, 3)
    ], metric='accuracy')()
    ```

    Args:
      step_ranks: A list of tuple (gating step, rank threshold, trigger histogram
        size).
          gating step - At which step this rule will be triggered.
          rank threshold - A float number in range [0.0, 1.0] indicating the
            rank percentage or an integer (> 0) indicating the absolute rank
            as the threshold for early stopping.
          trigger histogram size - The minimal number of historical trials
            reported at current step for this rule to trigger.
      metric: Based on which metric the rank will be computed.
        Use str for metric name or a callable object that takes a measurement
        object at a given step as input and returns a float value.
      maximize: If True, trials whose value ranks below the threshold are
        stopped; otherwise trials whose value ranks above the threshold are
        stopped.

    Returns:
      A `StepWise` early stopping policy.
    """
    assert isinstance(step_ranks, list), step_ranks
    for v in step_ranks:
        if (
            not isinstance(v, tuple)
            or len(v) != 3
            or not isinstance(v[0], int)
            or not isinstance(v[1], (int, float))
            or not isinstance(v[2], int)
        ):
            raise ValueError(
                f'Invalid definition in `step_ranks`: {v}. '
                f'Expect a tuple of 3 elements: '
                f'(step: int, rank: Union[float, int], trigger_size: int).'
            )
        if isinstance(v[1], float) and (v[1] > 1 or v[1] < 0):
            raise ValueError(
                f'Rank must be within range [0.0, 1.0] when it is percentage '
                f'(float). Encountered: {v[1]} in {v}.'
            )

    def _cmp(x, y) -> bool:
        if y is None:
            return False
        return x < y if maximize else x > y

    def _value(m: pg.tuning.Measurement) -> float:
        if isinstance(metric, str):
            return m.reward if metric == 'reward' else m.metrics[metric]
        assert callable(metric), metric
        return metric(m)

    def _value_by_rank(
        h: list[pg.tuning.Measurement], threshold: int | float
    ) -> float | None:
        if isinstance(threshold, float):
            assert 0.0 <= threshold <= 1.0, threshold
            k = int((len(h) - 1) * threshold)
        else:
            assert isinstance(threshold, int)
            k = threshold - 1
            if k < -len(h) or k >= len(h):
                return None
        return sorted([_value(r) for r in h], reverse=maximize)[k]

    def _make_predicate(rank: int | float, trigger_size: int):
        def _predicate(
            m: pg.tuning.Measurement, history: list[pg.tuning.Measurement]
        ):
            if not history or len(history) < trigger_size:
                return False
            return _cmp(_value(m), _value_by_rank(history, rank))

        return _predicate

    return StepWise(
        plan=[
            (step, _make_predicate(rank, trigger_size))
            for step, rank, trigger_size in step_ranks
        ]
    )

early_stop_by_value

early_stop_by_value(
    step_values: list[tuple[int, float]],
    metric: str | Callable[[Measurement], float] = "reward",
    maximize: bool = True,
) -> StepWise

Step-wise early stopping policy based on the value of reward/metric.

Example:

policy = early_stop_by_value([
  # Stop at step 1 if trial reward is less than 0.2.
  (1, 0.2),

  # Stop at step 2 if trial reward is less than 0.8.
  (2, 0.8),
])()

Parameters:

Name	Type	Description	Default
step_values	list[tuple[int, float]]	A list of tuple (gating step, value threshold). gating step - At which step this rule will be triggered. value threshold - A float number indicating the threshold value for early stopping.	required
metric	str | Callable[[Measurement], float]	Based on which metric the value should be compared against. Use str for metric name or a callable object that takes a measurement object at a given step as input and returns a float value.	'reward'
maximize	bool	If True, trials whose value is below the threshold are stopped; otherwise trials whose value is above the threshold are stopped.	True

Returns:

Type	Description
StepWise	A StepWise early stopping policy.

Source code in pygx/algo/early_stopping/_step_wise.py

@pg.symbolize
def early_stop_by_value(
    step_values: list[
        tuple[
            int,  # Gating step.
            float,
        ]
    ],  # Value threshold.
    metric: str | Callable[[pg.tuning.Measurement], float] = 'reward',
    maximize: bool = True,
) -> 'StepWise':
    """Step-wise early stopping policy based on the value of reward/metric.

    **Example:**

    ```python
    policy = early_stop_by_value([
      # Stop at step 1 if trial reward is less than 0.2.
      (1, 0.2),

      # Stop at step 2 if trial reward is less than 0.8.
      (2, 0.8),
    ])()
    ```

    Args:
      step_values: A list of tuple (gating step, value threshold).
          gating step - At which step this rule will be triggered.
          value threshold - A float number indicating the threshold value for
            early stopping.
      metric: Based on which metric the value should be compared against.
        Use str for metric name or a callable object that takes a measurement
        object at a given step as input and returns a float value.
      maximize: If True, trials whose value is below the threshold are
        stopped; otherwise trials whose value is above the threshold are
        stopped.

    Returns:
      A `StepWise` early stopping policy.
    """
    assert isinstance(step_values, list), step_values
    for v in step_values:
        if (
            not isinstance(v, tuple)
            or len(v) != 2
            or not isinstance(v[0], int)
            or not isinstance(v[1], numbers.Number)
        ):
            raise ValueError(
                f'Invalid definition in `step_values`: {v}. '
                f'Expect a tuple of 2 elements: '
                f'(step: int, threshold: float).'
            )

    def _cmp(x, y) -> bool:
        return x < y if maximize else x > y

    def _value(m: pg.tuning.Measurement) -> float:
        if isinstance(metric, str):
            return m.reward if metric == 'reward' else m.metrics[metric]
        assert callable(metric), metric
        return metric(m)

    def _make_predicate(threshold: float):
        def _predicate(m: pg.tuning.Measurement, unused_history):
            v = _value(m)
            ret = _cmp(v, threshold)
            return ret

        return _predicate

    return StepWise(
        plan=[
            (step, _make_predicate(threshold))
            for step, threshold in step_values
        ]
    )

options: show_root_heading: false show_root_toc_entry: false members_order: source heading_level: 2