dragon.data.zarr.Store

class Store

Bases: DDict, Store

This class implements a Zarr store using the Dragon distributed dictionary (DDict). This allows users to load a dataset fully into the memory of a set of nodes and keep it very close to Zarr clients. It can also be used as an initially empty store without data loaded from a file-based store.

When cloning from a file-based store, this class uses a process Pool to load data in parallel asychnronously from object construction. This allows the caller to interact with the store while it is still being loaded. Any missing key will be loaded directly by the client if not already loaded by the Pool. This object can be serialized (e.g., via Pickle) and shared with another process, who can then also access the store.

This version of the DDict-based Zarr store supports Zarr versions 2.8.X.

Example usage:

from dragon.data.zarr import Store
from dragon.native.machine import System
import zarr

dstore = Store(
            nloaders=32,
            managers_per_node=1,
            n_nodes=System().nnodes,
            total_mem=(2 * (1024**3)),
            path="/path/to/zarr/data",
        )

warm_zg = zarr.group(store=dstore)
print(f"Dragon-based Store: {warm_zg.tree()}", flush=True)

# access data while loading proceeds in the background
vals = warm_zg["/0"][100][:]

# wait for the load to complete and inspect the exact volume of data loaded
tot_bytes = dstore.wait_on_load()
print(f"Total bytes loaded={tot_bytes}", flush=True)

__init__(*args, path: str = None, nloaders: int = 2, dimension_separator: str = None, use_copy_store: bool = True, **kwargs)

Construct a new DDict-backed Zarr store that is either empty or a cached form of a Zarr file-based store. This class is a subclass of DDict and accepts all arguments it takes.

Parameters:

• path – Either None for an empty store, or a path to a file-based store to load into this store.

• nloaders – The number of processes used to load data in parallel from the file-based store.

• dimension_separator – (optional) Separator placed between the dimensions of a chunk. Used only when determining how to parallelize loading of a store from a path.

• use_copy_store – Use the copy_store to duplicate data from the base Zarr store if possible. This is faster when keys in the store are large and there are many of them.

Methods

__init__(*args[, path, nloaders, ...])	Construct a new DDict-backed Zarr store that is either empty or a cached form of a Zarr file-based store.
advance()	Advance to next available persisted checkpoint.
attach(serialized_dict, *[, timeout, trace])	Within Python you typically do not need to call this method explicitly.
bget(key)	Read the key written through bput.
bput(key, value)	Store a non-persistent key/value pair by brodcasting to all managers across the Distributed Dictionary.
checkpoint()	Calling checkpoint advances the checkpoint for this client.
clear()	Empty the distributed dictionary of all keys and values.
clone(clone_list)	Clone the current dictionary to the list of provided serialized dictionaries.
close()	Do nothing by default
copy([name])	Returns a copy of the Distributed Dictionary.
destroy([allow_restart])	Destroy a Distributed Dictionary instance, freeing all the resources that were allocated when it was created.
detach()	Detach from the Distributed Dictionary and free all local resources of this client.
end_batch_put()
filter(mgr_code, mgr_code_args, comparator)	Calling this instantiates a tree of process groups where mgr_code is expected to be a function that is invoked as mgr_code(args) where args are (dd, out_queue)+mgr_code_args.
freeze()	Freeze the DDict by placing it into read-only mode.
get(k[,d])
get_name()
getitems(keys, *, contexts)	Retrieve data from multiple keys.
is_erasable()
is_listable()
is_readable()
is_writeable()
items()	Returns a view of all key/value pairs in the Distributed Dictionary.
keys()	Return a list of the keys in the store.
listdir([path])
load_cmplt()	Check if loading is complete from another store is complete
local_items()	Returns a DDictItemsView of the key/value pairs that are local to the process invoking this method.
local_keys()	Returns a DDictKeysView of the keys that are local to the process invoking this method.
local_len()	Return the number of keys that are stored on managers that are colocated with this client.
local_values()	Returns a DDictValuesView of the values that are local to the process invoking this method.
manager(id)	Return a version of the current DDict that will always choose the given manager for storing and retrieving data.
persist()	Immediately persist the current checkpoint using the provided persister backend.
persisted_ids()	Get a list of persisted checkpoint IDs.
pickler([key_pickler, value_pickler])	Create a copy of the DDict which will utilize a specific key and value pickler.
pop(key[, default])	Pop the given key and its value from the distributed dictionary and return the associated value.
popitem()	Returns a random key/value pair from the Distributed Dictionary.
pput(key, value)	Persistently store a key/value pair within the Distributed Dictionary.
rename(src_path, dst_path)
reset_get_timer()	Reset the timer that accumulates with each getitem call.
restore(chkpt)	Restore a persisted checkpoint to the provided checkpoint ID.
rmdir([path])
rollback()	Calling rollback decrements the checkpoint id to its previous value.
serialize()	Returns a serialized, base64 encoded descriptor (i.e. string) that may be shared with other processes for attaching.
setdefault(k[,d])
setup_logging()
start_batch_put([persist])	Start a Batch Put operation.
sync_to_newest_checkpoint()	Advance the checkpoint identifier of this client to the newest checkpoint across all managers.
synchronize_ddicts(serialized_ddicts)	Synchronize managers across all parallel dictionaries.
unfreeze()	Unfreeze the DDict by resetting the read-only state to False.
update(dict2)	Adds all key/value pairs from dict2 into this Distributed Dictionary.
values()	When called this returns a view of all values in the Distributed Dictionary that can be iterated or otherwise inspected (i.e. for len) in an efficient manner.
wait_on_load([timeout])	Wait until the data is loaded into the DDict and then return the number of bytes loaded.
which_manager(key)	Return the manager id of the manager to which this key would be sent on a put/store operation.

Attributes

checkpoint_id	Returns the client's current checkpoint id.
dstats	Returns a dict of manager stats, one for each manager of the distributed dictionary.
empty_managers	Return a list of manager IDs that after a restart were empty because their persisted state could not be retrieved.
get_timer	Get the value of the timer that accumulates on getitem calls
is_frozen	Return a True of False value depending on the state of the DDict.
loaded_bytes	Return the number of bytes currently loaded
local_manager	Returns a local manager id if one exists.
local_managers	Returns manager ids of all managers that are local to this node.
main_manager	Returns the main manager id.
manager_nodes	For each manager, a dragon.native.machine.Node object where the manager resides is returned.
show_gets
stats	Returns a list of manager stats, one for each manager of the distributed dictionary.

__init__(*args, path: str = None, nloaders: int = 2, dimension_separator: str = None, use_copy_store: bool = True, **kwargs)

Construct a new DDict-backed Zarr store that is either empty or a cached form of a Zarr file-based store. This class is a subclass of DDict and accepts all arguments it takes.

Parameters:

• path – Either None for an empty store, or a path to a file-based store to load into this store.

• nloaders – The number of processes used to load data in parallel from the file-based store.

• dimension_separator – (optional) Separator placed between the dimensions of a chunk. Used only when determining how to parallelize loading of a store from a path.

• use_copy_store – Use the copy_store to duplicate data from the base Zarr store if possible. This is faster when keys in the store are large and there are many of them.

__len__()

Returns the number of keys stored in the entire Distributed Dictionary or just from the selected manager if this DDict Client was directed to a specific manager by calling the manager method.

Returns:

The number of stored keys in the current checkpoint plus any persistent keys.

Raises:

Various exceptions can be raised including TimeoutError.

keys() → list [object ]

Return a list of the keys in the store. If the store is acting as a cache of an existing store, this will be cached once loading is complete to speed up access.

unfreeze()

Unfreeze the DDict by resetting the read-only state to False.

Raises:

DDictError – If the DDict could not be unfrozen for some reason.

__contains__(key: object ) → bool

Returns True if key is in the Distributed Dictionary and False otherwise.

Parameters:

key – A possible key stored in the DDict.

Returns bool:

True or False depending on if the key is there or not.

Raises:

Various exceptions can be raised including TimeoutError.

load_cmplt() → bool

Check if loading is complete from another store is complete

wait_on_load(timeout: float = None) → int

Wait until the data is loaded into the DDict and then return the number of bytes loaded. If timeout is None (default) this call blocks until the load is done. It timeout is a postive numbers, it blocks at most timeout seconds before raising TimeoutError.

Parameters:

timeout – Either None to block or a postive number of seconds to wait for

property loaded_bytes: int

Return the number of bytes currently loaded

property show_gets

reset_get_timer()

Reset the timer that accumulates with each getitem call. This is diagnostic

property get_timer

Get the value of the timer that accumulates on getitem calls

__getitem__(key)

Get the value that is associated with the given key.

Parameters:

key – The key of a stored key/value pair.

Returns:

The value associated with the key.

Raises:

Exception – Various exceptions can be raised including TimeoutError and KeyError.

__delitem__(key: object ) → None

Deletes a key/value pair from the Distributed Dictionary if it exists.

Raises:

Various exceptions can be raised including TimeoutError and KeyError.

__setitem__(key: object , value: object ) → None

Store the key/value pair in the current checkpoint within the Distributed Dictionary. Due to the nature of a parallel, distributed dictionary, insertion order into the distributed dictionary is not maintained.

Parameters:

• key – The key of the pair. It must be serializable.

• value – the value of the pair. It also must be serializable.

Raises:

Exception – Various exceptions can be raised including TimeoutError.

advance() → None

Advance to next available persisted checkpoint. This operation is for read only mode and directs the DDict to load a next available persisted checkpoint. This can be useful in replaying checkpoints for provenance (i.e. watching how you arrived at a given checkpoint state).

classmethod attach(serialized_dict: str , *, timeout: float = None, trace: bool = False) → DDict

Within Python you typically do not need to call this method explicitly. It will be done automatically when you pass a Distributed Dictionary from one process to another. However, you can do this explicitly if desired/needed.

Parameters:

• serialized_dict – A serialized distributed dictionary.

• timeout – None or a float or int value. A value of None means to wait forever. Otherwise it is the number of seconds to wait while an operation is performed. This timeout is applied to all subsequent client operations that are performed by the process that is attaching this DDict.

• trace – If True, specifies that all operations on the distributed dictionary should be logged in detail within the client log.

Returns:

An attached serialized dictionary.

Raises:

• TimeoutError – If the timeout expires.

• Exception – Other exceptions are possible if for instance the serialized dictionary no longer exists.

bget(key: object ) → object

Read the key written through bput. Each manager has a copy of the key, the client should be able to request the key from its main manager. Clients request the key from the chosen manager if one has been set. Otherwise the client requests the key from its main manager.

Parameters:

key – The key of a stored key/value pair.

Returns:

The value associated with the key.

Raises:

Exception – Various exceptions can be raised including TimeoutError and KeyError.

bput(key: object , value: object ) → None

Store a non-persistent key/value pair by brodcasting to all managers across the Distributed Dictionary. This is useful when multiple clients requesting the same key. This should be used carefully as each manager holds a duplicate of the key/value pair, but it can be useful when a key/value pair is needed across all the nodes of the allocation by all worker processes. The broadcast put distributes the key/value pair in a tree fashion to maximize store performance and the corresponding bget operation can then retrieve it locally when the bgetting process is colocated with a manager and otherwise will get it from its randomly assigned main manager.

Parameters:

• key – A serializable object that will be stored as the key in the DDict.

• value – A serializable object that will be stored as the value.

checkpoint() → None

Calling checkpoint advances the checkpoint for this client. In subsequent calls to the distributed dictionary, like gets or puts, if the chosen manager does not have the current checkpoint in its working set, the get/put operations will advance the manager’s working set to the given checkpoint or block until the checkpoint becomes available. Calling this operation itself does not block.

property checkpoint_id: int

Returns the client’s current checkpoint id.

Returns:

The current checkpoint id of the client.

clear() → None

Empty the distributed dictionary of all keys and values.

clone(clone_list: list [str ]) → None

Clone the current dictionary to the list of provided serialized dictionaries.

Parameters:

clone_list – A list of serialized DDicts which will then be clones

of this DDict.

close() → None

Do nothing by default

copy(name: str = '') → DDict

Returns a copy of the Distributed Dictionary.

Returns:

A second DDict that is a copy of the first assuming that no other processes were concurrently using this DDict.

destroy(allow_restart=False) → None

Destroy a Distributed Dictionary instance, freeing all the resources that were allocated when it was created. Any clients that are still attached to the dictionary and try to do an operation on it will experience an exception if attempting subsequent operations.

detach() → None

Detach from the Distributed Dictionary and free all local resources of this client. But leave in place the DDict for other clients and processes.

property dstats: dict [int , DDictManagerStats]

Returns a dict of manager stats, one for each manager of the distributed dictionary. See the DDictManagerStats structure for a description of its contents.

property empty_managers: list [int ]

Return a list of manager IDs that after a restart were empty because their persisted state could not be retrieved.

end_batch_put() → None

filter(mgr_code: LambdaType, mgr_code_args: tuple , comparator: LambdaType, branching_factor: int = 5) → FilterContextManager

Calling this instantiates a tree of process groups where mgr_code is expected to be a function that is invoked as mgr_code(args) where args are (dd, out_queue)+mgr_code_args. For instance, if mgr_code_args are (x,) then mgr_code(dd, outqueue, x) is how mgr_code is invoked.

The dd of the mgr_code arguments is this distributed dictionary directed toward one manager in the collection of dd managers. In other words, dd is as if the manager method had been invoked on this distributed dictionary so mgr_code only interacts with the manager it was provided. In addition, mgr_code is executed on the same node where the manager it is directed toward is running. This means that mgr code will get the best possible performance while filtering data that is associated with its manager. The mgr_code can do whatever computation is desired, but its chosen output is put into the outqueue.

All data written to outqueue is aggregated with data coming from each manager in a tree-like fashion so as to be scalable to tens of thousands of nodes. All data put in the outqueue by mgr_code is assumed to be ordered from best to worst. When data is aggregated for sending up the tree, it is aggregated according to some kind of ordering which is determined by the comparator function. The comparator will be called as comparator(x,y) and should return True if x is better than y and False otherwise. If there is no ordering, or the ordering is not relevant to the filtering, then comparator(x,y) may return a constant value of False or True and there will be no ordering of the data.

The branching_factor of the filtering tree has a default value, but may be provided by the user to create a tree of whatever width is desired. Note that branching_factor is the max branching factor. Depending on the number of managers, some nodes in the tree may/will have smaller numbers of children.

The filter function returns a Context Manager that supplies an Iterator over which you can iterate on the filtered values. So you can write with dd.filter(...) as candidates: and then iterate over candidates inside the context to read the filtered values.

Assuming your distributed dictionary is called dd, this will get num_needed elements from the result of filtering the distributed dictionary by calling the function get_largest on each distributed dictionary manager.

Parameters:

mgr_code – A function taking arguments as described above that will

run on the same node as a distributed dictionary manager and will be directed toward that manager.

Parameters:

• mgr_code_args – A tuple of arguments to pass to the mgr_code as described above.

• comparator – A function taking two arguments that should return True if the first argument of the values being filtered is “better” than the second and False otherwise. Note that returning a constant value of True or False will result in the filtering imposing no order which may be fine in some use cases.

• branching_factor – The maximum branching factor of any interior node in the filtering tree (i.e. any aggregator).

Returns:

A Context Manager that supplies an iterator which you be used to iterate over the filtered values.

freeze() → None

Freeze the DDict by placing it into read-only mode.

Raises:

DDictError – If the DDict could not be frozen for some reason.

get(k[, d]) → D[k] if k in D, else d.  d defaults to None.

get_name() → str

getitems(keys: Sequence [str ], *, contexts: Mapping [str , Context]) → Mapping [str , Any ]

Retrieve data from multiple keys.

Parameters

keysIterable[str]

The keys to retrieve

contexts: Mapping[str, Context]

A mapping of keys to their context. Each context is a mapping of store specific information. E.g. a context could be a dict telling the store the preferred output array type: {"meta_array": cupy.empty(())}

Returns

Mapping

A collection mapping the input keys to their results.

Notes

This default implementation uses __getitem__() to read each key sequentially and ignores contexts. Overwrite this method to implement concurrent reads of multiple keys and/or to utilize the contexts.

is_erasable()

property is_frozen: bool

Return a True of False value depending on the state of the DDict.

Returns:

True or False to indicate if the DDict is currently frozen.

Raises:

DDictError – If the DDict cannot get this status from its main manager.

is_listable()

is_readable()

is_writeable()

items() → DDictItemsView

Returns a view of all key/value pairs in the Distributed Dictionary.

Returns:

A view of all key/value pairs.

listdir(path: str = '') → List [str ]

local_items() → DDictItemsView

Returns a DDictItemsView of the key/value pairs that are local to the process invoking this method.

Returns:

A view of the current DDict which has only the co-located node local items of the DDict in it.

local_keys() → DDictKeysView

Returns a DDictKeysView of the keys that are local to the process invoking this method. This is useful when a local process wants to work with data stored locally that will be transformed and then later requested by other processes globally.

Returns:

A DDictKeysView of the current DDict which has only the co-located node local keys of the DDict in it.

local_len() → int

Return the number of keys that are stored on managers that are colocated with this client.

Returns:

The number of keys stored on this node of the Dragon run-time.

Raises:

• DDictCheckpointSyncError – If the checkpoint the client is at has been retired.

• RuntimeError – Other errors are possible including TimeoutError.

property local_manager: int

Returns a local manager id if one exists. The manager designated as the main manager for the client if it is on the same node as its local manager. Otherwise, if no local manager exists, then None is returned.

property local_managers: list [int ]

Returns manager ids of all managers that are local to this node.

local_values() → DDictValuesView

Returns a DDictValuesView of the values that are local to the process invoking this method.

Returns:

A view of the current DDict which has only the co-located values of the DDict in it.

property main_manager: int

Returns the main manager id. This will always exist and will be the same as the local manager id if a local manager exists. Otherwise, it will be the id of a random manager from another node.

manager(id: int ) → DDict

Return a version of the current DDict that will always choose the given manager for storing and retrieving data. This is only useful when storing and/or retrieving data locally. If you need data to be globally available then you should only store data that would be globally stored there anyway. One way to accomplish this is to store data globally, but then work on locally stored keys. You can discover “local_keys” of a manager by calling getting a manager-directed handle to the DDict and iterating over its keys.

Parameters:

id – The manager id of the chosen manager.

Returns:

A version of the same DDict which will direct all gets and puts to the specified manager.

Raises:

Exception – If the manager id is not a valid id.

property manager_nodes: list [Node]

For each manager, a dragon.native.machine.Node object where the manager resides is returned.

persist() → None

Immediately persist the current checkpoint using the provided persister backend. Normally persistence occurs automatically when a checkpoint falls out of the working set. Calling this will cause a checkpoint to persist immediately.

persisted_ids() → list [int ]

Get a list of persisted checkpoint IDs.

Returns:

The list of persisted checkpoint IDs.

pickler(key_pickler=None, value_pickler=None) → DDict

Create a copy of the DDict which will utilize a specific key and value pickler.

Parameters:

• key_pickler – A pickler to de/serialize keys. Defaults to None.

• value_pickler – A pickler to de/serialize values. Defaults to None.

Returns:

The same DDict with the desired pickling attributes.

pop(key: object , default: object = None) → object

Pop the given key and its value from the distributed dictionary and return the associated value. If the given key is not found in the dictionary, then KeyError is raised unless a default value is provided, in which case the default value is returned if the key is not found in the dictionary.

Parameters:

• key – A key to be popped from the distributed dictionary.

• default – A default value to be returned if the key is not in the distributed dictionary.

Returns:

The associated value if key is popped and the default value otherwise.

popitem() → tuple [object , object ]

Returns a random key/value pair from the Distributed Dictionary.

Returns:

A random key/value pair.

Raises:

NotImplementedError – Not implemented.

pput(key: object , value: object ) → None

Persistently store a key/value pair within the Distributed Dictionary. This is useful when checkpointing is employed in the dictionary. A persistent put of a key/value pair means that the key/value pair persists across checkpoints. Persistent key/value pairs are useful when putting constant values or other values that don’t change across checkpoints.

Parameters:

• key – A serializable object that will be stored as the key in the DDict.

• value – A serializable object that will be stored as the value.

rename(src_path: str , dst_path: str ) → None

restore(chkpt: int ) → None

Restore a persisted checkpoint to the provided checkpoint ID.

Parameters:

chkpt – The checkpoint ID which should be restored.

rmdir(path: str = '') → None

rollback() → None

Calling rollback decrements the checkpoint id to its previous value. Again this call does not block. If rollback causes the checkpoint id to roll back to a checkpoint that a chosen manager no longer has in its working set, then subsequent operations may fail with a exception indicating the Checkpoint is no longer available, raising a DDictCheckpointSyncError exception.

serialize() → str

Returns a serialized, base64 encoded descriptor (i.e. string) that may be shared with other processes for attaching. This is especially useful when sharing with C or C++ code. Within Python you can pass the Distributed Dictionary to another process and it will be automatically serialized and attached so using this method is not needed when passing to another Python process.

Returns:

A serialized, base64 encoded string that may be used for attaching to the dictionary.

setdefault(k[, d]) → D.get(k,d), also set D[k]=d if k not in D

setup_logging()

start_batch_put(persist=False) → None

Start a Batch Put operation. This allows efficient data loading from a process or processes while multiple put operations are being performed. A start_batch_put should be followed by a series of put operations (i.e. __setitem__ or pput) and then concluded by a call to end_batch_put. The advantage of a batch put is the elimination of confirmation of each put operation thereby reducing the amount of communication and time spent waiting for put operations to complete. With batch put the put operations are streamed to each manager.

Parameters:

persist – If True, then the put operations should be persistent pput operations. Defaults to False.

property stats: list [DDictManagerStats]

Returns a list of manager stats, one for each manager of the distributed dictionary. See the DDictManagerStats structure for a description of its contents.

sync_to_newest_checkpoint() → None

Advance the checkpoint identifier of this client to the newest checkpoint across all managers. This does not guarantee that all managers have advanced to the same checkpoint. It does guarantee that the client that calls this will have advanced to the newest checkpoint across all the mangerrs. See the ddict_checkpoint_pi.py demo in ddict/ddict_checkpoint_pi.py for an example of an application that uses this method.

classmethod synchronize_ddicts(serialized_ddicts: list [str ]) → None

Synchronize managers across all parallel dictionaries. This is useful when you have two or more identical instances of a DDict and are using one to recover other instances. This method will look for any empty managers in the list of serialized dictionaries and fill them with their parallel counterpart from another non-empty dictionary manager.

Parameters:

serialized_ddicts – A list of serialized DDicts to synchronize.

update(dict2: DDict) → None

Adds all key/value pairs from dict2 into this Distributed Dictionary.

Parameters:

dict2 – Another distributed dictionary.

Raises:

NotImplementedError – Not implemented.

values() → DDictValuesView

When called this returns a view of all values in the Distributed Dictionary that can be iterated or otherwise inspected (i.e. for len) in an efficient manner.

Returns:

An view of the values in the DDict.

which_manager(key: object ) → int

Return the manager id of the manager to which this key would be sent on a put/store operation. This can be useful when wanting to minimize the movement of data.

Parameters:

key – A key that might be stored at some future time. It must be serializable.

Returns:

The manager id of the manager where this key would be stored.