diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml
index 6654e6fee..71be16006 100644
--- a/agent_context/MAP.yaml
+++ b/agent_context/MAP.yaml
@@ -430,6 +430,7 @@ topics:
DepthSidecarManager, depth_meta.json, fragment_id, finalize, expert action schema, qpos_slice, qvel_slice,
validate_dataset, create_split_manifest, recover_recording, episode_uuid, language_persistent,
SharedLanguageRegistry, resolve_language, configure_ods, ODSReuseProfileCfg,
+ ODSFlowController, ODSFlowControllerCfg, get_flow_status, ODS 供数控制,
episode_sample_budget, OnlineDataBudgetExhaustedError, ODS 自动配置, BatchSource, SequenceSpecCfg,
OnlineDataSource, TensorDataSource, LeRobotDataSource, OfflineDataset, SequenceDataset,
DataUnavailableError, validate_data_schema, data_schema_fingerprint, ChunkBC, SimpleBCPickPlace, 数据集加载, 离线加载]
@@ -438,6 +439,7 @@ topics:
- embodichain/data_pipeline/engine/data.py
- embodichain/data_pipeline/engine/language.py
- embodichain/data_pipeline/configuration.py
+ - embodichain/data_pipeline/runtime.py
- embodichain/data_pipeline/datasets/online_data.py
- embodichain/data_pipeline/datasets/sequence.py
- embodichain/data_pipeline/datasets/sources.py
diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md
index dceae966b..6d9352adf 100644
--- a/agent_context/topics/data-pipeline/data-pipeline.md
+++ b/agent_context/topics/data-pipeline/data-pipeline.md
@@ -11,6 +11,7 @@ belongs to [data-assets](../data-assets/data-assets.md).
|---|---|
| Online worker and sample selection | `embodichain/data_pipeline/engine/data.py`: `OnlineDataEngine` |
| Throughput-derived settings and protected reuse profiles | `embodichain/data_pipeline/configuration.py`: `configure_ods`; sampler enforcement in `engine/data.py` |
+| Explicit owner-side refill, waits and pacing | `embodichain/data_pipeline/runtime.py`: `ODSFlowController`; compact READY-only flow status in `engine/data.py` |
| PyTorch iteration | `embodichain/data_pipeline/datasets/`: `sequence_data.py` owns `SequenceDataset`; `online_data.py` preserves `OnlineDataset`; `OfflineDataset` is an alias |
| Shared sequence contract and sources | `embodichain/data_pipeline/datasets/sequence.py`, `sources.py`, `_sampling.py` |
| Bounded numeric LeRobot loading | `embodichain/data_pipeline/datasets/lerobot.py`: `LeRobotDataSource`, explicit field mapping |
@@ -82,7 +83,9 @@ action contracts independently of the expert-controller schema; see
Distinguish a terminal worker failure from a sampling request with no eligible
unlocked window. The latter raises `DataUnavailableError` (a `RuntimeError`)
immediately without itself setting `FAILED` or triggering a refill. Datasets
-propagate both kinds of exception without an automatic retry loop.
+propagate both kinds of exception without an automatic retry loop. Applications
+can opt into `ODSFlowController` around their manual engine for bounded retries,
+refill and plan-based pacing; start/stop and training stay application-owned.
## Persistence contract
diff --git a/agent_context/topics/data-pipeline/online-sampling.md b/agent_context/topics/data-pipeline/online-sampling.md
index bf2d47a85..f39fc0102 100644
--- a/agent_context/topics/data-pipeline/online-sampling.md
+++ b/agent_context/topics/data-pipeline/online-sampling.md
@@ -51,6 +51,39 @@
so DataLoader can collate; batch mode re-samples a dynamic chunk per yield. Multi-worker
readers sample shared memory independently.
+## Explicit owner-side flow control
+
+`runtime.py` owns opt-in `ODSFlowController` around a caller-managed manual
+engine. Its creating process/thread calls the controller; the application still
+starts/stops the engine and owns training, evaluation and resource accounting.
+An optional `ODSConfiguration` must match engine cache, producer count, quotas
+and batch size. A plan or production limit requires an explicit positive
+`gym_config.num_envs` dividing cache size; unplanned/unlimited legacy engines
+retain their existing layout support.
+With a plan, `prepare_batch()` schedules quota-threshold refill, bounded quota
+waits and batch-start pacing; a pause does not accumulate catch-up credits.
+Without a plan it only polls, leaving round/frozen refill decisions explicit.
+The optional accepted-episode limit covers initial fill and whole-cache refills;
+partial-cache remainder is not produced, and direct engine requests bypass it.
+Already pending generations are adopted without a new engine request or refill
+callback; only requests made through the controller receive request timestamps.
+
+Controller waits drain lifecycle events and check health/caller deadlines.
+`sample_context()` retries only live no-window/quota exceptions within a total
+bound; Source/Dataset behavior stays immediate. Generation waits can accept an
+absolute monotonic deadline that includes earlier work. Callback interval/refill
+times use absolute host monotonic seconds; the application aligns its own clock.
+The controller does not restore source or pacing state. The
+[configuration guide](../../../docs/source/features/ods/configuration.md#apply-the-plan)
+owns the integration example and callback details.
+
+`engine.check_health()` is READY-only and checks existing worker errors/liveness
+without sampling/refill. `get_flow_status()` returns compact generation and
+optional quota counters without per-row audit records or payload clones.
+Producer counters and quota totals use separate lock snapshots. Unlocked quota
+uses the same valid/continuity and write-window test as the full audit, but is a
+coarse scheduling bound; exact context/filter eligibility stays with sampling.
+
## Optional publication accounting
`OnlineDataEngineCfg.lifecycle_event_capacity` enables a bounded shared log;
@@ -79,8 +112,9 @@ quota from no valid windows and terminal worker failure. It is immediate and
does not partially spend a batch. `get_sampling_budget_snapshot()` returns
retained identities, budgets, delivered/remaining counts and lifetime delivery;
clean stopped engines remain auditable. The owner decides when to refill, wait
-with a bound, or stop. The fixed publication budget is not an adaptive replay
-retention policy and does not restore old sources after process restart.
+with a bound, or stop, directly or through the optional flow controller. The
+fixed publication budget is not an adaptive replay retention policy and does
+not restore old sources after process restart.
`configuration.py` derives bounded budgets, cache/refill settings and pacing
targets from measured accepted episodes/s and training chunks/s. It reads an
diff --git a/docs/source/api_reference/embodichain/embodichain.data_pipeline.engine.rst b/docs/source/api_reference/embodichain/embodichain.data_pipeline.engine.rst
index 78ecc03e1..3829f7eb2 100644
--- a/docs/source/api_reference/embodichain/embodichain.data_pipeline.engine.rst
+++ b/docs/source/api_reference/embodichain/embodichain.data_pipeline.engine.rst
@@ -44,6 +44,12 @@ require a DatasetManager.
``request_refill()`` and observe publication with ``wait_for_generation()``.
The default remains sample-threshold refill. ``get_stats()`` returns detached
counters and excludes the active write window from valid-frame counts.
+``check_health()`` requires READY and surfaces worker failures without sampling
+or refill. ``get_flow_status()`` supplies compact generation, accepted-episode
+and optional quota counters for frequent scheduling checks. Its producer and
+quota counters use separate lock snapshots; unlocked quota is a coarse bound
+that does not apply every sequence/filter requirement. Use
+``get_sampling_budget_snapshot()`` for the retained-source audit.
``shutdown_timeout`` is the finite, positive graceful-close wait in seconds
(default ``5.0``). Increasing it allows longer worker teardown; it does not
change forced terminate/kill waits or their terminal failure semantics.
diff --git a/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst
index cae6fc059..bfe2e2bf7 100644
--- a/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst
+++ b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst
@@ -22,6 +22,7 @@ storage for LeRobot datasets on Python 3.10--3.12).
depth_video
engine
recording
+ runtime
Datasets
--------
@@ -58,6 +59,31 @@ Online Data Engine
data
+Owner-side flow control
+-----------------------
+
+The application starts and stops its manual-refill engine. The optional
+controller applies a measured plan, drains publication events, checks health,
+and bounds generation/data/quota waits. A plan or production limit requires an
+explicit producer count dividing cache size. Sources and Datasets retain
+immediate sampling/error behavior. See :doc:`/features/ods/configuration` for
+matching engine settings, callback clocks and an integration example.
+
+.. automodule:: embodichain.data_pipeline.runtime
+
+.. autosummary::
+
+ embodichain.data_pipeline.runtime.ODSFlowControllerCfg
+ embodichain.data_pipeline.runtime.ODSFlowController
+
+.. autoclass:: embodichain.data_pipeline.runtime.ODSFlowControllerCfg
+ :members:
+ :show-inheritance:
+
+.. autoclass:: embodichain.data_pipeline.runtime.ODSFlowController
+ :members:
+ :show-inheritance:
+
Recording integrity
-------------------
diff --git a/docs/source/features/ods/architecture.md b/docs/source/features/ods/architecture.md
index d79f9dd7f..c23e74979 100644
--- a/docs/source/features/ods/architecture.md
+++ b/docs/source/features/ods/architecture.md
@@ -17,7 +17,7 @@ flowchart TB
Stored["Resident tensors / LeRobot files"] --> Source
Source --> Learner["Learner
Direct sampling or SequenceDataset"]
Learner --> Eval["Independent student evaluation"]
- Host["Application owner
configure_ods, pacing, refill and shutdown"]
+ Host["Application owner
Plan, optional flow control and shutdown"]
Host -.-> Engine
Host -.-> Learner
```
@@ -40,6 +40,7 @@ that requests targeted collection from the producer.
| Production and shared state | `engine/data.py`: lifecycle, worker health, publication, sampling, generation accounting, provenance and optional live delivery ledger. |
| Language metadata | `engine/language.py`: append-only engine-local task/subtask dictionaries used by copied numeric language IDs. |
| Planning | `configuration.py`: typed reuse/runtime inputs, measured-rate calculations and immutable `ODSConfiguration`. The host applies the plan. |
+| Owner-side flow control | `runtime.py`: opt-in `ODSFlowController` for manual refill, bounded waits, lifecycle draining and plan-based batch pacing. The application starts/stops the engine and owns training. |
| Contract and schema | `datasets/sequence.py`, `schema.py`: `BatchSource`, `SequenceSpecCfg`, typed sampling errors, semantic declaration validation and schema fingerprints. |
| Live/resident adaptation | `datasets/sources.py`: `OnlineDataSource` delegates to its caller-owned engine; `TensorDataSource` owns resident numeric data and its optional separate ledger. |
| Numeric SDK loading | `datasets/lerobot.py`: `LeRobotDataSource`, explicit field mapping, bounded reads and a resident allocation limit. |
@@ -58,7 +59,8 @@ origin IDs also do not establish portable dataset content identity.
**Live:** the owner starts the engine and waits for initial fill. The worker
accepts task rollouts and publishes complete episodes. `OnlineDataSource`
delegates to the engine's locked sampler and returns canonical owned contexts.
-The host schedules further fills and training; shutdown belongs to the creator.
+The host schedules further fills and training, optionally through
+`ODSFlowController`; shutdown belongs to the engine creator.
Refill publishes batches of rows incrementally: readable rows can temporarily
belong to different generations. A generation filter is explicit;
`wait_for_generation()` reports completion of the entire requested fill.
@@ -84,7 +86,9 @@ training-efficiency advantage.
- `configure_ods()` is a static planner. It does not run calibration, mutate a
live engine, choose a learning-safe profile or invoke the optimizer. The
- application enforces active/average pacing and planned refill thresholds.
+ application can explicitly select `ODSFlowController` to enforce the plan's
+ refill threshold and batch-start pacing with bounded waits. The controller
+ is creator-process/thread only and does not start or stop the engine.
- `request_refill()` is owner-only and coalesces an outstanding request.
`wait_for_generation()` waits for a complete fill. Optional bounded lifecycle
logging must be drained during production and waits.
@@ -94,8 +98,9 @@ training-efficiency advantage.
there is no distributed global quota across independent sources/runs.
- `DataUnavailableError` means no eligible window exists;
`DataBudgetExhaustedError` means eligible data cannot fund the requested batch.
- Worker/read failures remain terminal errors. Dataset iteration adds no hidden
- retry, refill or budget reset.
+ The opt-in controller retries only these transient live errors up to its
+ deadline while draining events and checking health. Worker/read failures
+ propagate. Dataset iteration adds no hidden retry, refill or budget reset.
- Model/optimizer/RNG recovery belongs to the application. The demo recreates a
live cache with a new run ID; exact replay-cache or protected offline-ledger
recovery is not implemented.
diff --git a/docs/source/features/ods/configuration.md b/docs/source/features/ods/configuration.md
index 41fdbccf4..42c656eae 100644
--- a/docs/source/features/ods/configuration.md
+++ b/docs/source/features/ods/configuration.md
@@ -2,8 +2,9 @@
Use `embodichain.data_pipeline.configuration` to plan bounded data consumption
from two rates: accepted trajectory production and active learner demand.
-The result is a static configuration plan. A consumer must apply its cache and
-sampling settings and implement the returned pacing/refill policy.
+The result is a static configuration plan. A consumer applies its cache and
+sampling settings, then can select `ODSFlowController` for explicit owner-side
+pacing/refill or provide its own scheduler.
This method applies across tasks and learners. The [Franka integration](https://github.com/DexForce/EmbodiChain/blob/main/examples/ods/simple_bc_demo/docs/CONFIGURATION.md)
provides a worked binding; its measurements and reuse bounds are provisional
@@ -151,15 +152,16 @@ error; they do not silently remove protection.
| --- | --- |
| `buffer_size`, `batch_size`, `producer_envs` | Host constructs the cache/consumer at these measured settings. |
| `episode_sample_budget`, `episode_sample_hard_limit` | Pass to `OnlineDataEngineCfg`; engine samplers enforce per-source delivery under the producer lock. Resident sources have their own optional ledger. |
-| `active_chunks_per_second_cap` | Consumer limits short-window service starts. The engine does not throttle optimizer calls. |
-| `training_chunks_per_second_target` | Consumer controls average delivery/completed-training rate and reports actual waits. |
-| `refill_remaining_chunks` | Host compares remaining usable quota and schedules manual refill. |
-| `publication_episodes_per_second_target` | Host paces production/refill requests. The engine has no finite accepted-episode production budget. |
+| `active_chunks_per_second_cap`, `training_chunks_per_second_target` | `ODSFlowController.prepare_batch()` spaces batch starts at the smaller rate. The host calls it before sampling/training; actual completed throughput still needs measurement. |
+| `refill_remaining_chunks` | With a plan, the controller compares coarse unlocked quota with this threshold and schedules manual refill. |
+| `publication_episodes_per_second_target` | Production target to compare with observed publication. The controller schedules whole-cache refills by quota, without a separate per-episode time gate. |
For example, after supplying the task's validated `gym_config` and state width:
```python
-from embodichain.data_pipeline.engine import OnlineDataEngineCfg
+from embodichain.data_pipeline.datasets import OnlineDataSource, SequenceSpecCfg
+from embodichain.data_pipeline.engine import OnlineDataEngine, OnlineDataEngineCfg
+from embodichain.data_pipeline.runtime import ODSFlowController, ODSFlowControllerCfg
engine_cfg = OnlineDataEngineCfg(
gym_config={**your_gym_config, "num_envs": plan.producer_envs},
@@ -172,18 +174,98 @@ engine_cfg = OnlineDataEngineCfg(
episode_sample_hard_limit=plan.episode_sample_hard_limit,
include_provenance=True,
)
+
+# The application supplies the truthful schema, feature transform and train_step.
+with OnlineDataEngine(engine_cfg) as engine:
+ source = OnlineDataSource(
+ engine, schema=your_schema, observation_transform=your_feature_transform
+ )
+ flow = ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(max_generated_episodes=3 * plan.buffer_size),
+ plan=plan,
+ )
+ sequence = SequenceSpecCfg(history_length=2, prediction_horizon=64)
+ for _ in range(1000):
+ if not flow.prepare_batch(plan.batch_size):
+ break # Final supply cannot fund another complete batch.
+ batch = flow.sample_context(source, plan.batch_size, sequence)
+ train_step(batch)
```
+The context manager starts and stops the engine. The controller requires manual
+refill and runs only in its creating process/thread; it creates no background
+scheduler or DataLoader worker. It validates the engine's cache, producer count
+and assigned/hard quotas against the plan; `prepare_batch()` validates batch
+size. With either a plan or a production limit, `engine.cfg.gym_config` must
+declare a positive integer `num_envs` that divides cache size, so each fill
+produces exactly one cache. Unplanned, unlimited legacy engines do not acquire
+this restriction. The host retains responsibility for workload and semantic
+schema compatibility. Constructing a controller does not apply engine settings.
+
+With a plan, `prepare_batch()` schedules refill, waits for enough coarse unlocked
+quota and spaces batch starts. It returns `False` when no further whole-cache
+refill fits the production limit and remaining quota cannot fund a full batch.
+Slow updates or evaluation pauses do not accumulate catch-up credits. Call it
+before each planned batch: `sample_context()` alone does not apply pacing.
+
+`max_generated_episodes` includes initial fill and must be at least cache size.
+It limits requests made through this controller to complete cache refills;
+unused partial-cache capacity remains unproduced. An already pending refill
+is adopted without another engine request or a fabricated request timestamp.
+Calls made directly to the engine bypass this controller policy.
+
+Without `plan`, `prepare_batch()` only polls health/events. Refill stays explicit
+through `flow.request_refill()`, which returns a target generation or `None`
+when no whole cache fits. This supports caller-defined round schedules without
+automatically changing their data distribution.
+
The planner's remaining-quota threshold is **not** engine `refill_threshold`.
The engine's default `sample_threshold` policy counts delivered samples against
`refill_threshold * buffer_size` after successful calls. Quota exhaustion does
not itself request refill. Use manual scheduling for the plan above.
-`request_refill()` returns the target generation and coalesces an outstanding
-request. `wait_for_generation()` waits for the complete requested generation,
-not first publication. Continue useful training while supply is active when
-appropriate, observe health, drain enabled lifecycle events and use bounded
-waits. Source errors remain failures, not permission to continue on stale data.
+### Bounded waits and callbacks
+
+`ODSFlowControllerCfg` defaults to 30 seconds for data/quota waits, 1800 seconds
+for a complete generation and 0.02 seconds between polling probes. A transient
+live `DataUnavailableError` or `DataBudgetExhaustedError` is retried only within
+the configured data deadline. Worker failures, invalid inputs and caller
+cancellation propagate. Callbacks run synchronously; an interval observer
+failure preserves an already active error with a diagnostic note. Resident
+sources are sampled directly; Source and Dataset APIs retain their immediate
+error semantics.
+
+`flow.wait_for_generation(generation, deadline=...)` waits for the whole fill,
+not its first publication. Its optional absolute `time.monotonic()` deadline
+can shorten the configured generation timeout. This lets a collector count time
+spent writing a previous generation against the same request deadline.
+
+All waiting/pacing loops poll caller progress, drain lifecycle events and check
+engine health. Optional callbacks let the application retain its own accounting:
+
+| Callback | Meaning |
+| --- | --- |
+| `progress()` | Check application deadlines or cancellation; raise to interrupt. |
+| `on_events(events)` | Receive drained publication/eviction records. Events are drained even without an observer; install one to retain accounting. |
+| `on_refill(generation, requested_at)` | Observe new requests made by this controller, once per generation, in absolute host monotonic seconds. Adopting an already pending generation emits no request callback. |
+| `on_interval(kind, started, ended)` | Receive `data_wait`, `quota_wait` and `usage_throttle` intervals in absolute host monotonic seconds, including zero-duration waits. |
+
+Convert callback times explicitly when an application reports seconds since
+launch; do not mix them with Unix acceptance timestamps. The application owns
+completed optimizer counts, telemetry and shutdown. `flow.drain_events()` can
+also drain final records after stop/failure. An enabled lifecycle log still
+needs capacity for initial fill before `start()` returns. The controller does
+not snapshot or restore cache, quotas or pacing state.
+
+For frequent probes, `engine.check_health()` checks readiness/liveness without
+sampling or refill. `get_flow_status()` returns compact detached generation,
+accepted-episode and optional quota counters without constructing per-row audit
+records or copying trajectory payloads. Producer and quota counters use separate
+lock snapshots. Unlocked quota requires a valid continuity-qualified anchor and
+excludes the write window; sequence/mode/generation filters may further restrict
+sampling. Use `get_sampling_budget_snapshot()` for the full retained-source
+audit rather than in every pacing poll.
## Strict budget and consumer accounting
@@ -201,8 +283,8 @@ the demo rejects protected offline resume without that state.
The initial update count is a ceiling, not a convergence recommendation.
After production stops, the host may consume only remaining quotas and must
-stop when a requested full batch cannot be funded. No generic background
-controller derives a new safe K or training budget automatically.
+stop when a requested full batch cannot be funded. Flow control applies the
+given profile and plan; it does not derive a new safe K or convergence budget.
## Verify the plan
diff --git a/docs/source/features/online_data.md b/docs/source/features/online_data.md
index 9e936966c..9d9c9a551 100644
--- a/docs/source/features/online_data.md
+++ b/docs/source/features/online_data.md
@@ -13,6 +13,7 @@ delivery budgets make data use auditable. Shared APIs live in
| `SequenceDataset` | Iterate causal contexts from a source; the caller owns production, loading, retry and refill. |
| `OnlineDataset`, `ChunkSizeSampler` | Preserve item/batch chunk iteration and dynamic chunk lengths. |
| `configure_ods()` | Calculate bounded reuse, cache and pacing/refill targets from measurements and a versioned profile; the host applies the plan. |
+| `ODSFlowController` | Apply an optional measured plan to a caller-managed manual engine: explicit refill, bounded waits, lifecycle draining and batch-start pacing. |
## Guides and ownership
@@ -108,6 +109,8 @@ identities.
| Interface | Information and limits |
| --- | --- |
+| `check_health()` | Require READY and surface worker failures without sampling or requesting refill. |
+| `get_flow_status()` | Compact READY-only generation/accepted-episode counters and optional unlocked/lifetime delivery totals. Counters use separate lock snapshots; quota is a coarse scheduling bound, not proof that a specific context/filter can be sampled. |
| `get_stats()` | Production counts/time, refill generations and sampled chunks. Readable-frame counts exclude the active write window. Sampled chunks count deliveries, not completed optimizer work. |
| `get_sampling_budget_snapshot()` | With quotas enabled, retained identities and assigned/delivered/remaining allowances, plus lifetime delivered chunks. Unlocked quota is a scheduling hint; context and generation filters can further restrict eligibility. |
| `drain_lifecycle_events()` | Optional publication/eviction records and monotonic timestamps. Enable `lifecycle_event_capacity` and drain during production and waits; a full log backpressures the producer. |
@@ -265,6 +268,11 @@ errors take precedence. Quota depletion does not automatically request refill.
Use the [configuration guide](ods/configuration.md) to connect strict sampling
protection to an explicit host pacing/refill policy.
+For a caller-managed manual engine, the optional
+[`ODSFlowController`](ods/configuration.md#apply-the-plan) provides reusable
+bounded retries, refill and pacing. It does not change Source or Dataset error
+semantics, own the optimizer, or start/stop the engine.
+
## Shared online and offline context sources
`OnlineDataSource`, `TensorDataSource` and `LeRobotDataSource` expose
diff --git a/embodichain/data_pipeline/__init__.py b/embodichain/data_pipeline/__init__.py
index cd0606b3e..8e3ea1615 100644
--- a/embodichain/data_pipeline/__init__.py
+++ b/embodichain/data_pipeline/__init__.py
@@ -27,9 +27,9 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
- from . import configuration, datasets, depth_video, engine, recording
+ from . import configuration, datasets, depth_video, engine, recording, runtime
-__all__ = ["configuration", "datasets", "depth_video", "engine", "recording"]
+__all__ = ["configuration", "datasets", "depth_video", "engine", "recording", "runtime"]
def __getattr__(name: str) -> ModuleType:
diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py
index a36d99e5c..d43968a0a 100644
--- a/embodichain/data_pipeline/engine/data.py
+++ b/embodichain/data_pipeline/engine/data.py
@@ -1602,6 +1602,19 @@ def is_init(self) -> bool:
"""Whether the engine is ready to serve initialized data."""
return self.state is OnlineDataEngineState.READY
+ def check_health(self) -> None:
+ """Require a ready engine and surface existing worker failures.
+
+ This check neither samples data nor requests a refill. It checks worker
+ liveness through the existing owner/consumer error channel.
+
+ Raises:
+ RuntimeError: If the engine is not READY or its worker has exited.
+ BaseException: The original reported worker or startup exception.
+ """
+ self._check_sampling_budget_health()
+ self._require_state(OnlineDataEngineState.READY, "check_health")
+
# -----------------------------------------------------------------------
# Sampling
# -----------------------------------------------------------------------
@@ -2035,6 +2048,70 @@ def _check_sampling_budget_health(self) -> None:
if self.state is OnlineDataEngineState.READY:
self._ensure_worker_alive()
+ def _unlocked_sampling_chunks(
+ self, quota_rows: list[int], write_start: int, write_end: int
+ ) -> int:
+ """Sum coarse eligible quota while the caller holds the copy lock."""
+ valid = self.shared_buffer["valid"].bool()
+ continuity = self.shared_buffer.get("continuity_id", None)
+ eligible = (
+ (valid if continuity is None else valid & (continuity >= 0))
+ .any(-1)
+ .tolist()
+ )
+ return sum(
+ quota_rows[slot * 4 + 2] - quota_rows[slot * 4 + 3]
+ for slot in range(self.buffer_size)
+ if quota_rows[slot * 4] >= 0
+ and not write_start <= slot < write_end
+ and eligible[slot]
+ )
+
+ def get_flow_status(self) -> dict[str, int | None]:
+ """Return detached counters for READY-only host flow control.
+
+ Producer counters and quota totals use separate lock snapshots. The
+ quota sum includes only published, unlocked rows with a valid anchor
+ and nonnegative continuity when continuity metadata is present. It is
+ a coarse bound; sequence lengths and filters may reduce eligibility.
+ This method creates no per-row audit records or payload copies.
+
+ Returns:
+ Requested/completed generations and successful episodes, plus
+ ``unlocked_remaining_chunks`` and lifetime ``total_delivered_chunks``.
+ Quota counters are None when protection is disabled.
+
+ Raises:
+ RuntimeError: If the engine is not READY or its worker has exited.
+ BaseException: The original reported worker or startup exception.
+ """
+ self.check_health()
+ with self._progress.values.get_lock():
+ values = self._progress.values[:3]
+ status: dict[str, int | None] = {
+ name: int(value)
+ for name, value in zip(
+ ("requested_generation", "completed_generation", "successful_episodes"),
+ values,
+ )
+ }
+ quota = getattr(self, "_sampling_budget", None)
+ if quota is None:
+ return {
+ **status,
+ "unlocked_remaining_chunks": None,
+ "total_delivered_chunks": None,
+ }
+ with self._lock_index.get_lock():
+ self.check_health()
+ start, end = self._lock_index[:]
+ quota_rows = quota.rows[:]
+ status["unlocked_remaining_chunks"] = self._unlocked_sampling_chunks(
+ quota_rows, start, end
+ )
+ status["total_delivered_chunks"] = quota.meta[1]
+ return status
+
def get_sampling_budget_snapshot(self) -> dict[str, object]:
"""Return a detached quota audit under the producer/sampling copy lock.
@@ -2060,17 +2137,11 @@ def get_sampling_budget_snapshot(self) -> dict[str, object]:
with self._lock_index.get_lock():
self._check_sampling_budget_health()
start, end = self._lock_index[:]
- valid = self.shared_buffer["valid"].bool()
- continuity = self.shared_buffer.get("continuity_id", None)
- eligible = (
- (valid if continuity is None else valid & (continuity >= 0))
- .any(-1)
- .tolist()
- )
+ quota_rows = quota.rows[:]
+ unlocked = self._unlocked_sampling_chunks(quota_rows, start, end)
rows = []
- unlocked = 0
for slot in range(self.buffer_size):
- episode_id, generation, assigned, delivered = quota.rows[
+ episode_id, generation, assigned, delivered = quota_rows[
slot * 4 : slot * 4 + 4
]
if episode_id < 0:
@@ -2089,8 +2160,6 @@ def get_sampling_budget_snapshot(self) -> dict[str, object]:
"locked": locked,
}
)
- if not locked and eligible[slot]:
- unlocked += remaining
return {
"enabled": True,
"rows": rows,
diff --git a/embodichain/data_pipeline/runtime.py b/embodichain/data_pipeline/runtime.py
new file mode 100644
index 000000000..358fd166c
--- /dev/null
+++ b/embodichain/data_pipeline/runtime.py
@@ -0,0 +1,375 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ----------------------------------------------------------------------------
+
+"""Owner-side flow control for an explicitly started online data engine."""
+
+from __future__ import annotations
+
+import math
+import os
+import sys
+import threading
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from embodichain.utils import configclass
+
+from .configuration import ODSConfiguration
+from .datasets.sequence import DataBudgetExhaustedError, DataUnavailableError
+from .engine.data import _add_exception_note
+
+if TYPE_CHECKING:
+ import torch
+ from tensordict import TensorDict
+
+ from .datasets import OnlineDataSource, SequenceSpecCfg
+ from .engine import OnlineDataEngine
+
+__all__ = ["ODSFlowControllerCfg", "ODSFlowController"]
+
+
+@configclass
+class ODSFlowControllerCfg:
+ """Bound owner-side waits and optional whole-cache production.
+
+ Args:
+ data_wait_timeout: Maximum seconds blocked on missing windows or quota.
+ generation_timeout: Maximum seconds waiting for a complete refill.
+ poll_interval: Maximum sleep/probe duration between health/event checks.
+ max_generated_episodes: Optional accepted-episode ceiling. A refill must
+ fit a complete cache; unused partial-cache capacity is not produced.
+ """
+
+ data_wait_timeout: float = 30.0
+ generation_timeout: float = 1800.0
+ poll_interval: float = 0.02
+ max_generated_episodes: int | None = None
+
+ def __post_init__(self) -> None:
+ for name in ("data_wait_timeout", "generation_timeout", "poll_interval"):
+ value = getattr(self, name)
+ if (
+ isinstance(value, bool)
+ or not isinstance(value, (int, float))
+ or not math.isfinite(value)
+ or value <= 0
+ ):
+ raise ValueError(f"{name} must be finite and positive")
+ limit = self.max_generated_episodes
+ if limit is not None and (type(limit) is not int or limit < 1):
+ raise ValueError("max_generated_episodes must be a positive integer")
+
+
+class ODSFlowController:
+ """Apply a measured plan without changing sampling or engine ownership.
+
+ The creating process/thread owns this controller. The application starts
+ and stops the engine and owns optimizer/evaluation deadlines. Sources and
+ Datasets retain their immediate sampling/error contracts.
+
+ Args:
+ engine: Creator-owned engine configured for manual refill.
+ cfg: Wait bounds and optional production limit.
+ plan: Optional configuration to enforce before each batch. Cache,
+ producer count, source allowances and batch size must match the plan.
+ progress: Called during polling; may raise for cancellation/deadlines.
+ on_events: Receives drained lifecycle records, including after stop.
+ on_refill: Receives a new requested generation and its absolute host
+ monotonic request time. Coalesced requests are reported once.
+ on_interval: Receives wait kind and absolute host monotonic start/end.
+ """
+
+ def __init__(
+ self,
+ engine: OnlineDataEngine,
+ cfg: ODSFlowControllerCfg | None = None,
+ *,
+ plan: ODSConfiguration | None = None,
+ progress: Callable[[], object] | None = None,
+ on_events: Callable[[list[dict[str, object]]], object] | None = None,
+ on_refill: Callable[[int, float], object] | None = None,
+ on_interval: Callable[[str, float, float], object] | None = None,
+ ) -> None:
+ engine._require_owner_process("ODSFlowController")
+ self.cfg = cfg if cfg is not None else ODSFlowControllerCfg()
+ self.cfg.validate()
+ if engine.cfg.refill_policy != "manual":
+ raise ValueError("ODSFlowController requires manual refill")
+ limit = self.cfg.max_generated_episodes
+ if limit is not None and limit < engine.buffer_size:
+ raise ValueError(
+ "Production limit cannot be smaller than the initial cache"
+ )
+ if plan is not None or limit is not None:
+ producers = engine.cfg.gym_config.get("num_envs")
+ if (
+ type(producers) is not int
+ or producers < 1
+ or engine.buffer_size % producers
+ ):
+ raise ValueError(
+ "Planned or limited flow requires explicit num_envs dividing the cache"
+ )
+ if plan is not None:
+ if not isinstance(plan, ODSConfiguration):
+ raise TypeError("plan must be an ODSConfiguration")
+ if (
+ engine.buffer_size != plan.buffer_size
+ or engine.cfg.episode_sample_budget != plan.episode_sample_budget
+ or engine.cfg.episode_sample_hard_limit
+ != plan.episode_sample_hard_limit
+ or producers != plan.producer_envs
+ ):
+ raise ValueError(
+ "Engine cache, producer count and quotas must match the plan"
+ )
+ rate = min(
+ plan.training_chunks_per_second_target,
+ plan.active_chunks_per_second_cap,
+ )
+ if not math.isfinite(rate) or rate <= 0:
+ raise ValueError("Plan training rate must be finite and positive")
+ self.engine, self.plan = engine, plan
+ self._progress, self._on_events = progress, on_events
+ self._on_refill, self._on_interval = on_refill, on_interval
+ self._owner = os.getpid(), threading.get_ident()
+ self._last_refill = 0
+ self._next_batch = 0.0
+
+ def _require_owner(self) -> None:
+ if self._owner != (os.getpid(), threading.get_ident()):
+ raise RuntimeError("ODS flow control is owner-process/thread only")
+
+ def drain_events(self) -> list[dict[str, object]]:
+ """Drain accounting records, also during failure/shutdown diagnostics.
+
+ Returns:
+ Detached lifecycle records delivered to the optional event observer.
+ """
+ self._require_owner()
+ events = self.engine.drain_lifecycle_events()
+ if events and self._on_events is not None:
+ self._on_events(events)
+ return events
+
+ def poll(self) -> None:
+ """Check caller deadlines, drain lifecycle events and verify worker health."""
+ self._require_owner()
+ if self._progress is not None:
+ self._progress()
+ self.drain_events()
+ self.engine.check_health()
+
+ def _record(self, kind: str, started: float | None) -> None:
+ ended = time.monotonic()
+ if self._on_interval is not None:
+ primary = sys.exc_info()[1]
+ try:
+ self._on_interval(kind, ended if started is None else started, ended)
+ except BaseException as error:
+ if primary is None:
+ raise
+ _add_exception_note(
+ primary, f"ODS interval observer also failed: {error}"
+ )
+
+ def _sleep(self, deadline: float, message: str) -> None:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ raise TimeoutError(message)
+ time.sleep(min(self.cfg.poll_interval, remaining))
+
+ def _room(self, state: dict[str, int | None]) -> bool:
+ limit = self.cfg.max_generated_episodes
+ return (
+ limit is None
+ or state["successful_episodes"] + self.engine.buffer_size <= limit
+ )
+
+ def request_refill(self) -> int | None:
+ """Request or coalesce a fill within the accepted-episode ceiling.
+
+ Returns:
+ Pending or newly requested generation, or None if no whole cache fits.
+ """
+ self.poll()
+ state = self.engine.get_flow_status()
+ pending = state["requested_generation"] > state["completed_generation"]
+ if pending:
+ return state["requested_generation"]
+ if not self._room(state):
+ return None
+ started = time.monotonic()
+ generation = self.engine.request_refill()
+ if generation != self._last_refill:
+ self._last_refill = generation
+ if self._on_refill is not None:
+ self._on_refill(generation, started)
+ return generation
+
+ def wait_for_generation(
+ self, generation: int, *, deadline: float | None = None
+ ) -> None:
+ """Wait with a total bound, optionally shortened by an absolute deadline.
+
+ Args:
+ generation: Target generation returned by a refill request.
+ deadline: Optional absolute host monotonic time limiting this wait.
+
+ Raises:
+ TimeoutError: If the generation does not complete within the bound.
+ """
+ self._require_owner()
+ if deadline is not None and (
+ isinstance(deadline, bool) or not math.isfinite(deadline)
+ ):
+ raise ValueError("deadline must be a finite host monotonic time")
+ limit = time.monotonic() + self.cfg.generation_timeout
+ if deadline is not None:
+ limit = min(limit, deadline)
+ started = time.monotonic()
+ try:
+ while True:
+ self.poll()
+ remaining = limit - time.monotonic()
+ try:
+ self.engine.wait_for_generation(
+ generation,
+ timeout=max(1e-9, min(self.cfg.poll_interval, remaining)),
+ )
+ except TimeoutError:
+ if time.monotonic() >= limit:
+ raise TimeoutError(
+ f"ODS generation {generation} missed its deadline"
+ ) from None
+ continue
+ self.drain_events()
+ return
+ finally:
+ self._record("data_wait", started)
+
+ def prepare_batch(self, batch_size: int) -> bool:
+ """Refill/pace a planned batch; False means final supply is exhausted.
+
+ Unplanned engines are only polled; their refill schedule stays explicit.
+ Quota readiness is a coarse bound. Sampling still validates the exact
+ context/filter request and can require a bounded data wait.
+
+ Args:
+ batch_size: Positive chunk count, matching the plan when supplied.
+
+ Returns:
+ False only when capped final supply cannot fund a complete batch.
+
+ Raises:
+ TimeoutError: If missing unlocked quota exceeds the data wait bound.
+ """
+ if type(batch_size) is not int or batch_size < 1:
+ raise ValueError("batch_size must be a positive integer")
+ self._require_owner()
+ if self.plan is None:
+ self.poll()
+ return True
+ if batch_size != self.plan.batch_size:
+ raise ValueError("batch_size must match the measured plan")
+ blocked = None
+ try:
+ while True:
+ self.poll()
+ state = self.engine.get_flow_status()
+ remaining = state["unlocked_remaining_chunks"]
+ pending = state["requested_generation"] > state["completed_generation"]
+ if not pending and remaining <= self.plan.refill_remaining_chunks:
+ pending = self.request_refill() is not None
+ if remaining >= batch_size:
+ break
+ if not pending and not self._room(state):
+ return False
+ if blocked is None:
+ blocked = time.monotonic()
+ self._sleep(
+ blocked + self.cfg.data_wait_timeout,
+ "No unlocked ODS quota before the data deadline",
+ )
+ finally:
+ self._record("quota_wait", blocked)
+ started = time.monotonic()
+ try:
+ while time.monotonic() < self._next_batch:
+ self.poll()
+ time.sleep(
+ max(
+ 0.0,
+ min(
+ self.cfg.poll_interval, self._next_batch - time.monotonic()
+ ),
+ )
+ )
+ # Slow updates/evaluation do not accumulate catch-up credits.
+ self._next_batch = time.monotonic() + batch_size / min(
+ self.plan.training_chunks_per_second_target,
+ self.plan.active_chunks_per_second_cap,
+ )
+ finally:
+ self._record("usage_throttle", started)
+ return True
+
+ def sample_context(
+ self,
+ source: OnlineDataSource,
+ batch_size: int,
+ sequence: SequenceSpecCfg,
+ *,
+ generator: torch.Generator | None = None,
+ ) -> TensorDict:
+ """Retry only transient live-sampling errors within the data deadline.
+
+ Call :meth:`prepare_batch` first to apply a plan's pacing/supply limit.
+ Resident sources are sampled directly and never retried by this class.
+
+ Args:
+ source: Live source wrapping this controller's engine.
+ batch_size: Number of chunks to sample.
+ sequence: Requested context and prediction window.
+ generator: Optional sampling random generator.
+
+ Returns:
+ The source's batch, including masks and provenance.
+
+ Raises:
+ TimeoutError: If transient unavailability exceeds the data wait bound.
+ """
+ self._require_owner()
+ if source.engine is not self.engine:
+ raise ValueError("source must use the controlled live engine")
+ blocked = None
+ try:
+ while True:
+ self.poll()
+ try:
+ return source.sample_context(
+ batch_size, sequence, generator=generator
+ )
+ except (DataUnavailableError, DataBudgetExhaustedError):
+ if blocked is None:
+ blocked = time.monotonic()
+ self._sleep(
+ blocked + self.cfg.data_wait_timeout,
+ "No eligible ODS context before the data deadline",
+ )
+ finally:
+ self._record("data_wait", blocked)
diff --git a/examples/ods/simple_bc_demo/README.md b/examples/ods/simple_bc_demo/README.md
index a3e1408f7..83923c470 100644
--- a/examples/ods/simple_bc_demo/README.md
+++ b/examples/ods/simple_bc_demo/README.md
@@ -8,7 +8,7 @@ robot and object state.
Run the commands below from the repository root. This directory owns the local
task, configuration, model, trainer, checkpoint tools and evaluation. The shared
-ODS engine, throughput configurator and numeric source APIs remain in
+ODS engine, throughput configurator, flow controller and numeric source APIs remain in
`embodichain.data_pipeline`.
## Layout
@@ -132,6 +132,12 @@ replacement can retire unused allowances. The historical K=2660 is provisional.
Throughput matching determines pace, while the learning protocol determines
the total production and update budgets.
+The trainer uses the public `ODSFlowController` for manual refill scheduling,
+bounded data waits, quota readiness and batch pacing. The example owns the
+optimizer, evaluation, resource accounting and engine shutdown. Sampling from a
+source or Dataset still returns immediately or raises; waiting is explicit in
+the controller.
+
Measure the actual learner, representation, horizons, batch size and producer
concurrency together. The public `measure_ods_throughput` helper preserves rate
units and observation windows. The generic CPU sampler-cost benchmark remains
diff --git a/examples/ods/simple_bc_demo/collection.py b/examples/ods/simple_bc_demo/collection.py
index d7487c02b..15f907a7b 100644
--- a/examples/ods/simple_bc_demo/collection.py
+++ b/examples/ods/simple_bc_demo/collection.py
@@ -19,26 +19,27 @@
from __future__ import annotations
import time
-import math
from pathlib import Path
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
+
+from embodichain.data_pipeline.datasets import DataBudgetExhaustedError
from examples.ods.simple_bc_demo.lerobot import DirectLeRobotWriter
+if TYPE_CHECKING:
+ from embodichain.data_pipeline.runtime import ODSFlowController
+
def collect_lerobot(
- engine: Any,
+ flow: ODSFlowController,
initial: dict[str, Any],
destination: Path,
*,
schema: dict[str, Any],
episodes: int,
snapshot: Callable[[Any], dict[str, Any]],
- request_refill: Callable[[], int],
- progress: Callable[[], None],
clock: Callable[[], float],
record_interval: Callable[[str, float, float], None],
- generation_timeout: float = 1800.0,
) -> dict[str, Any]:
"""Overlap generations with a bounded writer using cache-sized snapshots.
@@ -52,8 +53,8 @@ def collect_lerobot(
rows, steps = initial["valid"].shape
if episodes < rows or episodes % rows:
raise ValueError("Collection requires a positive whole-cache episode count")
- if not math.isfinite(generation_timeout) or generation_timeout <= 0:
- raise ValueError("Generation timeout must be finite and positive")
+ engine = flow.engine
+ progress = flow.poll
started = clock()
writer = None
try:
@@ -72,8 +73,13 @@ def collect_lerobot(
for index in range(episodes // rows):
# The current snapshot is an independent copy. Producing the next
# cache can safely overlap enqueue/write of these accepted rows.
- generation = request_refill() if (index + 1) * rows < episodes else None
- deadline = time.monotonic() + generation_timeout
+ more = (index + 1) * rows < episodes
+ generation = flow.request_refill() if more else None
+ if more and generation is None:
+ raise DataBudgetExhaustedError(
+ "Production limit cannot supply the requested offline corpus"
+ )
+ deadline = time.monotonic() + flow.cfg.generation_timeout
for row in range(rows):
progress()
writer.submit(
@@ -91,14 +97,7 @@ def collect_lerobot(
)
if generation is not None:
before = clock()
- while engine.get_stats()["completed_generation"] < generation:
- progress()
- if time.monotonic() >= deadline:
- raise TimeoutError(
- "Offline refill did not complete before its deadline"
- )
- time.sleep(0.02)
- progress()
+ flow.wait_for_generation(generation, deadline=deadline)
record_interval("offline_refill_wait", before, clock())
current = snapshot(engine)
before = clock()
diff --git a/examples/ods/simple_bc_demo/data.py b/examples/ods/simple_bc_demo/data.py
index f89e97012..2bef7f8da 100644
--- a/examples/ods/simple_bc_demo/data.py
+++ b/examples/ods/simple_bc_demo/data.py
@@ -27,10 +27,12 @@
import json
import math
import os
+import threading
import time
from collections.abc import Iterable
from typing import Any
+import psutil
import torch
from embodichain.data_pipeline.datasets import data_schema_fingerprint
@@ -542,3 +544,109 @@ def resource_snapshot(
),
"torch_cuda_this_process": cuda,
}
+
+
+class _ResourceSampler:
+ """Best-effort wall-clock resource samples, never an exact peak or CPU total.
+
+ A background thread covers owner-side waits and short runs independently of
+ evaluation frequency. Collection cost is measured; this instrumentation
+ requires its own performance validation before use in a timed comparison.
+ """
+
+ def __init__(
+ self, cuda_devices: tuple[int, ...], interval_seconds: float = 5.0
+ ) -> None:
+ if not math.isfinite(interval_seconds) or interval_seconds <= 0:
+ raise ValueError("Resource sampling interval must be finite and positive")
+ self.cuda_devices = cuda_devices
+ self.interval_seconds = interval_seconds
+ self._stop = threading.Event()
+ self._lock = threading.Lock()
+ self._thread: threading.Thread | None = None
+ self._records: list[dict[str, Any]] = []
+ self._errors: list[dict[str, Any]] = []
+ self._collection_seconds = 0.0
+ self._closed = False
+
+ def start(self) -> None:
+ if self._thread is not None or self._closed:
+ raise RuntimeError("Resource sampler cannot be restarted")
+ self.capture("before_worker_startup")
+ self._thread = threading.Thread(
+ target=self._run, name="ods-resource-sampler", daemon=True
+ )
+ self._thread.start()
+
+ def _run(self) -> None:
+ while not self._stop.wait(self.interval_seconds):
+ self.capture("periodic")
+
+ def capture(self, boundary: str) -> None:
+ with self._lock:
+ if self._closed:
+ return
+ started = time.monotonic()
+ snapshot = None
+ error_record = None
+ try:
+ process = psutil.Process()
+ pids = [process.pid] + [
+ child.pid for child in process.children(recursive=True)
+ ]
+ # resource_snapshot checks is_initialized before reading CUDA
+ # allocator counters; it never creates a CUDA context itself.
+ snapshot = resource_snapshot(pids, cuda_devices=self.cuda_devices)
+ snapshot["boundary"] = boundary
+ snapshot["collection_started_monotonic_seconds"] = started
+ snapshot["process_tree_scope"] = "best_effort_current_descendants"
+ except Exception as error:
+ error_record = {
+ "boundary": boundary,
+ "monotonic_seconds": started,
+ "error": f"{type(error).__name__}: {error}",
+ }
+ elapsed = time.monotonic() - started
+ with self._lock:
+ self._collection_seconds += elapsed
+ if snapshot is not None:
+ snapshot["collection_seconds"] = elapsed
+ self._records.append(snapshot)
+ if error_record is not None:
+ error_record["collection_seconds"] = elapsed
+ self._errors.append(error_record)
+
+ def stop(self) -> list[dict[str, Any]]:
+ self._stop.set()
+ if self._thread is not None:
+ self._thread.join(timeout=10.0)
+ if self._thread.is_alive():
+ with self._lock:
+ self._closed = True
+ raise RuntimeError("Resource sampler did not stop within 10 seconds")
+ with self._lock:
+ self._closed = True
+ return list(self._records)
+
+ def report(self) -> dict[str, Any]:
+ with self._lock:
+ return {
+ "started": self._thread is not None,
+ "stopped": self._closed
+ and (self._thread is None or not self._thread.is_alive()),
+ "interval_seconds": self.interval_seconds,
+ "samples": len(self._records),
+ "collection_wall_seconds": self._collection_seconds,
+ "errors": list(self._errors),
+ "missing_pid_observations": sum(
+ len(row.get("missing_pids", [])) for row in self._records
+ ),
+ "scope": "sampled process-tree PSS and observed CPU counters; Torch allocator counters belong to the learner only",
+ "limitations": [
+ "Memory values are sampled peaks, not exact peaks",
+ "CPU counters omit work after a process's last observation and may miss short-lived descendants",
+ "Missing processes or collection errors remain unknown, not zero",
+ "Collection wall time sums sample service durations, including any overlapping forced/periodic captures; it is not end-to-end overhead",
+ "End-to-end overhead and scheduling perturbation need workload-specific measurement",
+ ],
+ }
diff --git a/examples/ods/simple_bc_demo/trainer.py b/examples/ods/simple_bc_demo/trainer.py
index d6593057f..5c87fd0e0 100644
--- a/examples/ods/simple_bc_demo/trainer.py
+++ b/examples/ods/simple_bc_demo/trainer.py
@@ -27,7 +27,6 @@
import random
import subprocess
import sys
-import threading
import time
from collections import deque
from collections.abc import Callable
@@ -37,25 +36,24 @@
import numpy as np
import psutil
import torch
-from tensordict import TensorDict
from embodichain.data_pipeline.datasets import (
- BatchSource,
DataBudgetExhaustedError,
- DataUnavailableError,
OnlineDataSource,
SequenceSpecCfg,
)
+from embodichain.data_pipeline.configuration import ODSConfiguration
+from embodichain.data_pipeline.runtime import ODSFlowController, ODSFlowControllerCfg
+
from embodichain.data_pipeline.engine.data import (
OnlineDataEngine,
OnlineDataEngineCfg,
- OnlineDataEngineState,
)
from examples.ods.simple_bc_demo.data import (
_content_manifest,
EfficiencyTracker,
- resource_snapshot,
+ _ResourceSampler,
)
from examples.ods.simple_bc_demo.evaluation import _evaluate_worker, _receive
from examples.ods.simple_bc_demo.observations import (
@@ -73,112 +71,6 @@
)
-class _ResourceSampler:
- """Best-effort wall-clock resource samples, never an exact peak or CPU total.
-
- A background thread covers owner-side waits and short runs independently of
- evaluation frequency. Collection cost is measured; this instrumentation
- requires its own performance validation before use in a timed comparison.
- """
-
- def __init__(
- self, cuda_devices: tuple[int, ...], interval_seconds: float = 5.0
- ) -> None:
- if not np.isfinite(interval_seconds) or interval_seconds <= 0:
- raise ValueError("Resource sampling interval must be finite and positive")
- self.cuda_devices = cuda_devices
- self.interval_seconds = interval_seconds
- self._stop = threading.Event()
- self._lock = threading.Lock()
- self._thread: threading.Thread | None = None
- self._records: list[dict[str, Any]] = []
- self._errors: list[dict[str, Any]] = []
- self._collection_seconds = 0.0
- self._closed = False
-
- def start(self) -> None:
- if self._thread is not None or self._closed:
- raise RuntimeError("Resource sampler cannot be restarted")
- self.capture("before_worker_startup")
- self._thread = threading.Thread(
- target=self._run, name="ods-resource-sampler", daemon=True
- )
- self._thread.start()
-
- def _run(self) -> None:
- while not self._stop.wait(self.interval_seconds):
- self.capture("periodic")
-
- def capture(self, boundary: str) -> None:
- with self._lock:
- if self._closed:
- return
- started = time.monotonic()
- snapshot = None
- error_record = None
- try:
- process = psutil.Process()
- pids = [process.pid] + [
- child.pid for child in process.children(recursive=True)
- ]
- # resource_snapshot checks is_initialized before reading CUDA
- # allocator counters; it never creates a CUDA context itself.
- snapshot = resource_snapshot(pids, cuda_devices=self.cuda_devices)
- snapshot["boundary"] = boundary
- snapshot["collection_started_monotonic_seconds"] = started
- snapshot["process_tree_scope"] = "best_effort_current_descendants"
- except Exception as error:
- error_record = {
- "boundary": boundary,
- "monotonic_seconds": started,
- "error": f"{type(error).__name__}: {error}",
- }
- elapsed = time.monotonic() - started
- with self._lock:
- self._collection_seconds += elapsed
- if snapshot is not None:
- snapshot["collection_seconds"] = elapsed
- self._records.append(snapshot)
- if error_record is not None:
- error_record["collection_seconds"] = elapsed
- self._errors.append(error_record)
-
- def stop(self) -> list[dict[str, Any]]:
- self._stop.set()
- if self._thread is not None:
- self._thread.join(timeout=10.0)
- if self._thread.is_alive():
- with self._lock:
- self._closed = True
- raise RuntimeError("Resource sampler did not stop within 10 seconds")
- with self._lock:
- self._closed = True
- return list(self._records)
-
- def report(self) -> dict[str, Any]:
- with self._lock:
- return {
- "started": self._thread is not None,
- "stopped": self._closed
- and (self._thread is None or not self._thread.is_alive()),
- "interval_seconds": self.interval_seconds,
- "samples": len(self._records),
- "collection_wall_seconds": self._collection_seconds,
- "errors": list(self._errors),
- "missing_pid_observations": sum(
- len(row.get("missing_pids", [])) for row in self._records
- ),
- "scope": "sampled process-tree PSS and observed CPU counters; Torch allocator counters belong to the learner only",
- "limitations": [
- "Memory values are sampled peaks, not exact peaks",
- "CPU counters omit work after a process's last observation and may miss short-lived descendants",
- "Missing processes or collection errors remain unknown, not zero",
- "Collection wall time sums sample service durations, including any overlapping forced/periodic captures; it is not end-to-end overhead",
- "End-to-end overhead and scheduling perturbation need workload-specific measurement",
- ],
- }
-
-
def _write_evaluation_scalars(writer: Any, report: dict[str, Any]) -> None:
"""Index posthoc evaluation by the evaluated policy, not the live learner."""
policy_step = int(report["policy_step"])
@@ -228,39 +120,6 @@ def _result(
raise
-def _wait(
- engine: OnlineDataEngine,
- generation: int,
- args: argparse.Namespace,
- tracker: EfficiencyTracker,
- *,
- progress: Callable[[], Any] | None = None,
-) -> None:
- """Bound the full generation wait while draining events and checking health."""
- timeout = _timeout(args, tracker)
- deadline = time.monotonic() + timeout
- try:
- while True:
- if progress is not None:
- progress()
- remaining = min(deadline - time.monotonic(), _timeout(args, tracker))
- if remaining <= 0:
- raise TimeoutError(
- f"ODS generation {generation} was not published in {timeout}s"
- )
- try:
- engine.wait_for_generation(generation, timeout=min(0.1, remaining))
- except TimeoutError:
- continue
- if progress is not None:
- progress()
- return
- except TimeoutError:
- if args.wall_time_budget and tracker.elapsed_seconds() >= args.wall_time_budget:
- raise _BudgetExpired from None
- raise
-
-
def _selection(report: dict[str, Any]) -> tuple[float, float, float]:
return report["success_rate"], report["lift_rate"], -report["goal_error"]
@@ -276,11 +135,7 @@ def _state_digest(values: dict[str, torch.Tensor]) -> str:
def _normalization_values(model: torch.nn.Module) -> dict[str, torch.Tensor]:
"""Include every input and output scale that changes policy semantics."""
- names = (
- ("mean", "std", "target_mean", "target_std")
- if isinstance(model, ChunkBC)
- else ("mean", "std")
- )
+ names = ("mean", "std", "target_mean", "target_std")
return {name: getattr(model, name) for name in names}
@@ -439,32 +294,6 @@ def _offline_reset_positions(corpus: Any, count: int) -> list[list[float]]:
return corpus.data["obs"][:count, 0, 27:30].tolist()
-def _sample(
- source: BatchSource,
- count: int,
- tracker: EfficiencyTracker,
- sequence: SequenceSpecCfg,
-) -> TensorDict:
- """Bound live-data waits while propagating offline quota/read failures."""
- blocked = None
- engine = getattr(source, "engine", None)
- while True:
- try:
- batch = source.sample_context(count, sequence)
- break
- except (DataUnavailableError, DataBudgetExhaustedError) as error:
- if engine is None or engine.state is not OnlineDataEngineState.READY:
- raise
- if blocked is None:
- blocked = tracker.elapsed_seconds()
- if tracker.elapsed_seconds() - blocked >= 30:
- raise TimeoutError("No eligible ODS context for 30 seconds") from error
- time.sleep(0.02)
- ended = tracker.elapsed_seconds()
- tracker.record_interval("data_wait", ended if blocked is None else blocked, ended)
- return batch
-
-
def run(args: argparse.Namespace) -> None:
"""Execute one explicitly labeled diagnostic or efficiency experiment."""
started = time.monotonic() - (time.time() - psutil.Process().create_time())
@@ -592,6 +421,7 @@ def run(args: argparse.Namespace) -> None:
"embodichain/data_pipeline/datasets/lerobot.py",
"embodichain/data_pipeline/datasets/_sampling.py",
"embodichain/data_pipeline/configuration.py",
+ "embodichain/data_pipeline/runtime.py",
"embodichain/data_pipeline/engine/data.py",
"embodichain/lab/gym/envs/embodied_env.py",
"embodichain/lab/gym/envs/base_env.py",
@@ -672,6 +502,7 @@ def restart_evaluator() -> dict[str, Any]:
}
engine = None
+ flow = None
corpus = None
writer = None
tensorboard = None
@@ -713,12 +544,8 @@ def checkpoint_in_budget(path: Path) -> bool:
selection_budget is None or checkpoint_save_times[path] <= selection_budget
)
- def drain_usage() -> list[dict[str, Any]]:
- if engine is None or usage is None:
- return []
- events = engine.drain_lifecycle_events()
- usage.lifecycle(events)
- return events
+ def drain_usage() -> list[dict[str, object]]:
+ return [] if flow is None else flow.drain_events()
def save_usage(report: dict[str, Any] | None = None) -> None:
if usage is None:
@@ -738,15 +565,15 @@ def save_usage(report: dict[str, Any] | None = None) -> None:
)
)
- recorded_refills: set[int] = set()
+ def record_refill(generation: int, started: float) -> None:
+ tracker._event(
+ "refill_requested", started - tracker.start_monotonic, generation=generation
+ )
- def request_refill() -> int:
- before = tracker.elapsed_seconds()
- generation = engine.request_refill()
- if generation not in recorded_refills:
- tracker._event("refill_requested", before, generation=generation)
- recorded_refills.add(generation)
- return generation
+ def record_flow_interval(kind: str, started: float, ended: float) -> None:
+ tracker.record_interval(
+ kind, started - tracker.start_monotonic, ended - tracker.start_monotonic
+ )
end_reason = "update_budget"
generation_counts: dict[int, int] = {}
@@ -892,6 +719,17 @@ def finish_resource_sampling() -> None:
)["content_id"]
else:
engine = _engine(config, args)
+ flow = ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(
+ max_generated_episodes=args.max_generated_episodes
+ ),
+ plan=ODSConfiguration(**auto_plan) if auto_plan else None,
+ progress=lambda: _timeout(args, tracker),
+ on_events=usage.lifecycle if usage is not None else None,
+ on_refill=record_refill,
+ on_interval=record_flow_interval,
+ )
if auto_report is not None:
(args.output / "auto_configuration.json").write_text(
json.dumps(auto_report, indent=2)
@@ -935,23 +773,14 @@ def finish_resource_sampling() -> None:
load_lerobot_corpus,
)
- def collection_progress() -> None:
- drain_usage()
- _timeout(args, tracker)
- # Also checks process liveness when a hard exit has not
- # yet been reflected in the shared lifecycle state.
- engine.wait_for_generation(1, timeout=0.1)
-
path = args.output / "lerobot_dataset"
collection_report = collect_lerobot(
- engine,
+ flow,
initial,
path,
schema=schema,
episodes=offline_episodes,
snapshot=_snapshot,
- request_refill=request_refill,
- progress=collection_progress,
clock=tracker.elapsed_seconds,
record_interval=tracker.record_interval,
)
@@ -963,6 +792,7 @@ def collection_progress() -> None:
engine.stop()
drain_usage()
engine = None
+ flow = None
resource_sampler.capture("after_offline_producer_shutdown")
before = tracker.elapsed_seconds()
corpus = load_lerobot_corpus(
@@ -1225,24 +1055,23 @@ def evaluate(
),
)
loop_started = tracker.elapsed_seconds()
- next_quota_update = time.monotonic()
for round_index in range(args.rounds):
model.train()
generation = 1
barrier = None
if (
round_index
- and engine is not None
+ and flow is not None
and args.mode == "online"
and auto_plan is None
- and (
- args.max_generated_episodes is None
- or engine.get_stats()["successful_episodes"]
- < args.max_generated_episodes
- )
):
- generation = request_refill()
- barrier = args.updates_per_round // 4 if args.mode == "online" else 0
+ try:
+ generation = flow.request_refill()
+ except _BudgetExpired:
+ end_reason = "wall_time_budget"
+ break
+ if generation is not None:
+ barrier = args.updates_per_round // 4
for update_index in range(args.updates_per_round):
if (
args.wall_time_budget
@@ -1260,7 +1089,9 @@ def evaluate(
)
== "updates_and_generations"
):
- completed_generation = engine.get_stats()["completed_generation"]
+ completed_generation = engine.get_flow_status()[
+ "completed_generation"
+ ]
if completed_generation > last_saved_generation:
if pending_ids:
if args.train_device.startswith("cuda:"):
@@ -1283,76 +1114,15 @@ def evaluate(
f"generation_{completed_generation}_publication_dev",
)
)
- if auto_plan is not None:
- quota_wait_started = None
- while True:
- drain_usage()
- publication = engine.get_stats()
- quota = engine.get_sampling_budget_snapshot()
- pending = (
- publication["requested_generation"]
- > publication["completed_generation"]
- )
- room = (
- publication["successful_episodes"]
- < args.max_generated_episodes
- )
- if (
- not pending
- and room
- and quota["unlocked_remaining_chunks"]
- <= auto_plan["refill_remaining_chunks"]
- ):
- request_refill()
- pending = True
- if quota["unlocked_remaining_chunks"] >= args.batch_size:
- break
- if not pending and not room:
- end_reason = "usage_budget"
- break
- if quota_wait_started is None:
- quota_wait_started = tracker.elapsed_seconds()
- if tracker.elapsed_seconds() - quota_wait_started >= 30:
- raise TimeoutError(
- "No unlocked trajectory quota for 30 seconds"
- )
- _timeout(args, tracker)
- time.sleep(0.02)
- quota_wait_ended = tracker.elapsed_seconds()
- tracker.record_interval(
- "quota_wait",
- (
- quota_wait_ended
- if quota_wait_started is None
- else quota_wait_started
- ),
- quota_wait_ended,
- )
- if end_reason == "usage_budget":
+ try:
+ if flow is not None and not flow.prepare_batch(args.batch_size):
+ end_reason = "usage_budget"
break
- before_pacing = tracker.elapsed_seconds()
- while time.monotonic() < next_quota_update:
- drain_usage()
- _timeout(args, tracker)
- engine.get_sampling_budget_snapshot()
- time.sleep(
- min(0.02, max(0.0, next_quota_update - time.monotonic()))
- )
- tracker.record_interval(
- "usage_throttle", before_pacing, tracker.elapsed_seconds()
- )
- next_quota_update = time.monotonic() + args.batch_size / min(
- auto_plan["training_chunks_per_second_target"],
- auto_plan["active_chunks_per_second_cap"],
- )
- if barrier is not None and update_index == barrier:
- before = tracker.elapsed_seconds()
- try:
- _wait(engine, generation, args, tracker, progress=drain_usage)
- finally:
- tracker.record_interval(
- "data_wait", before, tracker.elapsed_seconds()
- )
+ if barrier is not None and update_index == barrier:
+ flow.wait_for_generation(generation)
+ except _BudgetExpired:
+ end_reason = "wall_time_budget"
+ break
if (
args.wall_time_budget
and tracker.elapsed_seconds() >= args.wall_time_budget
@@ -1362,7 +1132,17 @@ def evaluate(
service_started = tracker.elapsed_seconds()
try:
with sampling_rng.scope():
- batch = _sample(source, args.batch_size, tracker, sequence)
+ batch = (
+ source.sample_context(args.batch_size, sequence)
+ if flow is None
+ else flow.sample_context(source, args.batch_size, sequence)
+ )
+ if flow is None:
+ ended = tracker.elapsed_seconds()
+ tracker.record_interval("data_wait", ended, ended)
+ except _BudgetExpired:
+ end_reason = "wall_time_budget"
+ break
except DataBudgetExhaustedError:
if corpus is None:
raise
@@ -1399,15 +1179,11 @@ def evaluate(
)
step += 1
train_steps += 1
- target_mask = batch.get("loss_mask", None)
+ target_mask = batch["loss_mask"]
supervised_ids = (
completed_target_ids
if usage is not None
- else (
- batch["frame_ids"]
- if target_mask is None
- else batch["frame_ids"][target_mask]
- )
+ else batch["frame_ids"][target_mask]
)
if usage is None:
pending_ids.append(supervised_ids)
@@ -1429,11 +1205,9 @@ def evaluate(
tracker.record_interval(
"training_loop", update_started, tracker.elapsed_seconds()
)
- versions, counts = (
- batch["generations"]
- if target_mask is None
- else batch["generations"][target_mask]
- ).unique(return_counts=True)
+ versions, counts = batch["generations"][target_mask].unique(
+ return_counts=True
+ )
for version, count in zip(versions.tolist(), counts.tolist()):
generation_counts[version] = (
generation_counts.get(version, 0) + count
@@ -1549,7 +1323,7 @@ def evaluate(
if auto_plan is not None:
from examples.ods.simple_bc_demo.flow import summarize_flow
- flow = summarize_flow(
+ flow_measurements = summarize_flow(
monotonic_origin_seconds=tracker.start_monotonic,
intervals=tracker.intervals,
refill_requests=[
@@ -1563,7 +1337,7 @@ def evaluate(
buffer_size=args.buffer_size,
)
(args.output / "flow_measurements.json").write_text(
- json.dumps(flow, indent=2)
+ json.dumps(flow_measurements, indent=2)
)
save_usage(usage_report)
if args.defer_evaluation:
diff --git a/tests/data_pipeline/test_online_data_budget.py b/tests/data_pipeline/test_online_data_budget.py
index 1aaa0a53d..0eb4e26fa 100644
--- a/tests/data_pipeline/test_online_data_budget.py
+++ b/tests/data_pipeline/test_online_data_budget.py
@@ -29,7 +29,7 @@
from __future__ import annotations
import multiprocessing as mp
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
import pytest
import torch
@@ -492,3 +492,147 @@ def test_shared_quota_draw_preserves_explicit_generator_and_global_rng() -> None
)
assert torch.equal(first.get_state(), second.get_state())
assert torch.equal(global_before, torch.get_rng_state())
+
+
+def test_flow_status_matches_audit_across_ineligible_and_replaced_sources() -> None:
+ engine = _engine(rows=6, budget=3, hard=3)
+
+ def compare(expected_remaining: int, expected_delivered: int) -> None:
+ compact = engine.get_flow_status()
+ audit = engine.get_sampling_budget_snapshot()
+ assert compact["unlocked_remaining_chunks"] == expected_remaining
+ assert compact["total_delivered_chunks"] == expected_delivered
+ assert (
+ compact["unlocked_remaining_chunks"] == audit["unlocked_remaining_chunks"]
+ )
+ assert compact["total_delivered_chunks"] == audit["total_delivered_chunks"]
+
+ compare(18, 0)
+ with engine._lock_index.get_lock():
+ engine._sampling_budget.debit(torch.tensor([0, 0, 5]))
+ engine._lock_index[:] = [0, 1]
+ engine.shared_buffer["valid"][1].zero_()
+ engine.shared_buffer["continuity_id"][2].fill_(-1)
+ # An unpublished row with stale valid payload cannot lend credits.
+ engine._sampling_budget.rows[3 * 4] = -1
+ compare(5, 3)
+ with engine._lock_index.get_lock():
+ engine._sampling_budget.publish(5, 1, 6, 2)
+ engine.shared_buffer["ods_episode_id"][5].fill_(6)
+ # The replacement retires unused credits and resets only its own delivery.
+ compare(6, 3)
+ with engine._lock_index.get_lock():
+ engine._sampling_budget.debit(torch.tensor([5, 5, 5]))
+ compare(3, 6)
+
+
+@pytest.mark.parametrize("with_continuity", [False, True])
+def test_flow_status_requires_a_joint_valid_continuity_anchor(
+ with_continuity: bool,
+) -> None:
+ engine = _engine(rows=2, budget=3, hard=3)
+ with engine._lock_index.get_lock():
+ if with_continuity:
+ # Separate valid frames and continuity-qualified frames do not
+ # constitute one usable anchor.
+ engine.shared_buffer["valid"][0] = torch.tensor([True, False, False, False])
+ engine.shared_buffer["continuity_id"][0] = torch.tensor([-1, 0, 0, 0])
+ else:
+ engine.shared_buffer = engine.shared_buffer.exclude("continuity_id")
+ engine.shared_buffer["valid"][0].zero_()
+ assert engine.get_flow_status()["unlocked_remaining_chunks"] == 3
+ assert engine.get_sampling_budget_snapshot()["unlocked_remaining_chunks"] == 3
+
+
+@pytest.mark.parametrize("budget", [None, 3])
+def test_flow_status_is_detached_and_does_not_build_full_audits(
+ monkeypatch, budget: int | None
+) -> None:
+ engine = _engine(rows=2, budget=budget, hard=budget)
+ with engine._progress.values.get_lock():
+ engine._progress.values[:3] = [3, 2, 7]
+
+ def unexpected_audit():
+ pytest.fail("Compact flow status must not invoke full diagnostics")
+
+ monkeypatch.setattr(engine, "get_stats", unexpected_audit)
+ monkeypatch.setattr(engine, "get_sampling_budget_snapshot", unexpected_audit)
+ assert engine.check_health() is None
+ with torch.device("meta"):
+ status = engine.get_flow_status()
+ assert status == {
+ "requested_generation": 3,
+ "completed_generation": 2,
+ "successful_episodes": 7,
+ "unlocked_remaining_chunks": None if budget is None else 6,
+ "total_delivered_chunks": None if budget is None else 0,
+ }
+ assert all(value is None or type(value) is int for value in status.values())
+ status["requested_generation"] = -99
+ status["unlocked_remaining_chunks"] = -99
+ assert engine.get_flow_status()["requested_generation"] == 3
+ assert engine.get_flow_status()["unlocked_remaining_chunks"] == (
+ None if budget is None else 6
+ )
+ assert not engine._fill_signal.is_set()
+
+
+@pytest.mark.parametrize(
+ "state",
+ [
+ OnlineDataEngineState.CREATED,
+ OnlineDataEngineState.STARTING,
+ OnlineDataEngineState.STOPPED,
+ ],
+)
+@pytest.mark.parametrize("method", ["check_health", "get_flow_status"])
+def test_flow_apis_require_ready_without_removing_stopped_audits(
+ state: OnlineDataEngineState, method: str
+) -> None:
+ engine = _engine()
+ engine._set_state(state)
+ with pytest.raises(RuntimeError, match=f"current state is {state.value}"):
+ getattr(engine, method)()
+ if state is OnlineDataEngineState.STOPPED:
+ assert engine.get_sampling_budget_snapshot()["unlocked_remaining_chunks"] == 6
+
+
+@pytest.mark.parametrize("method", ["check_health", "get_flow_status"])
+@pytest.mark.parametrize("shared_error", [False, True])
+def test_flow_apis_preserve_worker_and_local_failure_exceptions(
+ method: str, shared_error: bool
+) -> None:
+ engine = _engine(rows=1, budget=1, hard=1)
+ engine.sample_batch(1, 1)
+ error = TimeoutError("producer refill failed")
+ if shared_error:
+ engine._record_worker_error(error)
+ else:
+ engine._worker_error = error
+ engine._set_state(OnlineDataEngineState.FAILED)
+ with pytest.raises(TimeoutError, match="producer refill failed") as captured:
+ getattr(engine, method)()
+ assert captured.value is error
+
+
+@pytest.mark.parametrize("method", ["check_health", "get_flow_status"])
+def test_flow_apis_detect_ready_worker_exit(monkeypatch, method: str) -> None:
+ engine = _engine()
+ process = MagicMock(exitcode=17)
+ process.is_alive.return_value = False
+ engine._sim_process = process
+ monkeypatch.setattr(engine, "_is_owner_process", lambda: True)
+ with pytest.raises(RuntimeError, match="exit code 17"):
+ getattr(engine, method)()
+ assert engine.state is OnlineDataEngineState.FAILED
+ process.join.assert_called_once_with(timeout=0)
+
+
+@pytest.mark.parametrize("method", ["check_health", "get_flow_status"])
+def test_flow_apis_reject_failed_engine_without_error_envelope(method: str) -> None:
+ engine = _engine()
+ engine._set_state(OnlineDataEngineState.FAILED)
+ with pytest.raises(
+ data_module.OnlineDataWorkerError, match="without an error envelope"
+ ):
+ getattr(engine, method)()
diff --git a/tests/data_pipeline/test_runtime.py b/tests/data_pipeline/test_runtime.py
new file mode 100644
index 000000000..48b651f40
--- /dev/null
+++ b/tests/data_pipeline/test_runtime.py
@@ -0,0 +1,795 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ----------------------------------------------------------------------------
+
+"""CPU flow-control contracts, including bounded child-process publication."""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import replace
+import os
+import threading
+import time
+from types import SimpleNamespace
+
+import pytest
+import torch
+from tensordict import TensorDict
+
+from embodichain.data_pipeline import runtime as runtime_module
+from embodichain.data_pipeline.configuration import (
+ ODSReuseProfileCfg,
+ ODSRuntimeCfg,
+ configure_ods,
+)
+from embodichain.data_pipeline.datasets import (
+ DataBudgetExhaustedError,
+ DataUnavailableError,
+ OnlineDataSource,
+ SequenceSpecCfg,
+)
+from embodichain.data_pipeline.engine import data as engine_module
+from embodichain.data_pipeline.engine.data import (
+ OnlineDataEngine,
+ OnlineDataEngineCfg,
+ OnlineDataEngineState,
+)
+from embodichain.data_pipeline.runtime import ODSFlowController, ODSFlowControllerCfg
+
+
+def _plan():
+ # Two chunks every quarter second; two trajectories with four credits each.
+ return configure_ods(
+ 2,
+ 8,
+ profile=ODSReuseProfileCfg(
+ profile_id="runtime-cpu",
+ scope="deterministic flow tests",
+ qualified=False,
+ min_chunks_per_episode=4,
+ max_chunks_per_episode=4,
+ hard_max_chunks_per_episode=8,
+ latency_safety_factor=1,
+ ),
+ runtime=ODSRuntimeCfg(
+ batch_size=2,
+ producer_envs=1,
+ first_publication_p95_seconds=0.5,
+ ),
+ )
+
+
+class _Clock:
+ def __init__(self):
+ self.now = 0.0
+ self.sleeps = []
+ self.scheduled = []
+
+ def monotonic(self):
+ return self.now
+
+ def sleep(self, seconds):
+ assert seconds >= 0, "Polling must never request a negative sleep"
+ self.sleeps.append(seconds)
+ self.now += seconds
+ due = [item for item in self.scheduled if item[0] <= self.now]
+ self.scheduled = [item for item in self.scheduled if item[0] > self.now]
+ for _, callback in due:
+ callback()
+
+
+@pytest.fixture
+def clock(monkeypatch):
+ result = _Clock()
+ monkeypatch.setattr(runtime_module, "time", result)
+ return result
+
+
+class _Engine:
+ def __init__(self, clock):
+ plan = _plan()
+ self.clock = clock
+ self.buffer_size = plan.buffer_size
+ self.cfg = SimpleNamespace(
+ refill_policy="manual",
+ episode_sample_budget=plan.episode_sample_budget,
+ episode_sample_hard_limit=plan.episode_sample_hard_limit,
+ gym_config={"num_envs": 1},
+ )
+ self.owner = os.getpid()
+ self.events = []
+ self.error = None
+ self.health_checks = 0
+ self.refills = []
+ self.waits = []
+ self.status = {
+ "requested_generation": 1,
+ "completed_generation": 1,
+ "successful_episodes": self.buffer_size,
+ "unlocked_remaining_chunks": 8,
+ "total_delivered_chunks": 0,
+ }
+
+ def _require_owner_process(self, operation):
+ if os.getpid() != self.owner:
+ raise RuntimeError(f"{operation} requires owner process")
+
+ def check_health(self):
+ self.health_checks += 1
+ if self.error is not None:
+ raise self.error
+
+ def drain_lifecycle_events(self):
+ events, self.events = self.events, []
+ return events
+
+ def get_flow_status(self):
+ self.check_health()
+ return dict(self.status)
+
+ def request_refill(self):
+ if self.status["requested_generation"] == self.status["completed_generation"]:
+ self.status["requested_generation"] += 1
+ generation = self.status["requested_generation"]
+ self.refills.append(generation)
+ return generation
+
+ def publish(self):
+ generation = self.status["requested_generation"]
+ self.status.update(
+ completed_generation=generation,
+ successful_episodes=self.status["successful_episodes"] + self.buffer_size,
+ unlocked_remaining_chunks=8,
+ )
+ self.events.append({"kind": "published", "generation": generation})
+
+ def wait_for_generation(self, generation, timeout):
+ self.waits.append((generation, timeout))
+ if generation > self.status["requested_generation"]:
+ raise ValueError("generation has not been requested")
+ if self.status["completed_generation"] >= generation:
+ return
+ self.clock.sleep(timeout)
+ if self.status["completed_generation"] < generation:
+ raise TimeoutError("publication pending")
+
+
+class _Source:
+ def __init__(self, engine, responses):
+ self.engine = engine
+ self.responses = deque(responses)
+ self.calls = 0
+
+ def sample_context(self, *args, **kwargs):
+ self.calls += 1
+ response = (
+ self.responses[0] if len(self.responses) == 1 else self.responses.popleft()
+ )
+ if isinstance(response, BaseException):
+ raise response
+ return response
+
+
+@pytest.mark.parametrize(
+ "field", ["data_wait_timeout", "generation_timeout", "poll_interval"]
+)
+@pytest.mark.parametrize("value", [0, -1, True, float("nan"), float("inf")])
+def test_wait_bounds_must_be_positive_finite_numbers(field, value):
+ with pytest.raises(ValueError, match=field):
+ ODSFlowControllerCfg(**{field: value})
+
+
+@pytest.mark.parametrize("value", [0, -1, True, 1.5])
+def test_production_limit_requires_a_positive_integer(value):
+ with pytest.raises(ValueError, match="max_generated_episodes"):
+ ODSFlowControllerCfg(max_generated_episodes=value)
+
+
+def test_controller_requires_creator_engine_manual_refill_and_matching_plan(clock):
+ engine = _Engine(clock)
+ engine.owner = -1
+ with pytest.raises(RuntimeError, match="owner"):
+ ODSFlowController(engine)
+ engine.owner = os.getpid()
+ engine.cfg.refill_policy = "sample_threshold"
+ with pytest.raises(ValueError, match="manual"):
+ ODSFlowController(engine)
+ engine.cfg.refill_policy = "manual"
+ with pytest.raises(ValueError, match="initial cache"):
+ ODSFlowController(engine, ODSFlowControllerCfg(max_generated_episodes=1))
+ with pytest.raises(ValueError, match="match the plan"):
+ ODSFlowController(engine, plan=replace(_plan(), buffer_size=3))
+ with pytest.raises(ValueError, match="match the plan"):
+ ODSFlowController(engine, plan=replace(_plan(), episode_sample_budget=3))
+
+
+@pytest.mark.parametrize("planned", [True, False])
+@pytest.mark.parametrize("producers", [None, 0, True, 1.5, 2])
+def test_protected_flow_requires_explicit_divisible_producer_count(
+ clock, planned, producers
+):
+ engine = _Engine(clock)
+ engine.buffer_size = 3 # A two-environment writer would overlap the final row.
+ engine.cfg.gym_config = {} if producers is None else {"num_envs": producers}
+ with pytest.raises(ValueError, match="explicit num_envs dividing the cache"):
+ ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(max_generated_episodes=None if planned else 6),
+ plan=replace(_plan(), buffer_size=3) if planned else None,
+ )
+
+
+def test_divisible_producer_count_must_still_match_measured_plan(clock):
+ engine = _Engine(clock)
+ # Both one and two producers divide this cache, but the plan measured one.
+ engine.cfg.gym_config["num_envs"] = 2
+ with pytest.raises(ValueError, match="producer count.*match the plan"):
+ ODSFlowController(engine, plan=_plan())
+
+
+@pytest.mark.parametrize(
+ "operation",
+ ["poll", "drain_events", "request_refill", "prepare_batch", "wait_for_generation"],
+)
+def test_control_operations_reject_other_threads_before_touching_engine(
+ clock, operation
+):
+ engine = _Engine(clock)
+ flow = ODSFlowController(engine)
+ errors = []
+
+ def other_thread():
+ try:
+ method = getattr(flow, operation)
+ (
+ method(1)
+ if operation in {"prepare_batch", "wait_for_generation"}
+ else method()
+ )
+ except BaseException as error:
+ errors.append(error)
+
+ thread = threading.Thread(target=other_thread)
+ thread.start()
+ thread.join(2)
+ assert not thread.is_alive()
+ assert len(errors) == 1 and isinstance(errors[0], RuntimeError)
+ assert "owner-process/thread" in str(errors[0])
+ assert engine.health_checks == 0 and engine.refills == [] and engine.waits == []
+
+
+def test_unplanned_prepare_does_not_implicitly_request_data(clock):
+ engine = _Engine(clock)
+ engine.status["unlocked_remaining_chunks"] = 0
+ flow = ODSFlowController(engine)
+ assert flow.prepare_batch(3)
+ assert engine.refills == [] and clock.sleeps == []
+
+
+def test_controller_rejects_inherited_process_ownership(clock, monkeypatch):
+ engine = _Engine(clock)
+ flow = ODSFlowController(engine)
+ monkeypatch.setattr(
+ runtime_module, "os", SimpleNamespace(getpid=lambda: os.getpid() + 1)
+ )
+ with pytest.raises(RuntimeError, match="owner-process/thread"):
+ flow.poll()
+ assert engine.health_checks == 0
+
+
+def test_planned_batch_size_mismatch_fails_before_requesting_data(clock):
+ engine = _Engine(clock)
+ flow = ODSFlowController(engine, plan=_plan())
+ with pytest.raises(ValueError, match="measured plan"):
+ flow.prepare_batch(3)
+ assert engine.refills == [] and engine.health_checks == 0
+
+
+def test_refill_observer_reports_each_coalesced_generation_once(clock):
+ engine = _Engine(clock)
+ observed = []
+ flow = ODSFlowController(engine, on_refill=lambda *event: observed.append(event))
+ assert flow.request_refill() == flow.request_refill() == 2
+ engine.publish()
+ clock.sleep(1)
+ assert flow.request_refill() == 3
+ assert observed == [(2, 0), (3, 1)]
+
+
+def test_pending_refill_completion_race_does_not_start_an_extra_generation(
+ clock, monkeypatch
+):
+ engine = _Engine(clock)
+ engine.request_refill()
+ engine.refills.clear()
+ observed = []
+ flow = ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(max_generated_episodes=4),
+ on_refill=lambda *event: observed.append(event),
+ )
+ status = engine.get_flow_status
+
+ def complete_after_snapshot():
+ snapshot = status()
+ if snapshot["requested_generation"] > snapshot["completed_generation"]:
+ engine.publish()
+ return snapshot
+
+ monkeypatch.setattr(engine, "get_flow_status", complete_after_snapshot)
+ assert flow.request_refill() == 2
+ assert engine.status["completed_generation"] == 2
+ assert engine.status["requested_generation"] == 2
+ assert engine.status["successful_episodes"] == 4
+ assert engine.refills == []
+ # Adopting another caller's earlier request must not fabricate its timestamp.
+ assert observed == []
+ assert flow.request_refill() is None
+
+
+def test_adopting_pending_refill_preserves_only_actual_request_timestamps(clock):
+ engine = _Engine(clock)
+ engine.request_refill()
+ observed = []
+ flow = ODSFlowController(engine, on_refill=lambda *event: observed.append(event))
+ clock.sleep(3)
+ assert flow.request_refill() == 2
+ assert flow.request_refill() == 2
+ assert observed == []
+ engine.publish()
+ clock.sleep(1)
+ assert flow.request_refill() == 3
+ assert observed == [(3, 4)]
+
+
+def test_partial_pending_fill_is_coalesced_even_when_another_cache_would_not_fit(clock):
+ engine = _Engine(clock)
+ flow = ODSFlowController(engine, ODSFlowControllerCfg(max_generated_episodes=4))
+ assert flow.request_refill() == 2
+ engine.status["successful_episodes"] = 3
+ assert flow.request_refill() == 2
+ assert engine.status["requested_generation"] == 2
+
+
+def test_production_ceiling_allows_only_complete_caches(clock):
+ engine = _Engine(clock)
+ # Initial2 + one refill2 fits5; the unused final capacity1 must not be filled.
+ flow = ODSFlowController(engine, ODSFlowControllerCfg(max_generated_episodes=5))
+ assert flow.request_refill() == 2
+ assert flow.request_refill() == 2
+ engine.publish()
+ assert flow.request_refill() is None
+ assert engine.status["successful_episodes"] == 4
+
+
+def test_remaining_final_supply_is_consumed_before_stopping(clock):
+ engine = _Engine(clock)
+ flow = ODSFlowController(
+ engine, ODSFlowControllerCfg(max_generated_episodes=2), plan=_plan()
+ )
+ engine.status["unlocked_remaining_chunks"] = 2
+ assert flow.prepare_batch(2)
+ engine.status["unlocked_remaining_chunks"] = 1
+ assert flow.prepare_batch(2) is False
+ assert engine.refills == [] and clock.sleeps == []
+
+
+def test_quota_wait_drains_publications_then_funds_a_batch(clock):
+ engine = _Engine(clock)
+ engine.status["unlocked_remaining_chunks"] = 0
+ publications, intervals = [], []
+ flow = ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(data_wait_timeout=1, poll_interval=0.1),
+ plan=_plan(),
+ on_events=publications.extend,
+ on_interval=lambda *interval: intervals.append(interval),
+ )
+ clock.scheduled.append((0.2, engine.publish))
+ assert flow.prepare_batch(2)
+ assert publications == [{"kind": "published", "generation": 2}]
+ assert intervals[0] == ("quota_wait", 0, pytest.approx(0.2))
+ assert engine.refills == [2]
+
+
+def test_quota_wait_has_one_total_deadline(clock):
+ engine = _Engine(clock)
+ engine.status["unlocked_remaining_chunks"] = 0
+ intervals = []
+ flow = ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(data_wait_timeout=0.25, poll_interval=0.1),
+ plan=_plan(),
+ on_interval=lambda *interval: intervals.append(interval),
+ )
+ with pytest.raises(TimeoutError, match="quota"):
+ flow.prepare_batch(2)
+ assert clock.now == pytest.approx(0.25)
+ assert intervals == [("quota_wait", 0, pytest.approx(0.25))]
+
+
+def test_pacing_uses_slower_cap_without_catch_up_after_slow_service(clock):
+ plan = replace(_plan(), active_chunks_per_second_cap=4)
+ flow = ODSFlowController(_Engine(clock), plan=plan)
+ assert flow.prepare_batch(2)
+ clock.sleep(10) # Long optimizer/evaluation pause earns no burst allowance.
+ assert flow.prepare_batch(2)
+ assert clock.now == 10
+ assert flow.prepare_batch(2)
+ assert clock.now == pytest.approx(10.5)
+
+
+def test_poll_callback_crossing_pacing_deadline_never_sleeps_negative(clock):
+ callbacks = []
+ armed = False
+
+ def progress():
+ if armed:
+ callbacks.append(clock.now)
+ if len(callbacks) == 2:
+ clock.sleep(0.3)
+
+ flow = ODSFlowController(_Engine(clock), plan=_plan(), progress=progress)
+ assert flow.prepare_batch(2)
+ armed = True
+ assert flow.prepare_batch(2)
+ assert clock.now == pytest.approx(0.3)
+ assert flow.prepare_batch(2)
+ assert clock.now == pytest.approx(0.55)
+
+
+@pytest.mark.parametrize("external_deadline,expected", [(None, 0.4), (0.15, 0.15)])
+def test_generation_wait_uses_total_or_shorter_absolute_deadline(
+ clock, external_deadline, expected
+):
+ engine = _Engine(clock)
+ engine.request_refill()
+ flow = ODSFlowController(
+ engine, ODSFlowControllerCfg(generation_timeout=0.4, poll_interval=0.1)
+ )
+ with pytest.raises(TimeoutError, match="deadline"):
+ flow.wait_for_generation(2, deadline=external_deadline)
+ assert clock.now == pytest.approx(expected)
+ assert all(0 < duration <= 0.1 for _, duration in engine.waits)
+
+
+def test_generation_wait_releases_events_and_propagates_invalid_generation(clock):
+ engine = _Engine(clock)
+ engine.request_refill()
+ events = []
+ flow = ODSFlowController(
+ engine, ODSFlowControllerCfg(poll_interval=0.1), on_events=events.extend
+ )
+ clock.scheduled.append((0.2, engine.publish))
+ flow.wait_for_generation(2)
+ assert events == [{"kind": "published", "generation": 2}]
+ with pytest.raises(ValueError, match="not been requested"):
+ flow.wait_for_generation(3)
+
+
+@pytest.mark.parametrize("wait", ["quota", "pacing", "generation", "sample"])
+def test_application_progress_can_interrupt_every_wait(clock, wait):
+ class Cancelled(Exception):
+ pass
+
+ def progress():
+ if clock.now >= 0.1:
+ raise Cancelled("application deadline")
+
+ engine = _Engine(clock)
+ flow = ODSFlowController(engine, plan=_plan(), progress=progress)
+ if wait == "quota":
+ engine.status["unlocked_remaining_chunks"] = 0
+ operation = lambda: flow.prepare_batch(2)
+ elif wait == "pacing":
+ flow.prepare_batch(2)
+ operation = lambda: flow.prepare_batch(2)
+ elif wait == "generation":
+ engine.request_refill()
+ operation = lambda: flow.wait_for_generation(2)
+ else:
+ source = _Source(engine, [DataUnavailableError("pending")])
+ operation = lambda: flow.sample_context(source, 2, SequenceSpecCfg())
+ with pytest.raises(Cancelled, match="application deadline"):
+ operation()
+ assert clock.now <= 0.12
+
+
+def test_transient_sampling_retries_drain_events_without_changing_request(clock):
+ engine = _Engine(clock)
+ events, intervals = [], []
+ batch = object()
+ source = _Source(
+ engine,
+ [
+ DataUnavailableError("locked"),
+ DataBudgetExhaustedError("pending quota"),
+ batch,
+ ],
+ )
+ flow = ODSFlowController(
+ engine,
+ ODSFlowControllerCfg(poll_interval=0.1),
+ on_events=events.extend,
+ on_interval=lambda *interval: intervals.append(interval),
+ )
+ clock.scheduled.append((0.1, lambda: engine.events.append({"kind": "published"})))
+ assert flow.sample_context(source, 2, SequenceSpecCfg()) is batch
+ assert source.calls == 3 and engine.refills == []
+ assert events == [{"kind": "published"}]
+ assert intervals == [("data_wait", 0, pytest.approx(0.2))]
+
+
+def test_sampling_retry_deadline_does_not_reset_after_each_transient_error(clock):
+ engine = _Engine(clock)
+ source = _Source(engine, [DataUnavailableError("no eligible window")])
+ flow = ODSFlowController(
+ engine, ODSFlowControllerCfg(data_wait_timeout=0.25, poll_interval=0.1)
+ )
+ with pytest.raises(TimeoutError, match="context"):
+ flow.sample_context(source, 2, SequenceSpecCfg())
+ assert clock.now == pytest.approx(0.25)
+ assert source.calls == 4
+
+
+@pytest.mark.parametrize("worker_failure", [True, False])
+def test_worker_and_schema_errors_are_propagated_without_retry(clock, worker_failure):
+ engine = _Engine(clock)
+ original = ValueError("worker failed" if worker_failure else "schema mismatch")
+ source = _Source(engine, [original])
+ if worker_failure:
+ engine.error = original
+ flow = ODSFlowController(engine)
+ with pytest.raises(ValueError) as caught:
+ flow.sample_context(source, 2, SequenceSpecCfg())
+ assert caught.value is original
+ assert source.calls == (0 if worker_failure else 1)
+ assert clock.sleeps == []
+
+
+def test_interval_observer_failure_retains_original_worker_error(clock):
+ engine = _Engine(clock)
+ original = ValueError("worker planner failed")
+ engine.error = original
+ source = _Source(engine, [object()])
+
+ def observer(*interval):
+ raise RuntimeError("interval export failed")
+
+ flow = ODSFlowController(engine, on_interval=observer)
+ with pytest.raises(ValueError) as caught:
+ flow.sample_context(source, 2, SequenceSpecCfg())
+ assert caught.value is original
+ assert source.calls == 0
+ assert any("interval export failed" in note for note in original.__notes__)
+
+
+def test_interval_observer_failure_without_primary_error_is_propagated(clock):
+ engine = _Engine(clock)
+ source = _Source(engine, [object()])
+ original = RuntimeError("interval export failed")
+
+ def observer(*interval):
+ raise original
+
+ flow = ODSFlowController(engine, on_interval=observer)
+ with pytest.raises(RuntimeError) as caught:
+ flow.sample_context(source, 2, SequenceSpecCfg())
+ assert caught.value is original
+ assert source.calls == 1
+ assert clock.sleeps == []
+
+
+def test_diagnostic_events_remain_drainable_after_worker_failure(clock):
+ engine = _Engine(clock)
+ original = RuntimeError("worker exited")
+ engine.error = original
+ engine.events.append({"kind": "published", "generation": 1})
+ events = []
+ flow = ODSFlowController(engine, on_events=events.extend)
+ with pytest.raises(RuntimeError) as caught:
+ flow.poll()
+ assert caught.value is original
+ engine.events.append({"kind": "evicted", "generation": 1})
+ assert flow.drain_events() == [{"kind": "evicted", "generation": 1}]
+ assert events == [
+ {"kind": "published", "generation": 1},
+ {"kind": "evicted", "generation": 1},
+ ]
+
+
+def test_sampling_rejects_an_uncontrolled_source(clock):
+ engine = _Engine(clock)
+ source = _Source(_Engine(clock), [object()])
+ with pytest.raises(ValueError, match="controlled live engine"):
+ ODSFlowController(engine).sample_context(source, 2, SequenceSpecCfg())
+ assert source.calls == 0
+
+
+def _buffer(engine):
+ rows, steps = engine.cfg.buffer_size, 4
+ shape = rows, steps
+ row_ids = torch.arange(rows, dtype=torch.int64)[:, None].expand(shape)
+ return TensorDict(
+ {
+ "obs": torch.arange(rows * steps * 2, dtype=torch.float32).reshape(
+ *shape, 2
+ ),
+ "actions": torch.zeros(*shape, 1),
+ "valid": torch.ones(shape, dtype=torch.bool),
+ "continuity_id": torch.zeros(shape, dtype=torch.int64),
+ "segment_id": torch.zeros(shape, dtype=torch.int64),
+ "segment_accepted": torch.ones(shape, dtype=torch.bool),
+ "ods_run_id": torch.full(shape, engine.run_id, dtype=torch.int64),
+ "ods_episode_id": row_ids.clone(),
+ "ods_generation": torch.ones(shape, dtype=torch.int64),
+ "ods_frame_index": torch.arange(steps).expand(shape).clone(),
+ },
+ batch_size=shape,
+ device="cpu",
+ ).share_memory_()
+
+
+_SCHEMA = {
+ "observation": {"names": ["position", "speed"]},
+ "action": {"names": ["request"], "semantics": "policy_request", "step_dt": 1 / 30},
+}
+
+
+def _engine_cfg(**changes):
+ values = dict(
+ buffer_size=2,
+ max_episode_steps=4,
+ state_dim=2,
+ action_dim=1,
+ action_semantics="policy_request",
+ action_contracts=("gripper.parallel@1",),
+ gym_config={"num_envs": 1},
+ include_provenance=True,
+ refill_policy="manual",
+ initialization_timeout=15,
+ shutdown_timeout=2,
+ )
+ return OnlineDataEngineCfg(**(values | changes))
+
+
+def test_flow_preserves_actual_source_batches_padding_and_independent_rng(monkeypatch):
+ monkeypatch.setattr(OnlineDataEngine, "_create_buffer", _buffer)
+ engine = OnlineDataEngine(_engine_cfg())
+ # This test exercises source sampling, while worker liveness is covered below.
+ monkeypatch.setattr(engine, "_ensure_worker_alive", lambda: None)
+ engine._set_state(OnlineDataEngineState.READY)
+ engine._init_signal.set()
+ with engine._lock_index.get_lock():
+ engine._lock_index[:] = [-1, -1]
+ try:
+ source = OnlineDataSource(engine, schema=_SCHEMA)
+ flow = ODSFlowController(engine)
+ sequence = SequenceSpecCfg(history_length=2, prediction_horizon=4)
+ direct_rng = torch.Generator().manual_seed(91)
+ flow_rng = torch.Generator().manual_seed(91)
+ global_rng = torch.get_rng_state().clone()
+ direct = source.sample_context(64, sequence, generator=direct_rng)
+ actual = flow.sample_context(source, 64, sequence, generator=flow_rng)
+ assert set(direct.keys()) == set(actual.keys())
+ assert all(torch.equal(direct[key], actual[key]) for key in direct.keys())
+ assert torch.equal(direct_rng.get_state(), flow_rng.get_state())
+ assert torch.equal(torch.get_rng_state(), global_rng)
+ assert not actual["loss_mask"][:, 0].any()
+ assert not actual["valid"].all()
+ finally:
+ engine._set_state(OnlineDataEngineState.CREATED)
+ engine.stop()
+
+
+def _numeric_worker(
+ cfg,
+ shared,
+ lock,
+ fill,
+ init,
+ close,
+ error_buffer,
+ error_length,
+ failed,
+ state,
+ language,
+ progress,
+ run_id,
+ events,
+ quota,
+):
+ """Publish numeric trajectories through real engine IPC, without a simulator."""
+ previous = {}
+ try:
+ while not close.is_set():
+ if not fill.wait(0.05):
+ continue
+ fill.clear()
+ with progress.values.get_lock():
+ generation = int(progress.values[0])
+ for slot in range(len(shared)):
+ required = 2 if slot in previous else 1
+ with events.condition:
+ if (
+ events.cursors[1] - events.cursors[0] + required
+ > events.capacity
+ ):
+ cfg.gym_config["backpressure_observed"].set()
+ if not events.wait_for_space(required, close):
+ return
+ with lock.get_lock():
+ lock[:] = [slot, slot + 1]
+ if slot in previous:
+ events.append(previous[slot] | {"kind": "evicted"})
+ shared["obs"][slot].fill_(generation)
+ shared["ods_episode_id"][slot].fill_(
+ generation * len(shared) + slot
+ )
+ shared["ods_generation"][slot].fill_(generation)
+ event = engine_module._publication_event(
+ shared[slot],
+ run_id=run_id,
+ episode_id=generation * len(shared) + slot,
+ generation=generation,
+ slot=slot,
+ monotonic_ns=time.monotonic_ns(),
+ )
+ events.append(event)
+ previous[slot] = event
+ lock[:] = [-1, -1]
+ with progress.values.get_lock():
+ progress.values[2] += 1
+ with progress.values.get_lock():
+ progress.values[1] = generation
+ init.set()
+ except BaseException as error:
+ engine_module._publish_worker_error(
+ error_buffer, error_length, failed, state, error
+ )
+
+
+def test_generation_wait_drains_a_full_bounded_log_from_a_real_child(monkeypatch):
+ import multiprocessing as mp
+
+ context = mp.get_context("forkserver")
+ observed = context.Event()
+ monkeypatch.setattr(OnlineDataEngine, "_create_buffer", _buffer)
+ monkeypatch.setattr(engine_module, "_sim_worker_fn", _numeric_worker)
+ engine = OnlineDataEngine(_engine_cfg(lifecycle_event_capacity=4))
+ # Test-only IPC handle: configclass copies ordinary authored configuration.
+ engine.cfg.gym_config["backpressure_observed"] = observed
+ events = []
+ try:
+ engine.start()
+ # Request directly so the initial publication records remain in the log.
+ generation = engine.request_refill()
+ assert observed.wait(10), "Producer did not reach the bounded-log barrier"
+ flow = ODSFlowController(
+ engine, ODSFlowControllerCfg(generation_timeout=5), on_events=events.extend
+ )
+ flow.wait_for_generation(generation)
+ assert [event["sequence"] for event in events] == list(range(6))
+ assert [event["kind"] for event in events].count("evicted") == 2
+ source = OnlineDataSource(engine, schema=_SCHEMA)
+ batch = flow.sample_context(source, 4, SequenceSpecCfg())
+ assert (batch["generations"][batch["valid"]] == generation).all()
+ assert (batch["obs"][batch["valid"]] == generation).all()
+ finally:
+ worker_pid = None if engine._sim_process is None else engine._sim_process.pid
+ engine.stop()
+ assert engine.state is OnlineDataEngineState.STOPPED
+ assert worker_pid not in {child.pid for child in mp.active_children()}