🐛 Compressed stash failures are silently swallowed

src/aiida/engine/daemon/execmanager.py

@@ -495,18 +495,27 @@ async def stash_calculation(calculation: CalcJobNode, transport: Transport) -> N
     ###
 
     if stash_mode == StashMode.COPY.value:
-        target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:]
+        # sharded by the source node, one directory per stash job
+        target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:] / calculation.uuid
 
-        async def _do_copy():
+        async def _do_copy() -> list[str]:
+            stashed_source_list: list[str] = []
             for source_filename in source_list:
                 if has_magic(source_filename):
+                    if fail_on_missing:
+                        msg = (
+                            'Stashing with glob patterns is not supported when fail_on_missing is True. '
+                            'Stashing failed.'
+                        )
+                        raise exceptions.StashingError(msg)
                     copy_instructions = []
                     for globbed_filename in await transport.glob_async(source_basepath / source_filename):
                         target_filepath = target_basepath / Path(globbed_filename).relative_to(source_basepath)
                         copy_instructions.append((globbed_filename, target_filepath))
                 else:
                     copy_instructions = [(source_basepath / source_filename, target_basepath / source_filename)]
 
+                stashed = False
                 for source_filepath, target_filepath in copy_instructions:
                     # If source is in a (nested) directory, create those directories first
                     target_dirname = target_filepath.parent
@@ -529,9 +538,13 @@ async def _do_copy():
                             f'Failed to copy {source_filepath} to {target_filepath}: {exc}'
                         ) from exc
                     EXEC_LOGGER.debug(f'Stashed from {source_filepath} to {target_filepath}')
+                    stashed = True
+                if stashed:
+                    stashed_source_list.append(source_filename)
+            return stashed_source_list
 
         try:
-            await _do_copy()
+            stashed_source_list = await _do_copy()
         except exceptions.StashingError as exception:
             await transport.rmtree_async(target_basepath)
             raise exception
@@ -542,7 +555,7 @@ async def _do_copy():
             computer=calculation.computer,
             target_basepath=str(target_basepath),
             stash_mode=StashMode(stash_mode),
-            source_list=source_list,
+            source_list=stashed_source_list,
             fail_on_missing=fail_on_missing,
         ).store()
 
@@ -555,52 +568,68 @@ async def _do_copy():
         # stash_mode values are identical with compression_format in transport plugin:
         # 'tar', 'tar.gz', 'tar.bz2', or 'tar.xz'
         compression_format = stash_mode
-        file_name = uuid
         authinfo = calculation.get_authinfo()
         aiida_remote_base = authinfo.get_workdir().format(username=transport.whoami())
 
-        target_destination = str(target_base / file_name) + '.' + compression_format
+        # sharded by the source node, one archive per stash job
+        target_destination = str(
+            target_base / uuid[:2] / uuid[2:4] / uuid[4:] / f'{calculation.uuid}.{compression_format}'
+        )
 
-        source_list_abs = [source_basepath / source for source in source_list]
+        # ``compress_async`` raises on any missing source, so resolve them here to honour ``fail_on_missing``.
+        # Only the entries that resolve are stashed and recorded on the node.
+        stashed_source_list: list[str] = []
+        source_list_abs: list[Path] = []
+        for source in source_list:
+            source_filepath = source_basepath / source
+            if has_magic(str(source_filepath)):
+                if fail_on_missing:
+                    msg = 'Stashing with glob patterns is not supported when fail_on_missing is True. Stashing failed.'
+                    raise exceptions.StashingError(msg)
+                if await transport.glob_async(source_filepath):
+                    stashed_source_list.append(source)
+                    source_list_abs.append(source_filepath)
+                else:
+                    EXEC_LOGGER.warning(f'No match for {source_filepath}. Skipping, because fail_on_missing=False')
+            elif await transport.path_exists_async(source_filepath):
+                stashed_source_list.append(source)
+                source_list_abs.append(source_filepath)
+            elif fail_on_missing:
+                msg = f'File {source_filepath} does not exist and fail_on_missing is True. Stashing failed.'
+                raise exceptions.StashingError(msg)
+            else:
+                EXEC_LOGGER.warning(f'File not found {source_filepath}. Skipping, because fail_on_missing=False')
 
-        # When fail_on_missing is True, check that all files exist before compressing
-        if fail_on_missing:
-            for source_filepath in source_list_abs:
-                if has_magic(str(source_filepath)):
-                    raise exceptions.StashingError(
-                        'Stashing with glob patterns is not supported when fail_on_missing is True. Stashing failed.'
-                    )
-                if not await transport.path_exists_async(source_filepath):
-                    raise exceptions.StashingError(
-                        f'File {source_filepath} does not exist and fail_on_missing is True. Stashing failed.'
-                    )
+        if not source_list_abs:
+            EXEC_LOGGER.warning(f'None of {source_list} exist in {source_basepath}. Nothing to stash.')
+            return
 
         remote_stash = RemoteStashCompressedData(
             computer=calculation.computer,
             target_basepath=target_destination,
             stash_mode=StashMode(stash_mode),
-            source_list=source_list,
+            source_list=stashed_source_list,
             dereference=dereference,
             fail_on_missing=fail_on_missing,
         )
 
+        # The path is unique to this job, so anything already there is a leftover of an earlier attempt, which
+        # no cleanup is guaranteed to have removed (the daemon may have died mid-stash): a retry must replace it.
         try:
             await transport.compress_async(
                 format=compression_format,
                 remotesources=source_list_abs,
                 remotedestination=target_destination,
                 root_dir=aiida_remote_base,
-                overwrite=False,
+                overwrite=True,
                 dereference=dereference,
             )
         except (OSError, ValueError) as exception:
-            EXEC_LOGGER.warning(f'Failed to stash {source_list} to {target_destination}: {exception}')
-            return
-            # note: if you raise here, you trigger the exponential backoff
-            # and if you don't raise, it appears as successful in verdi process list: Finished [0]
-            # An issue opened to investigate and fix this https://github.com/aiidateam/aiida-core/issues/6789
-            # raise exceptions.RemoteOperationError(f'failed '
-            # 'to compress {source_list} to {target_destination}: {exception}')
+            # remove our own partial leftover so a retry starts clean
+            if await transport.path_exists_async(target_destination):
+                await transport.remove_async(target_destination)
+            msg = f'Failed to stash {stashed_source_list} to {target_destination}: {exception}'
+            raise exceptions.StashingError(msg) from exception
         else:
             EXEC_LOGGER.debug(f'Stashed {source_list} to {target_destination}')
 

tests/calculations/test_stash.py

@@ -527,7 +527,7 @@ def test_all_modes(fixture_sandbox, aiida_localhost, generate_calc_job, tmp_path
 
     # Verify files were stashed
     if stash_mode == StashMode.COPY.value:
-        expected_base = Path(target_base) / source_node.uuid[0:2] / source_node.uuid[2:4] / source_node.uuid[4:]
+        expected_base = Path(stash_data_node.target_basepath)
     else:
         temp_for_extract = tmp_path / 'extract'
         temp_for_extract.mkdir()

tests/engine/daemon/test_execmanager.py

@@ -16,6 +16,7 @@
 from aiida.common.datastructures import CalcInfo, CodeInfo, FileCopyOperation, StashMode
 from aiida.common.exceptions import StashingError
 from aiida.common.folders import SandboxFolder
+from aiida.common.links import LinkType
 from aiida.engine.daemon import execmanager
 from aiida.orm import CalcJobNode, FolderData, PortableCode, RemoteData, SinglefileData
 from aiida.transports.plugins.local import LocalTransport
@@ -682,12 +683,8 @@ async def test_stashing(
     serialize_file_hierarchy,
     tmp_path,
     monkeypatch,
-    caplog,
 ):
     """Test `stash_calculation`"""
-
-    import logging
-
     computer_wdir = tmp_path / 'aiida'
     computer_wdir.mkdir()
     dest_path = tmp_path / 'stash_path'
@@ -739,10 +736,11 @@ def get_workdir(self, *args, **kwargs):
 
     if stash_mode != StashMode.COPY.value:
         # more detailed test on integrity of the zip file is in `test_all_plugins.py`
-        assert pathlib.Path(str(dest_path / node.uuid) + '.' + stash_mode).is_file()
+        archive = dest_path / uuid[:2] / uuid[2:4] / uuid[4:] / f'{uuid}.{stash_mode}'
+        assert archive.is_file()
 
         with LocalTransport() as transport:
-            transport.extract(str(dest_path / node.uuid) + '.' + stash_mode, dest_path / 'extracted')
+            transport.extract(archive, dest_path / 'extracted')
         base_path = dest_path / 'extracted'
 
     else:
@@ -758,7 +756,7 @@ def get_workdir(self, *args, **kwargs):
     # a calculation already stashed in the same shard, i.e. its UUID shares the first four characters
     if stash_mode == StashMode.COPY.value:
         other_uuid = uuid[:4] + 'beef-0000-0000-0000-000000000000'
-        other_stash = dest_path_error / other_uuid[:2] / other_uuid[2:4] / other_uuid[4:]
+        other_stash = dest_path_error / other_uuid[:2] / other_uuid[2:4] / other_uuid[4:] / other_uuid
         other_stash.mkdir(parents=True)
         (other_stash / 'aiida.out').write_text('other')
     else:
@@ -787,53 +785,153 @@ def get_workdir(self, *args, **kwargs):
             },
         )
 
+    async def mock_raise_oserror(*args, **kwargs):
+        raise OSError('mocked error')
+
+    async def mock_compress_partial(*args, **kwargs):
+        # Leave a partial archive behind, as an interrupted ``tar`` would
+        pathlib.Path(kwargs['remotedestination']).write_text('partial')
+        raise OSError('mocked error')
+
     with LocalTransport() as transport:
         if stash_mode == StashMode.COPY.value:
+            monkeypatch.setattr(transport, 'copy_async', mock_raise_oserror)
+            match = 'Failed to copy'
+        else:
+            monkeypatch.setattr(transport, 'compress_async', mock_compress_partial)
+            match = 'Failed to stash'
 
-            async def mock_copy_async(*args, **kwargs):
-                raise OSError('copy mocked error')
+        with pytest.raises(StashingError, match=match):
+            await execmanager.stash_calculation(node, transport)
 
-            monkeypatch.setattr(transport, 'copy_async', mock_copy_async)
+    # A failed stash must not remove anything that was already on disk, in any mode
+    removed = existing_paths - set(tmp_path.rglob('*'))
+    assert not removed, f'failed stash removed pre-existing paths: {sorted(str(path) for path in removed)}'
 
-            # StashingError should be raised for copy failures
-            with pytest.raises(StashingError, match='Failed to copy'):
-                await execmanager.stash_calculation(node, transport)
+    # Both modes also clean up their own half-written target
+    if stash_mode == StashMode.COPY.value:
+        assert not (dest_path_error / uuid[:2] / uuid[2:4] / uuid[4:] / uuid).exists()
+    else:
+        assert not (dest_path_error / uuid[:2] / uuid[2:4] / uuid[4:] / f'{uuid}.{stash_mode}').exists()
+
+
+@pytest.mark.parametrize('stash_mode', [StashMode.COPY.value, StashMode.COMPRESS_TARGZ.value])
+@pytest.mark.asyncio
+async def test_stashing_skips_missing(generate_calcjob_node, stash_mode, tmp_path, monkeypatch):
+    """With ``fail_on_missing=False`` missing sources are skipped and the stash node records only what was stashed."""
+    node = generate_calcjob_node()
+    workdir = tmp_path / 'workdir'
+    workdir.mkdir()
+    (workdir / 'present.out').write_text('present')
+    node.set_remote_workdir(str(workdir))
+    target_base = tmp_path / 'stash'
+    node.set_option(
+        'stash',
+        {
+            'source_list': ['present.out', 'missing.out', 'nomatch*'],
+            'target_base': str(target_base),
+            'stash_mode': stash_mode,
+            'fail_on_missing': False,
+        },
+    )
+
+    class MockAuthInfo:
+        def get_workdir(self, *args, **kwargs):
+            return str(workdir)
+
+    monkeypatch.setattr(node, 'get_authinfo', MockAuthInfo)
+    node.store()
+
+    with LocalTransport() as transport:
+        await execmanager.stash_calculation(node, transport)
+        remote_stash = node.base.links.get_outgoing(link_label_filter='remote_stash').one().node
+        if stash_mode == StashMode.COPY.value:
+            stashed = pathlib.Path(remote_stash.target_basepath)
         else:
+            stashed = tmp_path / 'extracted'
+            transport.extract(remote_stash.target_basepath, stashed)
 
-            async def mock_compress_async(*args, **kwargs):
-                raise OSError('compress mocked error')
+    assert [path.name for path in stashed.iterdir()] == ['present.out']
+    assert list(remote_stash.source_list) == ['present.out']
 
-            monkeypatch.setattr(transport, 'compress_async', mock_compress_async)
 
-            with caplog.at_level(logging.WARNING):
-                await execmanager.stash_calculation(node, transport)
-                assert any('Failed to stash' in message for message in caplog.messages)
+@pytest.mark.parametrize('stash_mode', [StashMode.COPY.value, StashMode.COMPRESS_TARGZ.value])
+@pytest.mark.asyncio
+async def test_stashing_fail_on_missing_rejects_glob(generate_calcjob_node, stash_mode, tmp_path, monkeypatch):
+    """With ``fail_on_missing=True`` glob patterns are rejected, since a non-matching one cannot be told apart."""
+    node = generate_calcjob_node(workdir=tmp_path)
+    node.set_option(
+        'stash',
+        {
+            'source_list': ['*.out'],
+            'target_base': str(tmp_path / 'stash'),
+            'stash_mode': stash_mode,
+            'fail_on_missing': True,
+        },
+    )
 
-    # A failed stash must not remove anything that was already on disk, in any mode
-    removed = existing_paths - set(tmp_path.rglob('*'))
-    assert not removed, f'failed stash removed pre-existing paths: {sorted(str(path) for path in removed)}'
+    class MockAuthInfo:
+        def get_workdir(self, *args, **kwargs):
+            return str(tmp_path)
 
-    # COPY also cleans up its own half-written target; the compress modes will do that in #7564
-    if stash_mode == StashMode.COPY.value:
-        assert not (dest_path_error / uuid[:2] / uuid[2:4] / uuid[4:]).exists()
+    monkeypatch.setattr(node, 'get_authinfo', MockAuthInfo)
 
-    ## 3) test that an existing stash target is never overwritten (see #7564)
-    if stash_mode != StashMode.COPY.value:
-        existing_archive = pathlib.Path(str(dest_path / uuid) + '.' + stash_mode)
-        existing_archive.write_text('tampered')
+    with LocalTransport() as transport, pytest.raises(StashingError, match='glob patterns'):
+        await execmanager.stash_calculation(node, transport)
 
+
+@pytest.mark.parametrize('stash_mode', [StashMode.COPY.value, StashMode.COMPRESS_TARGZ.value])
+@pytest.mark.asyncio
+async def test_stashing_same_source_twice(generate_calcjob_node, aiida_localhost, stash_mode, tmp_path, monkeypatch):
+    """Stash jobs of the same source node write distinct targets, each replacing only its own leftover."""
+    source_dir = tmp_path / 'source'
+    source_dir.mkdir()
+    (source_dir / 'aiida.out').write_text('out')
+    (source_dir / 'aiida.in').write_text('in')
+    remote = RemoteData(remote_path=str(source_dir), computer=aiida_localhost).store()
+
+    target_base = tmp_path / 'stash'
+    target_base.mkdir()
+
+    class MockAuthInfo:
+        def get_workdir(self, *args, **kwargs):
+            return str(source_dir)
+
+    targets = []
+    for source_list in (['aiida.out'], ['aiida.in']):
+        node = generate_calcjob_node(entry_point='aiida.calculations:core.stash')
         node.set_option(
             'stash',
             {
-                'source_list': ['*'],
-                'target_base': str(dest_path),
+                'source_list': source_list,
+                'target_base': str(target_base),
                 'stash_mode': stash_mode,
-                'dereference': True,
             },
         )
+        node.base.links.add_incoming(remote, link_type=LinkType.INPUT_CALC, link_label='source_node')
+        monkeypatch.setattr(node, 'get_authinfo', MockAuthInfo)
+        node.store()
+
+        # leftover of a previous attempt of this very job
+        if stash_mode == StashMode.COPY.value:
+            target = target_base / remote.uuid[:2] / remote.uuid[2:4] / remote.uuid[4:] / node.uuid
+            target.mkdir(parents=True)
+            (target / source_list[0]).write_text('stale')
+        else:
+            target = target_base / remote.uuid[:2] / remote.uuid[2:4] / remote.uuid[4:] / f'{node.uuid}.{stash_mode}'
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_text('stale')
 
-        with LocalTransport() as transport, caplog.at_level(logging.WARNING):
+        with LocalTransport() as transport:
             await execmanager.stash_calculation(node, transport)
+        targets.append(target)
 
-        assert existing_archive.read_text() == 'tampered'
-        assert any('already exists' in message for message in caplog.messages)
+    for target, filename, content in zip(targets, ('aiida.out', 'aiida.in'), ('out', 'in')):
+        if stash_mode == StashMode.COPY.value:
+            extracted = target
+        else:
+            extracted = tmp_path / f'extracted-{filename}'
+            with LocalTransport() as transport:
+                transport.extract(target, extracted)
+        assert [path.name for path in extracted.iterdir()] == [filename]
+        assert (extracted / filename).read_text() == content