Segmentation fault during orbax checkpoint restore for pi0_fast_droid

[hypercosmac]

I'm running into a segmentation fault when building a docker image using a slightly modified serve_policy.Dockerfile to run the pi0 fast droid model for inference on Modal GPUs, and the issue appears to be with orbax checkpoint restore -
I've double checked the paths (eg. metadata is read in successfully) - screenshot attached. Not sure if its a float 32/16 issue..has anyone faced anything similar?

File "/root/modal-inference.py", line 80, in run_openpi_inference
  policy = policy_config.create_trained_policy(config_obj, checkpoint_dir)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/app/openpi/policies/policy_config.py", line 61, in create_trained_policy
  model = train_config.model.load(_model.restore_params(checkpoint_dir / "params", dtype=jnp.float32))
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/app/openpi/models/model.py", line 312, in restore_params
  params = ckptr.restore(
           ^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 258, in restore
  restored = self._restore(directory, args=ckpt_args)

This is where the code stalls before running into the seg fault:

  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 482, in _read_array_index_and_device_put
    shard = await _read_shard(
            ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 425, in _read_shard
    await ts.array(out)[ts.d[:].translate_to[requested_domain.origin]][
asyncio.exceptions.CancelledError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/pkg/modal/_runtime/container_io_manager.py", line 728, in handle_input_exception
    yield
  File "/pkg/modal/_container_entrypoint.py", line 240, in run_input_sync
    res = io_context.call_finalized_function()
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/pkg/modal/_runtime/container_io_manager.py", line 180, in call_finalized_function
    res = self.finalized_function.callable(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/modal-inference.py", line 80, in run_openpi_inference
    policy = policy_config.create_trained_policy(config_obj, checkpoint_dir)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/openpi/policies/policy_config.py", line 61, in create_trained_policy
    model = train_config.model.load(_model.restore_params(checkpoint_dir / "params", dtype=jnp.float32))
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/openpi/models/model.py", line 312, in restore_params
    params = ckptr.restore(
             ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 258, in restore
    restored = self._restore(directory, args=ckpt_args)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/checkpointers/checkpointer.py", line 272, in _restore
    return self._handler.restore(directory, args=args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py", line 796, in restore
    return self._handler_impl.restore(directory, args=args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py", line 733, in restore
    tree_memory_size, restored_item = asyncio_utils.run_sync(
                                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/asyncio_utils.py", line 50, in run_sync
    return asyncio.run(coro)
           ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/base_events.py", line 653, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py", line 581, in _maybe_deserialize
    deserialized_batches += await asyncio.gather(*deserialized_batches_ops)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py", line 1224, in deserialize
    ret = await asyncio.gather(*deserialize_ops)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 585, in async_deserialize
    return await read_and_create_array(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 544, in read_and_create_array
    dbs = sum(await asyncio.gather(*read_array_coros), [])
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/serialization.py", line 490, in _read_array_index_and_device_put
    raise Exception(  # pylint: disable=broad-exception-raised
Exception: Encountered error while reading array index: (slice(0, 18, 1), slice(0, 2, 1), slice(0, 2048, 1), slice(0, 16384, 1)). See full TensorStore details: <bound method PyCapsule.spec of TensorStore({
  'base': {
    'driver': 'zarr',
    'dtype': 'float32',
    'fill_missing_data_reads': False,
    'kvstore': {
      'base': {
        'driver': 'file',
        'path': '/__modal/volumes/vo-3I8FqeT8KEQpvNZ4srXYED/openpi/checkpoints/pi0_fast_droid/params/',
      },
      'cache_pool': 'cache_pool#ocdbt',
      'config': {
        'compression': {'id': 'zstd'},
        'max_decoded_node_bytes': 100000000,
        'max_inline_value_bytes': 1024,
        'uuid': 'ca888d2fcf3b0d15d240930fdc9d4ce1',
        'version_tree_arity_log2': 4,
      },
      'driver': 'ocdbt',
      'experimental_read_coalescing_interval': '1ms',
      'experimental_read_coalescing_merged_bytes': 500000000000,
      'experimental_read_coalescing_threshold_bytes': 1000000,
      'path': 'params.PaliGemma.llm.layers.mlp.gating_einsum/',
    },
    'metadata': {
      'chunks': [18, 2, 2048, 4096],
      'compressor': {'id': 'zstd', 'level': 1},
      'dimension_separator': '.',
      'dtype': '<f4',
      'fill_value': None,
      'filters': None,
      'order': 'C',
      'shape': [18, 2, 2048, 16384],
      'zarr_format': 2,
    },
    'recheck_cached_data': False,
    'recheck_cached_metadata': False,
  },
  'context': {
    'cache_pool': {},
    'cache_pool#ocdbt': {'total_bytes_limit': 100000000},
    'data_copy_concurrency': {},
    'file_io_concurrency': {'limit': 128},
    'file_io_locking': {},
    'file_io_memmap': False,
    'file_io_sync': True,
    'ocdbt_coordinator': {},
  },
  'driver': 'cast',
  'dtype': 'float32',
  'transform': {
    'input_exclusive_max': [[18], [2], [2048], [16384]],
    'input_inclusive_min': [0, 0, 0, 0],
  },
})>.