LanceDB PR #3339 (lancedb/lancedb#3339) makes
both the Python BackgroundEventLoop and the Rust tokio runtime
fork-safe via pthread_atfork / os.register_at_fork hooks. This removes
the original reason LanceDataset forced spawn on Linux and lets the
default fork-mode DataLoader workers Just Work.

Three workarounds drop together:

1. LanceDataset._maybe_warn_fork_start_method (and the _fork_warning_emitted
   class flag): no longer force spawn. Workers fork cleanly; lancedb
   re-initialises its async layers in the child.

2. torch.multiprocessing.set_sharing_strategy('file_system') (was inside
   the same classmethod): only needed because spawn used /dev/shm fds
   that misbehaved on cloud instances. Fork inherits memory directly,
   so the default file_descriptor strategy is fine.

3. _ForwardWithCfg in lewm/pldm/prejepa.py: the wrapper around
   functools.partial(forward_fn, cfg=cfg) only existed because
   spawn pickles the bound method, which fails when the underlying
   callable lacks __name__. Fork doesn't pickle anything; a plain
   partial works.

Verified locally: built lancedb 0.31.0b11 from main, opened a Lance
table inside fork-mode DataLoader workers, confirmed Permutation
returns batches end-to-end (pre-#3339 this hangs indefinitely).

Requires lancedb release that includes #3339 (any 0.31.x once cut).

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

…n h5 (#214)

* HDF5Dataset: stream remote files via fsspec

Accept storage_options and recognise s3://, gs://, gcs://, azure://,
abfs://, http(s):// URIs. Open the file lazily per worker via
fsspec.filesystem(scheme, **storage_options) so DataLoader
multiprocessing keeps working — handle is stripped by __getstate__ and
re-opened inside _open() in each spawned worker.

Local paths still go through h5py with SWMR + 256MiB chunk cache; only
remote paths take the fsspec path. SWMR is local-only so it's omitted
for remote handles.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* add scripts/benchmark/: HDF5 vs Lance vs LeRobot throughput bench + LeRobot→{HDF5,Lance} converter

compare_h5_lance.py: self-contained bench. Auto-downloads local
HDF5/Lance from S3 on first run, exercises both local and S3 reads
(cached/uncached), optional LeRobot row. Creds blank by default — falls
through to AWS SDK chain (env vars, ~/.aws, EC2 instance role) so the
script is safe to commit.

convert.py: one-shot prep that reads lerobot/pusht_image via
LeRobotAdapter (with all native columns aliased through), writes the
same data to pusht.h5 and pusht.lance for apples-to-apples comparison.
Caches LeRobot's get_safe_version lookup to avoid one HuggingFace API
call per episode (otherwise times out around episode ~80).

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* update docstring

* fix pre-commit: ruff lint + format

- Add `# noqa: E402` on the post-monkey-patch imports in
  scripts/benchmark/convert.py (intentional placement so the
  get_safe_version cache is installed before SWM imports load lerobot).
- ruff format the three changed files.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* benchmark: add Video (per-episode mp4) format to the comparison

convert.py now also emits ./pusht.video alongside .h5 / .lance, using
the existing VideoWriter (one mp4 per episode + tabular .npz).

compare_h5_lance.py adds 'Video local (no cache)' and 'Video local
(cached)' rows. Reads via VideoDataset (decord-backed). Gated on
decord/imageio being installed; skipped cleanly if not. There is no
Video S3 row — the format is local-files only and S3 streaming would
require an extra fsspec download step that's outside the scope here.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* bench: print why Video row was skipped; previously failed silently

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* bench: collapse Video row to a single entry; FolderDataset has no keys_to_cache

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* VideoDataset: lazy-init decord in _reader, not just __init__

The class-level _decord attribute was set inside __init__, which
survives a forked process but not a spawned one — under
multiprocessing.set_start_method('spawn') (which the bench forces for
Lance fork-safety), each DataLoader worker is a fresh interpreter and
the unpickled dataset still has _decord=None on its class. _reader then
hit AttributeError on the first frame fetch in every worker.

Move the lazy import into a classmethod called from _reader, so workers
import decord on first use regardless of start method. __init__ still
calls it once for fail-fast behaviour when the dep is missing.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* convert: rescale float images to uint8 in _episode_to_step_lists

LeRobotAdapter returns pixels as float32 in [0, 1] because lerobot's
own pipeline runs torchvision.transforms.ToTensor() under the hood.
Downstream writers all assume uint8 HxWxC:
  - LanceWriter._encode_frame did frame.astype(np.uint8) on float [0,1]
    which truncated everything <1.0 to 0 — silently producing all-black
    JPEG blobs.
  - VideoWriter routed float-dtype pixels to .npz tabular instead of
    MP4 because is_image_column() requires uint8.
  - HDF5Writer accepted the float32 but inflated disk 4x.

Detect float HWC/NHWC arrays in _episode_to_step_lists, clip to [0, 1],
multiply by 255, cast to uint8. Single point fixes all three writer
paths. Existing uint8 sources (HDF5, Lance, Video) are unaffected
because the dtype check is float-only.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceDataset: batched JPEG decode (torchvision libjpeg-turbo + thread pool)

The old per-frame loop was ``[Image.open(BytesIO(b)).convert('RGB') ...]``
sequentially. Two replacements, in priority order, in a new
``_decode_images`` method:

  * torchvision.io.decode_jpeg when available — libjpeg-turbo backed,
    batch SIMD, GIL-released. Order-of-magnitude faster than PIL on CPU
    and supports CUDA decode if a GPU is present. Decoded tensors come
    back in (C, H, W) so we skip the post-decode permute.
  * PIL on a ThreadPoolExecutor as fallback — JPEG decode releases the
    GIL inside libjpeg, so threading still wins despite the GIL.

Malformed-JPEG decode failures fall through from torchvision to PIL
silently rather than crashing the batch. Optional dependency: torchvision
is already a core dep, so the fast path is on by default for everyone.
This narrows the small-frame gap where HDF5 was beating Lance because
PIL decode dominated the bytes-saved benefit at 96x96.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* benchmark: rewrite for tworoom + pusht across all formats with storage cost

Two scripts, paired:

scripts/benchmark/convert.py
  Converts both datasets to all three writable formats:
    pusht   (LeRobot HF) → h5, lance, video
    tworoom (HF .h5)     → h5 (symlink to HF cache), lance, video
  Idempotent — skips conversion when the local output already exists,
  pass --force to redo. Optional --no-upload to skip the S3 step.
  LeRobot import + get_safe_version cache is lazy so this script is
  importable even on machines without lerobot installed.

scripts/benchmark/compare_h5_lance.py
  One run, all rows. Per dataset/format/source/cache combination:
  throughput (samples/s, ms/step) and storage (local file/dir size or
  S3 prefix size). Auto-downloads any missing local copies from S3.
  Markdown table at the end. Includes a LeRobot HF row for pusht.

  Lance paths point directly at the .lance dir; LanceDataset
  auto-detects the parent + table-name split via _resolve_uri_and_table
  so we don't have to track table names separately.

Together: `python convert.py && python compare_h5_lance.py` reproduces
the full benchmark on a fresh EC2 with an IAM instance role attached.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceWriter image_codec: 'raw' (default) | 'jpeg' | 'both'; reader auto-detects

A small reshape is much cheaper than a JPEG decode. At small image sizes
(96×96) JPEG decode overhead outweighs the bytes-saved benefit on local
disk, so HDF5 was beating Lance in the bench despite Lance's columnar
projection win. This change lets users opt the storage form per dataset:

  * 'raw'  (default) — pixels stored as pa.list_(pa.uint8(), H*W*C) with
    field metadata b'image_shape': b'H,W,C'. No decode on read; the
    reader reshapes the flat fixed-size list back to (H, W, C) and
    permutes to (C, H, W).
  * 'jpeg'           — pixels stored as pa.binary, JPEG-encoded blobs.
    Same path that existed before. Cheaper on disk; pays decode cost.
  * 'both'           — write both: <X> as raw + <X>_jpeg as binary.
    Same data, two on-disk representations. The bench can read either.

Reader auto-detects the codec per column from schema type/metadata —
no flags needed, the same dataset handle handles legacy jpeg-only,
new raw-only, and both-codec tables transparently. Multi-camera +
'both' produces <cam>/<cam>_jpeg per camera with no special handling.

Output shape and dtype are unchanged: every image column produces
(num_steps, C, H, W) uint8 torch tensors regardless of codec.

Append to a legacy jpeg-only table still works — the codec is
inferred from the existing schema, not from the writer's default.

5 new tests: codec roundtrip for raw / jpeg / both, multi-camera with
codec='both', and append against a legacy jpeg table.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceDataset: drop ThreadPoolExecutor from _decode_images

The DataLoader already runs in N worker processes, each fetching a
batch independently. Adding our own thread pool inside `_decode_images`
oversubscribes when N≥cores and adds zero parallelism that
`torchvision.io.decode_jpeg` (libjpeg-turbo + SIMD) doesn't already
deliver inside its single C call. Drop the pool; the PIL fallback is
now a plain sequential list comprehension.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* convert.py: use 'hdf5' (registry key) not 'h5' (filename suffix)

PLAN was keyed by 'h5' but the format registry exposes the writer as
'hdf5' — get_format('h5') raised. Switch PLAN keys to 'hdf5'; local
filenames keep their .h5 suffix.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* convert.py: fix tworoom path; --force preserves the canonical source

Two issues with convert_tworoom:

1. `convert(str(h5_p), dest, dest_format=fmt, ...)` was failing because
   load_dataset() doesn't recognise a relative `.h5` filename as a local
   path — it routes through HF/registry resolution first. Use the same
   direct-write pattern convert_pusht uses: HDF5Dataset(path=...) plus a
   per-format writer with image_codec='both' for the lance target.

2. --force was wiping ./tworoom.h5 and re-downloading the canonical
   source from S3 on every run. The source is stable; --force should
   only redo the *derived* outputs (lance + video). Dropped the
   force-wipe of the source path.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceWriter: compress raw uint8 cols via large_binary + Lance 2.2 zstd

The previous raw codec used pa.list_(pa.uint8(), H*W*C) — a fixed-size
list. Lance v2's general compression (LZ4/Zstd via lance-encoding
metadata) does not apply to fixed_size_list; only to binary types.
Empirical: pa.list_<uint8> stays uncompressed at 1.0x; pa.large_binary
with the same content compresses 18-37x on tworoom frames.

Switch the raw codec on disk:
  * pa.large_binary instead of pa.list_(pa.uint8(), N).
  * Field metadata: image_shape='H,W,C' (used by reader for reshape) +
    lance-encoding:compression='zstd' (engages compression).
  * Writer passes data_storage_version='2.2' to db.create_table — the
    metadata-driven compression only takes effect under the v2.2 format.

Reader:
  * Auto-detect raw images by image_shape metadata — works for both
    binary (current) and fixed_size_list<uint8> (legacy) types, so
    older datasets remain readable without any user-facing change.
  * Decode raw binary via np.frombuffer + reshape per row; raw
    fixed_size_list path stays as a contiguous numpy reshape.

Output API is unchanged: every image col still produces (T, C, H, W)
uint8 torch tensors regardless of underlying storage.

convert.py: re-enable codec='both' for tworoom now that the raw col
compresses (~138GB uncompressed → ~4GB zstd on the synthetic env
content). pusht stays on 'both' as before.

Tests:
  * test_codec_raw_roundtrip rewritten for the new schema (large_binary
    + image_shape + lance-encoding:compression metadata).
  * New test_codec_raw_legacy_fixed_size_list_still_readable exercises
    the back-compat reader path against a hand-built legacy table.
  * test_codec_both_roundtrip schema assertion updated.
80 lance/format tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* training scripts: fix spawn-mode DataLoader pickle errors with Lance

Two unrelated bugs surfaced when LanceDataset (which forces spawn
start method on Linux for fork-safety) is used with the training
scripts. Both are unpickleable-callable failures during DataLoader
worker spawn. Both fixed in user-script land; the cleaner fix for #2
belongs in stable_pretraining.

#1 — closure-based normalizer (ALL 7 train scripts):

  def get_column_normalizer(dataset, source, target):
      mean, std = ...                              # captured by closure
      def norm_fn(x):                              # nested local function
          return ((x - mean) / std).float()
      return WrapTorchTransform(norm_fn, ...)

  Standard pickle can't serialise <locals>.norm_fn — it has no
  module-level qualified name.

  Fix: top-level ZScoreNormalizer class + column_normalizer helper
  added to stable_worldmodel.data.utils. Each training script now
  imports `column_normalizer as get_column_normalizer` and the local
  closure block is gone.

#2 — partial-as-Module-forward (lewm.py, pldm.py, prejepa.py):

  world_model = spt.Module(forward=partial(forward_fn, cfg=cfg), ...)

  spt.Module binds `forward=...` as a method on its instance.
  multiprocessing.reduction._reduce_method reduces bound methods via
  (getattr, (m.__self__, m.__func__.__name__)) — but partial has no
  __name__, so spawn-mode workers crash with AttributeError.

  Workaround: top-level _ForwardWithCfg callable class with
  __name__='forward' (the attribute name on Module, so the reducer's
  getattr round-trips). The proper fix is in stable_pretraining.Module
  to not subject `forward` to method-reduction semantics — comment in
  each script flags this.

Verified: lewm.py with `data.dataset.name=tworoom.lance num_workers=2`
now gets past DataLoader spawn (previously crashed with
'cannot pickle ... / partial has no __name__'). Hits an unrelated
pre-existing scheduler config issue in the next training-setup stage,
which is out of scope.

The other 4 scripts (gcbc/gciql/gcivl/hilp) define forward as a local
closure inside run() that captures more than just cfg; same bug class
but bigger surgery — flagged for a follow-up once spt itself is fixed.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* lewm/pldm: pass scheduler warmup_steps + max_steps explicitly

LinearWarmupCosineAnnealingLR requires both as constructor args. spt's
default factory normally derives them from `trainer.estimated_stepping_batches`,
but only when no scheduler kwargs are supplied at all — and Lance datasets
report a length that makes that derivation strip max_steps to None,
crashing the constructor. HDF5 worked by accident because its length
came back finite.

Compute total_steps = max_epochs * len(train), use 1% as warmup_steps
(mirrors what spt's factory would have produced for HDF5).

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* Revert raw/jpeg/both codec + zstd compression; keep JPEG-only path

Drops the codec selector and large_binary+zstd raw layout introduced in
23de71c and bcc0d62. The cross-codec branching in the writer/reader, the
codec test matrix, the LANCE_CODEC_PER_DATASET wiring in convert.py, and
the per-codec rows in the bench all go away. Net -396 lines.

The torchvision libjpeg-turbo fast path in `_decode_images` (from a45e413)
stays — that's the JPEG-side perf win, orthogonal to the codec choice.
Sequential PIL fallback (from 1d65253) stays too.

15/15 tests in tests/data/test_lance.py pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* benchmark: drop pusht; add LeRobot row backed by local tworoom dataset

96x96 pusht frames made per-format throughput differences noisy. Tworoom
(224x224) is now the bench's only dataset.

To keep a LeRobot row in the comparison, convert.py builds a local
LeRobot dataset via `LeRobotDataset.create()` from the tworoom HDF5
source (pixels → observation.images.main as AV1 mp4, proprio →
observation.state, action → action). The bench reads it via
`LeRobotAdapter(repo_id='local/tworoom', root='./tworoom.lerobot')`.
The dir gets uploaded to S3 alongside lance/video and synced down on
cold-start machines like the others.

Verified locally: 2-episode conversion (61+78 frames) produces a
readable LeRobot dataset; LeRobotAdapter returns the expected
(T, C, H, W) tensors.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* convert.py: enable LeRobot streaming_encoding for tworoom conversion

Without streaming, LeRobot's pipeline writes per-frame PNGs and then
spins up SVT-AV1 from scratch on every save_episode() — for tworoom's
~1500 episodes this means 1500 encoder boots, 1500 mp4 second-pass
finalizations, and a lot of disk thrash on the intermediate PNGs.

`streaming_encoding=True` keeps a single encoder warm across every
episode and feeds frames directly into it, skipping the PNG step
entirely. `VideoEncodingManager` flushes the streaming encoder on
exit. Read path is unchanged (LeRobotAdapter still seeks the per-frame
mp4 indices via parquet).

2-episode local smoke test went from ~5s (per-episode encoding) to
1.26s. Read-back via LeRobotAdapter returns identical (T, C, H, W)
tensors.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* convert.py: make LeRobot conversion opt-in via --lerobot

Even with streaming_encoding=True, lerobot's save_episode() finalizes
an mp4 chunk per episode (writing the moov atom, second-pass mux). On
a 1500+ episode dataset that per-episode tax dominates regardless of
how we configure the encoder, so the conversion is impractically slow
for the default workflow.

Default convert.py now does h5/lance/video only (fast). Users who
actually want the LeRobot bench row pass --lerobot. The bench script
already skips the LeRobot row gracefully when the local dir is absent.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* lewm.py: add SWM_PICKLE_DIAGNOSTIC callback for tracing _thread.lock leaks

Workers in spawn-mode DataLoaders crash with "cannot pickle
'_thread.lock' object" — pickle's default error doesn't tell us which
attribute path holds the lock. This callback walks the live train +
val DataLoaders' object graphs in `on_fit_start` (after spt's setup
callbacks have fired but before workers spawn) and prints every
attribute path whose pickle attempt fails on a thread lock. SystemExits
after to skip the rest of the run.

Activate with `SWM_PICKLE_DIAGNOSTIC=1`; default is off so normal runs
are unaffected. Will be reverted once the upstream lock leak is patched.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* cleanup: drop LeRobot from bench, remove pickle diagnostic, trim workaround docs

- convert.py + compare_h5_lance.py: remove all LeRobot wiring. Each
  episode triggered an mp4 finalization regardless of streaming_encoding,
  making conversion infeasibly slow on tworoom (~1500 episodes); the row
  isn't worth it.
- lewm.py: remove the _PickleDiagnostic callback + SWM_PICKLE_DIAGNOSTIC
  env-var hook used to trace the env_info _thread.lock leak. Bug is now
  identified and patched upstream-style.
- lewm/pldm/prejepa.py: shrink _ForwardWithCfg docstring to a 3-line
  comment that just points at the upstream fix location.

Net -257 lines. Ruff lint + format clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* tests/test_lance: drop unused torch import

Leftover from the codec-test removal — the deleted test_codec_* funcs
were the only consumers of torch in this file. Caught by ruff F401 on
pre-commit.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* gcbc/gciql/gcivl/hilp: hoist column_normalizer import to module level

These four were left with the import inside get_data() — leftover from
when the import replaced an inline `def get_column_normalizer(...)`
closure at the same indentation. Match the lewm/pldm/prejepa pattern
where the import lives at module top.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* benchmark: add lewm-pusht (224x224) alongside tworoom

Bring the LeWorldModel pusht (https://huggingface.co/datasets/quentinll/lewm-pusht)
into the bench so we have a second 224x224 dataset besides tworoom. Not
to be confused with lerobot/pusht_image, which is 96x96 and made the
per-format throughput differences too noisy to interpret.

convert.py:
- Refactor lance/video derivation into _derive_lance_video() helper
  shared by tworoom and pusht.
- Add convert_pusht(): tries S3 first (s3://.../pusht/pusht.h5) for
  a fast hop, falls back to downloading
  pusht_expert_train.h5.zst from HF and decompressing via the system
  `zstd` CLI on miss. After the first run + upload, future runs use S3.
- main() runs both convert_tworoom + convert_pusht; upload loop walks
  the whole PLAN.

compare_h5_lance.py:
- Add 'pusht' entry to DATASETS mirroring tworoom's local/s3 paths.
- The bench's per-dataset loop already iterates DATASETS, so pusht
  rows appear automatically.

* LanceDataset: switch torch sharing strategy to file_system on Linux

Spawn-mode DataLoader workers IPC tensors via PyTorch's shared-memory
layer. The default 'file_descriptor' strategy mmap()s files under
/dev/shm, and on cloud instances (e.g. EC2/Lambda) the ftruncate()
often returns EINVAL despite plenty of free shm space — kernel,
AppArmor, and AMI tmpfs flags all vary. Workers crash with:

    RuntimeError: unable to resize file <filename not specified>
    to the right size: Invalid argument (22)

'file_system' uses regular files and works everywhere. Throughput
delta is negligible for our batch sizes, and we already force spawn
on Linux for fork-safety, so this just plugs the matching IPC hole.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceDataset: drop _trainer back-reference on pickle

spt.Module calls `dataset.set_pl_trainer(trainer)` on every dataset to
inject `global_step` / `current_epoch` into samples. The trainer
transitively reaches `train_dataloader._iterator`
(`_MultiProcessingDataLoaderIter`), which raises NotImplementedError on
pickle. With `persistent_workers=True` the train iterator gets created
once and stays around, so when the val DataLoader spawns its workers,
the spawn closure traverses val_dataset._trainer -> ... -> train
iterator -> crash:

    NotImplementedError: ('{} cannot be pickled', '_MultiProcessingDataLoaderIter')

Drop `_trainer` in `__getstate__` alongside `_perm`. Workers see only a
snapshot of trainer state at spawn time anyway, so a missing trainer
on the worker side is fine — `process_sample` already handles it
(`if self._trainer is not None`).

This is also a latent bug in spt's DatasetMixin (the `_trainer`
back-reference is fragile across spawn for any wrapped dataset that
defines its own `__getstate__`); worth flagging upstream separately.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceDataset: silence torch.frombuffer non-writable warning

Lance returns immutable `bytes` for image blobs. `torch.frombuffer`
warns on non-writable input because mutating the resulting tensor
would silently mutate the underlying buffer. Wrap in `bytearray()`
to get a writable copy — same memory cost we already paid before via
the `bytes(b)` fallback, but unconditional and warning-free.

* LanceDataset: zero-copy JPEG decode + suppress non-writable warning

Reverts the bytearray() wrap in _decode_images. We pass the immutable
`bytes` from Lance straight into `torch.frombuffer` as a view (no
memcpy), then hand the view to torchvision's decoder which returns a
fresh output tensor. The shared memory between the view and the
source bytes is never written, so the underlying concern of the
warning ('writes via the view would corrupt source') doesn't apply.

Suppress the specific UserWarning at module load to keep logs clean.
~0.5 ms/batch saved at typical settings vs the bytearray() copy.

* formats: route s3:// (and other ://) URIs through to readers

Both Lance and HDF5 already accept remote URIs at the dataset class
level (Lance via object_store, HDF5 via fsspec). The format-dispatch
layer was the gap:

- HDF5.open_reader called Path(path).is_dir() on s3:// URIs, which
  silently mangled the URI (Path collapses `://` to `:/`).
- Neither format injected AWS region for s3:// — the bench passed it
  explicitly, but training scripts had no way to.

Fix: in both open_reader, branch on `://` in the path string. Pass the
URI through unchanged, and auto-inject `region` from
`AWS_DEFAULT_REGION` (default `us-east-1`) into the appropriate
storage-options dict when the caller didn't pass one.

After this:

  data.dataset.name=s3://bucket/path/foo.lance
  data.dataset.name=s3://bucket/path/foo.h5

just work in any training script, given AWS creds in env.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* LanceDataset: soften keys_to_cache warning message

The previous wording ('caching risks OOM') was overly alarming for what
is just a 'this knob is redundant' note. Reword to highlight the actual
property — Lance already has efficient random access — instead of
implying users will run out of memory.

* Drop spawn-mode workarounds: LanceDB main is now fork-safe

LanceDB PR #3339 (lancedb/lancedb#3339) makes
both the Python BackgroundEventLoop and the Rust tokio runtime
fork-safe via pthread_atfork / os.register_at_fork hooks. This removes
the original reason LanceDataset forced spawn on Linux and lets the
default fork-mode DataLoader workers Just Work.

Three workarounds drop together:

1. LanceDataset._maybe_warn_fork_start_method (and the _fork_warning_emitted
   class flag): no longer force spawn. Workers fork cleanly; lancedb
   re-initialises its async layers in the child.

2. torch.multiprocessing.set_sharing_strategy('file_system') (was inside
   the same classmethod): only needed because spawn used /dev/shm fds
   that misbehaved on cloud instances. Fork inherits memory directly,
   so the default file_descriptor strategy is fine.

3. _ForwardWithCfg in lewm/pldm/prejepa.py: the wrapper around
   functools.partial(forward_fn, cfg=cfg) only existed because
   spawn pickles the bound method, which fails when the underlying
   callable lacks __name__. Fork doesn't pickle anything; a plain
   partial works.

Verified locally: built lancedb 0.31.0b11 from main, opened a Lance
table inside fork-mode DataLoader workers, confirmed Permutation
returns batches end-to-end (pre-#3339 this hangs indefinitely).

Requires lancedb release that includes #3339 (any 0.31.x once cut).

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* Revert "Drop spawn-mode workarounds: LanceDB main is now fork-safe"

This reverts commit a0ee115.

* Consolidate spawn-mode workarounds into stable_worldmodel/_spawn_compat

Pulled the duplicated _ForwardWithCfg class (3 copies in lewm/pldm/
prejepa) and LanceDataset's _maybe_warn_fork_start_method classmethod
into a single module — stable_worldmodel/_spawn_compat — with one
clear docstring explaining WHY each piece exists, what bug it works
around, and when to drop it. Tagged at the top with a pointer to the
upstream ticket (lancedb_fork_segfault_ticket.md).

The module exports two helpers:

  - force_spawn_for_lancedb(): switches mp start method to spawn on
    Linux + flips torch sharing strategy to file_system. Called from
    LanceDataset.__init__.
  - ForwardWithCfg(fn, cfg): picklable replacement for partial that
    survives spawn-mode bound-method reduction. Used as the forward=
    arg to spt.Module.

No behaviour change vs the previous scattered version — same bug,
same workaround, just tagged and trackable in one place so future
cleanup is a single import-removal.

* training: pass ckpt_path only if it exists

spt.Manager (post a recent change) refuses to silently start training
from scratch when ckpt_path is set but the file doesn't exist. All
three of our training scripts always passed the path, even on fresh
runs, hitting:

    FileNotFoundError: ckpt_path was set to ... but no such file
    exists. Refusing to silently start training from scratch.

Guard each call with .exists() so fresh runs pass None and only
existing checkpoints are loaded for resume.

* refactor

* drop ForwardWithCfg hack; consolidate spawn workaround in lance.py

galilai-group/stable-pretraining#417 (PR #420) makes spt.Module accept
functools.partial as forward, so ForwardWithCfg is no longer needed.
Train scripts now pass forward=partial(fn, cfg=cfg) directly.

Move force_spawn_for_lancedb back into lance.py as a private helper
(the only caller). Delete _spawn_compat.py — both its inhabitants are
gone now.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* forward HF_TOKEN env var into lancedb storage_options for hf:// URIs

Lance picks up AWS creds automatically via the SDK chain, but not
HF_TOKEN — without this, hf://buckets/... reads silently fail or hit
auth retries that look like a hang. Inject only for hf:// URIs and
only when HF_TOKEN is set, so AWS/GS/Azure paths are untouched.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* update token injection

* fix(lance): scope HF_TOKEN injection to hf:// URIs only

object_store treats `token` as the AWS session token for S3
connections, so unconditionally setting it from HF_TOKEN broke
s3:// reads with `InvalidToken: The provided token is malformed`.
Gate on the hf:// prefix so AWS credential resolution is left
untouched on S3.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

* support string caching cols

---------

Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>