hpcflow 0.1.15__py3-none-any.whl → 0.2.0a271__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (275) hide show
  1. hpcflow/__init__.py +2 -11
  2. hpcflow/__pyinstaller/__init__.py +5 -0
  3. hpcflow/__pyinstaller/hook-hpcflow.py +40 -0
  4. hpcflow/_version.py +1 -1
  5. hpcflow/app.py +43 -0
  6. hpcflow/cli.py +2 -461
  7. hpcflow/data/demo_data_manifest/__init__.py +3 -0
  8. hpcflow/data/demo_data_manifest/demo_data_manifest.json +6 -0
  9. hpcflow/data/jinja_templates/test/test_template.txt +8 -0
  10. hpcflow/data/programs/hello_world/README.md +1 -0
  11. hpcflow/data/programs/hello_world/hello_world.c +87 -0
  12. hpcflow/data/programs/hello_world/linux/hello_world +0 -0
  13. hpcflow/data/programs/hello_world/macos/hello_world +0 -0
  14. hpcflow/data/programs/hello_world/win/hello_world.exe +0 -0
  15. hpcflow/data/scripts/__init__.py +1 -0
  16. hpcflow/data/scripts/bad_script.py +2 -0
  17. hpcflow/data/scripts/demo_task_1_generate_t1_infile_1.py +8 -0
  18. hpcflow/data/scripts/demo_task_1_generate_t1_infile_2.py +8 -0
  19. hpcflow/data/scripts/demo_task_1_parse_p3.py +7 -0
  20. hpcflow/data/scripts/do_nothing.py +2 -0
  21. hpcflow/data/scripts/env_specifier_test/input_file_generator_pass_env_spec.py +4 -0
  22. hpcflow/data/scripts/env_specifier_test/main_script_test_pass_env_spec.py +8 -0
  23. hpcflow/data/scripts/env_specifier_test/output_file_parser_pass_env_spec.py +4 -0
  24. hpcflow/data/scripts/env_specifier_test/v1/input_file_generator_basic.py +4 -0
  25. hpcflow/data/scripts/env_specifier_test/v1/main_script_test_direct_in_direct_out.py +7 -0
  26. hpcflow/data/scripts/env_specifier_test/v1/output_file_parser_basic.py +4 -0
  27. hpcflow/data/scripts/env_specifier_test/v2/main_script_test_direct_in_direct_out.py +7 -0
  28. hpcflow/data/scripts/generate_t1_file_01.py +7 -0
  29. hpcflow/data/scripts/import_future_script.py +7 -0
  30. hpcflow/data/scripts/input_file_generator_basic.py +3 -0
  31. hpcflow/data/scripts/input_file_generator_basic_FAIL.py +3 -0
  32. hpcflow/data/scripts/input_file_generator_test_stdout_stderr.py +8 -0
  33. hpcflow/data/scripts/main_script_test_direct_in.py +3 -0
  34. hpcflow/data/scripts/main_script_test_direct_in_direct_out.py +6 -0
  35. hpcflow/data/scripts/main_script_test_direct_in_direct_out_2.py +6 -0
  36. hpcflow/data/scripts/main_script_test_direct_in_direct_out_2_fail_allowed.py +6 -0
  37. hpcflow/data/scripts/main_script_test_direct_in_direct_out_2_fail_allowed_group.py +7 -0
  38. hpcflow/data/scripts/main_script_test_direct_in_direct_out_3.py +6 -0
  39. hpcflow/data/scripts/main_script_test_direct_in_direct_out_all_iters_test.py +15 -0
  40. hpcflow/data/scripts/main_script_test_direct_in_direct_out_env_spec.py +7 -0
  41. hpcflow/data/scripts/main_script_test_direct_in_direct_out_labels.py +8 -0
  42. hpcflow/data/scripts/main_script_test_direct_in_group_direct_out_3.py +6 -0
  43. hpcflow/data/scripts/main_script_test_direct_in_group_one_fail_direct_out_3.py +6 -0
  44. hpcflow/data/scripts/main_script_test_direct_sub_param_in_direct_out.py +6 -0
  45. hpcflow/data/scripts/main_script_test_hdf5_in_obj.py +12 -0
  46. hpcflow/data/scripts/main_script_test_hdf5_in_obj_2.py +12 -0
  47. hpcflow/data/scripts/main_script_test_hdf5_in_obj_group.py +12 -0
  48. hpcflow/data/scripts/main_script_test_hdf5_out_obj.py +11 -0
  49. hpcflow/data/scripts/main_script_test_json_and_direct_in_json_out.py +14 -0
  50. hpcflow/data/scripts/main_script_test_json_in_json_and_direct_out.py +17 -0
  51. hpcflow/data/scripts/main_script_test_json_in_json_out.py +14 -0
  52. hpcflow/data/scripts/main_script_test_json_in_json_out_labels.py +16 -0
  53. hpcflow/data/scripts/main_script_test_json_in_obj.py +12 -0
  54. hpcflow/data/scripts/main_script_test_json_out_FAIL.py +3 -0
  55. hpcflow/data/scripts/main_script_test_json_out_obj.py +10 -0
  56. hpcflow/data/scripts/main_script_test_json_sub_param_in_json_out_labels.py +16 -0
  57. hpcflow/data/scripts/main_script_test_shell_env_vars.py +12 -0
  58. hpcflow/data/scripts/main_script_test_std_out_std_err.py +6 -0
  59. hpcflow/data/scripts/output_file_parser_basic.py +3 -0
  60. hpcflow/data/scripts/output_file_parser_basic_FAIL.py +7 -0
  61. hpcflow/data/scripts/output_file_parser_test_stdout_stderr.py +8 -0
  62. hpcflow/data/scripts/parse_t1_file_01.py +4 -0
  63. hpcflow/data/scripts/script_exit_test.py +5 -0
  64. hpcflow/data/template_components/__init__.py +1 -0
  65. hpcflow/data/template_components/command_files.yaml +26 -0
  66. hpcflow/data/template_components/environments.yaml +13 -0
  67. hpcflow/data/template_components/parameters.yaml +14 -0
  68. hpcflow/data/template_components/task_schemas.yaml +139 -0
  69. hpcflow/data/workflows/workflow_1.yaml +5 -0
  70. hpcflow/examples.ipynb +1037 -0
  71. hpcflow/sdk/__init__.py +149 -0
  72. hpcflow/sdk/app.py +4266 -0
  73. hpcflow/sdk/cli.py +1479 -0
  74. hpcflow/sdk/cli_common.py +385 -0
  75. hpcflow/sdk/config/__init__.py +5 -0
  76. hpcflow/sdk/config/callbacks.py +246 -0
  77. hpcflow/sdk/config/cli.py +388 -0
  78. hpcflow/sdk/config/config.py +1410 -0
  79. hpcflow/sdk/config/config_file.py +501 -0
  80. hpcflow/sdk/config/errors.py +272 -0
  81. hpcflow/sdk/config/types.py +150 -0
  82. hpcflow/sdk/core/__init__.py +38 -0
  83. hpcflow/sdk/core/actions.py +3857 -0
  84. hpcflow/sdk/core/app_aware.py +25 -0
  85. hpcflow/sdk/core/cache.py +224 -0
  86. hpcflow/sdk/core/command_files.py +814 -0
  87. hpcflow/sdk/core/commands.py +424 -0
  88. hpcflow/sdk/core/element.py +2071 -0
  89. hpcflow/sdk/core/enums.py +221 -0
  90. hpcflow/sdk/core/environment.py +256 -0
  91. hpcflow/sdk/core/errors.py +1043 -0
  92. hpcflow/sdk/core/execute.py +207 -0
  93. hpcflow/sdk/core/json_like.py +809 -0
  94. hpcflow/sdk/core/loop.py +1320 -0
  95. hpcflow/sdk/core/loop_cache.py +282 -0
  96. hpcflow/sdk/core/object_list.py +933 -0
  97. hpcflow/sdk/core/parameters.py +3371 -0
  98. hpcflow/sdk/core/rule.py +196 -0
  99. hpcflow/sdk/core/run_dir_files.py +57 -0
  100. hpcflow/sdk/core/skip_reason.py +7 -0
  101. hpcflow/sdk/core/task.py +3792 -0
  102. hpcflow/sdk/core/task_schema.py +993 -0
  103. hpcflow/sdk/core/test_utils.py +538 -0
  104. hpcflow/sdk/core/types.py +447 -0
  105. hpcflow/sdk/core/utils.py +1207 -0
  106. hpcflow/sdk/core/validation.py +87 -0
  107. hpcflow/sdk/core/values.py +477 -0
  108. hpcflow/sdk/core/workflow.py +4820 -0
  109. hpcflow/sdk/core/zarr_io.py +206 -0
  110. hpcflow/sdk/data/__init__.py +13 -0
  111. hpcflow/sdk/data/config_file_schema.yaml +34 -0
  112. hpcflow/sdk/data/config_schema.yaml +260 -0
  113. hpcflow/sdk/data/environments_spec_schema.yaml +21 -0
  114. hpcflow/sdk/data/files_spec_schema.yaml +5 -0
  115. hpcflow/sdk/data/parameters_spec_schema.yaml +7 -0
  116. hpcflow/sdk/data/task_schema_spec_schema.yaml +3 -0
  117. hpcflow/sdk/data/workflow_spec_schema.yaml +22 -0
  118. hpcflow/sdk/demo/__init__.py +3 -0
  119. hpcflow/sdk/demo/cli.py +242 -0
  120. hpcflow/sdk/helper/__init__.py +3 -0
  121. hpcflow/sdk/helper/cli.py +137 -0
  122. hpcflow/sdk/helper/helper.py +300 -0
  123. hpcflow/sdk/helper/watcher.py +192 -0
  124. hpcflow/sdk/log.py +288 -0
  125. hpcflow/sdk/persistence/__init__.py +18 -0
  126. hpcflow/sdk/persistence/base.py +2817 -0
  127. hpcflow/sdk/persistence/defaults.py +6 -0
  128. hpcflow/sdk/persistence/discovery.py +39 -0
  129. hpcflow/sdk/persistence/json.py +954 -0
  130. hpcflow/sdk/persistence/pending.py +948 -0
  131. hpcflow/sdk/persistence/store_resource.py +203 -0
  132. hpcflow/sdk/persistence/types.py +309 -0
  133. hpcflow/sdk/persistence/utils.py +73 -0
  134. hpcflow/sdk/persistence/zarr.py +2388 -0
  135. hpcflow/sdk/runtime.py +320 -0
  136. hpcflow/sdk/submission/__init__.py +3 -0
  137. hpcflow/sdk/submission/enums.py +70 -0
  138. hpcflow/sdk/submission/jobscript.py +2379 -0
  139. hpcflow/sdk/submission/schedulers/__init__.py +281 -0
  140. hpcflow/sdk/submission/schedulers/direct.py +233 -0
  141. hpcflow/sdk/submission/schedulers/sge.py +376 -0
  142. hpcflow/sdk/submission/schedulers/slurm.py +598 -0
  143. hpcflow/sdk/submission/schedulers/utils.py +25 -0
  144. hpcflow/sdk/submission/shells/__init__.py +52 -0
  145. hpcflow/sdk/submission/shells/base.py +229 -0
  146. hpcflow/sdk/submission/shells/bash.py +504 -0
  147. hpcflow/sdk/submission/shells/os_version.py +115 -0
  148. hpcflow/sdk/submission/shells/powershell.py +352 -0
  149. hpcflow/sdk/submission/submission.py +1402 -0
  150. hpcflow/sdk/submission/types.py +140 -0
  151. hpcflow/sdk/typing.py +194 -0
  152. hpcflow/sdk/utils/arrays.py +69 -0
  153. hpcflow/sdk/utils/deferred_file.py +55 -0
  154. hpcflow/sdk/utils/hashing.py +16 -0
  155. hpcflow/sdk/utils/patches.py +31 -0
  156. hpcflow/sdk/utils/strings.py +69 -0
  157. hpcflow/tests/api/test_api.py +32 -0
  158. hpcflow/tests/conftest.py +123 -0
  159. hpcflow/tests/data/__init__.py +0 -0
  160. hpcflow/tests/data/benchmark_N_elements.yaml +6 -0
  161. hpcflow/tests/data/benchmark_script_runner.yaml +26 -0
  162. hpcflow/tests/data/multi_path_sequences.yaml +29 -0
  163. hpcflow/tests/data/workflow_1.json +10 -0
  164. hpcflow/tests/data/workflow_1.yaml +5 -0
  165. hpcflow/tests/data/workflow_1_slurm.yaml +8 -0
  166. hpcflow/tests/data/workflow_1_wsl.yaml +8 -0
  167. hpcflow/tests/data/workflow_test_run_abort.yaml +42 -0
  168. hpcflow/tests/jinja_templates/test_jinja_templates.py +161 -0
  169. hpcflow/tests/programs/test_programs.py +180 -0
  170. hpcflow/tests/schedulers/direct_linux/test_direct_linux_submission.py +12 -0
  171. hpcflow/tests/schedulers/sge/test_sge_submission.py +36 -0
  172. hpcflow/tests/schedulers/slurm/test_slurm_submission.py +14 -0
  173. hpcflow/tests/scripts/test_input_file_generators.py +282 -0
  174. hpcflow/tests/scripts/test_main_scripts.py +1361 -0
  175. hpcflow/tests/scripts/test_non_snippet_script.py +46 -0
  176. hpcflow/tests/scripts/test_ouput_file_parsers.py +353 -0
  177. hpcflow/tests/shells/wsl/test_wsl_submission.py +14 -0
  178. hpcflow/tests/unit/test_action.py +1066 -0
  179. hpcflow/tests/unit/test_action_rule.py +24 -0
  180. hpcflow/tests/unit/test_app.py +132 -0
  181. hpcflow/tests/unit/test_cache.py +46 -0
  182. hpcflow/tests/unit/test_cli.py +172 -0
  183. hpcflow/tests/unit/test_command.py +377 -0
  184. hpcflow/tests/unit/test_config.py +195 -0
  185. hpcflow/tests/unit/test_config_file.py +162 -0
  186. hpcflow/tests/unit/test_element.py +666 -0
  187. hpcflow/tests/unit/test_element_iteration.py +88 -0
  188. hpcflow/tests/unit/test_element_set.py +158 -0
  189. hpcflow/tests/unit/test_group.py +115 -0
  190. hpcflow/tests/unit/test_input_source.py +1479 -0
  191. hpcflow/tests/unit/test_input_value.py +398 -0
  192. hpcflow/tests/unit/test_jobscript_unit.py +757 -0
  193. hpcflow/tests/unit/test_json_like.py +1247 -0
  194. hpcflow/tests/unit/test_loop.py +2674 -0
  195. hpcflow/tests/unit/test_meta_task.py +325 -0
  196. hpcflow/tests/unit/test_multi_path_sequences.py +259 -0
  197. hpcflow/tests/unit/test_object_list.py +116 -0
  198. hpcflow/tests/unit/test_parameter.py +243 -0
  199. hpcflow/tests/unit/test_persistence.py +664 -0
  200. hpcflow/tests/unit/test_resources.py +243 -0
  201. hpcflow/tests/unit/test_run.py +286 -0
  202. hpcflow/tests/unit/test_run_directories.py +29 -0
  203. hpcflow/tests/unit/test_runtime.py +9 -0
  204. hpcflow/tests/unit/test_schema_input.py +372 -0
  205. hpcflow/tests/unit/test_shell.py +129 -0
  206. hpcflow/tests/unit/test_slurm.py +39 -0
  207. hpcflow/tests/unit/test_submission.py +502 -0
  208. hpcflow/tests/unit/test_task.py +2560 -0
  209. hpcflow/tests/unit/test_task_schema.py +182 -0
  210. hpcflow/tests/unit/test_utils.py +616 -0
  211. hpcflow/tests/unit/test_value_sequence.py +549 -0
  212. hpcflow/tests/unit/test_values.py +91 -0
  213. hpcflow/tests/unit/test_workflow.py +827 -0
  214. hpcflow/tests/unit/test_workflow_template.py +186 -0
  215. hpcflow/tests/unit/utils/test_arrays.py +40 -0
  216. hpcflow/tests/unit/utils/test_deferred_file_writer.py +34 -0
  217. hpcflow/tests/unit/utils/test_hashing.py +65 -0
  218. hpcflow/tests/unit/utils/test_patches.py +5 -0
  219. hpcflow/tests/unit/utils/test_redirect_std.py +50 -0
  220. hpcflow/tests/unit/utils/test_strings.py +97 -0
  221. hpcflow/tests/workflows/__init__.py +0 -0
  222. hpcflow/tests/workflows/test_directory_structure.py +31 -0
  223. hpcflow/tests/workflows/test_jobscript.py +355 -0
  224. hpcflow/tests/workflows/test_run_status.py +198 -0
  225. hpcflow/tests/workflows/test_skip_downstream.py +696 -0
  226. hpcflow/tests/workflows/test_submission.py +140 -0
  227. hpcflow/tests/workflows/test_workflows.py +564 -0
  228. hpcflow/tests/workflows/test_zip.py +18 -0
  229. hpcflow/viz_demo.ipynb +6794 -0
  230. hpcflow-0.2.0a271.dist-info/LICENSE +375 -0
  231. hpcflow-0.2.0a271.dist-info/METADATA +65 -0
  232. hpcflow-0.2.0a271.dist-info/RECORD +237 -0
  233. {hpcflow-0.1.15.dist-info → hpcflow-0.2.0a271.dist-info}/WHEEL +4 -5
  234. hpcflow-0.2.0a271.dist-info/entry_points.txt +6 -0
  235. hpcflow/api.py +0 -490
  236. hpcflow/archive/archive.py +0 -307
  237. hpcflow/archive/cloud/cloud.py +0 -45
  238. hpcflow/archive/cloud/errors.py +0 -9
  239. hpcflow/archive/cloud/providers/dropbox.py +0 -427
  240. hpcflow/archive/errors.py +0 -5
  241. hpcflow/base_db.py +0 -4
  242. hpcflow/config.py +0 -233
  243. hpcflow/copytree.py +0 -66
  244. hpcflow/data/examples/_config.yml +0 -14
  245. hpcflow/data/examples/damask/demo/1.run.yml +0 -4
  246. hpcflow/data/examples/damask/demo/2.process.yml +0 -29
  247. hpcflow/data/examples/damask/demo/geom.geom +0 -2052
  248. hpcflow/data/examples/damask/demo/load.load +0 -1
  249. hpcflow/data/examples/damask/demo/material.config +0 -185
  250. hpcflow/data/examples/damask/inputs/geom.geom +0 -2052
  251. hpcflow/data/examples/damask/inputs/load.load +0 -1
  252. hpcflow/data/examples/damask/inputs/material.config +0 -185
  253. hpcflow/data/examples/damask/profiles/_variable_lookup.yml +0 -21
  254. hpcflow/data/examples/damask/profiles/damask.yml +0 -4
  255. hpcflow/data/examples/damask/profiles/damask_process.yml +0 -8
  256. hpcflow/data/examples/damask/profiles/damask_run.yml +0 -5
  257. hpcflow/data/examples/damask/profiles/default.yml +0 -6
  258. hpcflow/data/examples/thinking.yml +0 -177
  259. hpcflow/errors.py +0 -2
  260. hpcflow/init_db.py +0 -37
  261. hpcflow/models.py +0 -2595
  262. hpcflow/nesting.py +0 -9
  263. hpcflow/profiles.py +0 -455
  264. hpcflow/project.py +0 -81
  265. hpcflow/scheduler.py +0 -322
  266. hpcflow/utils.py +0 -103
  267. hpcflow/validation.py +0 -166
  268. hpcflow/variables.py +0 -543
  269. hpcflow-0.1.15.dist-info/METADATA +0 -168
  270. hpcflow-0.1.15.dist-info/RECORD +0 -45
  271. hpcflow-0.1.15.dist-info/entry_points.txt +0 -8
  272. hpcflow-0.1.15.dist-info/top_level.txt +0 -1
  273. /hpcflow/{archive → data/jinja_templates}/__init__.py +0 -0
  274. /hpcflow/{archive/cloud → data/programs}/__init__.py +0 -0
  275. /hpcflow/{archive/cloud/providers → data/workflows}/__init__.py +0 -0
@@ -0,0 +1,954 @@
1
+ """
2
+ Persistence model based on writing JSON documents.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ from contextlib import contextmanager
8
+ import copy
9
+ import json
10
+ from pathlib import Path
11
+ from typing import cast, TYPE_CHECKING
12
+ from typing_extensions import override
13
+
14
+ from fsspec import filesystem, AbstractFileSystem # type: ignore
15
+ import numpy as np
16
+ from hpcflow.sdk.core import RUN_DIR_ARR_DTYPE, RUN_DIR_ARR_FILL
17
+ from hpcflow.sdk.core.errors import (
18
+ MissingParameterData,
19
+ MissingStoreEARError,
20
+ MissingStoreElementError,
21
+ MissingStoreElementIterationError,
22
+ )
23
+ from hpcflow.sdk.persistence.base import (
24
+ PersistentStoreFeatures,
25
+ PersistentStore,
26
+ StoreEAR,
27
+ StoreElement,
28
+ StoreElementIter,
29
+ StoreParameter,
30
+ StoreTask,
31
+ update_param_source_dict,
32
+ )
33
+ from hpcflow.sdk.submission.submission import JOBSCRIPT_SUBMIT_TIME_KEYS
34
+ from hpcflow.sdk.persistence.pending import CommitResourceMap
35
+ from hpcflow.sdk.persistence.store_resource import JSONFileStoreResource
36
+ from hpcflow.sdk.typing import DataIndex
37
+
38
+ if TYPE_CHECKING:
39
+ from collections.abc import Iterable, Iterator, Mapping, Sequence
40
+ from datetime import datetime
41
+ from typing import Any, ClassVar, Literal
42
+ from typing_extensions import Self
43
+ from numpy.typing import NDArray
44
+ from ..app import BaseApp
45
+ from ..core.json_like import JSONed, JSONDocument
46
+ from ..core.workflow import Workflow
47
+ from ..typing import ParamSource
48
+ from .types import (
49
+ ElemMeta,
50
+ IterMeta,
51
+ LoopDescriptor,
52
+ Metadata,
53
+ RunMeta,
54
+ StoreCreationInfo,
55
+ TaskMeta,
56
+ TemplateMeta,
57
+ )
58
+
59
+
60
+ class JsonStoreTask(StoreTask["TaskMeta"]):
61
+ """
62
+ Persisted task that is serialized using JSON.
63
+ """
64
+
65
+ @override
66
+ def encode(self) -> tuple[int, TaskMeta, dict[str, Any]]:
67
+ """Prepare store task data for the persistent store."""
68
+ assert self.task_template is not None
69
+ wk_task: TaskMeta = {
70
+ "id_": self.id_,
71
+ "element_IDs": self.element_IDs,
72
+ "index": self.index,
73
+ }
74
+ task = {"id_": self.id_, **self.task_template}
75
+ return self.index, wk_task, task
76
+
77
+ @override
78
+ @classmethod
79
+ def decode(cls, task_dat: TaskMeta) -> Self:
80
+ """Initialise a `StoreTask` from store task data
81
+
82
+ Note: the `task_template` is only needed for encoding because it is retrieved as
83
+ part of the `WorkflowTemplate` so we don't need to load it when decoding.
84
+
85
+ """
86
+ return cls(is_pending=False, **task_dat)
87
+
88
+
89
+ class JsonStoreElement(StoreElement["ElemMeta", None]):
90
+ """
91
+ Persisted element that is serialized using JSON.
92
+ """
93
+
94
+ @override
95
+ def encode(self, context: None) -> ElemMeta:
96
+ """Prepare store element data for the persistent store."""
97
+ dct = self.__dict__
98
+ del dct["is_pending"]
99
+ return cast("ElemMeta", dct)
100
+
101
+ @override
102
+ @classmethod
103
+ def decode(cls, elem_dat: ElemMeta, context: None) -> Self:
104
+ """Initialise a `JsonStoreElement` from store element data"""
105
+ return cls(is_pending=False, **elem_dat)
106
+
107
+
108
+ class JsonStoreElementIter(StoreElementIter["IterMeta", None]):
109
+ """
110
+ Persisted element iteration that is serialized using JSON.
111
+ """
112
+
113
+ @override
114
+ def encode(self, context: None) -> IterMeta:
115
+ """Prepare store element iteration data for the persistent store."""
116
+ dct = self.__dict__
117
+ del dct["is_pending"]
118
+ return cast("IterMeta", dct)
119
+
120
+ @override
121
+ @classmethod
122
+ def decode(cls, iter_dat: IterMeta, context: None) -> Self:
123
+ """Initialise a `JsonStoreElementIter` from persistent store element iteration data"""
124
+
125
+ iter_dat = copy.deepcopy(iter_dat) # to avoid mutating; can we avoid this?
126
+
127
+ # cast JSON string keys to integers:
128
+ if EAR_IDs := iter_dat["EAR_IDs"]:
129
+ for act_idx in list(EAR_IDs):
130
+ EAR_IDs[int(act_idx)] = EAR_IDs.pop(act_idx)
131
+
132
+ return cls(is_pending=False, **cast("dict", iter_dat))
133
+
134
+
135
+ class JsonStoreEAR(StoreEAR["RunMeta", None]):
136
+ """
137
+ Persisted element action run that is serialized using JSON.
138
+ """
139
+
140
+ @override
141
+ def encode(self, ts_fmt: str, context: None) -> RunMeta:
142
+ """Prepare store EAR data for the persistent store."""
143
+ return {
144
+ "id_": self.id_,
145
+ "elem_iter_ID": self.elem_iter_ID,
146
+ "action_idx": self.action_idx,
147
+ "commands_idx": self.commands_idx,
148
+ "data_idx": self.data_idx,
149
+ "submission_idx": self.submission_idx,
150
+ "commands_file_ID": self.commands_file_ID,
151
+ "success": self.success,
152
+ "skip": self.skip,
153
+ "start_time": self._encode_datetime(self.start_time, ts_fmt),
154
+ "end_time": self._encode_datetime(self.end_time, ts_fmt),
155
+ "snapshot_start": self.snapshot_start,
156
+ "snapshot_end": self.snapshot_end,
157
+ "exit_code": self.exit_code,
158
+ "metadata": self.metadata,
159
+ "run_hostname": self.run_hostname,
160
+ "port_number": self.port_number,
161
+ }
162
+
163
+ @override
164
+ @classmethod
165
+ def decode(cls, EAR_dat: RunMeta, ts_fmt: str, context: None) -> Self:
166
+ """Initialise a `JsonStoreEAR` from persistent store EAR data"""
167
+ # don't want to mutate EAR_dat:
168
+ EAR_dat = copy.deepcopy(EAR_dat)
169
+ start_time = cls._decode_datetime(EAR_dat.pop("start_time"), ts_fmt)
170
+ end_time = cls._decode_datetime(EAR_dat.pop("end_time"), ts_fmt)
171
+ return cls(
172
+ is_pending=False,
173
+ **cast("dict", EAR_dat),
174
+ start_time=start_time,
175
+ end_time=end_time,
176
+ )
177
+
178
+
179
+ class JSONPersistentStore(
180
+ PersistentStore[
181
+ JsonStoreTask,
182
+ JsonStoreElement,
183
+ JsonStoreElementIter,
184
+ JsonStoreEAR,
185
+ StoreParameter,
186
+ ]
187
+ ):
188
+ """
189
+ A store that writes JSON files for all its state serialization.
190
+ """
191
+
192
+ _name: ClassVar[str] = "json"
193
+ _features: ClassVar[PersistentStoreFeatures] = PersistentStoreFeatures(
194
+ create=True,
195
+ edit=True,
196
+ jobscript_parallelism=False,
197
+ EAR_parallelism=False,
198
+ schedulers=True,
199
+ submission=True,
200
+ )
201
+
202
+ _meta_res: ClassVar[str] = "metadata"
203
+ _params_res: ClassVar[str] = "parameters"
204
+ _subs_res: ClassVar[str] = "submissions"
205
+ _runs_res: ClassVar[str] = "runs"
206
+
207
+ _res_file_names: ClassVar[Mapping[str, str]] = {
208
+ _meta_res: "metadata.json",
209
+ _params_res: "parameters.json",
210
+ _subs_res: "submissions.json",
211
+ _runs_res: "runs.json",
212
+ }
213
+
214
+ _res_map: ClassVar[CommitResourceMap] = CommitResourceMap(
215
+ commit_tasks=(_meta_res,),
216
+ commit_loops=(_meta_res,),
217
+ commit_loop_num_iters=(_meta_res,),
218
+ commit_loop_parents=(_meta_res,),
219
+ commit_submissions=(_subs_res,),
220
+ commit_at_submit_metadata=(_subs_res,),
221
+ commit_js_metadata=(_subs_res,),
222
+ commit_elem_IDs=(_meta_res,),
223
+ commit_elements=(_meta_res,),
224
+ commit_element_sets=(_meta_res,),
225
+ commit_elem_iter_IDs=(_meta_res,),
226
+ commit_elem_iters=(_meta_res,),
227
+ commit_loop_indices=(_meta_res,),
228
+ commit_elem_iter_EAR_IDs=(_meta_res,),
229
+ commit_EARs_initialised=(_meta_res,),
230
+ commit_EARs=(_runs_res,),
231
+ commit_EAR_submission_indices=(_runs_res,),
232
+ commit_EAR_skips=(_runs_res,),
233
+ commit_EAR_starts=(_runs_res,),
234
+ commit_EAR_ends=(_runs_res,),
235
+ commit_template_components=(_meta_res,),
236
+ commit_parameters=(_params_res,),
237
+ commit_param_sources=(_params_res,),
238
+ commit_set_run_dirs=(_runs_res,),
239
+ commit_iter_data_idx=(_meta_res,),
240
+ commit_run_data_idx=(_runs_res,),
241
+ )
242
+
243
+ @classmethod
244
+ def _store_task_cls(cls) -> type[JsonStoreTask]:
245
+ return JsonStoreTask
246
+
247
+ @classmethod
248
+ def _store_elem_cls(cls) -> type[JsonStoreElement]:
249
+ return JsonStoreElement
250
+
251
+ @classmethod
252
+ def _store_iter_cls(cls) -> type[JsonStoreElementIter]:
253
+ return JsonStoreElementIter
254
+
255
+ @classmethod
256
+ def _store_EAR_cls(cls) -> type[JsonStoreEAR]:
257
+ return JsonStoreEAR
258
+
259
+ @classmethod
260
+ def _store_param_cls(cls) -> type[StoreParameter]:
261
+ return StoreParameter
262
+
263
+ def __init__(
264
+ self, app, workflow: Workflow | None, path: Path, fs: AbstractFileSystem
265
+ ):
266
+ self._resources = {
267
+ self._meta_res: self._get_store_resource(app, "metadata", path, fs),
268
+ self._params_res: self._get_store_resource(app, "parameters", path, fs),
269
+ self._subs_res: self._get_store_resource(app, "submissions", path, fs),
270
+ self._runs_res: self._get_store_resource(app, "runs", path, fs),
271
+ }
272
+ super().__init__(app, workflow, path, fs)
273
+
274
+ # store-specific cache data, assigned in `using_resource()` when
275
+ # `_use_parameters_metadata_cache` is True, and set back to None when exiting the
276
+ # `parameters_metadata_cache` context manager.
277
+ self._parameters_file_dat: dict[str, dict[str, Any]] | None = None
278
+
279
+ @contextmanager
280
+ def cached_load(self) -> Iterator[None]:
281
+ """Context manager to cache the metadata."""
282
+ with self.using_resource("metadata", "read"):
283
+ with self.using_resource("runs", "read"):
284
+ yield
285
+
286
+ @contextmanager
287
+ def using_resource(
288
+ self,
289
+ res_label: Literal["metadata", "submissions", "parameters", "attrs", "runs"],
290
+ action: str,
291
+ ) -> Iterator[Any]:
292
+ """Context manager for managing `StoreResource` objects associated with the store.
293
+
294
+ Notes
295
+ -----
296
+ This overridden method facilitates easier use of the
297
+ `JSONPersistentStore`-specific implementation of the `parameters_metadata_cache`,
298
+ which in this case is just a copy of the `parameters.json` file data.
299
+
300
+ """
301
+
302
+ if (
303
+ self._use_parameters_metadata_cache
304
+ and res_label == "parameters"
305
+ and action == "read"
306
+ ):
307
+ if not self._parameters_file_dat:
308
+ with super().using_resource(
309
+ cast("Literal['parameters']", res_label), action
310
+ ) as res:
311
+ self._parameters_file_dat = res
312
+ yield self._parameters_file_dat
313
+
314
+ else:
315
+ with super().using_resource(res_label, action) as res:
316
+ yield res
317
+
318
+ @contextmanager
319
+ def parameters_metadata_cache(self) -> Iterator[None]:
320
+ """Context manager for using the parameters-metadata cache."""
321
+ self._use_parameters_metadata_cache = True
322
+ try:
323
+ yield
324
+ finally:
325
+ self._use_parameters_metadata_cache = False
326
+ self._parameters_file_dat = None # clear cache data
327
+
328
+ def remove_replaced_dir(self) -> None:
329
+ """
330
+ Remove the directory containing replaced workflow details.
331
+ """
332
+ with self.using_resource("metadata", "update") as md:
333
+ if "replaced_workflow" in md:
334
+ assert self.fs is not None
335
+ self.remove_path(md["replaced_workflow"])
336
+ self.logger.debug("removing temporarily renamed pre-existing workflow.")
337
+ del md["replaced_workflow"]
338
+
339
+ def reinstate_replaced_dir(self) -> None:
340
+ """
341
+ Reinstate the directory containing replaced workflow details.
342
+ """
343
+ with self.using_resource("metadata", "read") as md:
344
+ if "replaced_workflow" in md:
345
+ assert self.fs is not None
346
+ self.logger.debug(
347
+ "reinstating temporarily renamed pre-existing workflow."
348
+ )
349
+ self.rename_path(md["replaced_workflow"], self.path)
350
+
351
+ @classmethod
352
+ def _get_store_resource(
353
+ cls, app: BaseApp, name: str, path: str | Path, fs: AbstractFileSystem
354
+ ) -> JSONFileStoreResource:
355
+ return JSONFileStoreResource(
356
+ app=app,
357
+ name=name,
358
+ path=path,
359
+ fs=fs,
360
+ filename=cls._res_file_names[name],
361
+ )
362
+
363
+ @classmethod
364
+ def write_empty_workflow(
365
+ cls,
366
+ app: BaseApp,
367
+ *,
368
+ template_js: TemplateMeta,
369
+ template_components_js: dict[str, Any],
370
+ wk_path: str,
371
+ fs: AbstractFileSystem,
372
+ name: str,
373
+ replaced_wk: str | None,
374
+ creation_info: StoreCreationInfo,
375
+ ts_fmt: str,
376
+ ts_name_fmt: str,
377
+ ) -> None:
378
+ """
379
+ Write an empty persistent workflow.
380
+ """
381
+ fs.mkdir(wk_path)
382
+ submissions: list[None] = []
383
+ parameters: dict[str, dict[None, None]] = {
384
+ "data": {},
385
+ "sources": {},
386
+ }
387
+ metadata: Metadata = {
388
+ "name": name,
389
+ "ts_fmt": ts_fmt,
390
+ "ts_name_fmt": ts_name_fmt,
391
+ "creation_info": creation_info,
392
+ "template_components": template_components_js,
393
+ "template": template_js,
394
+ "tasks": [],
395
+ "elements": [],
396
+ "iters": [],
397
+ "num_added_tasks": 0,
398
+ "loops": [],
399
+ }
400
+ runs: dict[str, list] = {
401
+ "runs": [],
402
+ "run_dirs": [],
403
+ }
404
+ if replaced_wk:
405
+ metadata["replaced_workflow"] = replaced_wk
406
+
407
+ cls._get_store_resource(app, "metadata", wk_path, fs)._dump(metadata)
408
+ cls._get_store_resource(app, "parameters", wk_path, fs)._dump(parameters)
409
+ cls._get_store_resource(app, "submissions", wk_path, fs)._dump(submissions)
410
+ cls._get_store_resource(app, "runs", wk_path, fs)._dump(runs)
411
+
412
+ def _append_tasks(self, tasks: Iterable[StoreTask]):
413
+ with self.using_resource("metadata", action="update") as md:
414
+ assert "tasks" in md and "template" in md and "num_added_tasks" in md
415
+ for task in tasks:
416
+ idx, wk_task_i, task_i = task.encode()
417
+ md["tasks"].insert(idx, cast("TaskMeta", wk_task_i))
418
+ md["template"]["tasks"].insert(idx, task_i)
419
+ md["num_added_tasks"] += 1
420
+
421
+ def _append_loops(self, loops: dict[int, LoopDescriptor]):
422
+ with self.using_resource("metadata", action="update") as md:
423
+ assert "loops" in md and "template" in md
424
+ for _, loop in loops.items():
425
+ md["loops"].append(
426
+ {
427
+ "num_added_iterations": loop["num_added_iterations"],
428
+ "iterable_parameters": loop["iterable_parameters"],
429
+ "output_parameters": loop["output_parameters"],
430
+ "parents": loop["parents"],
431
+ }
432
+ )
433
+ md["template"]["loops"].append(loop["loop_template"])
434
+
435
+ def _append_submissions(self, subs: dict[int, Mapping[str, JSONed]]):
436
+ with self.using_resource("submissions", action="update") as subs_res:
437
+ subs_res.extend(subs.values())
438
+
439
+ def _append_task_element_IDs(self, task_ID: int, elem_IDs: list[int]):
440
+ with self.using_resource("metadata", action="update") as md:
441
+ assert "tasks" in md
442
+ md["tasks"][task_ID]["element_IDs"].extend(elem_IDs)
443
+
444
+ def _append_elements(self, elems: Sequence[JsonStoreElement]):
445
+ with self.using_resource("metadata", action="update") as md:
446
+ assert "elements" in md
447
+ md["elements"].extend(elem.encode(None) for elem in elems)
448
+
449
+ def _append_element_sets(self, task_id: int, es_js: Sequence[Mapping]):
450
+ task_idx = self._get_task_id_to_idx_map()[task_id]
451
+ with self.using_resource("metadata", "update") as md:
452
+ assert "template" in md
453
+ md["template"]["tasks"][task_idx]["element_sets"].extend(es_js)
454
+
455
+ def _append_elem_iter_IDs(self, elem_ID: int, iter_IDs: Iterable[int]):
456
+ with self.using_resource("metadata", action="update") as md:
457
+ assert "elements" in md
458
+ md["elements"][elem_ID]["iteration_IDs"].extend(iter_IDs)
459
+
460
+ def _append_elem_iters(self, iters: Sequence[JsonStoreElementIter]):
461
+ with self.using_resource("metadata", action="update") as md:
462
+ assert "iters" in md
463
+ md["iters"].extend(it.encode(None) for it in iters)
464
+
465
+ def _append_elem_iter_EAR_IDs(
466
+ self, iter_ID: int, act_idx: int, EAR_IDs: Sequence[int]
467
+ ):
468
+ with self.using_resource("metadata", action="update") as md:
469
+ assert "iters" in md
470
+ md["iters"][iter_ID].setdefault("EAR_IDs", {}).setdefault(act_idx, []).extend(
471
+ EAR_IDs
472
+ )
473
+
474
+ def _update_elem_iter_EARs_initialised(self, iter_ID: int):
475
+ with self.using_resource("metadata", action="update") as md:
476
+ assert "iters" in md
477
+ md["iters"][iter_ID]["EARs_initialised"] = True
478
+
479
+ def _update_at_submit_metadata(self, at_submit_metadata: dict[int, dict[str, Any]]):
480
+ with self.using_resource("submissions", action="update") as subs_res:
481
+ for sub_idx, metadata_i in at_submit_metadata.items():
482
+ sub = subs_res[sub_idx]
483
+ assert isinstance(sub, dict)
484
+ for dt_str, parts_j in metadata_i["submission_parts"].items():
485
+ sub["at_submit_metadata"]["submission_parts"][dt_str] = parts_j
486
+
487
+ def _update_loop_index(self, loop_indices: dict[int, dict[str, int]]):
488
+ with self.using_resource("metadata", action="update") as md:
489
+ assert "iters" in md
490
+ for iter_ID, loop_idx in loop_indices.items():
491
+ md["iters"][iter_ID]["loop_idx"].update(loop_idx)
492
+
493
+ def _update_loop_num_iters(self, index: int, num_iters: list[list[list[int] | int]]):
494
+ with self.using_resource("metadata", action="update") as md:
495
+ assert "loops" in md
496
+ md["loops"][index]["num_added_iterations"] = num_iters
497
+
498
+ def _update_loop_parents(self, index: int, parents: list[str]):
499
+ with self.using_resource("metadata", action="update") as md:
500
+ assert "loops" in md
501
+ md["loops"][index]["parents"] = parents
502
+
503
+ def _update_iter_data_indices(self, iter_data_indices: dict[int, DataIndex]):
504
+ with self.using_resource("metadata", action="update") as md:
505
+ assert "iters" in md
506
+ for iter_ID, dat_idx in iter_data_indices.items():
507
+ md["iters"][iter_ID]["data_idx"].update(dat_idx)
508
+
509
+ def _update_run_data_indices(self, run_data_indices: dict[int, DataIndex]):
510
+ with self.using_resource("runs", action="update") as md:
511
+ assert "runs" in md
512
+ for run_ID, dat_idx in run_data_indices.items():
513
+ md["runs"][run_ID]["data_idx"].update(dat_idx)
514
+
515
+ def _append_EARs(self, EARs: Sequence[JsonStoreEAR]):
516
+ with self.using_resource("runs", action="update") as md:
517
+ assert "runs" in md
518
+ assert "run_dirs" in md
519
+ md["runs"].extend(i.encode(self.ts_fmt, None) for i in EARs)
520
+ md["run_dirs"].extend([None] * len(EARs))
521
+
522
+ def _set_run_dirs(self, run_dir_arr: np.ndarray, run_idx: np.ndarray):
523
+ with self.using_resource("runs", action="update") as md:
524
+ assert "run_dirs" in md
525
+ dirs_lst = md["run_dirs"]
526
+ for idx, r_idx in enumerate(run_idx):
527
+ dirs_lst[r_idx] = run_dir_arr[idx].item()
528
+ md["run_dirs"] = dirs_lst
529
+
530
+ def _update_EAR_submission_data(self, sub_data: Mapping[int, tuple[int, int | None]]):
531
+ with self.using_resource("runs", action="update") as md:
532
+ assert "runs" in md
533
+ for EAR_ID_i, (sub_idx_i, cmd_file_ID) in sub_data.items():
534
+ md["runs"][EAR_ID_i]["submission_idx"] = sub_idx_i
535
+ md["runs"][EAR_ID_i]["commands_file_ID"] = cmd_file_ID
536
+
537
+ def _update_EAR_start(
538
+ self,
539
+ run_starts: dict[int, tuple[datetime, dict[str, Any] | None, str, int | None]],
540
+ ):
541
+ with self.using_resource("runs", action="update") as md:
542
+ assert "runs" in md
543
+ for run_id, (s_time, s_snap, s_hn, port_number) in run_starts.items():
544
+ md["runs"][run_id]["start_time"] = s_time.strftime(self.ts_fmt)
545
+ md["runs"][run_id]["snapshot_start"] = s_snap
546
+ md["runs"][run_id]["run_hostname"] = s_hn
547
+ md["runs"][run_id]["port_number"] = port_number
548
+
549
+ def _update_EAR_end(
550
+ self, run_ends: dict[int, tuple[datetime, dict[str, Any] | None, int, bool]]
551
+ ):
552
+ with self.using_resource("runs", action="update") as md:
553
+ assert "runs" in md
554
+ for run_id, (e_time, e_snap, ext_code, success) in run_ends.items():
555
+ md["runs"][run_id]["end_time"] = e_time.strftime(self.ts_fmt)
556
+ md["runs"][run_id]["snapshot_end"] = e_snap
557
+ md["runs"][run_id]["exit_code"] = ext_code
558
+ md["runs"][run_id]["success"] = success
559
+
560
+ def _update_EAR_skip(self, skips: dict[int, int]):
561
+ with self.using_resource("runs", action="update") as md:
562
+ assert "runs" in md
563
+ for run_ID, reason in skips.items():
564
+ md["runs"][run_ID]["skip"] = reason
565
+
566
+ def _update_js_metadata(self, js_meta: dict[int, dict[int, dict[str, Any]]]):
567
+ with self.using_resource("submissions", action="update") as sub_res:
568
+ for sub_idx, all_js_md in js_meta.items():
569
+ sub = cast("dict[str, list[dict[str, Any]]]", sub_res[sub_idx])
570
+ for js_idx, js_meta_i in all_js_md.items():
571
+ self.logger.info(
572
+ f"updating jobscript metadata for (sub={sub_idx}, js={js_idx}): "
573
+ f"{js_meta_i!r}."
574
+ )
575
+ _at_submit_md = {
576
+ k: js_meta_i.pop(k)
577
+ for k in JOBSCRIPT_SUBMIT_TIME_KEYS
578
+ if k in js_meta_i
579
+ }
580
+ sub["jobscripts"][js_idx].update(**js_meta_i)
581
+ sub["jobscripts"][js_idx]["at_submit_metadata"].update(
582
+ **_at_submit_md
583
+ )
584
+
585
+ def _append_parameters(self, params: Sequence[StoreParameter]):
586
+ self._ensure_all_encoders()
587
+ with self.using_resource("parameters", "update") as params_u:
588
+ for param_i in params:
589
+ params_u["data"][str(param_i.id_)] = param_i.encode()
590
+ params_u["sources"][str(param_i.id_)] = param_i.source
591
+
592
+ def _set_parameter_values(self, set_parameters: dict[int, tuple[Any, bool]]):
593
+ """Set multiple unset persistent parameters."""
594
+ self._ensure_all_encoders()
595
+ param_objs = self._get_persistent_parameters(set_parameters)
596
+ with self.using_resource("parameters", "update") as params:
597
+ for param_id, (value, is_file) in set_parameters.items():
598
+ param_i = param_objs[param_id]
599
+ if is_file:
600
+ param_i = param_i.set_file(value)
601
+ else:
602
+ param_i = param_i.set_data(value)
603
+ params["data"][str(param_id)] = param_i.encode()
604
+
605
+ def _update_parameter_sources(self, sources: Mapping[int, ParamSource]):
606
+ """Update the sources of multiple persistent parameters."""
607
+ param_objs = self._get_persistent_parameters(sources)
608
+ with self.using_resource("parameters", "update") as params:
609
+ # no need to update data array:
610
+ for p_id, src_i in sources.items():
611
+ param_i = param_objs[p_id]
612
+ new_src_i = update_param_source_dict(param_i.source, src_i)
613
+ params["sources"][str(p_id)] = new_src_i
614
+
615
+ def _update_template_components(self, tc: dict[str, Any]):
616
+ with self.using_resource("metadata", "update") as md:
617
+ md["template_components"] = tc
618
+
619
+ def _get_num_persistent_tasks(self) -> int:
620
+ """Get the number of persistent tasks."""
621
+ if self.use_cache and self.num_tasks_cache is not None:
622
+ num = self.num_tasks_cache
623
+ else:
624
+ with self.using_resource("metadata", action="read") as md:
625
+ assert "tasks" in md
626
+ num = len(md["tasks"])
627
+ if self.use_cache and self.num_tasks_cache is None:
628
+ self.num_tasks_cache = num
629
+ return num
630
+
631
+ def _get_num_persistent_loops(self) -> int:
632
+ """Get the number of persistent loops."""
633
+ with self.using_resource("metadata", action="read") as md:
634
+ assert "loops" in md
635
+ return len(md["loops"])
636
+
637
+ def _get_num_persistent_submissions(self) -> int:
638
+ """Get the number of persistent submissions."""
639
+ with self.using_resource("submissions", "read") as subs_res:
640
+ return len(subs_res)
641
+
642
+ def _get_num_persistent_elements(self) -> int:
643
+ """Get the number of persistent elements."""
644
+ with self.using_resource("metadata", action="read") as md:
645
+ assert "elements" in md
646
+ return len(md["elements"])
647
+
648
+ def _get_num_persistent_elem_iters(self) -> int:
649
+ """Get the number of persistent element iterations."""
650
+ with self.using_resource("metadata", action="read") as md:
651
+ assert "iters" in md
652
+ return len(md["iters"])
653
+
654
+ def _get_num_persistent_EARs(self) -> int:
655
+ """Get the number of persistent EARs."""
656
+ if self.use_cache and self.num_EARs_cache is not None:
657
+ num = self.num_EARs_cache
658
+ else:
659
+ with self.using_resource("runs", action="read") as md:
660
+ assert "runs" in md
661
+ num = len(md["runs"])
662
+ if self.use_cache and self.num_EARs_cache is None:
663
+ self.num_EARs_cache = num
664
+ return num
665
+
666
+ def _get_num_persistent_parameters(self) -> int:
667
+ if self.use_cache and self.num_params_cache is not None:
668
+ num = self.num_params_cache
669
+ else:
670
+ with self.using_resource("parameters", "read") as params:
671
+ assert "data" in params
672
+ num = len(params["data"])
673
+ if self.use_cache and self.num_params_cache is None:
674
+ self.num_params_cache = num
675
+ return num
676
+
677
+ def _get_num_persistent_added_tasks(self) -> int:
678
+ with self.using_resource("metadata", "read") as md:
679
+ assert "num_added_tasks" in md
680
+ return md["num_added_tasks"]
681
+
682
+ @classmethod
683
+ def make_test_store_from_spec(
684
+ cls,
685
+ app: BaseApp,
686
+ spec,
687
+ dir=None,
688
+ path="test_store.json",
689
+ overwrite=False,
690
+ ts_fmt="%d/%m/%Y, %H:%M:%S", # FIXME: use the right default timestamp format
691
+ ):
692
+ """Generate an store for testing purposes."""
693
+
694
+ tasks_, elems, elem_iters, EARs = super().prepare_test_store_from_spec(spec)
695
+
696
+ path_ = Path(path).resolve()
697
+ tasks = [JsonStoreTask(**task_info).encode() for task_info in tasks_]
698
+ elements_ = [JsonStoreElement(**elem_info).encode(None) for elem_info in elems]
699
+ elem_iters_ = [
700
+ JsonStoreElementIter(**it_info).encode(None) for it_info in elem_iters
701
+ ]
702
+ EARs_ = [JsonStoreEAR(**ear_info).encode(ts_fmt, None) for ear_info in EARs]
703
+
704
+ persistent_data = {
705
+ "tasks": tasks,
706
+ "elements": elements_,
707
+ "iters": elem_iters_,
708
+ "runs": EARs_,
709
+ }
710
+
711
+ path_ = Path(dir or "", path_)
712
+ with path_.open("wt") as fp:
713
+ json.dump(persistent_data, fp, indent=2)
714
+
715
+ return cls(app=app, workflow=None, path=path_, fs=filesystem("file"))
716
+
717
+ def _get_persistent_template_components(self) -> dict[str, Any]:
718
+ with self.using_resource("metadata", "read") as md:
719
+ assert "template_components" in md
720
+ return md["template_components"]
721
+
722
+ def _get_persistent_template(self) -> dict[str, JSONed]:
723
+ with self.using_resource("metadata", "read") as md:
724
+ assert "template" in md
725
+ return cast("dict[str, JSONed]", md["template"])
726
+
727
+ def _get_persistent_tasks(self, id_lst: Iterable[int]) -> dict[int, JsonStoreTask]:
728
+ tasks, id_lst = self._get_cached_persistent_tasks(id_lst)
729
+ if id_lst:
730
+ with self.using_resource("metadata", action="read") as md:
731
+ assert "tasks" in md
732
+ new_tasks = {
733
+ i["id_"]: JsonStoreTask.decode({**i, "index": idx})
734
+ for idx, i in enumerate(cast("Sequence[TaskMeta]", md["tasks"]))
735
+ if id_lst is None or i["id_"] in id_lst
736
+ }
737
+ self.task_cache.update(new_tasks)
738
+ tasks.update(new_tasks)
739
+ return tasks
740
+
741
+ def _get_persistent_loops(
742
+ self, id_lst: Iterable[int] | None = None
743
+ ) -> dict[int, LoopDescriptor]:
744
+ with self.using_resource("metadata", "read") as md:
745
+ assert "loops" in md
746
+ return {
747
+ idx: cast("LoopDescriptor", i)
748
+ for idx, i in enumerate(md["loops"])
749
+ if id_lst is None or idx in id_lst
750
+ }
751
+
752
+ def _get_persistent_submissions(
753
+ self, id_lst: Iterable[int] | None = None
754
+ ) -> dict[int, Mapping[str, JSONed]]:
755
+ with self.using_resource("submissions", "read") as sub_res:
756
+ subs_dat = copy.deepcopy(
757
+ {
758
+ idx: i
759
+ for idx, i in enumerate(sub_res)
760
+ if id_lst is None or idx in id_lst
761
+ }
762
+ )
763
+ # cast jobscript `task_elements` keys:
764
+ for sub in subs_dat.values():
765
+ js: dict[str, Any]
766
+ assert isinstance(sub, dict)
767
+ for js in sub["jobscripts"]:
768
+ blk: dict[str, Any]
769
+ assert isinstance(js, dict)
770
+ for blk in js["blocks"]:
771
+ for key in list(te := blk["task_elements"]):
772
+ te[int(key)] = te.pop(key)
773
+
774
+ return subs_dat
775
+
776
+ def _get_persistent_elements(
777
+ self, id_lst: Iterable[int]
778
+ ) -> dict[int, JsonStoreElement]:
779
+ elems, id_lst_ = self._get_cached_persistent_elements(id_lst)
780
+ if id_lst_:
781
+ # could convert `id_lst` to e.g. slices if more efficient for a given store
782
+ with self.using_resource("metadata", action="read") as md:
783
+ try:
784
+ if "elements" not in md:
785
+ raise KeyError
786
+ elem_dat = {id_: md["elements"][id_] for id_ in id_lst_}
787
+ except KeyError:
788
+ raise MissingStoreElementError(id_lst_)
789
+ new_elems = {
790
+ k: JsonStoreElement.decode(v, None) for k, v in elem_dat.items()
791
+ }
792
+ self.element_cache.update(new_elems)
793
+ elems.update(new_elems)
794
+ return elems
795
+
796
+ def _get_persistent_element_iters(
797
+ self, id_lst: Iterable[int]
798
+ ) -> dict[int, JsonStoreElementIter]:
799
+ iters, id_lst_ = self._get_cached_persistent_element_iters(id_lst)
800
+ if id_lst_:
801
+ with self.using_resource("metadata", action="read") as md:
802
+ try:
803
+ if "iters" not in md:
804
+ raise KeyError
805
+ iter_dat = {id_: md["iters"][id_] for id_ in id_lst_}
806
+ except KeyError:
807
+ raise MissingStoreElementIterationError(id_lst_)
808
+ new_iters = {
809
+ k: JsonStoreElementIter.decode(v, None) for k, v in iter_dat.items()
810
+ }
811
+ self.element_iter_cache.update(new_iters)
812
+ iters.update(new_iters)
813
+ return iters
814
+
815
+ def _get_persistent_EARs(self, id_lst: Iterable[int]) -> dict[int, JsonStoreEAR]:
816
+ runs, id_lst_ = self._get_cached_persistent_EARs(id_lst)
817
+ if id_lst_:
818
+ with self.using_resource("runs", action="read") as md:
819
+ try:
820
+ if "runs" not in md:
821
+ raise KeyError
822
+ EAR_dat = {id_: md["runs"][id_] for id_ in id_lst_}
823
+ except KeyError:
824
+ raise MissingStoreEARError(id_lst_)
825
+ new_runs = {
826
+ k: JsonStoreEAR.decode(v, self.ts_fmt, None)
827
+ for k, v in EAR_dat.items()
828
+ }
829
+ self.EAR_cache.update(new_runs)
830
+ runs.update(new_runs)
831
+ return runs
832
+
833
+ def _get_persistent_parameters(
834
+ self, id_lst: Iterable[int], **kwargs
835
+ ) -> Mapping[int, StoreParameter]:
836
+ self._ensure_all_decoders()
837
+ params, id_lst_ = self._get_cached_persistent_parameters(id_lst)
838
+ if id_lst_:
839
+ with self.using_resource("parameters", "read") as params_:
840
+ try:
841
+ param_dat = {id_: params_["data"][str(id_)] for id_ in id_lst_}
842
+ src_dat = {id_: params_["sources"][str(id_)] for id_ in id_lst_}
843
+ except KeyError:
844
+ raise MissingParameterData(id_lst_)
845
+
846
+ new_params = {
847
+ k: StoreParameter.decode(id_=k, data=v, source=src_dat[k])
848
+ for k, v in param_dat.items()
849
+ }
850
+ self.parameter_cache.update(new_params)
851
+ params.update(new_params)
852
+ return params
853
+
854
+ def _get_persistent_param_sources(
855
+ self, id_lst: Iterable[int]
856
+ ) -> dict[int, ParamSource]:
857
+ sources, id_lst_ = self._get_cached_persistent_param_sources(id_lst)
858
+ if id_lst_:
859
+ with self.using_resource("parameters", "read") as params:
860
+ try:
861
+ new_sources = {id_: params["sources"][str(id_)] for id_ in id_lst_}
862
+ except KeyError:
863
+ raise MissingParameterData(id_lst_)
864
+ self.param_sources_cache.update(new_sources)
865
+ sources.update(new_sources)
866
+ return sources
867
+
868
+ def _get_persistent_parameter_set_status(
869
+ self, id_lst: Iterable[int]
870
+ ) -> dict[int, bool]:
871
+ with self.using_resource("parameters", "read") as params:
872
+ try:
873
+ param_dat = {id_: params["data"][str(id_)] for id_ in id_lst}
874
+ except KeyError:
875
+ raise MissingParameterData(id_lst)
876
+ return {k: v is not None for k, v in param_dat.items()}
877
+
878
+ def _get_persistent_parameter_IDs(self) -> list[int]:
879
+ with self.using_resource("parameters", "read") as params:
880
+ return [int(i) for i in params["data"]]
881
+
882
+ def get_ts_fmt(self) -> str:
883
+ """
884
+ Get the format for timestamps.
885
+ """
886
+ with self.using_resource("metadata", action="read") as md:
887
+ assert "ts_fmt" in md
888
+ return md["ts_fmt"]
889
+
890
+ def get_ts_name_fmt(self) -> str:
891
+ """
892
+ Get the format for timestamps to use in names.
893
+ """
894
+ with self.using_resource("metadata", action="read") as md:
895
+ assert "ts_name_fmt" in md
896
+ return md["ts_name_fmt"]
897
+
898
+ def get_creation_info(self) -> StoreCreationInfo:
899
+ """
900
+ Get information about the creation of the workflow.
901
+ """
902
+ with self.using_resource("metadata", action="read") as md:
903
+ assert "creation_info" in md
904
+ return copy.deepcopy(md["creation_info"])
905
+
906
+ def get_name(self) -> str:
907
+ """
908
+ Get the name of the workflow.
909
+ """
910
+ with self.using_resource("metadata", action="read") as md:
911
+ assert "name" in md
912
+ return md["name"]
913
+
914
+ def zip(
915
+ self,
916
+ path: str = ".",
917
+ log: str | None = None,
918
+ overwrite=False,
919
+ include_execute=False,
920
+ include_rechunk_backups=False,
921
+ ) -> str:
922
+ raise TypeError("unsupported operation: zipping-json")
923
+
924
+ def unzip(self, path: str = ".", log: str | None = None) -> str:
925
+ raise TypeError("unsupported operation: unzipping-json")
926
+
927
+ def rechunk_parameter_base(
928
+ self,
929
+ chunk_size: int | None = None,
930
+ backup: bool = True,
931
+ status: bool = True,
932
+ ) -> Any:
933
+ raise TypeError("unsupported operation: rechunk-json")
934
+
935
+ def rechunk_runs(
936
+ self,
937
+ chunk_size: int | None = None,
938
+ backup: bool = True,
939
+ status: bool = True,
940
+ ) -> Any:
941
+ raise TypeError("unsupported operation: rechunk-json")
942
+
943
+ def get_dirs_array(self) -> NDArray:
944
+ """
945
+ Retrieve the run directories array.
946
+ """
947
+ with self.using_resource("runs", action="read") as md:
948
+ dirs_lst = md["run_dirs"]
949
+ dirs_arr = np.zeros(len(dirs_lst), dtype=RUN_DIR_ARR_DTYPE)
950
+ dirs_arr[:] = RUN_DIR_ARR_FILL
951
+ for idx, i in enumerate(dirs_lst):
952
+ if i is not None:
953
+ dirs_arr[idx] = tuple(i)
954
+ return dirs_arr