def setUp(self):
super().setUp()
self.cpp_writer_mock = mock.Mock()
self.cpp_writer_mock.Append.side_effect = _mock_append
self.cpp_writer_mock.AppendPartial.side_effect = _mock_append
self.writer = trajectory_writer.TrajectoryWriter(self.cpp_writer_mock)
def test_numpy_squeeze(self):
# No data will ever be sent to the server so it doesn't matter that we use
# an invalid address.
client = client_lib.Client('localhost:1234')
writer = trajectory_writer.TrajectoryWriter(client, 5, 10)
for i in range(10):
writer.append({'a': i})
self.assertEqual(writer.history['a'][-1].numpy(), i)
def setUp(self):
super().setUp()
self.cpp_writer_mock = mock.Mock()
self.cpp_writer_mock.Append.side_effect = \
lambda x: [FakeWeakCellRef(y) if y is not None else None for y in x]
self.cpp_writer_mock.AppendPartial.side_effect = \
lambda x: [FakeWeakCellRef(y) if y is not None else None for y in x]
self.writer = trajectory_writer.TrajectoryWriter(self.cpp_writer_mock)
def setUp(self):
super().setUp()
self.cpp_writer_mock = mock.Mock()
self.cpp_writer_mock.Append.side_effect = \
lambda x: [FakeWeakCellRef(y) if y else None for y in x]
client_mock = mock.Mock()
client_mock._client.NewTrajectoryWriter.return_value = self.cpp_writer_mock
self.writer = trajectory_writer.TrajectoryWriter(client_mock, 1, 1)
def test_numpy(self):
# No data will ever be sent to the server so it doesn't matter that we use
# an invalid address.
client = client_lib.Client('localhost:1234')
writer = trajectory_writer.TrajectoryWriter(client, 5, 10)
for i in range(10):
writer.append({'a': i, 'b': np.ones([3, 3], np.float) * i})
np.testing.assert_array_equal(writer.history['a'][:].numpy(),
np.arange(i + 1, dtype=np.int64))
np.testing.assert_array_equal(
writer.history['b'][:].numpy(),
np.stack([np.ones([3, 3], np.float) * x for x in range(i + 1)]))
def trajectory_writer(self,
num_keep_alive_refs: int,
*,
get_signature_timeout_ms: Optional[int] = 3000):
"""Constructs a new `TrajectoryWriter`.
Note: The documentation is minimal as this is just a draft proposal to give
alpha testers something tangible to play around with.
Note: The chunk length is auto tuned by default. Use
`TrajectoryWriter.configure` to override this behaviour.
TODO(b/179978457): Add documentation and examples.
Args:
num_keep_alive_refs: The size of the circular buffer which each column
maintains for the most recent data appended to it. When a data reference
popped from the buffer it can no longer be referenced by new items. The
value `num_keep_alive_refs` can therefore be interpreted as maximum
number of steps which a trajectory can span.
get_signature_timeout_ms: The number of milliesconds to wait to pull table
signatures (if any) from the server. These signatures are used to
validate new items before they are sent to the server. Signatures are
only pulled once and cached. If set to None then the signature will not
fetched from the server. Default wait time is 3 seconds.
Returns:
A `TrajectoryWriter` with auto tuned chunk lengths in each column.
Raises:
ValueError: If num_keep_alive_refs < 1.
"""
if num_keep_alive_refs < 1:
raise ValueError(
f'num_keep_alive_refs ({num_keep_alive_refs}) must be a positive '
f'integer')
chunker_options = pybind.AutoTunedChunkerOptions(
num_keep_alive_refs, 1.0)
cpp_writer = self._client.NewTrajectoryWriter(
chunker_options, get_signature_timeout_ms)
return trajectory_writer_lib.TrajectoryWriter(cpp_writer)
def setUp(self):
super().setUp()
self.items = []
def _mock_create_item(unused_table, unused_priority, refs_per_col,
squeeze_per_col):
item = []
for refs, squeeze in zip(refs_per_col, squeeze_per_col):
if squeeze:
item.append(refs[0].data)
else:
item.append(tuple([ref.data for ref in refs]))
self.items.append(tuple(item))
self.cpp_writer_mock = mock.Mock()
self.cpp_writer_mock.Append.side_effect = _mock_append
self.cpp_writer_mock.AppendPartial.side_effect = _mock_append
self.cpp_writer_mock.CreateItem.side_effect = _mock_create_item
self.cpp_writer_mock.max_num_keep_alive_refs = 10
self.writer = trajectory_writer.TrajectoryWriter(self.cpp_writer_mock)
def trajectory_writer(self,
num_keep_alive_refs: int,
*,
validate_items: bool = True):
"""Constructs a new `TrajectoryWriter`.
Note: The chunk length is auto tuned by default. Use
`TrajectoryWriter.configure` to override this behaviour.
See `TrajectoryWriter` for more detailed documentation about the writer
itself.
Args:
num_keep_alive_refs: The size of the circular buffer which each column
maintains for the most recent data appended to it. When a data reference
popped from the buffer it can no longer be referenced by new items. The
value `num_keep_alive_refs` can therefore be interpreted as maximum
number of steps which a trajectory can span.
validate_items: Whether to validate items against the table signature
before they are sent to the server. This requires table signature to be
fetched from the server and cached locally.
Returns:
A `TrajectoryWriter` with auto tuned chunk lengths in each column.
Raises:
ValueError: If num_keep_alive_refs < 1.
"""
if num_keep_alive_refs < 1:
raise ValueError(
f'num_keep_alive_refs ({num_keep_alive_refs}) must be a positive '
f'integer'
)
chunker_options = pybind.AutoTunedChunkerOptions(num_keep_alive_refs, 1.0)
cpp_writer = self._client.NewTrajectoryWriter(chunker_options,
validate_items)
return trajectory_writer_lib.TrajectoryWriter(cpp_writer)
