`StashCalculation`: a new `CalcJob` plugin by khsrali · Pull Request #6772 · aiidateam/aiida-core

khsrali · 2025-02-26T15:36:05Z

Historically, stashing was only possible, if it was instructed before running a generic calcjob. The instruction had to be "attached" to the original calcjob, like this for example:

inputs = {
    'MyInputs': <MyInputs>,
    'metadata': {
        'computer': Computer.collection.get(label="localhost"),
        'options': {
            'resources': {'num_machines': 1}, 
            'stash': {
                'stash_mode': StashMode.COPY.value,
                'target_base': '/scratch/',
                 'source_list': ['heavy_data.xyz'],
            },
        },
    },
}
run(MyCalculation, **inputs)

However, if a user would realize they need to stash something only after running th calcjob, this would not be possible.

This PR, defines a new calcjob, that is able to perform a stashing operation even after a calculation is finished.
The usage is very similar, and for consistency and user-friendliness, we keep the instruction as part of the metadata. The only main input is obviously a source node, for example:

StashCalculation_ = CalculationFactory('core.stash')


MyCalculation = orm.load_node(pk=<PK>)
inputs = {
    'metadata': {
        'computer': Computer.collection.get(label="localhost"),
        'options': {
            'resources': {'num_machines': 1}, 
            'stash': {
                'stash_mode': StashMode.COPY.value,
                'target_base': '/scratch/',
                 'source_list': ['heavy_data.xyz'],
            },
        },
    },
    'source_node': MyCalculation,
}

result = run(StashCalculation_, **inputs)

P.S. This PR is part of the upcoming changes, reference in #6764

codecov · 2025-02-26T15:54:58Z

Codecov Report

❌ Patch coverage is 92.10526% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.31%. Comparing base (eb34b06) to head (6c44e6b).
⚠️ Report is 60 commits behind head on main.

Files with missing lines	Patch %	Lines
src/aiida/engine/daemon/execmanager.py	72.73%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6772      +/-   ##
==========================================
+ Coverage   78.29%   78.31%   +0.02%     
==========================================
  Files         566      567       +1     
  Lines       42766    42796      +30     
==========================================
+ Hits        33481    33510      +29     
- Misses       9285     9286       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

agoscinski

First quick review.

agoscinski · 2025-03-04T12:16:47Z

+        spec.inputs.pop('code', None)
+
+        # Ideally one could use the same computer as the one of the `source_node`.
+        # However, if another computer has access to the directory, we don't want to restrict.`


I am not sure how would I pass another computer here? Maybe you could also explain more scenario more detailed.

Suggested change

# However, if another computer has access to the directory, we don't want to restrict.`

# However if you cannot access the stash storage from the same computer anymore but you have access to it from another computer, you can can specify the computer with this commented block below

Thanks I updated the comment, and moved it to docstring.
Passing computer appears to be a very routing procedure it's done via metadata.computer.
It should be very clear, now, even for the users.

agoscinski · 2025-03-04T12:49:21Z


+    if calculation.process_type == 'aiida.calculations:core.stash':
+        remote_node = load_node(calculation.inputs.source_node.pk)
+        uuid = calculation.inputs.source_node.uuid


might increase readability

Suggested change

uuid = calculation.inputs.source_node.uuid

uuid = remote_node.uuid

khsrali · 2025-03-05T11:47:06Z

First quick review.

applied!

agoscinski · 2025-04-07T10:03:37Z

+
+        spec.input(
+            'source_node',
+            valid_type=(orm.RemoteData, orm.SinglefileData),


FolderData?

🤔 yeah, you're probably right

did you change it?

Alright, I added orm.FolderData to the list.

did you? I still don't see it here

agoscinski · 2025-04-07T10:36:04Z


    if stash_mode == StashMode.COPY.value:
-        target_basepath = Path(stash_options['target_base']) / uuid[:2] / uuid[2:4] / uuid[4:]
+        target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:]


Can you make the construction of thpa path uuid[:2] / uuid[2:4] / uuid[4:] part of the CalcJob or CalcInfo (not sure which makes more sense but I see that also the the uuid is used from the CalcInfo? This should make this path less magical and better understandable. Then you also need to use that function in engine.daemon.execmanager:upload_calculation (AFAIK this is where the calculation folder is created)

Yeah, but that would only work if stashing is done as a separate calcjob.
If stashing is performed as part as a generic calcjob, e.g. qe , then this is not possible, And making the folder should be done in execmanager.

Then do a util function that is used in both cases, that takes a uuid as input and in the docstring you can explain how it should be used. The current state is highly suboptimal.

I understand, but this has been like this even before this PR.
And it's used only here, and in execmanager.py::upload_calculation --but not in the exact same way--

So I'd leave that for another PR, if we have to functionalize that.

I understand but you are continuing in this PR the bad design. It should be simply fixable by creating an util function. Please explain then why this change needs a significant higher amount of effort beyond just creating an util function that would justifies a fix in a separate PR.

So I made a function that is being used only once?

Ok, I'm convinced. It's actually two places. :) also here:

workdir = Path(remote_working_directory).joinpath(calc_info.uuid[:2], calc_info.uuid[2:4])

Line 135 execmanager

is it ok, if I do this in a separate PR? I'll open a sub-issue for this.

yes issue is also okay or directly PR might be easier since it should be easy fix (if I am not mistaken)

Just had the look, the logic in upload_calculation is a bit not straight forward. I'd leave this for another PR.

agoscinski · 2025-04-07T10:37:09Z


            elif self._command == STASH_COMMAND:
-                if node.get_option('stash') is not None:
+                if (node.get_option('stash') is not None) or (type(self.process).__name__ == 'StashCalculation'):


why not isinstance(self.process, StashCalculation)?

don't remember why that didn't work,
I have to try it again 🤔

still open, It could be that you want to explicitly check that it is a StashCalculation and not a child class but that would indicate to me also something being wrong

Actually, now I realize the second condition is not needed at all.
The first condition would also serve for stashing via a calcjob. The instruction is similar:

inputs = { 'metadata': { 'computer': Computer.collection.get(label="localhost"), 'options': { 'resources': {'num_machines': 1}, 'stash': { 'stash_mode': StashMode.COPY.value, 'target_base': '/scratch/my_stashing/', 'source_list': ['aiida.in', '_aiidasubmit.sh'], }, }, }, 'source_node': node_1, }

Therefore node.get_option('stash') is not None will be True also in this case.

agoscinski · 2025-04-11T10:23:40Z


            elif self._command == STASH_COMMAND:
-                if node.get_option('stash') is not None:
+                if (node.get_option('stash') is not None) or (type(self.process).__name__ == 'StashCalculation'):


still open, It could be that you want to explicitly check that it is a StashCalculation and not a child class but that would indicate to me also something being wrong

agoscinski · 2025-04-11T10:53:32Z


    if stash_mode == StashMode.COPY.value:
-        target_basepath = Path(stash_options['target_base']) / uuid[:2] / uuid[2:4] / uuid[4:]
+        target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:]


Then do a util function that is used in both cases, that takes a uuid as input and in the docstring you can explain how it should be used. The current state is highly suboptimal.

agoscinski · 2025-04-11T10:54:08Z

+
+        spec.input(
+            'source_node',
+            valid_type=(orm.RemoteData, orm.SinglefileData),


did you change it?

khsrali · 2025-04-14T15:11:30Z

Hi @agoscinski

Thanks a lot for the nice review.
Especially, the ASCII art 😉 I applied your review.

agoscinski · 2025-04-15T13:34:24Z

+
+        spec.input(
+            'source_node',
+            valid_type=(orm.RemoteData, orm.SinglefileData),


did you? I still don't see it here

agoscinski · 2025-04-15T13:42:14Z


    if stash_mode == StashMode.COPY.value:
-        target_basepath = Path(stash_options['target_base']) / uuid[:2] / uuid[2:4] / uuid[4:]
+        target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:]


I understand but you are continuing in this PR the bad design. It should be simply fixable by creating an util function. Please explain then why this change needs a significant higher amount of effort beyond just creating an util function that would justifies a fix in a separate PR.

agoscinski · 2025-04-15T13:45:14Z

+
+class StashCalculation(CalcJob):
+    """
+    Utility to stash files from a remote folder.


not only remote since you also support SinglefileData (and hopefully FolderData)

agoscinski · 2025-04-15T13:56:39Z

+                    'target_base': str(target_base),
+                    'source_list': ['*'],
+                },
+            },


what happens when I forget to add stash options for the StashCalculation? Would like to know if the error message is clear to the user in this case.

very good point!!!
I totally, missed that. Now made those fields compulsory via:

spec.inputs['metadata']['options']['stash'].required = True spec.inputs['metadata']['options']['stash']['stash_mode'].required = True spec.inputs['metadata']['options']['stash']['target_base'].required = True spec.inputs['metadata']['options']['stash']['source_list'].required = True

agoscinski

This point is important and needs short discussion #6772 (comment) (maybe in person is easier) otherwise is almost there

agoscinski · 2025-04-17T13:14:12Z


    if stash_mode == StashMode.COPY.value:
-        target_basepath = Path(stash_options['target_base']) / uuid[:2] / uuid[2:4] / uuid[4:]
+        target_basepath = target_base / uuid[:2] / uuid[2:4] / uuid[4:]


yes issue is also okay or directly PR might be easier since it should be easy fix (if I am not mistaken)

khsrali · 2025-04-24T10:31:49Z

yes issue is also okay or directly PR might be easier since it should be easy fix (if I am not mistaken)

Just had the look, the logic in upload_calculation is a bit not straight forward. I'd leave this for another PR.

Opened an issue:
#6836

khsrali · 2025-04-25T07:53:08Z

@agoscinski has opened a PR #6837 that would allow a CalcJobNode be an input for another CalcJobNode.

It's better and more intuitive if StashCalculation would accept directly a CalcJobNode as an input.
Therefore, this PR has to wait until that change is merged.

Update:
That change is no longer justified, check this #6837 (comment)

khsrali · 2025-04-29T14:58:41Z

@agoscinski please see #6837 (comment)
Given that, I think this PR is ready to go.

agoscinski

In PR #6837 we explored the possibility to input a CalcJobNode as input to a CalcJob, so we could use its path for stashing. In the discussion there it was concluded that a CalcJob already returns a RemoteData node that contains the calcjob working directory. This RemoteData node can be then passed to the StashCalculation implemented in this PR and allows a same usage of the stash options as one would do it when stashing in the first place when running the CalcJob. Therefore the usage of RemoteData as input is also a user friendly. This needs however to be documented in the readthedocs in a future PR as this is not clear for a user (and developers^^)

khsrali requested review from agoscinski and unkcpz as code owners February 26, 2025 15:36

khsrali mentioned this pull request Feb 26, 2025

Upcoming changes on stashing #6764

Closed

7 tasks

khsrali mentioned this pull request Feb 28, 2025

RemoteStashCustomData: stashing via custom script, made possible #6777

Merged

agoscinski reviewed Mar 4, 2025

View reviewed changes

khsrali closed this Mar 5, 2025

khsrali force-pushed the stashing/calcjob branch from 239ff8e to c535928 Compare March 5, 2025 11:40

khsrali reopened this Mar 5, 2025

rebased

4c1504b

khsrali force-pushed the stashing/calcjob branch from 84f49a7 to 4c1504b Compare March 7, 2025 14:03

khsrali requested a review from agoscinski March 7, 2025 14:21

khsrali changed the title ~~StashCalculation: new CalcJob plugin~~ StashCalculation: a new CalcJob plugin Mar 7, 2025

Merge branch 'main' into stashing/calcjob

1d1a82e

agoscinski requested changes Apr 7, 2025

View reviewed changes

agoscinski requested changes Apr 11, 2025

View reviewed changes

revie applied

3bdfa64

khsrali requested a review from agoscinski April 14, 2025 15:11

agoscinski requested changes Apr 15, 2025

View reviewed changes

khsrali commented Apr 16, 2025

View reviewed changes

Comment thread src/aiida/calculations/stash.py Outdated

r

9d021a6

khsrali requested a review from agoscinski April 16, 2025 08:10

agoscinski requested changes Apr 17, 2025

View reviewed changes

Merge branch 'main' into stashing/calcjob

ed351e9

khsrali mentioned this pull request Apr 24, 2025

Make a util function for handling uuid directories for more readability #6836

Open

khsrali added the pr/blocked PR is blocked by another PR that should be merged first label Apr 25, 2025

khsrali mentioned this pull request Apr 29, 2025

Allow CalcJob as an input node #6837

Closed

Merge branch 'main' into stashing/calcjob

a3c51ca

khsrali removed the pr/blocked PR is blocked by another PR that should be merged first label Apr 29, 2025

Update src/aiida/calculations/stash.py

f0a0536

khsrali requested a review from agoscinski April 29, 2025 14:58

Merge branch 'main' into stashing/calcjob

6c44e6b

agoscinski approved these changes Apr 30, 2025

View reviewed changes

khsrali merged commit bc25323 into aiidateam:main Apr 30, 2025
37 of 38 checks passed

agoscinski mentioned this pull request Apr 30, 2025

Update/add models for RemoteData and RemoteStashCompressedData #6844

Merged

khsrali deleted the stashing/calcjob branch May 2, 2025 13:27

	# However, if another computer has access to the directory, we don't want to restrict.`
	# However if you cannot access the stash storage from the same computer anymore but you have access to it from another computer, you can can specify the computer with this commented block below

	uuid = calculation.inputs.source_node.uuid
	uuid = remote_node.uuid

Conversation

khsrali commented Feb 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Feb 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

agoscinski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

khsrali commented Mar 5, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

khsrali commented Apr 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

khsrali commented Feb 26, 2025 •

edited

Loading

codecov Bot commented Feb 26, 2025 •

edited

Loading

agoscinski left a comment •

edited

Loading

khsrali commented Apr 24, 2025 •

edited

Loading

khsrali commented Apr 25, 2025 •

edited

Loading

agoscinski left a comment •

edited

Loading