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)