hpcflow-new2 0.2.0a176__py3-none-any.whl → 0.2.0a178__py3-none-any.whl

hpcflow/sdk/core/workflow.py

@@ -25,6 +25,7 @@ from hpcflow.sdk.core import (
     ABORT_EXIT_CODE,
 )
 from hpcflow.sdk.core.actions import EARStatus
+from hpcflow.sdk.core.loop_cache import LoopCache
 from hpcflow.sdk.log import TimeIt
 from hpcflow.sdk.persistence import store_cls_from_str, DEFAULT_STORE_FORMAT
 from hpcflow.sdk.persistence.base import TEMPLATE_COMP_TYPES, AnySEAR
@@ -41,6 +42,7 @@ from hpcflow.sdk.submission.schedulers.direct import DirectScheduler
 from hpcflow.sdk.typing import PathLike
 from hpcflow.sdk.core.json_like import ChildObjectSpec, JSONLike
 from .utils import (
+    nth_key,
     read_JSON_file,
     read_JSON_string,
     read_YAML_str,
@@ -625,19 +627,28 @@ class Workflow:
             )
             with wk._store.cached_load():
                 with wk.batch_update(is_workflow_creation=True):
-                    for idx, task in enumerate(template.tasks):
+                    with wk._store.cache_ctx():
+                        for idx, task in enumerate(template.tasks):
+                            if status:
+                                status.update(
+                                    f"Adding task {idx + 1}/{len(template.tasks)} "
+                                    f"({task.name!r})..."
+                                )
+                            wk._add_task(task)
                         if status:
                             status.update(
-                                f"Adding task {idx + 1}/{len(template.tasks)} "
-                                f"({task.name!r})..."
+                                f"Preparing to add {len(template.loops)} loops..."
                             )
-                        wk._add_task(task)
-                    for idx, loop in enumerate(template.loops):
-                        if status:
-                            status.update(
-                                f"Adding loop {idx + 1}/" f"{len(template.loops)}..."
-                            )
-                        wk._add_loop(loop)
+                        if template.loops:
+                            # TODO: if loop with non-initialisable actions, will fail
+                            cache = LoopCache.build(workflow=wk, loops=template.loops)
+                            for idx, loop in enumerate(template.loops):
+                                if status:
+                                    status.update(
+                                        f"Adding loop {idx + 1}/"
+                                        f"{len(template.loops)} ({loop.name!r})"
+                                    )
+                                wk._add_loop(loop, cache=cache, status=status)
         except Exception:
             if status:
                 status.stop()
@@ -1101,7 +1112,7 @@ class Workflow:
 
     @TimeIt.decorator
     def _add_empty_loop(
-        self, loop: app.Loop
+        self, loop: app.Loop, cache: LoopCache
     ) -> Tuple[app.WorkflowLoop, List[app.ElementIteration]]:
         """Add a new loop (zeroth iterations only) to the workflow."""
 
@@ -1114,15 +1125,15 @@ class Workflow:
         self.template._add_empty_loop(loop_c)
 
         # all these element iterations will be initialised for the new loop:
-        iters = self.get_element_iterations_of_tasks(loop_c.task_insert_IDs)
-        iter_IDs = [i.id_ for i in iters]
+        iter_IDs = cache.get_iter_IDs(loop_c)
+        iter_loop_idx = cache.get_iter_loop_indices(iter_IDs)
 
         # create and insert a new WorkflowLoop:
         new_loop = self.app.WorkflowLoop.new_empty_loop(
             index=new_index,
             workflow=self,
             template=loop_c,
-            iterations=iters,
+            iter_loop_idx=iter_loop_idx,
         )
         self.loops.add_object(new_loop)
         wk_loop = self.loops[new_index]
@@ -1144,15 +1155,28 @@ class Workflow:
 
         self._pending["loops"].append(new_index)
 
+        # update cache loop indices:
+        cache.update_loop_indices(new_loop_name=loop_c.name, iter_IDs=iter_IDs)
+
         return wk_loop
 
     @TimeIt.decorator
-    def _add_loop(self, loop: app.Loop) -> None:
-        new_wk_loop = self._add_empty_loop(loop)
+    def _add_loop(
+        self, loop: app.Loop, cache: Optional[Dict] = None, status: Optional[Any] = None
+    ) -> None:
+        if not cache:
+            cache = LoopCache.build(workflow=self, loops=[loop])
+        new_wk_loop = self._add_empty_loop(loop, cache)
         if loop.num_iterations is not None:
             # fixed number of iterations, so add remaining N > 0 iterations:
-            for _ in range(loop.num_iterations - 1):
-                new_wk_loop.add_iteration()
+            if status:
+                status_prev = status.status
+            for iter_idx in range(loop.num_iterations - 1):
+                if status:
+                    status.update(
+                        f"{status_prev}: iteration {iter_idx + 2}/{loop.num_iterations}."
+                    )
+                new_wk_loop.add_iteration(cache=cache)
 
     def add_loop(self, loop: app.Loop) -> None:
         """Add a loop to a subset of workflow tasks."""
@@ -1326,6 +1350,7 @@ class Workflow:
                     iters.append(iter_i)
         return iters
 
+    @TimeIt.decorator
     def get_elements_from_IDs(self, id_lst: Iterable[int]) -> List[app.Element]:
         """Return element objects from a list of IDs."""
 
@@ -1334,6 +1359,7 @@ class Workflow:
         task_IDs = [i.task_ID for i in store_elems]
         store_tasks = self._store.get_tasks_by_IDs(task_IDs)
 
+        element_idx_by_task = defaultdict(set)
         index_paths = []
         for el, tk in zip(store_elems, store_tasks):
             elem_idx = tk.element_IDs.index(el.id_)
@@ -1343,15 +1369,23 @@ class Workflow:
                     "task_idx": tk.index,
                 }
             )
+            element_idx_by_task[tk.index].add(elem_idx)
+
+        elements_by_task = {}
+        for task_idx, elem_idx in element_idx_by_task.items():
+            task = self.tasks[task_idx]
+            elements_by_task[task_idx] = dict(
+                zip(elem_idx, task.elements[list(elem_idx)])
+            )
 
         objs = []
         for idx_dat in index_paths:
-            task = self.tasks[idx_dat["task_idx"]]
-            elem = task.elements[idx_dat["elem_idx"]]
+            elem = elements_by_task[idx_dat["task_idx"]][idx_dat["elem_idx"]]
             objs.append(elem)
 
         return objs
 
+    @TimeIt.decorator
     def get_element_iterations_from_IDs(
         self, id_lst: Iterable[int]
     ) -> List[app.ElementIteration]:
@@ -1365,6 +1399,8 @@ class Workflow:
         task_IDs = [i.task_ID for i in store_elems]
         store_tasks = self._store.get_tasks_by_IDs(task_IDs)
 
+        element_idx_by_task = defaultdict(set)
+
         index_paths = []
         for it, el, tk in zip(store_iters, store_elems, store_tasks):
             iter_idx = el.iteration_IDs.index(it.id_)
@@ -1376,11 +1412,18 @@ class Workflow:
                     "task_idx": tk.index,
                 }
             )
+            element_idx_by_task[tk.index].add(elem_idx)
+
+        elements_by_task = {}
+        for task_idx, elem_idx in element_idx_by_task.items():
+            task = self.tasks[task_idx]
+            elements_by_task[task_idx] = dict(
+                zip(elem_idx, task.elements[list(elem_idx)])
+            )
 
         objs = []
         for idx_dat in index_paths:
-            task = self.tasks[idx_dat["task_idx"]]
-            elem = task.elements[idx_dat["elem_idx"]]
+            elem = elements_by_task[idx_dat["task_idx"]][idx_dat["elem_idx"]]
             iter_ = elem.iterations[idx_dat["iter_idx"]]
             objs.append(iter_)
 
@@ -1653,7 +1696,14 @@ class Workflow:
 
         return wk
 
-    def zip(self, path=".", log=None, overwrite=False) -> str:
+    def zip(
+        self,
+        path=".",
+        log=None,
+        overwrite=False,
+        include_execute=False,
+        include_rechunk_backups=False,
+    ) -> str:
         """
         Parameters
         ----------
@@ -1662,7 +1712,13 @@ class Workflow:
             directory, the zip file will be created within this directory. Otherwise,
             this path is assumed to be the full file path to the new zip file.
         """
-        return self._store.zip(path=path, log=log, overwrite=overwrite)
+        return self._store.zip(
+            path=path,
+            log=log,
+            overwrite=overwrite,
+            include_execute=include_execute,
+            include_rechunk_backups=include_rechunk_backups,
+        )
 
     def unzip(self, path=".", log=None) -> str:
         """
@@ -2900,6 +2956,34 @@ class Workflow:
                     final_runs[loop_name].append(final[0])
         return dict(final_runs)
 
+    def rechunk_runs(
+        self,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        self._store.rechunk_runs(chunk_size=chunk_size, backup=backup, status=status)
+
+    def rechunk_parameter_base(
+        self,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        self._store.rechunk_parameter_base(
+            chunk_size=chunk_size, backup=backup, status=status
+        )
+
+    def rechunk(
+        self,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        """Rechunk metadata/runs and parameters/base arrays."""
+        self.rechunk_runs(chunk_size=chunk_size, backup=backup, status=status)
+        self.rechunk_parameter_base(chunk_size=chunk_size, backup=backup, status=status)
+
 
 @dataclass
 class WorkflowBlueprint:

hpcflow/sdk/persistence/base.py

@@ -716,6 +716,11 @@ class PersistentStore(ABC):
         """Cache for number of persistent tasks."""
         return self._cache["num_tasks"]
 
+    @property
+    def num_EARs_cache(self):
+        """Cache for total number of persistent EARs."""
+        return self._cache["num_EARs"]
+
     @property
     def param_sources_cache(self):
         """Cache for persistent parameter sources."""
@@ -730,6 +735,10 @@ class PersistentStore(ABC):
     def num_tasks_cache(self, value):
         self._cache["num_tasks"] = value
 
+    @num_EARs_cache.setter
+    def num_EARs_cache(self, value):
+        self._cache["num_EARs"] = value
+
     def _reset_cache(self):
         self._cache = {
             "tasks": {},
@@ -739,6 +748,7 @@ class PersistentStore(ABC):
             "param_sources": {},
             "num_tasks": None,
             "parameters": {},
+            "num_EARs": None,
         }
 
     @contextlib.contextmanager
@@ -873,6 +883,7 @@ class PersistentStore(ABC):
         """Get the total number of persistent and pending element iterations."""
         return self._get_num_persistent_elem_iters() + len(self._pending.add_elem_iters)
 
+    @TimeIt.decorator
     def _get_num_total_EARs(self):
         """Get the total number of persistent and pending EARs."""
         return self._get_num_persistent_EARs() + len(self._pending.add_EARs)
@@ -1296,9 +1307,11 @@ class PersistentStore(ABC):
             self.save()
 
     @TimeIt.decorator
-    def update_param_source(self, param_id: int, source: Dict, save: bool = True) -> None:
-        self.logger.debug(f"Updating parameter ID {param_id!r} source to {source!r}.")
-        self._pending.update_param_sources[param_id] = source
+    def update_param_source(
+        self, param_sources: Dict[int, Dict], save: bool = True
+    ) -> None:
+        self.logger.debug(f"Updating parameter sources with {param_sources!r}.")
+        self._pending.update_param_sources.update(param_sources)
         if save:
             self.save()
 

hpcflow/sdk/persistence/json.py

@@ -303,12 +303,13 @@ class JSONPersistentStore(PersistentStore):
 
     def _get_num_persistent_tasks(self) -> int:
         """Get the number of persistent tasks."""
-        if self.num_tasks_cache is not None:
+        if self.use_cache and self.num_tasks_cache is not None:
             num = self.num_tasks_cache
         else:
             with self.using_resource("metadata", action="read") as md:
                 num = len(md["tasks"])
-                self.num_tasks_cache = num
+        if self.use_cache and self.num_tasks_cache is None:
+            self.num_tasks_cache = num
         return num
 
     def _get_num_persistent_loops(self) -> int:
@@ -333,8 +334,14 @@ class JSONPersistentStore(PersistentStore):
 
     def _get_num_persistent_EARs(self) -> int:
         """Get the number of persistent EARs."""
-        with self.using_resource("metadata", action="read") as md:
-            return len(md["runs"])
+        if self.use_cache and self.num_EARs_cache is not None:
+            num = self.num_EARs_cache
+        else:
+            with self.using_resource("metadata", action="read") as md:
+                num = len(md["runs"])
+        if self.use_cache and self.num_EARs_cache is None:
+            self.num_EARs_cache = num
+        return num
 
     def _get_num_persistent_parameters(self):
         with self.using_resource("parameters", "read") as params:

hpcflow/sdk/persistence/pending.py

@@ -275,6 +275,7 @@ class PendingChanges:
             EAR_ids = list(self.add_EARs.keys())
             self.logger.debug(f"commit: adding pending EARs with IDs: {EAR_ids!r}")
             self.store._append_EARs(EARs)
+            self.store.num_EARs_cache = None  # invalidate cache
             # pending start/end times/snapshots, submission indices, and skips that belong
             # to pending EARs are now committed (accounted for in `get_EARs` above):
             self.set_EAR_submission_indices = {
@@ -408,6 +409,7 @@ class PendingChanges:
     @TimeIt.decorator
     def commit_loop_indices(self) -> None:
         """Make pending update to element iteration loop indices persistent."""
+        # TODO: batch up
         for iter_ID, loop_idx in self.update_loop_indices.items():
             self.logger.debug(
                 f"commit: updating loop indices of iteration ID {iter_ID!r} with "

hpcflow/sdk/persistence/zarr.py

@@ -5,6 +5,8 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
+import shutil
+import time
 from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union
 
 import numpy as np
@@ -774,9 +776,16 @@ class ZarrPersistentStore(PersistentStore):
         """Get the number of persistent element iterations."""
         return len(self._get_iters_arr())
 
+    @TimeIt.decorator
     def _get_num_persistent_EARs(self) -> int:
         """Get the number of persistent EARs."""
-        return len(self._get_EARs_arr())
+        if self.use_cache and self.num_EARs_cache is not None:
+            num = self.num_EARs_cache
+        else:
+            num = len(self._get_EARs_arr())
+        if self.use_cache and self.num_EARs_cache is None:
+            self.num_EARs_cache = num
+        return num
 
     def _get_num_persistent_parameters(self):
         return len(self._get_parameter_base_array())
@@ -1145,7 +1154,14 @@ class ZarrPersistentStore(PersistentStore):
         with self.using_resource("attrs", action="read") as attrs:
             return attrs["name"]
 
-    def zip(self, path=".", log=None, overwrite=False):
+    def zip(
+        self,
+        path=".",
+        log=None,
+        overwrite=False,
+        include_execute=False,
+        include_rechunk_backups=False,
+    ):
         """
         Parameters
         ----------
@@ -1181,16 +1197,120 @@ class ZarrPersistentStore(PersistentStore):
             add_pw_to="target_options",
         )
         dst_zarr_store = zarr.storage.FSStore(url="", fs=zfs)
+        excludes = []
+        if not include_execute:
+            excludes.append("execute")
+        if not include_rechunk_backups:
+            excludes.append("runs.bak")
+            excludes.append("base.bak")
+
         zarr.convenience.copy_store(
             src_zarr_store,
             dst_zarr_store,
-            excludes="execute",
+            excludes=excludes or None,
             log=log,
         )
         del zfs  # ZipFileSystem remains open for instance lifetime
         status.stop()
         return dst_path
 
+    def _rechunk_arr(
+        self,
+        arr,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        arr_path = Path(self.workflow.path) / arr.path
+        arr_name = arr.path.split("/")[-1]
+
+        if status:
+            console = Console()
+            status = console.status("Rechunking...")
+            status.start()
+        backup_time = None
+
+        if backup:
+            if status:
+                status.update("Backing up...")
+            backup_path = arr_path.with_suffix(".bak")
+            if backup_path.is_dir():
+                pass
+            else:
+                tic = time.perf_counter()
+                shutil.copytree(arr_path, backup_path)
+                toc = time.perf_counter()
+                backup_time = toc - tic
+
+        tic = time.perf_counter()
+        arr_rc_path = arr_path.with_suffix(".rechunked")
+        arr = zarr.open(arr_path)
+        if status:
+            status.update("Creating new array...")
+        arr_rc = zarr.create(
+            store=arr_rc_path,
+            shape=arr.shape,
+            chunks=arr.shape if chunk_size is None else chunk_size,
+            dtype=object,
+            object_codec=MsgPack(),
+        )
+        if status:
+            status.update("Copying data...")
+        data = np.empty(shape=arr.shape, dtype=object)
+        bad_data = []
+        for idx in range(len(arr)):
+            try:
+                data[idx] = arr[idx]
+            except RuntimeError:
+                # blosc decompression errors
+                bad_data.append(idx)
+                pass
+        arr_rc[:] = data
+
+        arr_rc.attrs.put(arr.attrs.asdict())
+
+        if status:
+            status.update("Deleting old array...")
+        shutil.rmtree(arr_path)
+
+        if status:
+            status.update("Moving new array into place...")
+        shutil.move(arr_rc_path, arr_path)
+
+        toc = time.perf_counter()
+        rechunk_time = toc - tic
+
+        if status:
+            status.stop()
+
+        if backup_time:
+            print(f"Time to backup {arr_name}: {backup_time:.1f} s")
+
+        print(f"Time to rechunk and move {arr_name}: {rechunk_time:.1f} s")
+
+        if bad_data:
+            print(f"Bad data at {arr_name} indices: {bad_data}.")
+
+        return arr_rc
+
+    def rechunk_parameter_base(
+        self,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        arr = self._get_parameter_base_array()
+        return self._rechunk_arr(arr, chunk_size, backup, status)
+
+    def rechunk_runs(
+        self,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        arr = self._get_EARs_arr()
+        return self._rechunk_arr(arr, chunk_size, backup, status)
+
 
 class ZarrZipPersistentStore(ZarrPersistentStore):
     """A store designed mainly as an archive format that can be uploaded to data
@@ -1250,3 +1370,12 @@ class ZarrZipPersistentStore(ZarrPersistentStore):
     def delete_no_confirm(self) -> None:
         # `ZipFileSystem.rm()` does not seem to be implemented.
         raise NotImplementedError()
+
+    def _rechunk_arr(
+        self,
+        arr,
+        chunk_size: Optional[int] = None,
+        backup: Optional[bool] = True,
+        status: Optional[bool] = True,
+    ):
+        raise NotImplementedError

hpcflow/tests/unit/test_persistence.py

@@ -1,4 +1,6 @@
-from importlib import resources
+from pathlib import Path
+import numpy as np
+import zarr
 import pytest
 from hpcflow.sdk.core.test_utils import make_test_data_YAML_workflow
 from hpcflow.sdk.persistence.base import StoreEAR, StoreElement, StoreElementIter
@@ -239,3 +241,118 @@ def test_make_zarr_store_no_compressor(null_config, tmp_path):
         store="zarr",
         store_kwargs={"compressor": None},
     )
+
+
+@pytest.mark.integration
+def test_zarr_rechunk_data_equivalent(null_config, tmp_path):
+    t1 = hf.Task(
+        schema=hf.task_schemas.test_t1_conditional_OS,
+        inputs={"p1": 101},
+        repeats=3,
+    )
+    wk = hf.Workflow.from_template_data(
+        tasks=[t1],
+        template_name="test_run_rechunk",
+        workflow_name="test_run_rechunk",
+        path=tmp_path,
+    )
+    wk.submit(wait=True, status=False, add_to_known=False)
+    wk.rechunk_runs(backup=True, status=False, chunk_size=None)  # None -> one chunk
+
+    arr = wk._store._get_EARs_arr()
+    assert arr.chunks == arr.shape
+
+    bak_path = (Path(wk.path) / arr.path).with_suffix(".bak")
+    arr_bak = zarr.open(bak_path)
+
+    assert arr_bak.chunks == (1,)
+
+    # check backup and new runs data are equal:
+    assert np.all(arr[:] == arr_bak[:])
+
+    # check attributes are equal:
+    assert arr.attrs.asdict() == arr_bak.attrs.asdict()
+
+
+@pytest.mark.integration
+def test_zarr_rechunk_data_equivalent_custom_chunk_size(null_config, tmp_path):
+    t1 = hf.Task(
+        schema=hf.task_schemas.test_t1_conditional_OS,
+        inputs={"p1": 101},
+        repeats=3,
+    )
+    wk = hf.Workflow.from_template_data(
+        tasks=[t1],
+        template_name="test_run_rechunk",
+        workflow_name="test_run_rechunk",
+        path=tmp_path,
+    )
+    wk.submit(wait=True, status=False, add_to_known=False)
+    wk.rechunk_runs(backup=True, status=False, chunk_size=2)
+
+    arr = wk._store._get_EARs_arr()
+    assert arr.chunks == (2,)
+
+    bak_path = (Path(wk.path) / arr.path).with_suffix(".bak")
+    arr_bak = zarr.open(bak_path)
+
+    assert arr_bak.chunks == (1,)
+
+    # check backup and new runs data are equal:
+    assert np.all(arr[:] == arr_bak[:])
+
+
+@pytest.mark.integration
+def test_zarr_rechunk_data_no_backup_load_runs(null_config, tmp_path):
+    t1 = hf.Task(
+        schema=hf.task_schemas.test_t1_conditional_OS,
+        inputs={"p1": 101},
+        repeats=3,
+    )
+    wk = hf.Workflow.from_template_data(
+        tasks=[t1],
+        template_name="test_run_rechunk",
+        workflow_name="test_run_rechunk",
+        path=tmp_path,
+    )
+    wk.submit(wait=True, status=False, add_to_known=False)
+    wk.rechunk_runs(backup=False, status=False)
+
+    arr = wk._store._get_EARs_arr()
+
+    bak_path = (Path(wk.path) / arr.path).with_suffix(".bak")
+    assert not bak_path.is_file()
+
+    # check we can load runs:
+    runs = wk._store._get_persistent_EARs(id_lst=list(range(wk.num_EARs)))
+    run_ID = []
+    for i in runs.values():
+        run_ID.append(i.id_)
+
+
+@pytest.mark.integration
+def test_zarr_rechunk_data_no_backup_load_parameter_base(null_config, tmp_path):
+    t1 = hf.Task(
+        schema=hf.task_schemas.test_t1_conditional_OS,
+        inputs={"p1": 101},
+        repeats=3,
+    )
+    wk = hf.Workflow.from_template_data(
+        tasks=[t1],
+        template_name="test_run_rechunk",
+        workflow_name="test_run_rechunk",
+        path=tmp_path,
+    )
+    wk.submit(wait=True, status=False, add_to_known=False)
+    wk.rechunk_parameter_base(backup=False, status=False)
+
+    arr = wk._store._get_parameter_base_array()
+
+    bak_path = (Path(wk.path) / arr.path).with_suffix(".bak")
+    assert not bak_path.is_file()
+
+    # check we can load parameters:
+    params = wk.get_all_parameters()
+    param_IDs = []
+    for i in params:
+        param_IDs.append(i.id_)

hpcflow/tests/unit/test_utils.py

@@ -13,6 +13,8 @@ from hpcflow.sdk.core.utils import (
     get_nested_indices,
     is_fsspec_url,
     linspace_rect,
+    nth_key,
+    nth_value,
     process_string_nodes,
     replace_items,
     check_valid_py_identifier,
@@ -556,3 +558,22 @@ def test_dict_values_process_flat_single_item_lists():
         "b": [4],
         "c": [5],
     }
+
+
+def test_nth_key():
+    dct = {"a": 1, "b": 2}
+    assert [nth_key(dct, i) for i in range(len(dct))] == ["a", "b"]
+
+
+def test_nth_value():
+    dct = {"a": 1, "b": 2}
+    assert [nth_value(dct, i) for i in range(len(dct))] == [1, 2]
+
+
+def test_nth_key_raises():
+    dct = {"a": 1, "b": 2}
+    with pytest.raises(Exception):
+        nth_key(dct, 2)
+
+    with pytest.raises(Exception):
+        nth_key(dct, -1)