async def create(self, name, doc_fields=None):
"""
Create an experiment document.
Args:
name (str): Name of the experiment.
doc_fields: Other fields of the experiment document.
The following fields will be set by default if absent:
* start_time: set to ``datetime.utcnow()``.
* heartbeat: set to `start_time`.
* status: set to "RUNNING".
Returns:
ObjectId: The ID of the inserted document.
"""
doc_fields = validate_experiment_doc(
pop_experiment_id(dict(doc_fields or ())))
doc_fields['name'] = name
if 'start_time' not in doc_fields:
doc_fields['start_time'] = datetime.utcnow()
if 'heartbeat' not in doc_fields:
doc_fields['heartbeat'] = doc_fields['start_time']
if 'status' not in doc_fields:
doc_fields['status'] = 'RUNNING'
await self.ensure_indexes()
return (await self.collection.insert_one(doc_fields)).inserted_id
async def set_heartbeat(self, id, doc_fields=None):
"""
Set the heartbeat time of an experiment.
Args:
id (str or ObjectId): ID of the experiment.
doc_fields: Other fields to be updated, optional.
Raises:
KeyError: If the experiment with `id` does not exist.
"""
doc_fields = validate_experiment_doc(
pop_experiment_id(dict(doc_fields or ())))
doc_fields['heartbeat'] = datetime.utcnow()
await self.ensure_indexes()
return await self._update(id, doc_fields)
async def update(self, id, doc_fields):
"""
Update an experiment document in ``[coll_name].runs``.
Args:
id (str or ObjectId): ID of the experiment.
doc_fields: Fields of the experiment document to be updated.
Raises:
KeyError: If the experiment with `id` does not exist.
"""
id = validate_experiment_id(id)
doc_fields = validate_experiment_doc(
pop_experiment_id(dict(doc_fields or ())))
await self.ensure_indexes()
if doc_fields:
return await self._update(id, doc_fields)
async def set_finished(self, id, status, doc_fields=None):
"""
Set the status of an experiment to "COMPLETED" or "FAILED".
The "stop_time" and the "heartbeat" will be set to current time.
Args:
id (str or ObjectId): ID of the experiment.
status ({"COMPLETED", "FAILED"}): The final experiment status.
doc_fields: Other fields to be updated, optional.
Raises:
KeyError: If the experiment with `id` does not exist.
"""
if status not in ('COMPLETED', 'FAILED'):
raise ValueError('Invalid `status`: {!r}'.format(status))
doc_fields = validate_experiment_doc(
pop_experiment_id(dict(doc_fields or ())))
doc_fields['stop_time'] = doc_fields['heartbeat'] = datetime.utcnow()
doc_fields['status'] = status
await self.ensure_indexes()
return await self._update(id, doc_fields)
async def iter_docs(self, filter=None, skip=None, limit=None,
sort_by=None, include_deleted=False):
"""
Iterate through experiment documents.
Args:
filter: The filter for querying experiment documents.
If `None`, all experiments will be queried.
skip: Number of documents to skip at front.
limit: Limiting the returned document by this number.
sort_by: The sort ordering. If not specified, will sort
by the DESCENDING order of "heartbeat".
include_deleted (bool): Whether or not to include deleted
documents? (default :obj:`False`)
Yields:
The matched documents, in DESCENDING order of "heartbeat".
"""
# assemble the query
filter_ = to_database_experiment_doc(
validate_experiment_doc(dict(filter or ())))
if not include_deleted:
filter_['deleted'] = {'$ne': True}
if sort_by is None:
sort_by = [('heartbeat', pymongo.DESCENDING)]
# open the cursor and fetch documents
cursor = self.collection.find(
filter_,
sort=sort_by
)
if skip:
cursor = cursor.skip(skip)
if limit:
cursor = cursor.limit(limit)
async for doc in cursor:
yield from_database_experiment_doc(doc)