diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 0c1f35e6bc..1b24b0f7b0 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -441,6 +441,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, 数据集加载, 离线加载] @@ -449,6 +450,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 dceae966bd..f795d24b6f 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 | @@ -70,10 +71,14 @@ shared memory: `CREATED → STARTING → READY`; `FAILED` and `STOPPED` are term Startup waits for initial fill with a timeout and demonstration generation has bounded attempts. Worker errors are published to all consumers; failed engines are replaced with new instances rather than restarted. +The monitor waits for sentinel readiness without consuming the forkserver's +single-reader exit-status pipe. Owner and monitor serialize status reads and +process-handle closure with a local lock; shutdown releases it before joining +the monitor. Consumer deserialization recreates that process-local lock. `sample_batch()` is READY-only. Selection excludes the currently written range, -checks validity/continuity/segment boundaries, and clones the result while -holding the shared write lock. Moving selection or cloning outside that scope +checks validity/continuity/segment boundaries, and copies the selected payload +while holding the shared write lock. Moving selection or copying outside that scope can expose partially overwritten rows. Refill follows the configured automatic or manual policy. Explicit raw-policy recording validates its ordered live action contracts independently of the expert-controller schema; see @@ -82,7 +87,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 bf2d47a85e..0caed0c023 100644 --- a/agent_context/topics/data-pipeline/online-sampling.md +++ b/agent_context/topics/data-pipeline/online-sampling.md @@ -22,7 +22,7 @@ samplers, writes the rollout, then clears the range after a full fill. 5. `sample_batch()` is READY-only. While holding the lock it excludes the current write range, checks every row in a sampled window, enforces `continuity_id`, applies - segment/boundary/episode mode rules, and clones the final batch before releasing the lock + segment/boundary/episode mode rules, and copies the final batch before releasing the lock (`engine/data.py`). `sample_batch_with_context()` samples every valid anchor, including episode ends, with a fixed causal-history/future extent. The current frame is at @@ -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 @@ -117,7 +151,7 @@ request schema remain independent. Configuration examples live in the within the stored episode. The timestamp is Unix nanoseconds at successful acceptance/publication, stamped while holding the publication lock before the row window becomes readable. It is not a monotonic or simulation timestamp. -Samplers clone provenance with observations/actions under that same lock. +Samplers copy provenance with observations/actions under that same lock. Invalid frames, including history padding, retain `-1` in every provenance field. `get_stats()` returns detached counters and excludes the active write window from valid-frame counts. These IDs do not provide shared-buffer checkpoint @@ -185,7 +219,7 @@ separates acquisition-inclusive and data-ready budgets. before selecting a window. Invalid arguments raise `ValueError`. If no unlocked valid window satisfies `chunk_size` and `sampling_mode`, it raises `DataUnavailableError` immediately; that eligibility failure does not itself change the engine state. -Refill accounting runs only after a successful clone. `OnlineDataset.__iter__()` +Refill accounting runs only after a successful copy. `OnlineDataset.__iter__()` and `SequenceDataset` propagate these exceptions, so retry policy belongs to the consumer. `DataUnavailableError` subclasses `RuntimeError`; quota exhaustion and terminal worker/read failures remain distinct and must not be retried by @@ -196,7 +230,7 @@ matching exception messages. | Symptom | Inspect | |---|---| | Engine times out or all readers fail | Inspect the shared worker error envelope and demo retry limit; do not catch only in one DataLoader consumer. | -| Samples contain rows being replaced | Selection or clone moved outside the `lock_index` critical section; restore the shared lock boundary. | +| Samples contain rows being replaced | Selection or copying moved outside the `lock_index` critical section; restore the shared lock boundary. | | Sequence crosses a reset or fragment | Fix validity/`continuity_id`/segment eligibility in `sample_batch()`, not downstream training code. | | Engine cannot restart | Expected after FAILED/STOPPED; create a new engine instance. | | Non-CPU online buffer is not truly process-shared | `_create_buffer()` currently does not reject it but skips `share_memory_()`; keep `buffer_device="cpu"` unless the implementation adds another explicit interprocess transport. | 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 78ecc03e1b..3829f7eb2e 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 cae6fc0597..bfe2e2bf72 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 d79f9dd7f7..c23e749790 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 41fdbccf43..0cac31ed09 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,101 @@ 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. Its `sample_context()` accepts only an +`OnlineDataSource` wrapping the controlled engine. Sample resident sources +directly through their own `sample_context()`; do not pass them to the controller. +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 +286,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 9e936966c8..9d9c9a5515 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 cd0606b3e4..8e3ea1615f 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/datasets/sources.py b/embodichain/data_pipeline/datasets/sources.py index 78e910be45..f01e738db5 100644 --- a/embodichain/data_pipeline/datasets/sources.py +++ b/embodichain/data_pipeline/datasets/sources.py @@ -165,7 +165,11 @@ def sample_context( if native in raw.keys(): result[canonical] = raw[native] _check_features(result, self.schema) - return _pad_context(result, raw["valid"], raw["loss_mask"]) + if self.observation_transform is not None: + # The engine already padded every copied leaf. Only a transform + # can make invalid observations nonzero again. + _pad_context(result.select("obs"), raw["valid"], raw["loss_mask"]) + return result def resolve_language(self, batch: TensorDict) -> dict[str, object]: """Resolve copied task/subtask IDs through the engine's text registry. diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py index a36d99e5c2..4304584eaa 100644 --- a/embodichain/data_pipeline/engine/data.py +++ b/embodichain/data_pipeline/engine/data.py @@ -31,6 +31,7 @@ from typing import TYPE_CHECKING, Literal from multiprocessing.sharedctypes import Synchronized, SynchronizedArray from multiprocessing.synchronize import Event as MpEvent +from multiprocessing.connection import wait as wait_for_process_events import torch from tensordict import TensorDict @@ -188,6 +189,7 @@ def _publish_worker_error( def _monitor_worker_process( process: mp.Process, + process_status_lock: threading.RLock, close_signal: MpEvent, error_buffer: SynchronizedArray, error_length: Synchronized, @@ -196,13 +198,19 @@ def _monitor_worker_process( ) -> None: """Broadcast hard worker exits without retaining the engine instance.""" try: - process.join() - if close_signal.is_set() or failed_signal.is_set(): - return - error = RuntimeError( - "OnlineDataEngine simulation worker exited unexpectedly " - f"(exit code {process.exitcode})." - ) + # Forkserver status lives in a single-consumer pipe: concurrent + # Process.join()/is_alive()/exitcode reads can replace a valid status + # with 255 after the second reader gets EOF. Readiness waiting does not + # consume that pipe and must not hold the owner-side shutdown lock. + wait_for_process_events([process.sentinel]) + with process_status_lock: + if close_signal.is_set() or failed_signal.is_set(): + return + process.join(timeout=0) + error = RuntimeError( + "OnlineDataEngine simulation worker exited unexpectedly " + f"(exit code {process.exitcode})." + ) except BaseException as caught_error: if close_signal.is_set(): return @@ -1276,6 +1284,7 @@ def __init__(self, cfg: OnlineDataEngineCfg) -> None: # Handle to the simulation subprocess, set in start() and used in stop(). self._sim_process: mp.Process | None = None self._monitor_thread: threading.Thread | None = None + self._process_status_lock = threading.RLock() self._channel_error: BaseException | None = None self._worker_error: BaseException | None = None @@ -1324,6 +1333,7 @@ def start(self) -> None: target=_monitor_worker_process, args=( process, + self._process_status_lock, self._close_signal, self._worker_error_buffer, self._worker_error_length, @@ -1427,19 +1437,23 @@ def _ensure_worker_alive(self) -> None: ) return - if self._sim_process is None: - error = RuntimeError("OnlineDataEngine simulation worker was not created.") - self._record_worker_error(error) - raise error - if self._sim_process.is_alive(): - return - - self._sim_process.join(timeout=0) + with self._process_status_lock: + process = self._sim_process + if process is None: + error = RuntimeError( + "OnlineDataEngine simulation worker was not created." + ) + self._record_worker_error(error) + raise error + if process.is_alive(): + return + process.join(timeout=0) + exit_code = process.exitcode worker_error = self._receive_worker_error() if worker_error is None: worker_error = RuntimeError( "OnlineDataEngine simulation worker exited unexpectedly " - f"(exit code {self._sim_process.exitcode})." + f"(exit code {exit_code})." ) self._record_worker_error(worker_error) self._set_state(OnlineDataEngineState.FAILED) @@ -1602,6 +1616,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 # ----------------------------------------------------------------------- @@ -1743,7 +1770,7 @@ def _sample_batch( if sampling_mode == "boundary" and chunk_size < 2: raise ValueError("boundary sampling requires chunk_size >= 2.") - # Hold the producer's window lock through the final clone. The worker + # Hold the producer's window lock through the indexed copy. The worker # may continue writing its already-advertised window, which is excluded # below, but cannot switch to one of our selected rows until the copy is # complete. @@ -1880,15 +1907,17 @@ def _sample_batch( future, segment_ids if sampling_mode == "segment" else None, ) - result = self.shared_buffer[row_indices[:, None], indices].clone() + # Two tensor indices gather independent storage for every leaf; + # cloning that owned result would copy the complete batch twice. + result = self.shared_buffer[row_indices[:, None], indices] _pad_context(result, mask, loss_mask) else: time_offsets = torch.arange(chunk_size, device=start_indices.device) time_indices = start_indices[:, None] + time_offsets[None, :] - result = self.shared_buffer[row_indices[:, None], time_indices].clone() + result = self.shared_buffer[row_indices[:, None], time_indices] if sampling_budget is not None: - # Clone and padding must succeed before any quota is spent. + # Copy and padding must succeed before any quota is spent. # A later optimizer/consumer failure does not refund delivery. error = self._receive_worker_error() if error is not None: @@ -2035,6 +2064,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 +2153,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 +2176,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, @@ -2180,10 +2265,11 @@ def _detect_preexisting_worker_exit(self) -> BaseException | None: return error try: - if process.is_alive(): - return None - process.join(timeout=0) - exit_code = process.exitcode + with self._process_status_lock: + if process.is_alive(): + return None + process.join(timeout=0) + exit_code = process.exitcode except BaseException as inspection_error: error = RuntimeError( "OnlineDataEngine could not inspect its simulation worker before " @@ -2215,46 +2301,49 @@ def _shutdown_worker(self) -> bool: forced_shutdown = False - try: - is_alive = process.is_alive() - except ValueError: - is_alive = False + with self._process_status_lock: + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False - try: - process.join(timeout=self.cfg.shutdown_timeout if is_alive else 0) - except (AssertionError, ValueError): - pass + try: + process.join(timeout=self.cfg.shutdown_timeout if is_alive else 0) + except (AssertionError, ValueError): + pass - try: - is_alive = process.is_alive() - except ValueError: - is_alive = False - if is_alive: - forced_shutdown = True - self._forced_shutdown_attempted = True - process.terminate() - process.join(timeout=3.0) + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + if is_alive: + forced_shutdown = True + self._forced_shutdown_attempted = True + process.terminate() + process.join(timeout=3.0) - try: - is_alive = process.is_alive() - except ValueError: - is_alive = False - if is_alive and hasattr(process, "kill"): - forced_shutdown = True - self._forced_shutdown_attempted = True - process.kill() - process.join(timeout=1.0) + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + if is_alive and hasattr(process, "kill"): + forced_shutdown = True + self._forced_shutdown_attempted = True + process.kill() + process.join(timeout=1.0) - try: - is_alive = process.is_alive() - except ValueError: - is_alive = False - if is_alive: - raise RuntimeError( - "OnlineDataEngine simulation worker remained alive after " - "graceful shutdown, terminate(), and kill()." - ) + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + if is_alive: + raise RuntimeError( + "OnlineDataEngine simulation worker remained alive after " + "graceful shutdown, terminate(), and kill()." + ) + # The monitor may be waiting for the status lock after its sentinel + # became ready. Never join that thread while holding the same lock. monitor_thread = self._monitor_thread if ( monitor_thread is not None @@ -2266,9 +2355,10 @@ def _shutdown_worker(self) -> bool: "OnlineDataEngine worker monitor did not stop after the worker exited." ) - if hasattr(process, "close"): - process.close() - self._sim_process = None + with self._process_status_lock: + if hasattr(process, "close"): + process.close() + self._sim_process = None self._monitor_thread = None return forced_shutdown @@ -2388,6 +2478,7 @@ def __getstate__(self) -> dict: state["_sim_process"] = None state["_monitor_thread"] = None state["_lifecycle_condition"] = None + state["_process_status_lock"] = None state["_channel_error"] = None state["_worker_error"] = None return state @@ -2396,6 +2487,7 @@ def __setstate__(self, state: dict) -> None: """Restore process-local synchronization after consumer deserialization.""" self.__dict__.update(state) self._lifecycle_condition = threading.Condition(threading.RLock()) + self._process_status_lock = threading.RLock() def __del__(self) -> None: try: diff --git a/embodichain/data_pipeline/runtime.py b/embodichain/data_pipeline/runtime.py new file mode 100644 index 0000000000..834a92ba17 --- /dev/null +++ b/embodichain/data_pipeline/runtime.py @@ -0,0 +1,377 @@ +# ---------------------------------------------------------------------------- +# 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. + This helper accepts only a live source wrapping the controlled engine. + Sample resident sources through their own ``sample_context()`` method; + they are not accepted by this controller. + + 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 a3e1408f77..83923c470d 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 d7487c02b5..15f907a7bd 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 f89e97012e..2bef7f8da7 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 3673dd54f4..16cbaaacb9 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()) @@ -616,6 +445,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", @@ -695,6 +525,7 @@ def restart_evaluator() -> dict[str, Any]: } engine = None + flow = None corpus = None writer = None tensorboard = None @@ -736,12 +567,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: @@ -761,15 +588,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] = {} @@ -915,6 +742,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) @@ -958,23 +796,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, ) @@ -986,6 +815,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( @@ -1248,24 +1078,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 @@ -1283,7 +1112,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:"): @@ -1306,76 +1137,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 @@ -1385,7 +1155,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 @@ -1422,15 +1202,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) @@ -1452,11 +1228,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 @@ -1572,7 +1346,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=[ @@ -1586,7 +1360,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/scripts/benchmark/data_pipeline/benchmark_ods_context_copy.py b/scripts/benchmark/data_pipeline/benchmark_ods_context_copy.py new file mode 100644 index 0000000000..7a7858e70e --- /dev/null +++ b/scripts/benchmark/data_pipeline/benchmark_ods_context_copy.py @@ -0,0 +1,515 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Measure owned ODS CPU sampling copies without simulation or a learner. + +Run: python scripts/benchmark/data_pipeline/benchmark_ods_context_copy.py \ + --implementation baseline=/tmp/embodichain-baseline \ + --implementation optimized=. --output-dir outputs/benchmarks/ods-copy + +Each checkout is imported in a separate subprocess. Input construction, +correctness hashes, output release and a separate sampled-RSS pass are outside +latency samples. This isolates synchronous sampling, not training throughput. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +from pathlib import Path +import platform +import statistics +import subprocess +import sys +import threading +import time +from typing import Any + +__all__: list[str] = [] + + +def _hash_tensor(value: Any) -> str: + """Hash every byte without introducing a second full image copy.""" + return hashlib.sha256(memoryview(value.contiguous().numpy())).hexdigest() + + +def _digest(batch: Any) -> dict[str, str]: + """Describe and hash all returned TensorDict leaves outside timing.""" + return { + str(key): f"{tuple(value.shape)}:{value.dtype}:{_hash_tensor(value)}" + for key, value in batch.items(include_nested=True, leaves_only=True) + } + + +def _fixture(cameras: int) -> Any: + """Construct seeded numeric and RGBA payloads with dirty tails and boundaries.""" + import torch + from tensordict import TensorDict + + generator = torch.Generator(device="cpu").manual_seed(317) + shape = (8, 600) + frames = torch.arange(shape[1])[None].expand(shape) + rows = torch.arange(shape[0])[:, None].expand(shape) + lengths = 600 - torch.arange(shape[0]) * 7 + valid = frames < lengths[:, None] + continuity = (frames >= 300).to(torch.int64) + continuity[~valid] = -1 + obs = {"state": torch.randn(*shape, 59, generator=generator)} + for camera in range(cameras): + obs[f"rgba_{camera}"] = torch.randint( + 0, 256, (*shape, 256, 256, 4), dtype=torch.uint8, generator=generator + ) + data = TensorDict( + { + "obs": TensorDict(obs, batch_size=shape), + "actions": torch.randn(*shape, 7, generator=generator), + "rewards": torch.randn(*shape, 1, generator=generator), + "valid": valid, + "continuity_id": continuity, + "segment_id": (frames >= 150).to(torch.int64), + "segment_accepted": valid.clone(), + "task_index": torch.zeros(shape, dtype=torch.int64), + "subtask_index": torch.ones(shape, dtype=torch.int64), + "ods_run_id": torch.full(shape, 17, dtype=torch.int64), + "ods_episode_id": rows.clone(), + "ods_frame_index": frames.clone(), + "ods_generation": torch.ones(shape, dtype=torch.int64), + "ods_accepted_at_ns": torch.full(shape, 12345, dtype=torch.int64), + "ods_published_at_monotonic_ns": torch.full( + shape, 23456, dtype=torch.int64 + ), + }, + batch_size=shape, + device="cpu", + ) + # The actual sampler owns these independent CPU shared-memory tensors. + return data.share_memory_() + + +def _run_worker(args: argparse.Namespace) -> None: + """Time one checkout and vision shape with real engine/source methods.""" + import psutil + import torch + + from embodichain.data_pipeline.datasets import OnlineDataSource, SequenceSpecCfg + from embodichain.data_pipeline.engine.data import ( + OnlineDataEngine, + OnlineDataEngineCfg, + OnlineDataEngineState, + ) + + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + imported = Path(sys.modules[OnlineDataEngine.__module__].__file__).resolve() + if not imported.is_relative_to(Path(args.root).resolve()): + raise RuntimeError(f"Wrong checkout imported: {imported}") + prepared = _fixture(args.cameras) + + class FixtureEngine(OnlineDataEngine): + """Retain normal synchronization and sampling, substituting only allocation.""" + + def _create_buffer(self) -> Any: + """Supply the prefilled shared fixture instead of constructing a simulator.""" + return prepared + + engine = FixtureEngine( + OnlineDataEngineCfg( + gym_config={"num_envs": 1}, + buffer_size=8, + max_episode_steps=600, + state_dim=59, + action_dim=7, + include_provenance=True, + refill_policy="manual", + ) + ) + # Consumer identity bypasses owner-only simulation liveness, not copy locks. + engine._owner_pid = -1 + engine.run_id = 17 + engine._set_state(OnlineDataEngineState.READY) + engine._init_signal.set() + with engine._lock_index.get_lock(): + engine._lock_index[:] = [-1, -1] + source = OnlineDataSource( + engine, + schema={ + "action": { + "names": [f"joint_{index}" for index in range(7)], + "semantics": "expert_controller_command", + "step_dt": 1 / 30, + } + }, + ) + spec = SequenceSpecCfg(2, 32) + rng = torch.Generator(device="cpu") + process = psutil.Process() + rows: list[dict[str, Any]] = [] + for batch_size in (1, 32): + for api in ("chunk", "context", "source"): + + def reset(seed: int) -> None: + """Restore equivalent local/global sampling randomness outside timing.""" + rng.manual_seed(seed) + torch.manual_seed(seed) + + def sample() -> Any: + """Call one public synchronous production sampling method.""" + if api == "chunk": + return engine.sample_batch(batch_size, 33) + if api == "context": + return engine.sample_batch_with_context( + batch_size, 2, 32, generator=rng + ) + return source.sample_context(batch_size, spec, generator=rng) + + for repeat in range(args.warmup): + reset(1000 + repeat) + result = sample() + del result + samples: list[float] = [] + for repeat in range(args.repeats): + reset(2000 + repeat) + start = time.perf_counter() + result = sample() + samples.append((time.perf_counter() - start) * 1000) + del result + + # Separate warmed memory observation; sampling is approximate, not a + # true allocator peak. Retain precisely one returned batch in all arms. + gc.collect() + reset(3000) + before = process.memory_info().rss + memory_samples = [before] + finished = threading.Event() + + def poll_rss() -> None: + """Observe this subprocess RSS every millisecond in the memory pass.""" + while not finished.wait(0.001): + memory_samples.append(process.memory_info().rss) + + observer = threading.Thread(target=poll_rss) + observer.start() + result = sample() + after = process.memory_info().rss + memory_samples.append(after) + finished.set() + observer.join() + del result + + # Ensure context checks actually cover padding, including B=1. + correctness_seed = 4000 + while True: + reset(correctness_seed) + result = sample() + if api == "chunk" or not result["valid"].all(): + break + del result + correctness_seed += 1 + digest = _digest(result) + rng_hash = _hash_tensor( + torch.get_rng_state() if api == "chunk" else rng.get_state() + ) + global_preserved = api == "chunk" or torch.equal( + torch.get_rng_state(), + torch.Generator().manual_seed(correctness_seed).get_state(), + ) + invalid = ~result["valid"] + padding_ok = True + for key, value in result.items(include_nested=True, leaves_only=True): + if key in {"valid", "loss_mask"}: + continue + identifier = isinstance(key, str) and ( + key.startswith("ods_") + or key + in { + "frame_ids", + "generations", + "task_index", + "subtask_index", + "accepted_at_ns", + "published_at_monotonic_ns", + } + ) + padding_ok &= bool((value[invalid] == (-1 if identifier else 0)).all()) + # Replacement and later sampling must leave every returned leaf owned. + with engine._lock_index.get_lock(): + for value in prepared["obs"].values(): + ( + value.bitwise_xor_(255) + if value.dtype == torch.uint8 + else value.neg_() + ) + prepared["actions"].neg_() + prepared["ods_episode_id"].add_(8) + prepared["ods_generation"].add_(1) + reset(correctness_seed) + replacement = sample() + owned = _digest(result) == digest + del replacement + with engine._lock_index.get_lock(): + for value in prepared["obs"].values(): + ( + value.bitwise_xor_(255) + if value.dtype == torch.uint8 + else value.neg_() + ) + prepared["actions"].neg_() + prepared["ods_episode_id"].sub_(8) + prepared["ods_generation"].sub_(1) + rows.append( + { + "api": api, + "batch_size": batch_size, + "cameras": args.cameras, + "samples_ms": samples, + "cost_time_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "cpu_delta_mb": (after - before) / 1024**2, + "sampled_peak_delta_mb": (max(memory_samples) - before) / 1024**2, + "rss_baseline_mb": before / 1024**2, + "output_bytes": sum( + v.numel() * v.element_size() + for _, v in result.items(include_nested=True, leaves_only=True) + ), + "correctness_seed": correctness_seed, + "padding_frames": int(invalid.sum()), + "padding_ok": padding_ok, + "owned_after_replacement": owned, + "global_rng_preserved": global_preserved, + "output_hashes": digest, + "rng_hash": rng_hash, + } + ) + del result + print( + f"{args.label}: cameras={args.cameras} B={batch_size} {api}: {rows[-1]['cost_time_ms']:.3f} ms", + flush=True, + ) + engine._set_state(OnlineDataEngineState.CREATED) + engine._owner_pid = os.getpid() + engine.stop() + packet = { + "label": args.label, + "root": args.root, + "commit": subprocess.check_output( + ["git", "-C", args.root, "rev-parse", "HEAD"], text=True + ).strip(), + "production_source_sha256": { + str(path.relative_to(Path(args.root))): hashlib.sha256( + path.read_bytes() + ).hexdigest() + for path in ( + Path(args.root) / "embodichain/data_pipeline/engine/data.py", + Path(args.root) / "embodichain/data_pipeline/datasets/sources.py", + Path(args.root) / "embodichain/data_pipeline/datasets/_sampling.py", + ) + }, + "hardware": { + "platform": platform.platform(), + "cpu": Path("/proc/cpuinfo") + .read_text() + .split("model name\t: ")[1] + .splitlines()[0], + "affinity": process.cpu_affinity(), + "total_memory_bytes": psutil.virtual_memory().total, + }, + "software": { + "python": sys.version, + "torch": torch.__version__, + "tensordict": __import__("tensordict").__version__, + }, + "method": { + "warmup": args.warmup, + "repeats": args.repeats, + "torch_threads": torch.get_num_threads(), + "interop_threads": torch.get_num_interop_threads(), + "device": "cpu", + "cache_shape": [8, 600], + "state_width": 59, + "action_width": 7, + "image_shape": [256, 256, 4], + "image_dtype": "uint8", + "history": 2, + "future": 32, + "persistent_cache_bytes": sum( + value.numel() * value.element_size() + for _, value in prepared.items(include_nested=True, leaves_only=True) + ), + }, + "rows": rows, + } + Path(args.worker_output).write_text(json.dumps(packet, indent=2) + "\n") + + +def _write_report(packets: list[dict[str, Any]], output: Path) -> None: + """Write exactly three tables with correctness and scoped timing evidence.""" + reference = { + (r["cameras"], r["batch_size"], r["api"]): r + for p in packets + if p["label"] == packets[0]["label"] + for r in p["rows"] + } + lines = [ + "# ODS CPU sampling-copy benchmark", + "", + "Synchronous production methods; no simulation, LeRobot reader, learner or GPU. Returned output release, correctness and memory observation are excluded from latency. Repeated samples support median and observed range; no percentile-tail claim is made.", + "", + "## Time & Memory", + "", + "| workload | implementation | cost_time_ms | min–max ms | cpu_delta_mb | sampled_peak_delta_mb | gpu_delta_mb | peak_gpu_mb |", + "| --- | --- | --- | --- | --- | --- | --- | --- |", + ] + quality: list[tuple[str, dict[str, Any], bool]] = [] + for packet in packets: + for row in packet["rows"]: + baseline = reference[(row["cameras"], row["batch_size"], row["api"])] + equivalent = ( + row["output_hashes"] == baseline["output_hashes"] + and row["rng_hash"] == baseline["rng_hash"] + ) + quality.append((packet["label"], row, equivalent)) + lines.append( + f"| camera={row['cameras']} B={row['batch_size']} {row['api']} | {packet['label']} | {row['cost_time_ms']:.3f} | {row['min_ms']:.3f}–{row['max_ms']:.3f} | {row['cpu_delta_mb']:.1f} | {row['sampled_peak_delta_mb']:.1f} | N/A | N/A |" + ) + lines += [ + "", + "## Success & Other Metrics", + "", + "| workload | implementation | success_rate | all-leaf and RNG equivalence | padding frames | padding correct | owned after replacement |", + "| --- | --- | --- | --- | --- | --- | --- |", + ] + for label, row, equivalent in quality: + passed = ( + equivalent + and row["padding_ok"] + and row["owned_after_replacement"] + and row["global_rng_preserved"] + ) + lines.append( + f"| camera={row['cameras']} B={row['batch_size']} {row['api']} | {label} | {int(passed):.0%} | {equivalent} | {row['padding_frames']} | {row['padding_ok']} | {row['owned_after_replacement']} |" + ) + ranking = [] + for label in dict.fromkeys(p["label"] for p in packets): + checks = [ + equivalent + and row["padding_ok"] + and row["owned_after_replacement"] + and row["global_rng_preserved"] + for name, row, equivalent in quality + if name == label + ] + ranking.append((sum(checks) / len(checks), label, len(checks))) + lines += [ + "", + "## Leaderboard", + "", + "| rank | algorithm | overall_success_rate | cases |", + "| --- | --- | --- | --- |", + ] + for rank, (success, label, count) in enumerate( + sorted(ranking, key=lambda entry: -entry[0]), 1 + ): + lines.append(f"| {rank} | {label} | {success:.0%} | {count} |") + lines += [ + "", + "Memory is a separate warmed pass: CPU retained RSS delta and maximum observed increment from 1 ms process-RSS polling, not an allocator peak. Persistent cache bytes and software/hardware/affinity are recorded in the JSON. Shared-memory fixture construction and dirty-tail/continuity setup are excluded. No worker writes during timing, so producer contention and training overlap remain unmeasured. Implementations use the same script and seeds; checkout order alternates by camera count. First-implementation hashes are the comparison reference, not independent semantic proof.", + "", + ] + output.write_text("\n".join(lines)) + + +def run_all_benchmarks() -> None: + """Run isolated checkout comparisons and emit raw JSON plus one Markdown report.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--implementation", action="append", default=[], metavar="LABEL=PATH" + ) + parser.add_argument( + "--output-dir", type=Path, default=Path("outputs/benchmarks/ods-copy") + ) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--repeats", type=int, default=7) + parser.add_argument("--cameras", type=int, default=0) + parser.add_argument("--root") + parser.add_argument("--label") + parser.add_argument("--worker-output") + args = parser.parse_args() + if args.warmup < 1 or args.repeats < 1: + parser.error("warmup and repeats must be positive") + if args.worker_output: + _run_worker(args) + return + if not args.implementation: + parser.error("at least one --implementation LABEL=PATH is required") + implementations = [ + (label, str(Path(root).resolve())) + for label, root in (item.split("=", 1) for item in args.implementation) + ] + args.output_dir.mkdir(parents=True, exist_ok=True) + packets = [] + for cameras in (0, 1, 2): + for label, root in ( + implementations if cameras % 2 == 0 else implementations[::-1] + ): + destination = args.output_dir / f"{label}-camera-{cameras}.json" + environment = { + **os.environ, + "PYTHONPATH": root, + "OMP_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "NUMEXPR_NUM_THREADS": "1", + } + subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "--root", + root, + "--label", + label, + "--cameras", + str(cameras), + "--warmup", + str(args.warmup), + "--repeats", + str(args.repeats), + "--worker-output", + str(destination.resolve()), + ], + env=environment, + check=True, + ) + packets.append(json.loads(destination.read_text())) + # Stable baseline-first packet order, independent of alternating execution. + packets.sort( + key=lambda packet: ( + [label for label, _ in implementations].index(packet["label"]), + packet["rows"][0]["cameras"], + ) + ) + (args.output_dir / "raw.json").write_text(json.dumps(packets, indent=2) + "\n") + report = args.output_dir / "report.md" + _write_report(packets, report) + print(f"Markdown report saved: {report.resolve()}") + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/tests/data_pipeline/test_ods_runtime.py b/tests/data_pipeline/test_ods_runtime.py new file mode 100644 index 0000000000..48b651f407 --- /dev/null +++ b/tests/data_pipeline/test_ods_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()} diff --git a/tests/data_pipeline/test_online_data.py b/tests/data_pipeline/test_online_data.py index 95661abd8b..652b3cf52b 100644 --- a/tests/data_pipeline/test_online_data.py +++ b/tests/data_pipeline/test_online_data.py @@ -117,6 +117,100 @@ def test_action_context_rejects_invalid_lengths(history, future) -> None: _make_fake_engine().sample_batch_with_context(2, history, future) +@pytest.mark.parametrize("context", [False, True]) +def test_sample_owns_nested_noncontiguous_images_across_source_replacement( + context: bool, +) -> None: + """Both sampling APIs detach every leaf, including repeated image indices.""" + steps = 6 + batch_size = 16 # One source row guarantees repeated anchors within the batch. + engine = _make_fake_engine(buffer_size=1, max_episode_steps=steps) + source = engine.shared_buffer + camera = torch.arange(steps * 2 * 3 * 4, dtype=torch.uint8).reshape( + 1, steps, 2, 3, 4 + ) + source["obs"] = TensorDict( + { + "joint": source["obs"], + "cameras": TensorDict( + {"front": camera.transpose(-3, -2), "wrist": camera.flip(-1)}, + batch_size=source.batch_size, + ), + }, + batch_size=source.batch_size, + ) + for key, value in ( + ("ods_run_id", 17), + ("ods_episode_id", 23), + ("ods_generation", 1), + ("ods_accepted_at_ns", 100), + ("ods_published_at_monotonic_ns", 200), + ): + source[key] = torch.full(source.batch_size, value, dtype=torch.int64) + source["ods_frame_index"] = torch.arange(steps).reshape(1, steps) + source.share_memory_() + assert not source["obs", "cameras", "front"].is_contiguous() + original = source.clone() + + def sample() -> TensorDict: + if context: + return engine.sample_batch_with_context( + batch_size, 2, 3, generator=torch.Generator().manual_seed(39) + ) + return engine.sample_batch(batch_size, 1) + + try: + batch = sample() + saved = batch.clone() + valid = batch["valid"] + frames = batch["ods_frame_index"].clamp(min=0) + anchor_index = 1 if context else 0 + anchors = batch["ods_frame_index"][:, anchor_index] + assert anchors.unique().numel() < batch_size + for key, value in original.items(include_nested=True, leaves_only=True): + actual = batch[key] + assert ( + actual.untyped_storage().data_ptr() + != source[key].untyped_storage().data_ptr() + ) + expected = value[0, frames] + identifier = isinstance(key, str) and ( + key.startswith("ods_") or key in {"task_index", "subtask_index"} + ) + mask = valid.reshape(*valid.shape, *((1,) * (expected.ndim - valid.ndim))) + expected = expected.masked_fill(~mask, -1 if identifier else 0) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(source[key], value) + + # Model a producer reusing the same physical slot, then another consumer read. + replacement = original.clone() + replacement["obs", "joint"].add_(10) + replacement["obs", "cameras", "front"].fill_(251) + replacement["obs", "cameras", "wrist"].fill_(252) + replacement["actions"].add_(10) + replacement["ods_episode_id"].fill_(24) + replacement["ods_generation"].fill_(2) + with engine._lock_index.get_lock(): + source.copy_(replacement) + later = sample() + assert (later["ods_episode_id"][later["valid"]] == 24).all() + for key, value in saved.items(include_nested=True, leaves_only=True): + torch.testing.assert_close(batch[key], value) + + # Duplicate gather positions own distinct output elements as well. + unique_anchors, counts = torch.unique(anchors, return_counts=True) + matches = (anchors == unique_anchors[counts.argmax()]).nonzero().flatten() + first, second = matches[:2].tolist() + for key, value in batch.items(include_nested=True, leaves_only=True): + second_before = value[second].clone() + value[first].zero_() + torch.testing.assert_close(value[second], second_before) + if key in source.keys(include_nested=True, leaves_only=True): + torch.testing.assert_close(source[key], replacement[key]) + finally: + engine.stop() + + # --------------------------------------------------------------------------- # Helper # --------------------------------------------------------------------------- @@ -241,6 +335,7 @@ def _make_fake_engine( engine._worker_error = None engine._owner_pid = os.getpid() engine._lifecycle_condition = threading.Condition(threading.RLock()) + engine._process_status_lock = threading.RLock() engine._stop_requested = False engine._cleanup_complete = False engine._forced_shutdown_attempted = False @@ -494,14 +589,22 @@ def _install_fake_starting_worker( """Install a process mock that optionally completes the initial fill.""" process = MagicMock(pid=1234, exitcode=None) process.is_alive.return_value = True + sentinel_read, sentinel_write = os.pipe() + process.sentinel = sentinel_read + write_closed = False def join(timeout=None) -> None: + nonlocal write_closed if timeout is None: engine._close_signal.wait(timeout=5.0) if engine._close_signal.is_set(): process.is_alive.return_value = False + if not write_closed: + os.close(sentinel_write) + write_closed = True process.join.side_effect = join + process.close.side_effect = lambda: os.close(sentinel_read) if initialize: process.start.side_effect = engine._init_signal.set process_context = MagicMock() @@ -510,6 +613,165 @@ def join(timeout=None) -> None: return process +def _forkserver_lifecycle_worker( + started: mp.Event, + finish: mp.Event, + exit_code: int, + error_channels: tuple | None, +) -> None: + """Exercise parent status handling without importing a simulator.""" + started.set() + if not finish.wait(CONSUMER_RESULT_TIMEOUT): + os._exit(124) + if error_channels is not None: + engine_module._publish_worker_error( + *error_channels, ValueError("original recorder close failure") + ) + os._exit(exit_code) + + +def _start_lifecycle_worker( + engine: OnlineDataEngine, + *, + finish: mp.Event | None = None, + exit_code: int = 0, + publish_close_error: bool = False, +) -> mp.Process: + started = engine._mp_ctx.Event() + channels = ( + engine._worker_error_buffer, + engine._worker_error_length, + engine._worker_failed_signal, + engine._state_value, + ) + process = engine._mp_ctx.Process( + target=_forkserver_lifecycle_worker, + args=( + started, + engine._close_signal if finish is None else finish, + exit_code, + channels if publish_close_error else None, + ), + ) + process.start() + engine._sim_process = process + assert started.wait(CONSUMER_RESULT_TIMEOUT), "Lifecycle worker did not start" + engine._monitor_thread = threading.Thread( + target=engine_module._monitor_worker_process, + args=( + process, + engine._process_status_lock, + engine._close_signal, + *channels, + ), + daemon=True, + ) + engine._monitor_thread.start() + return process + + +def test_forkserver_monitor_preserves_normal_shutdown_status() -> None: + """A monitor and owner must preserve the real zero status after close().""" + engine = _make_fake_engine() + process = _start_lifecycle_worker(engine) + popen = process._popen + monitor = engine._monitor_thread + try: + engine.stop() + assert popen.returncode == 0 + assert monitor is not None and not monitor.is_alive() + assert engine.state is OnlineDataEngineState.STOPPED + assert not engine._worker_failed_signal.is_set() + finally: + engine.stop() + + +def test_forkserver_status_pipe_has_one_reader_during_shutdown( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Pause the real status read where competing joins formerly produced 255.""" + import multiprocessing.forkserver as forkserver + + engine = _make_fake_engine() + process = _start_lifecycle_worker(engine) + popen = process._popen + original_read = forkserver.read_signed + first_read, second_read, release_read = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + readers: list[int] = [] + errors: list[BaseException] = [] + + def paused_read(sentinel: int) -> int: + readers.append(threading.get_ident()) + first_read.set() + if len(readers) > 1: + second_read.set() + assert release_read.wait(PROCESS_CLEANUP_TIMEOUT) + return original_read(sentinel) + + def stop() -> None: + try: + engine.stop() + except BaseException as error: + errors.append(error) + + monkeypatch.setattr(forkserver, "read_signed", paused_read) + owner = threading.Thread(target=stop) + owner.start() + try: + assert first_read.wait(PROCESS_CLEANUP_TIMEOUT) + # An old monitor.join() enters the same pipe read while the first + # reader is paused. The status lock now permits only one consumer. + competing_read = second_read.wait(0.2) + finally: + release_read.set() + owner.join(PROCESS_CLEANUP_TIMEOUT) + monkeypatch.setattr(forkserver, "read_signed", original_read) + if owner.is_alive(): + engine._close_signal.set() + assert not owner.is_alive() + assert errors == [] + assert not competing_read and len(readers) == 1 + assert popen.returncode == 0 + assert engine._cleanup_complete + + +def test_forkserver_monitor_broadcasts_real_hard_exit_code() -> None: + """Sentinel monitoring still publishes an unexpected nonzero worker exit.""" + engine = _make_fake_engine() + finish = engine._mp_ctx.Event() + process = _start_lifecycle_worker(engine, finish=finish, exit_code=23) + popen = process._popen + finish.set() + try: + assert engine._worker_failed_signal.wait(PROCESS_CLEANUP_TIMEOUT) + with pytest.raises(RuntimeError, match="exit code 23"): + engine.stop() + assert popen.returncode == 23 + assert engine.state is OnlineDataEngineState.FAILED + assert engine._cleanup_complete + finally: + engine.stop() + + +def test_forkserver_shutdown_preserves_original_worker_exception() -> None: + """A real close-time exception survives monitor/owner status serialization.""" + engine = _make_fake_engine() + process = _start_lifecycle_worker(engine, publish_close_error=True) + popen = process._popen + try: + with pytest.raises(ValueError, match="original recorder close failure"): + engine.stop() + assert popen.returncode == 0 + assert engine.state is OnlineDataEngineState.FAILED + assert engine._cleanup_complete + finally: + engine.stop() + + def _sample_engine_in_subprocess( engine: OnlineDataEngine, result_queue: mp.Queue ) -> None: diff --git a/tests/data_pipeline/test_online_data_budget.py b/tests/data_pipeline/test_online_data_budget.py index 1aaa0a53de..ca60b2ae57 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 @@ -158,12 +158,19 @@ def test_every_sampling_api_obeys_shared_episode_budget(method: str) -> None: def test_failed_payload_copy_does_not_spend_quota(monkeypatch) -> None: engine = _engine() before = engine.get_sampling_budget_snapshot() + original_getitem = TensorDict.__getitem__ + + def fail_payload_gather(buffer, index): + if ( + buffer is engine.shared_buffer + and isinstance(index, tuple) + and all(isinstance(value, torch.Tensor) for value in index) + ): + raise RuntimeError("copy failed") + return original_getitem(buffer, index) + with monkeypatch.context() as context: - context.setattr( - TensorDict, - "clone", - lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("copy failed")), - ) + context.setattr(TensorDict, "__getitem__", fail_payload_gather) with pytest.raises(RuntimeError, match="copy failed"): engine.sample_batch(2, 2) assert engine.get_sampling_budget_snapshot() == before @@ -492,3 +499,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_sources.py b/tests/data_pipeline/test_sources.py index 0e6dc7ec6d..9fae59b7a3 100644 --- a/tests/data_pipeline/test_sources.py +++ b/tests/data_pipeline/test_sources.py @@ -156,6 +156,100 @@ def test_live_and_resident_sources_share_windows_masks_and_rng( assert not (resident.data["obs"] == -100).any() +def test_nested_camera_context_padding_and_provenance_match_resident(paired_sources): + live, resident = paired_sources + shape = live.engine.shared_buffer.batch_size + camera = torch.arange(2 * 5 * 3 * 2 * 4, dtype=torch.uint8).reshape(2, 5, 3, 2, 4) + nested = TensorDict( + { + "joint": resident.data["obs"], + "cameras": TensorDict( + {"front": camera, "wrist": camera.flip(-1)}, batch_size=shape + ), + }, + batch_size=shape, + ) + raw = live.engine.shared_buffer.clone().unlock_() + raw["obs"] = nested.clone() + live.engine.shared_buffer = raw.share_memory_() + data = resident.data.clone() + data["obs"] = nested.clone() + schema = {"action": _SCHEMA["action"]} + online = OnlineDataSource(live.engine, schema=schema) + offline = TensorDataSource(data, schema=schema, source_id="nested-camera") + spec = SequenceSpecCfg(3, 5) + a = online.sample_context(32, spec, generator=torch.Generator().manual_seed(72)) + b = offline.sample_context(32, spec, generator=torch.Generator().manual_seed(72)) + assert set(a.keys(include_nested=True, leaves_only=True)) == set( + b.keys(include_nested=True, leaves_only=True) + ) + assert (~a["valid"]).any() + assert not a["loss_mask"][:, : spec.history_length - 1].any() + for key, value in a.items(include_nested=True, leaves_only=True): + torch.testing.assert_close(value, b[key]) + if isinstance(key, tuple) and key[0] == "obs": + assert (value[~a["valid"]] == 0).all() + for key in ( + "frame_ids", + "generations", + "accepted_at_ns", + "published_at_monotonic_ns", + "task_index", + "subtask_index", + ): + assert (a[key][~a["valid"]] == -1).all() + + +@pytest.mark.parametrize("nested", [False, True]) +def test_online_observation_transform_zeroes_only_invalid_observations( + paired_sources, nested +): + live, _ = paired_sources + schema = _SCHEMA if not nested else {"action": _SCHEMA["action"]} + original = live.engine.shared_buffer.clone() + observed_inputs = [] + + def transform(observation): + observed_inputs.append(observation.clone()) + transformed = observation + 7 + if nested: + return TensorDict( + {"joint": transformed, "offset": transformed[..., :1] + 3}, + batch_size=observation.shape[:2], + ) + return transformed + + source = OnlineDataSource( + live.engine, schema=schema, observation_transform=transform + ) + spec = SequenceSpecCfg(3, 5) + actual = source.sample_context( + 32, spec, generator=torch.Generator().manual_seed(29) + ) + reference = live.sample_context( + 32, spec, generator=torch.Generator().manual_seed(29) + ) + valid = actual["valid"] + assert (~valid).any() + torch.testing.assert_close(observed_inputs[0], reference["obs"]) + leaves = ( + {"joint": actual["obs", "joint"], "offset": actual["obs", "offset"]} + if nested + else {"joint": actual["obs"]} + ) + for key, value in leaves.items(): + expected = ( + reference["obs"] + 7 if key == "joint" else reference["obs"][..., :1] + 10 + ) + torch.testing.assert_close(value[valid], expected[valid]) + assert (value[~valid] == 0).all() + for key in actual.keys(): + if key != "obs": + torch.testing.assert_close(actual[key], reference[key]) + for key, value in original.items(include_nested=True, leaves_only=True): + torch.testing.assert_close(live.engine.shared_buffer[key], value) + + @pytest.mark.parametrize("with_continuity", [False, True]) @pytest.mark.parametrize("generations", [None, (1,)]) def test_segment_filter_preserves_source_validity_and_episode_anchors(