Skip to content

Commit

Permalink
First version of renaming group paths and relabeling the Logger on gr…
Browse files Browse the repository at this point in the history
…oup relabel.
  • Loading branch information
GeigerJ2 committed Feb 17, 2025
1 parent 59f6fb2 commit b795eda
Show file tree
Hide file tree
Showing 5 changed files with 192 additions and 45 deletions.
7 changes: 7 additions & 0 deletions src/aiida/cmdline/commands/cmd_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -319,12 +319,19 @@ def _dry_run_callback(pks):
@with_dbenv()
def group_relabel(group, label):
"""Change the label of a group."""
# TODO: Add a message here that if one has the profile mirrored, they should also run the command `verdi profile
# TODO: mirror relabel-group ` to update the mirrored profile.
try:
group.label = label
except UniquenessError as exception:
echo.echo_critical(str(exception))
else:
echo.echo_success(f"Label changed to '{label}'")
msg = (
'Note that if you are mirroring your profile data to disk, to reflect the relabeling of the group, '
'run the command: `verdi profile mirror --update-groups.'
)
echo.echo_report(msg)


@verdi_group.command('description')
Expand Down
47 changes: 38 additions & 9 deletions src/aiida/cmdline/commands/cmd_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,14 @@ def profile_delete(force, delete_data, profiles):
show_default=True,
help='If a top-level process calls sub-processes, create a designated directory only for the top-level process.',
)
# TODO: Implement this...
# TODO: Possibly
@click.option(
'--update-groups/--no-update-groups',
default=False,
show_default=True,
help='Update directories if nodes have been added to other groups, or organized differently in terms of groups.',
)
@options.INCLUDE_INPUTS()
@options.INCLUDE_OUTPUTS()
@options.INCLUDE_ATTRIBUTES()
Expand All @@ -319,6 +327,7 @@ def profile_mirror(
organize_by_groups,
symlink_duplicates,
delete_missing,
update_groups,
only_top_level_calcs,
only_top_level_workflows,
include_inputs,
Expand All @@ -343,6 +352,11 @@ def profile_mirror(
# ? Does it even make sense to provide both options, as they are mutually exclusive?
incremental = not overwrite

if not organize_by_groups and update_groups:
# Add check outside in cmd_profile?
msg = '`--update-groups` selected, even though `--organize-by-groups` is set to False.'
echo.echo_critical(msg)

if path is None:
path = Path.cwd() / f'{profile.name}-mirror'

Expand Down Expand Up @@ -421,6 +435,9 @@ def profile_mirror(
num_processes_to_dump = len(profile_dumper.processes_to_dump)
num_processes_to_delete = len(profile_dumper.processes_to_delete)

# num_groups_to_dump = len(profile_dumper.groups_to_dump)
num_groups_to_delete = len(profile_dumper.groups_to_delete)

if dry_run:
dry_run_message = (
f'Dry run for mirroring of profile `{profile.name}`. '
Expand All @@ -438,23 +455,35 @@ def profile_mirror(
# TODO: Maybe add y/n confirmation here?
echo.echo_report(msg)

if num_processes_to_dump == 0:
msg = 'No processes to dump.'
echo.echo_success(msg)
else:
profile_dumper.dump_processes()
msg = f'Dumped {num_processes_to_dump} new nodes.'
echo.echo_success(msg)
if dump_processes:
if num_processes_to_dump == 0:
msg = 'No processes to dump.'
echo.echo_success(msg)
else:
profile_dumper.dump_processes()
msg = f'Dumped {num_processes_to_dump} new nodes.'
echo.echo_success(msg)

if delete_missing:
# breakpoint()
if num_processes_to_delete == 0:
echo.echo_success('No processes to delete.')
else:
profile_dumper.delete_processes()

echo.echo_success(f'Deleted {num_processes_to_delete} node directories.')

if num_groups_to_delete == 0:
echo.echo_success('No groups to delete.')
else:
profile_dumper.delete_groups()
echo.echo_success(f'Deleted {num_groups_to_delete} group directories.')

if update_groups:
relabeled_paths = profile_dumper.update_groups()

msg = 'Mirrored directories and '
echo.echo_success(msg)
print(relabeled_paths)

# Append the current dump time to dumping safeguard file
with safeguard_file_path.open('a') as fhandle:
fhandle.write(f'Last profile mirror time: {last_dump_time.isoformat()}\n')
Expand Down
45 changes: 35 additions & 10 deletions src/aiida/tools/dumping/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,9 @@

from aiida.common.exceptions import NotExistent

# TODO: Possibly mirror hierarchy of mirrored directory inside json file
# TODO: Currently, json file has only top-level "groups", "workflows", and "calculations"


@dataclass
class DumpLog:
Expand Down Expand Up @@ -94,10 +97,12 @@ class DumpLogStoreCollection:

calculations: DumpLogStore
workflows: DumpLogStore
groups: DumpLogStore
# data: DumpLogStore


class DumpLogger:
"""Main logger class using dataclasses for better structure."""
"""Main dumping logger singleton."""

DUMP_LOG_FILE: str = '.dump_log.json'
_instance: 'DumpLogger | None' = None # Class-level singleton instance
Expand All @@ -116,7 +121,8 @@ def __init__(
dump_parent_path: Path | None = None,
calculations: DumpLogStore | None = None,
workflows: DumpLogStore | None = None,
# counter: int = 0,
groups: DumpLogStore | None = None,
# data: DumpLogStore | None = None,
) -> None:
# Ensure __init__ is only called once
if hasattr(self, '_initialized') and self._initialized:
Expand All @@ -125,7 +131,8 @@ def __init__(
self.dump_parent_path = dump_parent_path or Path.cwd()
self.calculations = calculations or DumpLogStore()
self.workflows = workflows or DumpLogStore()
# self.counter = counter
self.groups = groups or DumpLogStore()
# self.dat = data or DumpLogStore()

# Mark the object as initialized
self._initialized = True
Expand All @@ -144,20 +151,26 @@ def del_entry(self, store: DumpLogStore, uuid: str) -> bool:
@property
def log(self) -> DumpLogStoreCollection:
"""Retrieve the current state of the log as a dataclass."""
return DumpLogStoreCollection(calculations=self.calculations, workflows=self.workflows)
return DumpLogStoreCollection(calculations=self.calculations, workflows=self.workflows, groups=self.groups)

def save_log(self) -> None:
"""Save the log to a JSON file."""

def serialize_logs(container: DumpLogStore) -> dict:
serialized = {}
for uuid, entry in container.entries.items():
serialized[uuid] = {'path': str(entry.path), 'time': entry.time.isoformat()}
serialized[uuid] = {
'path': str(entry.path),
'time': entry.time.isoformat(),
'links': [str(link) for link in entry.links],
}
return serialized

log_dict = {
'calculations': serialize_logs(self.calculations),
'workflows': serialize_logs(self.workflows),
'groups': serialize_logs(self.groups),
# 'data': serialize_logs(self.data),
}

with self.log_file_path.open('w', encoding='utf-8') as f:
Expand Down Expand Up @@ -185,12 +198,20 @@ def deserialize_logs(category_data: dict) -> DumpLogStore:
container = DumpLogStore()
for uuid, entry in category_data.items():
container.add_entry(
uuid, DumpLog(path=Path(entry['path']), time=datetime.fromisoformat(entry['time']))
uuid,
DumpLog(
path=Path(entry['path']),
time=datetime.fromisoformat(entry['time']),
links=[Path(p) for p in entry['links']],
),
)

return container

instance.calculations = deserialize_logs(data['calculations'])
instance.workflows = deserialize_logs(data['workflows'])
instance.groups = deserialize_logs(data['groups'])
# instance.data = deserialize_logs(data['data'])

except (json.JSONDecodeError, OSError):
raise
Expand All @@ -206,7 +227,7 @@ def get_store_by_uuid(self, uuid: str) -> DumpLogStore:
if uuid in store.entries:
return store

msg = f"No corresponding `DumpLogStore` found for UUID: `{uuid}`."
msg = f'No corresponding `DumpLogStore` found for UUID: `{uuid}`.'
raise NotExistent(msg)

def get_path_by_uuid(self, uuid: str) -> Path | None:
Expand All @@ -215,13 +236,17 @@ def get_path_by_uuid(self, uuid: str) -> Path | None:

try:
current_store = self.get_store_by_uuid(uuid=uuid)
path = current_store.entries[uuid].path
return path
except NotExistent as exc:
raise NotExistent(exc.args[0]) from exc
try:
path = current_store.entries[uuid].path
return path
except KeyError as exc:
msg = f"UUID: `{uuid}` not contained in store `{current_store}`."
msg = f'UUID: `{uuid}` not contained in store `{current_store}`.'
raise KeyError(msg) from exc
except:
# For debugging
import ipdb

ipdb.set_trace()
raise
Loading

0 comments on commit b795eda

Please sign in to comment.