From 0a200c32c8efa03d5d81e31aa689713d32c6e532 Mon Sep 17 00:00:00 2001 From: Ali Date: Fri, 14 Feb 2025 19:29:56 +0100 Subject: [PATCH 1/3] verdi node delete --clean workdir --- src/aiida/cmdline/commands/cmd_node.py | 66 +++++- tests/cmdline/commands/test_node.py | 269 ++++++++++++++++++++++--- 2 files changed, 306 insertions(+), 29 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_node.py b/src/aiida/cmdline/commands/cmd_node.py index 79efcebcef..d988eb5a6b 100644 --- a/src/aiida/cmdline/commands/cmd_node.py +++ b/src/aiida/cmdline/commands/cmd_node.py @@ -328,9 +328,14 @@ def extras(nodes, keys, fmt, identifier, raw): @click.argument('identifier', nargs=-1, metavar='NODES') @options.DRY_RUN() @options.FORCE() +@click.option( + '--clean-workdir', + is_flag=True, + help='Also clean the remote work directory. (only applies to workflows and CalcJobNodes)', +) @options.graph_traversal_rules(GraphTraversalRules.DELETE.value) @with_dbenv() -def node_delete(identifier, dry_run, force, **traversal_rules): +def node_delete(identifier, dry_run, force, clean_workdir, **traversal_rules): """Delete nodes from the provenance graph. This will not only delete the nodes explicitly provided via the command line, but will also include @@ -349,6 +354,65 @@ def node_delete(identifier, dry_run, force, **traversal_rules): except ValueError: pks.append(NodeEntityLoader.load_entity(obj).pk) + if clean_workdir: + from aiida.manage import get_manager + from aiida.orm import AuthInfo, CalcJobNode, QueryBuilder, User, load_computer + from aiida.orm.utils.remote import get_calcjob_remote_paths + from aiida.tools.graph.graph_traversers import get_nodes_delete + + backend = get_manager().get_backend() + # For here we ignore missing nodes will be raised via func:``delete_nodes`` in the next block + pks_set_to_delete = get_nodes_delete( + pks, get_links=False, missing_callback=lambda missing_pks: None, backend=backend, **traversal_rules + )['nodes'] + + qb = QueryBuilder() + qb.append(CalcJobNode, filters={'id': {'in': pks_set_to_delete}}, project='id') + calcjobs_pks = [result[0] for result in qb.all()] + + if not calcjobs_pks: + echo.echo_report('--clean-workdir ignored. No CalcJobNode associated with the given node, found.') + else: + path_mapping = get_calcjob_remote_paths( + calcjobs_pks, + only_not_cleaned=True, + ) + + if not path_mapping: + echo.echo_report('--clean-workdir ignored. CalcJobNode work directories are already cleaned.') + else: + descendant_pks = [remote_folder.pk for paths in path_mapping.values() for remote_folder in paths] + + if not force and not dry_run: + echo.echo_warning( + f'YOU ARE ABOUT TO CLEAN {len(descendant_pks)} REMOTE DIRECTORIES! ' 'THIS CANNOT BE UNDONE!' + ) + echo.echo_info( + 'Remote directories of nodes with the following pks would be cleaned: ' + + ' '.join(map(str, descendant_pks)) + ) + click.confirm('Shall I continue?', abort=True) + + if dry_run: + echo.echo_report( + 'Remote folders of these node are marked for deletion: ' + ' '.join(map(str, descendant_pks)) + ) + else: + user = User.collection.get_default() + counter = 0 + for computer_uuid, paths in path_mapping.items(): + computer = load_computer(uuid=computer_uuid) + transport = AuthInfo.collection.get( + dbcomputer_id=computer.pk, aiidauser_id=user.pk + ).get_transport() + + with transport: + for remote_folder in paths: + remote_folder._clean(transport=transport) + counter += 1 + + echo.echo_report(f'{counter} remote folders cleaned on {computer.label}') + def _dry_run_callback(pks): if not pks or force: return False diff --git a/tests/cmdline/commands/test_node.py b/tests/cmdline/commands/test_node.py index 0e52708051..010e5447d2 100644 --- a/tests/cmdline/commands/test_node.py +++ b/tests/cmdline/commands/test_node.py @@ -13,7 +13,9 @@ import gzip import io import os +import uuid import warnings +from pathlib import Path import pytest @@ -21,6 +23,9 @@ from aiida.cmdline.commands import cmd_node from aiida.cmdline.utils.echo import ExitCode from aiida.common import timezone +from aiida.common.exceptions import NotExistent +from aiida.common.links import LinkType +from aiida.orm import CalcJobNode, RemoteData, WorkflowNode def get_result_lines(result): @@ -290,7 +295,6 @@ def test_catch_bad_pk(self, run_cli_command): """Test that an invalid root_node pk (non-numeric, negative, or decimal), or non-existent pk will produce an error """ - from aiida.common.exceptions import NotExistent from aiida.orm import load_node # Forbidden pk @@ -557,39 +561,248 @@ def test_rehash_invalid_entry_point(self, run_cli_command): run_cli_command(cmd_node.rehash, options, raises=True) -@pytest.mark.parametrize( - 'options', - ( - ['--verbosity', 'info'], - ['--verbosity', 'info', '--force'], - ['--create-forward'], - ['--call-calc-forward'], - ['--call-work-forward'], - ['--force'], - ), -) -def test_node_delete_basics(run_cli_command, options): - """Testing the correct translation for the `--force` and `--verbosity` options. - This just checks that the calls do not except and that in all cases with the - force flag there is no messages. - """ - from aiida.common.exceptions import NotExistent +class TestVerdiDelete: + """Tests for the ``verdi node delete`` command.""" + + @pytest.fixture(autouse=True) + def init__(self): + self.workflow_nodes = [] + self.calcjob_nodes = [] + self.remote_nodes = [] + self.remote_folders = [] + + @pytest.fixture() + def setup_node_hierarchy(self, aiida_localhost, tmp_path): + """Set up a WorkflowNode with multiple CalcJobNodes and RemoteData nodes. + + :param aiida_localhost: the localhost computer + :param tmp_path: the temporary directory + :param n_workflows: the number of WorkflowNodes to create + :param n_calcjobs: the number of CalcJobNodes per WorkflowNode to create + :param n_remotes: the number of RemoteData nodes per CalcJobNode to create + + the WorkflowNode will have the following structure: + + WorkflowNode + ├── CalcJobNode + │ ├── RemoteData + │ └── RemoteData + └── CalcJobNode + ├── RemoteData + └── RemoteData + + :return: a tuple of the WorkflowNode, the list of CalcJobNodes, and the list of RemoteData nodes.""" + + def _setup(n_workflows=1, n_calcjobs=1, n_remotes=1): + for _ in range(n_workflows): + workflow_node = WorkflowNode() + workflow_node.store() + self.workflow_nodes.append(workflow_node) + + for i in range(n_calcjobs): + calcjob_node = CalcJobNode(computer=aiida_localhost) + calcjob_node.add_incoming(workflow_node, link_type=LinkType.CALL_CALC, link_label='call') + + workdir = tmp_path / f'calcjob_{uuid.uuid4()}' + workdir.mkdir() + Path(workdir / 'fileA.txt').write_text('test stringA') + self.remote_folders.append(workdir) + + calcjob_node.set_remote_workdir(str(workdir)) + calcjob_node.set_option('output_filename', 'fileA.txt') + calcjob_node.store() + self.calcjob_nodes.append(calcjob_node) + + for _ in range(n_remotes): + remote_node = RemoteData(remote_path=str(workdir), computer=aiida_localhost) + remote_node.base.links.add_incoming(calcjob_node, LinkType.CREATE, link_label='remote_folder') + remote_node.store() + self.remote_nodes.append(remote_node) + + # this is equal to (n_workflows + n_workflows * n_calcjobs + n_workflows * n_calcjobs * n_remotes) + self.total_nodes = len(self.workflow_nodes) + len(self.calcjob_nodes) + len(self.remote_nodes) + + return _setup + + def verify_deletion(self, nodes_deleted=True, folders_deleted=True): + """Verify that the nodes and remote folders are deleted or not deleted.""" + + if nodes_deleted: + for workflow_node in self.workflow_nodes: + with pytest.raises(Exception, match='deleted'): + WorkflowNode.objects.get(pk=workflow_node.pk) + + for calcjob_node in self.calcjob_nodes: + with pytest.raises(Exception, match='deleted'): + CalcJobNode.objects.get(pk=calcjob_node.pk) + + for remote_node in self.remote_nodes: + with pytest.raises(Exception, match='deleted'): + RemoteData.objects.get(pk=remote_node.pk) + else: + for workflow_node in self.workflow_nodes: + WorkflowNode.objects.get(pk=workflow_node.pk) + for calcjob_node in self.calcjob_nodes: + CalcJobNode.objects.get(pk=calcjob_node.pk) + for remote_node in self.remote_nodes: + RemoteData.objects.get(pk=remote_node.pk) + + for remote_folder in self.remote_folders: + if folders_deleted: + assert not remote_folder.exists() + else: + assert remote_folder.exists() + + def test_setup_node_hierarchy(self, setup_node_hierarchy): + """Test the `setup_node_hierarchy` and `verify_deletion` fixtures.""" + # Guard the guardians + setup_node_hierarchy() + self.verify_deletion(nodes_deleted=False, folders_deleted=False) + + def test_node_delete_dry_run(self, run_cli_command, setup_node_hierarchy): + """Test the `--dry-run` option. + Nothing should be deleted and the proper message should be printed without prompting y/n. + + Note: To speed up the test, I do all parameters in one tests instead of using `@pytest.mark.parametrize`. + This way I can reuse the setup_node_hierarchy fixture. This works because nothing should be deleted! + """ + setup_node_hierarchy(1, 1, 1) + all_workflow_pks = [str(workflow_node.pk) for workflow_node in self.workflow_nodes] + + # 1) + options = ['--dry-run'] + result = run_cli_command(cmd_node.node_delete, options + all_workflow_pks) + assert f'Report: {self.total_nodes} Node(s) marked for deletion' in str(result.stdout_bytes) + assert 'Report: This was a dry run, exiting without deleting anything' in str(result.stdout_bytes) + self.verify_deletion(False, False) + + # 2) + options = ['--dry-run', '--clean-workdir'] + result = run_cli_command(cmd_node.node_delete, options + all_workflow_pks) + assert ( + 'Report: Remote folders of these node are marked for ' + f"deletion: {' '.join(str(remote_node.pk) for remote_node in self.remote_nodes)}" + in str(result.stdout_bytes) + ) + assert f'Report: {self.total_nodes} Node(s) marked for deletion' in str(result.stdout_bytes) + assert 'Report: This was a dry run, exiting without deleting anything' in str(result.stdout_bytes) + self.verify_deletion(False, False) + + # 3) This is important! Should not delete! + options = ['--dry-run', '--force'] + result = run_cli_command(cmd_node.node_delete, options + all_workflow_pks) + assert f'Report: {self.total_nodes} Node(s) marked for deletion' in str(result.stdout_bytes) + assert 'Report: This was a dry run, exiting without deleting anything' in str(result.stdout_bytes) + self.verify_deletion(False, False) + + # 4) This is important! Should not delete! + options = ['--dry-run', '--force', '--clean-workdir'] + result = run_cli_command(cmd_node.node_delete, options + all_workflow_pks) + assert ( + 'Report: Remote folders of these node are marked for ' + f"deletion: {' '.join(str(remote_node.pk) for remote_node in self.remote_nodes)}" + in str(result.stdout_bytes) + ) + assert f'Report: {self.total_nodes} Node(s) marked for deletion' in str(result.stdout_bytes) + assert 'Report: This was a dry run, exiting without deleting anything' in str(result.stdout_bytes) + self.verify_deletion(False, False) + + @pytest.mark.parametrize( + 'options, user_input, nodes_deleted, folders_deleted', + [ + ([], 'n', False, False), + ([], 'y', True, False), + (['--force'], '', True, False), + (['--force', '--clean-workdir'], '', True, True), + (['--clean-workdir'], 'y\ny', True, True), + (['--clean-workdir'], 'y\nn', False, True), + ], + ) + def test_node_delete_prompt_flow( + self, run_cli_command, setup_node_hierarchy, options, user_input, nodes_deleted, folders_deleted + ): + """Test the prompt flow with various options.""" + setup_node_hierarchy(1, 1, 1) + all_workflow_pks = [str(workflow_node.pk) for workflow_node in self.workflow_nodes] + + result = run_cli_command( + cmd_node.node_delete, + options + all_workflow_pks, + user_input=user_input, + raises=True if 'n' in user_input.split('\n') else False, + ) + self.verify_deletion(nodes_deleted=nodes_deleted, folders_deleted=folders_deleted) + + if options == [] and user_input == 'y': + assert ( + f'YOU ARE ABOUT TO DELETE {self.total_nodes} NODES! THIS CANNOT BE UNDONE!\\nShall I continue? [y/N]' + in str(result.stdout_bytes) + ) + elif options == [] and user_input == 'n': + assert 'Aborted!' in str(result.stderr_bytes) + elif options == ['--force']: + assert 'YOU ARE ABOUT TO DELETE' not in str(result.stdout_bytes) + assert '[y/N]' not in str(result.stdout_bytes) + elif options == ['--force', '--clean-workdir']: + assert 'YOU ARE ABOUT TO DELETE' not in str(result.stdout_bytes) + assert '[y/N]' not in str(result.stdout_bytes) + elif options == ['--clean-workdir'] and user_input == 'y\ny': + assert ( + f'YOU ARE ABOUT TO CLEAN {len(self.remote_nodes)} REMOTE DIRECTORIES! THIS CANNOT BE UNDONE!' + '\\nShall I continue? [y/N]' in str(result.stdout_bytes) + ) + assert ( + f'YOU ARE ABOUT TO DELETE {self.total_nodes} NODES! THIS CANNOT BE UNDONE!' + '\\nShall I continue? [y/N]' in str(result.stdout_bytes) + ) + + elif options == ['--clean-workdir'] and user_input == 'y\nn': + # This is a special case, if the user's imagination came invented a "hacky" use with --clean-workdir + # To only delete the folders, but not the nodes. + assert 'Aborted!' in str(result.stderr_bytes) + # And later if decided to delete the nodes, as well, while the nodes are already deleted, + # no error should be raised, and it should proceed with printing a message only + result = run_cli_command(cmd_node.node_delete, options + all_workflow_pks, user_input='y\ny') + self.verify_deletion(nodes_deleted=True, folders_deleted=True) + assert '--clean-workdir ignored. CalcJobNode work directories are already cleaned.' in str( + result.stdout_bytes + ) + assert ( + f'YOU ARE ABOUT TO DELETE {self.total_nodes} NODES! THIS CANNOT BE UNDONE!\\nShall I continue? [y/N]' + in str(result.stdout_bytes) + ) - node = orm.Data().store() - pk = node.pk + @pytest.mark.parametrize( + 'options', + ( + ['--verbosity', 'info'], + ['--verbosity', 'info', '--force'], + ['--create-forward'], + ['--call-calc-forward'], + ['--call-work-forward'], + ['--force'], + ), + ) + def test_node_delete_basics(self, run_cli_command, options): + # Legacy test, this can somehow get merged with the more extensive tests above. + """This just checks that the calls do not except and that in all cases with the + force flag there is no messages. + """ - run_cli_command(cmd_node.node_delete, options + [str(pk), '--dry-run'], use_subprocess=True) + node = orm.Data().store() + pk = node.pk - # To delete the created node - run_cli_command(cmd_node.node_delete, [str(pk), '--force'], use_subprocess=True) + run_cli_command(cmd_node.node_delete, options + [str(pk), '--dry-run'], use_subprocess=True) - with pytest.raises(NotExistent): - orm.load_node(pk) + # To delete the created node + run_cli_command(cmd_node.node_delete, [str(pk), '--force'], use_subprocess=True) + with pytest.raises(NotExistent): + orm.load_node(pk) -def test_node_delete_missing_pk(run_cli_command): - """Check that no exception is raised when a non-existent pk is given (just warns).""" - run_cli_command(cmd_node.node_delete, ['999']) + def test_node_delete_missing_pk(self, run_cli_command): + """Check that no exception is raised when a non-existent pk is given (just warns).""" + run_cli_command(cmd_node.node_delete, ['999']) @pytest.fixture(scope='class') From 337458057328f4168f5f8d0e305b32cb806c1825 Mon Sep 17 00:00:00 2001 From: Ali Date: Fri, 14 Feb 2025 20:13:30 +0100 Subject: [PATCH 2/3] . --- tests/cmdline/commands/test_node.py | 49 ++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/tests/cmdline/commands/test_node.py b/tests/cmdline/commands/test_node.py index 010e5447d2..8b7eb55d6e 100644 --- a/tests/cmdline/commands/test_node.py +++ b/tests/cmdline/commands/test_node.py @@ -570,6 +570,10 @@ def init__(self): self.calcjob_nodes = [] self.remote_nodes = [] self.remote_folders = [] + # We keep track of the PKs to verify deletion + self.workflow_pks = [] + self.calcjob_pks = [] + self.remote_pks = [] @pytest.fixture() def setup_node_hierarchy(self, aiida_localhost, tmp_path): @@ -598,6 +602,7 @@ def _setup(n_workflows=1, n_calcjobs=1, n_remotes=1): workflow_node = WorkflowNode() workflow_node.store() self.workflow_nodes.append(workflow_node) + self.workflow_pks.append(workflow_node.pk) for i in range(n_calcjobs): calcjob_node = CalcJobNode(computer=aiida_localhost) @@ -612,12 +617,14 @@ def _setup(n_workflows=1, n_calcjobs=1, n_remotes=1): calcjob_node.set_option('output_filename', 'fileA.txt') calcjob_node.store() self.calcjob_nodes.append(calcjob_node) + self.calcjob_pks.append(calcjob_node.pk) for _ in range(n_remotes): remote_node = RemoteData(remote_path=str(workdir), computer=aiida_localhost) remote_node.base.links.add_incoming(calcjob_node, LinkType.CREATE, link_label='remote_folder') remote_node.store() self.remote_nodes.append(remote_node) + self.remote_pks.append(remote_node.pk) # this is equal to (n_workflows + n_workflows * n_calcjobs + n_workflows * n_calcjobs * n_remotes) self.total_nodes = len(self.workflow_nodes) + len(self.calcjob_nodes) + len(self.remote_nodes) @@ -628,24 +635,24 @@ def verify_deletion(self, nodes_deleted=True, folders_deleted=True): """Verify that the nodes and remote folders are deleted or not deleted.""" if nodes_deleted: - for workflow_node in self.workflow_nodes: - with pytest.raises(Exception, match='deleted'): - WorkflowNode.objects.get(pk=workflow_node.pk) + for workflow_pk in self.workflow_pks: + with pytest.raises(NotExistent): + WorkflowNode.objects.get(pk=workflow_pk) - for calcjob_node in self.calcjob_nodes: - with pytest.raises(Exception, match='deleted'): - CalcJobNode.objects.get(pk=calcjob_node.pk) + for calcjob_pk in self.calcjob_pks: + with pytest.raises(NotExistent): + CalcJobNode.objects.get(pk=calcjob_pk) - for remote_node in self.remote_nodes: - with pytest.raises(Exception, match='deleted'): - RemoteData.objects.get(pk=remote_node.pk) + for remote_pk in self.remote_pks: + with pytest.raises(NotExistent): + RemoteData.objects.get(pk=remote_pk) else: - for workflow_node in self.workflow_nodes: - WorkflowNode.objects.get(pk=workflow_node.pk) - for calcjob_node in self.calcjob_nodes: - CalcJobNode.objects.get(pk=calcjob_node.pk) - for remote_node in self.remote_nodes: - RemoteData.objects.get(pk=remote_node.pk) + for workflow_pk in self.workflow_pks: + WorkflowNode.objects.get(pk=workflow_pk) + for calcjob_pk in self.calcjob_pks: + CalcJobNode.objects.get(pk=calcjob_pk) + for remote_pk in self.remote_pks: + RemoteData.objects.get(pk=remote_pk) for remote_folder in self.remote_folders: if folders_deleted: @@ -804,6 +811,18 @@ def test_node_delete_missing_pk(self, run_cli_command): """Check that no exception is raised when a non-existent pk is given (just warns).""" run_cli_command(cmd_node.node_delete, ['999']) + def test_node_delete_no_calcjob_to_cleandir(self, run_cli_command): + """Check that no exception is raised when a node is deleted with no calcjob to clean.""" + node = orm.Data().store() + pk = node.pk + result = run_cli_command(cmd_node.node_delete, ['--clean-workdir', '--force', str(pk)]) + assert '--clean-workdir ignored. No CalcJobNode associated with the given node, found.' in str( + result.stdout_bytes + ) + + with pytest.raises(NotExistent): + orm.load_node(pk) + @pytest.fixture(scope='class') def create_nodes_verdi_node_list(aiida_profile_clean_class): From 85553701d052353e8e7298c2db3d0d76f3ddaeb1 Mon Sep 17 00:00:00 2001 From: Ali Date: Tue, 18 Feb 2025 20:31:45 +0100 Subject: [PATCH 3/3] review applied --- src/aiida/cmdline/commands/cmd_calcjob.py | 17 +--- src/aiida/cmdline/commands/cmd_node.py | 95 +++++++++++------------ src/aiida/orm/utils/remote.py | 32 +++++++- tests/cmdline/commands/test_node.py | 2 +- 4 files changed, 77 insertions(+), 69 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_calcjob.py b/src/aiida/cmdline/commands/cmd_calcjob.py index 376301247e..f49dc9c4d4 100644 --- a/src/aiida/cmdline/commands/cmd_calcjob.py +++ b/src/aiida/cmdline/commands/cmd_calcjob.py @@ -259,8 +259,7 @@ def calcjob_cleanworkdir(calcjobs, past_days, older_than, computers, force, exit If both are specified, a logical AND is done between the two, i.e. the calcjobs that will be cleaned have been modified AFTER [-p option] days from now, but BEFORE [-o option] days from now. """ - from aiida import orm - from aiida.orm.utils.remote import get_calcjob_remote_paths + from aiida.orm.utils.remote import clean_mapping_remote_paths, get_calcjob_remote_paths if calcjobs: if past_days is not None and older_than is not None: @@ -286,19 +285,7 @@ def calcjob_cleanworkdir(calcjobs, past_days, older_than, computers, force, exit warning = f'Are you sure you want to clean the work directory of {path_count} calcjobs?' click.confirm(warning, abort=True) - user = orm.User.collection.get_default() - - for computer_uuid, paths in path_mapping.items(): - counter = 0 - computer = orm.load_computer(uuid=computer_uuid) - transport = orm.AuthInfo.collection.get(dbcomputer_id=computer.pk, aiidauser_id=user.pk).get_transport() - - with transport: - for remote_folder in paths: - remote_folder._clean(transport=transport) - counter += 1 - - echo.echo_success(f'{counter} remote folders cleaned on {computer.label}') + clean_mapping_remote_paths(path_mapping) def get_remote_and_path(calcjob, path=None): diff --git a/src/aiida/cmdline/commands/cmd_node.py b/src/aiida/cmdline/commands/cmd_node.py index d988eb5a6b..21357e054b 100644 --- a/src/aiida/cmdline/commands/cmd_node.py +++ b/src/aiida/cmdline/commands/cmd_node.py @@ -331,7 +331,7 @@ def extras(nodes, keys, fmt, identifier, raw): @click.option( '--clean-workdir', is_flag=True, - help='Also clean the remote work directory. (only applies to workflows and CalcJobNodes)', + help='Also clean the remote work directory, if applicable.', ) @options.graph_traversal_rules(GraphTraversalRules.DELETE.value) @with_dbenv() @@ -354,10 +354,22 @@ def node_delete(identifier, dry_run, force, clean_workdir, **traversal_rules): except ValueError: pks.append(NodeEntityLoader.load_entity(obj).pk) + def _dry_run_callback(pks): + if not pks or force: + return False + echo.echo_warning(f'YOU ARE ABOUT TO DELETE {len(pks)} NODES! THIS CANNOT BE UNDONE!') + echo.echo_info('The nodes with the following pks would be deleted: ' + ' '.join(map(str, pks))) + return not click.confirm('Shall I continue?', abort=True) + + def _perform_delete(): + _, was_deleted = delete_nodes(pks, dry_run=dry_run or _dry_run_callback, **traversal_rules) + if was_deleted: + echo.echo_success('Finished deletion.') + if clean_workdir: from aiida.manage import get_manager - from aiida.orm import AuthInfo, CalcJobNode, QueryBuilder, User, load_computer - from aiida.orm.utils.remote import get_calcjob_remote_paths + from aiida.orm import CalcJobNode, QueryBuilder + from aiida.orm.utils.remote import clean_mapping_remote_paths, get_calcjob_remote_paths from aiida.tools.graph.graph_traversers import get_nodes_delete backend = get_manager().get_backend() @@ -372,58 +384,39 @@ def node_delete(identifier, dry_run, force, clean_workdir, **traversal_rules): if not calcjobs_pks: echo.echo_report('--clean-workdir ignored. No CalcJobNode associated with the given node, found.') - else: - path_mapping = get_calcjob_remote_paths( - calcjobs_pks, - only_not_cleaned=True, - ) + _perform_delete() + return - if not path_mapping: - echo.echo_report('--clean-workdir ignored. CalcJobNode work directories are already cleaned.') - else: - descendant_pks = [remote_folder.pk for paths in path_mapping.values() for remote_folder in paths] - - if not force and not dry_run: - echo.echo_warning( - f'YOU ARE ABOUT TO CLEAN {len(descendant_pks)} REMOTE DIRECTORIES! ' 'THIS CANNOT BE UNDONE!' - ) - echo.echo_info( - 'Remote directories of nodes with the following pks would be cleaned: ' - + ' '.join(map(str, descendant_pks)) - ) - click.confirm('Shall I continue?', abort=True) - - if dry_run: - echo.echo_report( - 'Remote folders of these node are marked for deletion: ' + ' '.join(map(str, descendant_pks)) - ) - else: - user = User.collection.get_default() - counter = 0 - for computer_uuid, paths in path_mapping.items(): - computer = load_computer(uuid=computer_uuid) - transport = AuthInfo.collection.get( - dbcomputer_id=computer.pk, aiidauser_id=user.pk - ).get_transport() - - with transport: - for remote_folder in paths: - remote_folder._clean(transport=transport) - counter += 1 - - echo.echo_report(f'{counter} remote folders cleaned on {computer.label}') + path_mapping = get_calcjob_remote_paths( + calcjobs_pks, + only_not_cleaned=True, + ) - def _dry_run_callback(pks): - if not pks or force: - return False - echo.echo_warning(f'YOU ARE ABOUT TO DELETE {len(pks)} NODES! THIS CANNOT BE UNDONE!') - echo.echo_info('The nodes with the following pks would be deleted: ' + ' '.join(map(str, pks))) - return not click.confirm('Shall I continue?', abort=True) + if not path_mapping: + echo.echo_report('--clean-workdir ignored. CalcJobNode work directories are already cleaned.') + _perform_delete() + return - _, was_deleted = delete_nodes(pks, dry_run=dry_run or _dry_run_callback, **traversal_rules) + descendant_pks = [remote_folder.pk for paths in path_mapping.values() for remote_folder in paths] + + if not force and not dry_run: + echo.echo_warning( + f'YOU ARE ABOUT TO CLEAN {len(descendant_pks)} REMOTE DIRECTORIES! ' 'THIS CANNOT BE UNDONE!' + ) + echo.echo_info( + 'Remote directories of nodes with the following pks would be cleaned: ' + + ' '.join(map(str, descendant_pks)) + ) + click.confirm('Shall I continue?', abort=True) + + if dry_run: + echo.echo_report( + 'Remote folders of these node are marked for deletion: ' + ' '.join(map(str, descendant_pks)) + ) + else: + clean_mapping_remote_paths(path_mapping) - if was_deleted: - echo.echo_success('Finished deletion.') + _perform_delete() @verdi_node.command('rehash') diff --git a/src/aiida/orm/utils/remote.py b/src/aiida/orm/utils/remote.py index f55cedc35a..1e7727d68f 100644 --- a/src/aiida/orm/utils/remote.py +++ b/src/aiida/orm/utils/remote.py @@ -13,12 +13,13 @@ import os import typing as t +from aiida import orm +from aiida.cmdline.utils import echo from aiida.orm.nodes.data.remote.base import RemoteData if t.TYPE_CHECKING: from collections.abc import Sequence - from aiida import orm from aiida.orm.implementation import StorageBackend from aiida.transports import Transport @@ -45,6 +46,34 @@ def clean_remote(transport: Transport, path: str) -> None: pass +def clean_mapping_remote_paths(path_mapping, silent=False): + """Clean the remote folders for a given mapping of computer UUIDs to a list of remote folders. + + :param path_mapping: a dictionary where the keys are the computer UUIDs and the values are lists of remote folders + It's designed to accept the output of `get_calcjob_remote_paths` + :param transport: the transport to use to clean the remote folders + :param silent: if True, the `echo` output will be suppressed + """ + + user = orm.User.collection.get_default() + + if not user: + raise ValueError('No default user found') + + for computer_uuid, paths in path_mapping.items(): + counter = 0 + computer = orm.load_computer(uuid=computer_uuid) + transport = orm.AuthInfo.collection.get(dbcomputer_id=computer.pk, aiidauser_id=user.pk).get_transport() + + with transport: + for remote_folder in paths: + remote_folder._clean(transport=transport) + counter += 1 + + if not silent: + echo.echo_success(f'{counter} remote folders cleaned on {computer.label}') + + def get_calcjob_remote_paths( pks: list[int] | None = None, past_days: int | None = None, @@ -70,7 +99,6 @@ def get_calcjob_remote_paths( """ from datetime import timedelta - from aiida import orm from aiida.common import timezone from aiida.orm import CalcJobNode diff --git a/tests/cmdline/commands/test_node.py b/tests/cmdline/commands/test_node.py index 8b7eb55d6e..1b88e9cc07 100644 --- a/tests/cmdline/commands/test_node.py +++ b/tests/cmdline/commands/test_node.py @@ -764,7 +764,7 @@ def test_node_delete_prompt_flow( ) elif options == ['--clean-workdir'] and user_input == 'y\nn': - # This is a special case, if the user's imagination came invented a "hacky" use with --clean-workdir + # This is a special case, the user's imagination may invent a "hacky" solution with --clean-workdir # To only delete the folders, but not the nodes. assert 'Aborted!' in str(result.stderr_bytes) # And later if decided to delete the nodes, as well, while the nodes are already deleted,