Assertion fails in unpack_helper()

[qzed]

🐛 Bug

When remote logging, resuming/continuing a run occasionally fails, triggering an assertion. See log/trace excerpt below for details.

Trace

``` │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/sdk/run.py:864 in __init__ │ │ │ │ 861 │ │ log_system_params: Optional[bool] = False, │ │ 862 │ │ capture_terminal_logs: Optional[bool] = True, │ │ 863 │ ): │ │ ❱ 864 │ │ super().__init__(run_hash, repo=repo, read_only=read_only, experiment=experiment │ │ 865 │ │ │ │ 866 │ │ self._system_resource_tracker: ResourceTracker = None │ │ 867 │ │ if not read_only: │ │ │ │ ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │ │ │ capture_terminal_logs = True │ │ │ │ experiment = None │ │ │ │ force_resume = True │ │ │ │ log_system_params = False │ │ │ │ read_only = False │ │ │ │ repo = 'aim://10.4.28.205:53800' │ │ │ │ run_hash = '147ff01f895d4e828fceb6f0' │ │ │ │ self = > │ │ │ │ system_tracking_interval = 10 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/sdk/run.py:280 in __init__ │ │ │ │ 277 │ │ if not read_only: │ │ 278 │ │ │ logger.debug(f'Opening Run {self.hash} in write mode') │ │ 279 │ │ │ │ │ ❱ 280 │ │ │ if self.check_metrics_version(): │ │ 281 │ │ │ │ if self.repo.is_remote_repo: │ │ 282 │ │ │ │ │ logger.warning( │ │ 283 │ │ │ │ │ │ f'Cannot track Run with remote repo {self.repo.path}. Please upg │ │ │ │ ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │ │ │ experiment = None │ │ │ │ force_resume = True │ │ │ │ read_only = False │ │ │ │ repo = 'aim://10.4.28.205:53800' │ │ │ │ run_hash = '147ff01f895d4e828fceb6f0' │ │ │ │ self = > │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/sdk/base_run.py:81 in │ │ check_metrics_version │ │ │ │ 78 │ │ │ 79 │ def check_metrics_version(self) -> bool: │ │ 80 │ │ metric_dtypes = ('float', 'float64', 'int') │ │ ❱ 81 │ │ traces_tree = self.meta_run_tree.get('traces', {}) │ │ 82 │ │ │ │ 83 │ │ v1_metric_found = False │ │ 84 │ │ for ctx_metadata in traces_tree.values(): │ │ │ │ ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │ │ │ metric_dtypes = ('float', 'float64', 'int') │ │ │ │ self = > │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ in aim.storage.treeview.TreeView.get:32 │ │ │ │ in aim.storage.treeview.TreeView.__getitem__:28 │ │ │ │ in aim.storage.treeview.TreeView.collect:23 │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/storage/treeviewproxy.py:154 in │ │ collect │ │ │ │ 151 │ def collect( │ │ 152 │ │ self, path: Union[AimObjectKey, AimObjectPath] = (), strict: bool = True, resolv │ │ 153 │ ) -> AimObject: │ │ ❱ 154 │ │ return self.tree.collect(self.absolute_path(path), strict, resolve_objects) │ │ 155 │ │ │ 156 │ def __delitem__(self, path: Union[AimObjectKey, AimObjectPath]): │ │ 157 │ │ del self.tree[self.absolute_path(path)] │ │ │ │ ╭────────────────────────────────────── locals ──────────────────────────────────────╮ │ │ │ path = 'traces' │ │ │ │ resolve_objects = False │ │ │ │ self = │ │ │ │ strict = True │ │ │ ╰────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/storage/treeviewproxy.py:73 in │ │ collect │ │ │ │ 70 │ def collect( │ │ 71 │ │ self, path: Union[AimObjectKey, AimObjectPath] = (), strict: bool = True, resolv │ │ 72 │ ) -> AimObject: │ │ ❱ 73 │ │ return self._rpc_client.run_instruction(self._hash, self._handler, 'collect', (p │ │ 74 │ │ │ 75 │ def __delitem__(self, path: Union[AimObjectKey, AimObjectPath]): │ │ 76 │ │ self._rpc_client.run_instruction(self._hash, self._handler, '__delitem__', (path │ │ │ │ ╭───────────────────────────────────── locals ─────────────────────────────────────╮ │ │ │ path = ('meta', 'chunks', '147ff01f895d4e828fceb6f0', 'traces') │ │ │ │ resolve_objects = False │ │ │ │ self = │ │ │ │ strict = True │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/ext/transport/client.py:253 in │ │ run_instruction │ │ │ │ 250 │ │ │ ) │ │ 251 │ │ │ return │ │ 252 │ │ │ │ ❱ 253 │ │ return self._run_read_instructions(queue_id, resource, method, args) │ │ 254 │ │ │ 255 │ def _run_read_instructions(self, queue_id, resource, method, args): │ │ 256 │ │ endpoint = f'{self._http_protocol}{self._tracking_endpoint}/{self.uri}/read-inst │ │ │ │ ╭──────────────────────────────────────── locals ─────────────────────────────────────────╮ │ │ │ args = (('meta', 'chunks', '147ff01f895d4e828fceb6f0', 'traces'), True, False) │ │ │ │ is_write_only = False │ │ │ │ method = 'collect' │ │ │ │ queue_id = '147ff01f895d4e828fceb6f0' │ │ │ │ resource = '5519f717-f1b9-4def-9bed-53307197ffa7' │ │ │ │ self = │ │ │ ╰─────────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/ext/transport/client.py:273 in │ │ _run_read_instructions │ │ │ │ 270 │ │ │ │ 271 │ │ if response.status_code == 400: │ │ 272 │ │ │ raise_exception(response.json().get('exception')) │ │ ❱ 273 │ │ return decode_tree(unpack_stream(response.iter_content(chunk_size=None))) │ │ 274 │ │ │ 275 │ def _run_write_instructions(self, instructions: [Tuple[bytes, bytes]]): │ │ 276 │ │ msg = pack_args(iter(instructions)) │ │ │ │ ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │ │ │ args = (('meta', 'chunks', '147ff01f895d4e828fceb6f0', 'traces'), True, False) │ │ │ │ endpoint = 'http://10.4.28.205:53800/tracking/74e8733e-ebb5-4c36-a6af-01291cf1f782/read-inst'+8 │ │ │ │ method = 'collect' │ │ │ │ queue_id = '147ff01f895d4e828fceb6f0' │ │ │ │ request_data = { │ │ │ │ │ 'resource_handler': '5519f717-f1b9-4def-9bed-53307197ffa7', │ │ │ │ │ 'method_name': 'collect', │ │ │ │ │ 'args': │ │ │ │ 'AAAAAAABAAAABgoAAAD+AAAAAAAAAAD+AAEAAAAGFAAAAP4AAAAAAAAAAP7+AAAAAAAAAAD+AAUAAAAE'+232 │ │ │ │ } │ │ │ │ resource = '5519f717-f1b9-4def-9bed-53307197ffa7' │ │ │ │ response = │ │ │ │ self = │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ in aim.storage.treeutils.decode_tree:270 │ │ │ │ in aim.storage.treeutils.fold_tree:95 │ │ │ │ in iter_fold_tree:126 │ │ │ │ in aim.storage.treeutils.DecodePathsVals.__next__:201 │ │ │ │ in aim.storage.treeutils.DecodePathsVals._next:214 │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/ext/transport/message_utils.py:63 in │ │ unpack_stream │ │ │ │ 60 │ │ 61 def unpack_stream(stream) -> Tuple[bytes, bytes]: │ │ 62 │ for msg in stream: │ │ ❱ 63 │ │ yield from unpack_helper(msg) │ │ 64 │ │ 65 │ │ 66 def raise_exception(server_exception): │ │ │ │ ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │ │ │ msg = b'#\x00\x00\x00\xfeH\x07\xfcy\xc3\xdf\xaf5\xfeamotp/bicycle\xfefirst_step\xfe\x00\t\x00\x00\x00\x02\xb… │ │ │ │ stream = .generate at 0x7672f91298a0> │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │ │ │ │ /home/luz/micromamba/envs/bosch-tracking-cu118/lib/python3.11/site-packages/aim/ext/transport/message_utils.py:54 in │ │ unpack_helper │ │ │ │ 51 │ (is_blob,), tail = struct.unpack('?', tail[:1]), tail[1:] │ │ 52 │ (value_size,), tail = struct.unpack('I', tail[:4]), tail[4:] │ │ 53 │ value, tail = tail[:value_size], tail[value_size:] │ │ ❱ 54 │ assert len(tail) == 0 │ │ 55 │ if is_blob: │ │ 56 │ │ yield key, BLOB(data=value) │ │ 57 │ else: │ │ │ │ ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │ │ │ is_blob = False │ │ │ │ key = b'\xfeH\x07\xfcy\xc3\xdf\xaf5\xfeamotp/bicycle\xfefirst_step\xfe' │ │ │ │ key_size = 35 │ │ │ │ msg = b'#\x00\x00\x00\xfeH\x07\xfcy\xc3\xdf\xaf5\xfeamotp/bicycle\xfefirst_step\xfe\x00\t\x00\x00\x00\x0… │ │ │ │ tail = b'\x1e\x00\x00\x00\xfeH\x07\xfcy\xc3\xdf\xaf5\xfeamotp/bicycle\xfefirst\xfe\x00\t\x00\x00\x00\x03\… │ │ │ │ value = b'\x02\xbc\r\x00\x00\x00\x00\x00\x00' │ │ │ │ value_size = 9 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ```

It looks like unpack_helper() expects a single message as input. But I think using response.iter_content(chunk_size=None) does not guarantee this, e.g., I think it might receive two or more messages, causing the assertion to fail.

To reproduce

Random. Resume/continue stopped runs in a remote logging setup.

Expected behavior

Runs can be resumed without errors.

Environment

• Aim Version: v3.27.0
• Python version: 3.11.9
• pip version: 24.3.1
• OS: Linux/Ubuntu 22.04