Add TemplateCodeInfo for code info in job template only

[unkcpz]

This is another piece required for containerized code implementation. In the current design, the code_info of calc_job is read from the code setup and plugin then pass to the job template to create the bash script.
However, job template needs more flexibility to control the different part of script runline where currently all the part (1. exec_name from code uuid, 2. code_info.cmdline_params, 3. mpi parameters from computer setting) are stacked together to the job template's code_info.

In this PR, the class TemplateCodeInfo is created to handle the elements needed to generate a job script. Where the code_uuid and withmpi fields are not used in job script generation.

[codecov]

Codecov Report

Merging #5350 (3b5914a) into develop (ffedc8b) will increase coverage by 2.57%.
The diff coverage is 96.00%.

❗ Current head 3b5914a differs from pull request most recent head 6c3701c. Consider uploading reports for the commit 6c3701c to get more accurate results

@@             Coverage Diff             @@
##           develop    #5350      +/-   ##
===========================================
+ Coverage    79.57%   82.13%   +2.57%     
===========================================
  Files          517      533      +16     
  Lines        36982    38504    +1522     
===========================================
+ Hits         29424    31623    +2199     
+ Misses        7558     6881     -677     

Flag	Coverage Δ
django	77.22% <96.00%> (?)
sqlalchemy	76.52% <96.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
aiida/schedulers/scheduler.py	85.46% <ø> (ø)
aiida/engine/processes/calcjobs/calcjob.py	90.20% <93.75%> (-0.05%)	⬇️
aiida/schedulers/datastructures.py	94.83% <100.00%> (+0.32%)	⬆️
aiida/orm/nodes/data/array/xy.py	18.58% <0.00%> (-41.42%)	⬇️
aiida/cmdline/params/options/commands/setup.py	54.55% <0.00%> (-38.45%)	⬇️
aiida/cmdline/commands/cmd_setup.py	56.87% <0.00%> (-38.03%)	⬇️
aiida/orm/nodes/data/array/projection.py	16.13% <0.00%> (-32.25%)	⬇️
aiida/manage/external/postgres.py	63.10% <0.00%> (-20.23%)	⬇️
aiida/tools/archive/exceptions.py	90.48% <0.00%> (-9.52%)	⬇️
aiida/cmdline/params/types/profile.py	64.45% <0.00%> (-7.29%)	⬇️
... and 448 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update ffedc8b...6c3701c. Read the comment docs.

[sphuber]

Thanks @unkcpz . Have a few suggestions. The tests are failing because the JobTemplate that is generated by the engine is written to the repository of the CalcJobNode by dumping it to JSON. The old CodeInfo was a simple dict and so serializable, but the new dataclasses are not JSON serializable.

In [1]: import json

In [2]: from dataclasses import dataclass

In [4]: @dataclass
   ...: class Test:
   ...:     a: int = 1
   ...: 

In [5]: a = Test()

In [6]: a.a
Out[6]: 1

In [7]: json.dumps(a)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-2f50cf32d976> in <module>
----> 1 json.dumps(a)

/usr/lib/python3.9/json/__init__.py in dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls, indent, separators, default, sort_keys, **kw)
    229         cls is None and indent is None and separators is None and
    230         default is None and not sort_keys and not kw):
--> 231         return _default_encoder.encode(obj)
    232     if cls is None:
    233         cls = JSONEncoder

/usr/lib/python3.9/json/encoder.py in encode(self, o)
    197         # exceptions aren't as detailed.  The list call should be roughly
    198         # equivalent to the PySequence_Fast that ''.join() would do.
--> 199         chunks = self.iterencode(o, _one_shot=True)
    200         if not isinstance(chunks, (list, tuple)):
    201             chunks = list(chunks)

/usr/lib/python3.9/json/encoder.py in iterencode(self, o, _one_shot)
    255                 self.key_separator, self.item_separator, self.sort_keys,
    256                 self.skipkeys, _one_shot)
--> 257         return _iterencode(o, 0)
    258 
    259 def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,

/usr/lib/python3.9/json/encoder.py in default(self, o)
    177 
    178         """
--> 179         raise TypeError(f'Object of type {o.__class__.__name__} '
    180                         f'is not JSON serializable')
    181 

TypeError: Object of type Test is not JSON serializable

However, you can turn it into a dictionary first. So in line 763, of CalcJob.presubmit turn

subfolder.create_file_from_filelike(io.StringIO(json.dumps(job_tmpl)), 'job_tmpl.json', 'w', encoding='utf8')

into

from dataclasses import asdict
subfolder.create_file_from_filelike(io.StringIO(json.dumps(asdict(job_tmpl))), 'job_tmpl.json', 'w', encoding='utf8')

[unkcpz]

Thanks @sphuber. I check the failed test caused by dataclass. I think we have to work around it with AttributeDict(asdict(tmpl_code_info)). It is tmpl_code_info who is the dataclass and can be escaped to dict not job_tmpl. Converting to dict is not enough since the initial benifite from using dataclass is we have a built-in way to use dot to access attributes. Therefore need to convert to AttributeDict in order to do so. It seems not very concise but I don't know a more elegant solution.

[chrisjsewell]

Therefore need to convert to AttributeDict in order to do so.

Err, I feel the better solution is just to handle data classes at JSON dump time:

import dataclasses
def encoder(obj):
     if dataclasses.is_dataclass(obj):
        return  dataclasses.as_dict(obj) 
    raise TypeError(f" {obj!r} is not JSON serializable")
     
subfolder.create_file_from_filelike(io.StringIO(json.dumps(job_tmpl, default=encoder)), 'job_tmpl.json', 'w', encoding='utf8')

But, I don't want to overcomplicate the issue, so I'll leave it up to you guys, and maybe open this as a more general issue:

I'm very much not a fan of these AttributeDict, having personally made the "mistake" of implementing them in another one of my packages and then later removing them: executablebooks/markdown-it-py#66.
One should use either dicts or dataclasses, but not mix the concepts.

@sphuber I think one "argument" for using them was dynamic auto-completion but, just out of interest, I found that since ipython/ipython#5304, you can anyhow auto-complete dict keys:

[sphuber]

I agree with @chrisjsewell that we should avoid using the custom dict classes and use dataclasses instead. These things were in aiida-core when I joined so I am not sure what the original motivation was, but you may well be right that the main one was auto-completion. I think we should deprecate all of the classes in aiida.common.extendeddicts and replace those that use them by dataclasses. But that is for another time.

@unkcpz My suggestion was not to convert the JobTemplateCodeInfo to a dict in the JobTemplate instance itself, but rather just before it is being dumped to a JSON file. After that it is not being "used" anymore by plugins. My suggestion was to simply call asdict manually since there is just one thing we need to convert.

But @chrisjsewell solution is more elegant and will automatically capture all instances of dataclasses if we were to add more in the future. I do wonder what the reasoning was for not adding a default support in JSON serializer for dataclasses in the standard library. Seems like a very common use case and in most cases where there are simple base types in the dataclass, it should be trivially serializable.

[unkcpz]

Thanks @chrisjsewell @sphuber. I update using the way Chris suggests.

For the extendeddicts classes, I think AttributeDict can quickly apply to the nested dict and make it dot visitable which is used in the NodeLinksManager. Is this also supported by dataclasses?

[sphuber]

For the extendeddicts classes, I think AttributeDict can quickly apply to the nested dict and make it dot visitable which is used in the NodeLinksManager. Is this also supported by dataclasses?

No, dataclasses do not recursively add attribute dereferencing for nested dicts that it contains. This is not what dataclasss are designed for though. I also don't think that it is recommended that its attributes are defined dynamically, as is the case for the NodeLinksManager, so we cannot or at least shouldn't use dataclasses for that.

The dataclasses should simply be used whenever we have a simple class that essentially just wraps some data through attributes. This is the case for CodeInfo, CalcInfo etc.

[sphuber]

@giovannipizzi @chrisjsewell even though the chances are small, this could potentially be breaking if there are scheduler plugins that rely on the type of the codes_info argument passed to the _get_run_line method. Do we accept these chances and put this in 2.0 still?

[ltalirz]

@giovannipizzi @chrisjsewell even though the chances are small, this could potentially be breaking if there are scheduler plugins that rely on the type of the codes_info argument passed to the _get_run_line method. Do we accept these chances and put this in 2.0 still?

The AiiDA registry currently contains four packages that add a scheduler plugin:

https://github.com/zhubonan/aiida-fireworks-scheduler
https://github.com/pzarabadip/aiida-metavo-scheduler
https://github.com/aiidateam/aiida-firecrest
https://github.com/dev-zero/aiida-statefile-schedulers

@unkcpz Can you check whether any of those would be impacted?
If not, I think we can assume it is safe to proceed here (of course there can be private plugins that can be affected; we will just document the change in API).

[unkcpz]

@ltalirz thanks!

@unkcpz Can you check whether any of those would be impacted?

Sure, I will go through them this week, if nothing particular is messed up by this PR, we can give it merged.

[unkcpz]

The AiiDA registry currently contains four packages that add a scheduler plugin:
https://github.com/zhubonan/aiida-fireworks-scheduler
https://github.com/pzarabadip/aiida-metavo-scheduler
https://github.com/aiidateam/aiida-firecrest
https://github.com/dev-zero/aiida-statefile-schedulers

Just check all scheduler implementations. There are no codes that override the part of generating exec command of script. So it should be pretty safe to go forward. 😄

[ltalirz]

Great, any objections @sphuber ?

[sphuber]

No, if you rebase the PR and verify that it won't break the schedulers on the plugin registry, then I will approve and merge this.

[unkcpz]

Hi @sphuber, I rebase PR and check the aiida-hyperqueue is not impacted.

[ltalirz]

Thanks @unkcpz , my understanding is this can move forward !