pygsti.protocols.protocol.Protocol

pygsti.protocols.protocol.Protocol#

class Protocol(name=None)#

Bases: MongoSerializable

An analysis routine that is run on experimental data. A generalized notion of a QCVV protocol.

A Protocol object represents things like, but not strictly limited to, QCVV protocols. This class is essentially a serializable run function that takes as input a ProtocolData object and returns a ProtocolResults object. This function describes the working of the “protocol”.

Parameters:

name (str, optional) – The name of this protocol, also used to (by default) name the results produced by this protocol. If None, the class name will be used.

Create a new Protocol object.

Parameters:

name (str, optional) – The name of this protocol, also used to (by default) name the results produced by this protocol. If None, the class name will be used.

Return type:

Protocol

Methods

__init__([name])

Create a new Protocol object.

add_mongodb_write_ops(write_ops, mongodb[, ...])

Accumulate write and update operations for writing this object to a MongoDB database.

from_dir(dirname[, quick_load])

Initialize a new Protocol object from dirname.

from_mongodb(mongodb, doc_id, **kwargs)

Create and initialize an object from a MongoDB instance.

from_mongodb_doc(mongodb, collection_name, ...)

Create and initialize an object from a MongoDB instance and pre-loaded primary document.

remove_from_mongodb(mongodb, doc_id[, ...])

Remove the documents corresponding to an instance of this class from a MongoDB database.

remove_me_from_mongodb(mongodb[, session, ...])

run(data[, memlimit, comm])

Run this protocol on data.

run_mpi(data, num_ranks, *[, mpiexec, ...])

Run this protocol in parallel using MPI workers launched as a subprocess.

setup_nameddict(final_dict)

Initializes a set of nested NamedDict dictionaries describing this protocol.

stage_slurm(data, num_ranks, slurm, work_dir, *)

Write all working files to work_dir and generate a SLURM batch script ready for sbatch submission.

write(dirname)

Write this protocol to a directory.

write_to_mongodb(mongodb[, session, ...])

Write this object to a MongoDB database.

Attributes

collection_name

add_mongodb_write_ops(write_ops, mongodb, overwrite_existing=False, **kwargs)#

Accumulate write and update operations for writing this object to a MongoDB database.

Similar to write_to_mongodb() but collects write operations instead of actually executing any write operations on the database. This function may be preferred to write_to_mongodb() when this object is being written as a part of a larger entity and executing write operations is saved until the end.

As in write_to_mongodb(), self.collection_name is the collection name and _id is either: 1) the ID used by a previous write or initial read-in, if one exists, OR 2) a new random bson.objectid.ObjectId

Parameters:

  • write_ops (WriteOpsByCollection) – An object that keeps track of pymongo write operations on a per-collection basis. This object accumulates write operations to be performed at some point in the future.

  • mongodb (pymongo.database.Database) – The MongoDB instance to write data to.

  • overwrite_existing (bool, optional) – Whether existing documents should be overwritten. The default of False causes a ValueError to be raised if a document with the given _id already exists and is different from what is being written.

  • **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _add_auxiliary_write_ops_and_update_doc method is allowed here.

Returns:

The identifier (_id value) of the main document that was written.

Return type:

bson.objectid.ObjectId

classmethod from_dir(dirname, quick_load=False)#

Initialize a new Protocol object from dirname.

quick_loadbool, optional

Setting this to True skips the loading of components that may take a long time to load.

Parameters:

  • dirname (str) – The directory name.

  • quick_load (bool, optional) – Setting this to True skips the loading of components that may take a long time to load.

Return type:

Protocol

classmethod from_mongodb(mongodb, doc_id, **kwargs)#

Create and initialize an object from a MongoDB instance.

Parameters:

  • mongodb (pymongo.database.Database) – The MongoDB instance to load from.

  • doc_id (bson.objecctid.ObjectId or dict) – The object ID or filter used to find a single object ID within the database. This document is loaded from the collection given by the collection_name attribute of this class.

  • **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _create_obj_from_doc_and_mongodb method is allowed here.

Return type:

object

classmethod from_mongodb_doc(mongodb, collection_name, doc, **kwargs)#

Create and initialize an object from a MongoDB instance and pre-loaded primary document.

Parameters:

  • mongodb (pymongo.database.Database) – The MongoDB instance to load from.

  • collection_name (str) – The collection name within mongodb that doc was loaded from. This is needed for the sole purpose of setting the created (returned) object’s database “coordinates”.

  • doc (dict) – The already-retrieved main document for the object being loaded. This takes the place of giving an identifier for this object.

  • **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _create_obj_from_doc_and_mongodb method is allowed here.

Return type:

object

classmethod remove_from_mongodb(mongodb, doc_id, collection_name=None, session=None, recursive='default')#

Remove the documents corresponding to an instance of this class from a MongoDB database.

Parameters:

  • mongodb (pymongo.database.Database) – The MongoDB instance to remove documents from.

  • doc_id (bson.objectid.ObjectId) – The identifier of the root document stored in the database.

  • collection_name (str, optional) – the MongoDB collection within mongodb where the main document resides. If None, then <this_class>.collection_name is used (which is usually what you want).

  • session (pymongo.client_session.ClientSession, optional) – MongoDB session object to use when interacting with the MongoDB database. This can be used to implement transactions among other things.

  • recursive (RecursiveRemovalSpecification, optional) – An object that filters the type of documents that are removed. Used when working with inter-related experiment designs, data, and results objects to only remove the types of documents you know aren’t being shared with other documents.

Return type:

None

run(data, memlimit=None, comm=None)#

Run this protocol on data.

Parameters:

  • data (ProtocolData) – The input data.

  • memlimit (int, optional) – A rough per-processor memory limit in bytes.

  • comm (mpi4py.MPI.Comm, optional) – When not None, an MPI communicator used to run this protocol in parallel.

Return type:

ProtocolResults

run_mpi(data, num_ranks, *, mpiexec='auto', extra_mpi_args=None, ranks_per_host=None, env=None, persistent_dir=None, dry_run=False, blas_threads_per_rank=0, **run_kwargs)#

Run this protocol in parallel using MPI workers launched as a subprocess. The subprocess environment variables will be set as

worker_env = {**_os.environ, **_blas_env, **(env or {})}

where _blas_env is inferred from blas_threads_per_rank (see Notes).

This method can be called from anywhere (e.g., Jupyter notebooks or scripts) without requiring the caller to manage MPI communicators or write launcher scripts manually.

Parameters:

  • data (ProtocolData) – The input data.

  • num_ranks (int) – Number of MPI worker processes to launch. When 1, falls back to a plain run() call with no MPI overhead.

  • mpiexec (str, keyword-only) – MPI launcher executable name or path. 'auto' (default) searches PATH for mpiexec, mpirun, or mpiexec.hydra in that order.

  • extra_mpi_args (list of str, keyword-only) – Extra arguments inserted between the launcher and the Python executable, e.g. ['--oversubscribe'] or ['--hostfile', 'hosts.txt'].

  • ranks_per_host (int, keyword-only) – Number of ranks per virtual host. Sets PYGSTI_MAX_HOST_PROCS in the worker environment. None (default) uses actual hostnames.

  • env (dict, keyword-only) – Extra environment variables forwarded to worker processes, merged on top of the current os.environ. ranks_per_host takes precedence over PYGSTI_MAX_HOST_PROCS supplied here.

  • persistent_dir (str or Path, keyword-only) – Permanent directory for data and results. Created if it does not already exist. When None (default), the existing on-disk data path is reused if available; otherwise a temporary directory is used and deleted on return. Required when dry_run=True.

  • dry_run (bool, keyword-only) –

    When True, we write working files to persistent_dir and return None. The recommended launch command is printed to console if self.verbosity is unset or positive.

    Requires persistent_dir to be set.

  • blas_threads_per_rank (int, keyword-only) – Number of threads each worker rank allows BLAS libraries to use. 0 (default) auto-detects an appropriate value. See Notes.

  • **run_kwargs – Keyword arguments forwarded to Protocol.run for each worker. Examples of arguments you might want to forward include simulator, optimizers, or disable_checkpointing. These arguments will be serialized via pickle.

Returns:

None when dry_run=True.

Return type:

ProtocolResults or None

Notes

Shared memory and host topology. pyGSTi uses a two-level communicator hierarchy: a host communicator (host_comm) groups ranks that share a physical host and can exchange data via shared memory, while an inter-host communicator (interhost_comm) connects one rank per host for cross-host MPI. Collective operations (gather, allreduce, broadcast) use a two-phase pattern — first within the host, then across hosts — when shared memory is enabled.

Setting PYGSTI_USE_SHARED_MEMORY=0 (via env) disables shared memory entirely. In this case host_comm is None and all communication goes directly through the full communicator; the host grouping has no effect. ranks_per_host is therefore only meaningful when shared memory is enabled.

Choosing ranks_per_host. On a single compute node all ranks share one hostname, so by default pyGSTi places them all in one host group and uses shared memory across all of them. Setting ranks_per_host=k splits the ranks into num_ranks // k virtual hosts of k ranks each, producing k-rank shared-memory groups connected by MPI across groups. This is useful for testing multi-node behavior on a single machine, or for matching a known NUMA or hardware partition boundary.

Thread oversubscription. When blas_threads_per_rank > 0, that value is set for OMP_NUM_THREADS, OPENBLAS_NUM_THREADS, MKL_NUM_THREADS, NUMEXPR_NUM_THREADS, and BLIS_NUM_THREADS in the worker environment. When blas_threads_per_rank == 0 (the default), the value is computed as max(1, num_cpus // num_ranks) where num_cpus is the physical CPU core count of the current machine (falling back to the logical count if the physical count is unavailable). These variables are applied before env, so explicit entries in env override them.

Job schedulers. To generate a SLURM batch script, use stage_slurm() instead. For other schedulers (PBS, LSF, etc.), use dry_run=True together with persistent_dir to write the data and runner script to a permanent location. The recommended MPI launch command will be written to console if self.verbosity is unset or positive.

setup_nameddict(final_dict)#

Initializes a set of nested NamedDict dictionaries describing this protocol.

This function is used by ProtocolResults objects when they’re creating nested dictionaries of their contents. This function returns a set of nested, single (key,val)-pair named-dictionaries which describe the particular attributes of this Protocol object named within its self._nameddict_attributes tuple. The final nested dictionary is set to be final_dict, which allows additional result quantities to easily be added.

Parameters:

final_dict (NamedDict) – the final-level (innermost-nested) NamedDict in the returned nested dictionary.

Return type:

NamedDict

stage_slurm(data, num_ranks, slurm, work_dir, *, ranks_per_host=None, blas_threads_per_rank=0, **run_kwargs)#

Write all working files to work_dir and generate a SLURM batch script ready for sbatch submission.

This is the SLURM-specific counterpart to run_mpi(). It does not launch any processes; call sbatch <slurm.script_path> from a terminal or batch system to submit the job.

Parameters:

  • data (ProtocolData) – The input data.

  • num_ranks (int) – Total number of MPI worker processes.

  • slurm (SlurmSettings) – SLURM script options. See SlurmSettings for details.

  • work_dir (str or Path) – Permanent directory where data, the protocol, the pickled kwargs, and the runner script are written. All paths embedded in the generated batch script point here.

  • ranks_per_host (int, keyword-only) – Number of ranks per shared-memory group. When set, a commented export PYGSTI_MAX_HOST_PROCS=<value> line is emitted in the batch script (uncomment to activate). None (default) uses a suggested value of num_ranks // slurm.nodes.

  • blas_threads_per_rank (int, keyword-only) –

    Number of BLAS threads per rank. 0 (default) auto-detects based on the current machine’s CPU count. This value sets both --cpus-per-task in the batch script and the BLAS environment variable exports.

    Note

    Auto-detection measures the machine on which you call stage_slurm (typically a login node), which may differ from the HPC compute nodes. Pass an explicit value if your compute nodes have a different core count.

  • **run_kwargs – Forwarded to run() in each worker. Serialized via pickle.

Returns:

script_path – The file path of the generated SLURM script.

Return type:

str

Notes

The generated script uses srun as the MPI launcher. srun is SLURM-native and inherits the job allocation automatically, so no -n flag is needed.

The script also contains a commented block of optional #SBATCH directives (--account, --qos, --constraint, etc.) that are commonly needed but too site-specific to set automatically. Uncomment and edit the relevant lines before submitting.

write(dirname)#

Write this protocol to a directory.

Parameters:

dirname (str) – The directory name to write. This directory will be created if needed, and the files in an existing directory will be overwritten.

Return type:

None

write_to_mongodb(mongodb, session=None, overwrite_existing=False, **kwargs)#

Write this object to a MongoDB database.

The collection name used is self.collection_name, and the _id is either: 1) the ID used by a previous write or initial read-in, if one exists, OR 2) a new random bson.objectid.ObjectId

Parameters:

  • mongodb (pymongo.database.Database) – The MongoDB instance to write data to.

  • session (pymongo.client_session.ClientSession, optional) – MongoDB session object to use when interacting with the MongoDB database. This can be used to implement transactions among other things.

  • overwrite_existing (bool, optional) – Whether existing documents should be overwritten. The default of False causes a ValueError to be raised if a document with the given _id already exists and is different from what is being written.

  • **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _add_auxiliary_write_ops_and_update_doc method is allowed here.

Returns:

The identifier (_id value) of the main document that was written.

Return type:

bson.objectid.ObjectId