from __future__ import annotations
import os
import sys
from datetime import datetime
from enum import Enum
from typing import Optional, Dict, Any, List, Union, Iterable, Tuple, Type, Sequence, ClassVar, TypeVar
from dateutil.parser import parse as date_parser
try:
from pydantic.v1 import BaseModel, Extra, constr, validator, PrivateAttr, Field, parse_obj_as
except ImportError:
from pydantic import BaseModel, Extra, constr, validator, PrivateAttr, Field, parse_obj_as
from qcelemental.models.results import Provenance
from qcportal.base_models import (
RestModelBase,
QueryModelBase,
QueryIteratorBase,
)
from qcportal.cache import RecordCache, get_records_with_cache
from qcportal.compression import CompressionEnum, decompress, get_compressed_ext
_T = TypeVar("_T")
[docs]
class PriorityEnum(int, Enum):
"""
The priority of a Task. Higher priority will be pulled first.
"""
high = 2
normal = 1
low = 0
@classmethod
def _missing_(cls, priority):
"""Attempts to find the correct priority in a case-insensitive way
If a string being converted to a PriorityEnum is missing, then this function
will convert the case and try to find the appropriate priority.
"""
if isinstance(priority, int):
# An integer that is outside the range of valid priorities
return
priority = priority.lower()
# Search this way rather than doing 'in' since we are comparing
# a string to an enum
for p in cls:
if priority == p.name:
return p
[docs]
class RecordStatusEnum(str, Enum):
"""
The state of a record object. The states which are available are a finite set.
"""
# This ordering shouldn't change in the near future, as it conflicts
# a bit with some migration testing
complete = "complete"
invalid = "invalid"
running = "running"
error = "error"
waiting = "waiting"
cancelled = "cancelled"
deleted = "deleted"
@classmethod
def _missing_(cls, name):
"""Attempts to find the correct status in a case-insensitive way
If a string being converted to a RecordStatusEnum is missing, then this function
will convert the case and try to find the appropriate status.
"""
name = name.lower()
# Search this way rather than doing 'in' since we are comparing
# a string to an enum
for status in cls:
if name == status:
return status
[docs]
@classmethod
def make_ordered_status(cls, statuses: Iterable[RecordStatusEnum]) -> List[RecordStatusEnum]:
"""Returns a list of the given statuses but in a defined order"""
order = [cls.complete, cls.error, cls.running, cls.waiting, cls.cancelled, cls.invalid, cls.deleted]
return sorted(statuses, key=lambda x: order.index(x))
class OutputTypeEnum(str, Enum):
"""
What type of data is stored
"""
stdout = "stdout"
stderr = "stderr"
error = "error"
class OutputStore(BaseModel):
"""
Storage of outputs and error messages, with optional compression
"""
class Config:
extra = Extra.forbid
output_type: OutputTypeEnum = Field(..., description="The type of output this is (stdout, error, etc)")
compression_type: CompressionEnum = Field(CompressionEnum.none, description="Compression method (such as lzma)")
data_: Optional[bytes] = Field(None, alias="data")
_data_url: Optional[str] = PrivateAttr(None)
_client: Any = PrivateAttr(None)
def propagate_client(self, client, history_base_url):
self._client = client
self._data_url = f"{history_base_url}/outputs/{self.output_type.value}/data"
def _fetch_raw_data(self):
if self.data_ is not None:
return
if self._client is None:
raise RuntimeError("No client to fetch output data from")
cdata, ctype = self._client.make_request(
"get",
self._data_url,
Tuple[bytes, CompressionEnum],
)
assert self.compression_type == ctype
self.data_ = cdata
@property
def data(self) -> Any:
self._fetch_raw_data()
return decompress(self.data_, self.compression_type)
class ComputeHistory(BaseModel):
class Config:
extra = Extra.forbid
id: int
record_id: int
status: RecordStatusEnum
manager_name: Optional[str]
modified_on: datetime
provenance: Optional[Provenance]
outputs_: Optional[Dict[str, OutputStore]] = Field(None, alias="outputs")
_client: Any = PrivateAttr(None)
_base_url: Optional[str] = PrivateAttr(None)
def propagate_client(self, client, record_base_url):
self._client = client
self._base_url = f"{record_base_url}/compute_history/{self.id}"
if self.outputs_ is not None:
for o in self.outputs_.values():
o.propagate_client(self._client, self._base_url)
def fetch_all(self):
self._fetch_outputs()
def _fetch_outputs(self):
if self._client is None:
raise RuntimeError("This compute history is not connected to a client")
self.outputs_ = self._client.make_request(
"get",
f"{self._base_url}/outputs",
Dict[str, OutputStore],
)
for o in self.outputs_.values():
o.propagate_client(self._client, self._base_url)
o._fetch_raw_data()
@property
def outputs(self) -> Dict[str, OutputStore]:
if self.outputs_ is None:
self._fetch_outputs()
return self.outputs_
def get_output(self, output_type: OutputTypeEnum) -> Any:
if not self.outputs:
return None
o = self.outputs.get(output_type, None)
if o is None:
return None
else:
return o.data
@property
def stdout(self) -> Any:
return self.get_output("stdout")
@property
def stderr(self) -> Any:
return self.get_output("stderr")
@property
def error(self) -> Any:
return self.get_output("error")
class NativeFile(BaseModel):
"""
Storage of native files, with compression
"""
class Config:
extra = Extra.forbid
name: str = Field(..., description="Name of the file")
compression_type: CompressionEnum = Field(..., description="Compression method (such as lzma)")
data_: Optional[bytes] = Field(None, alias="data")
_data_url: Optional[str] = PrivateAttr(None)
_client: Any = PrivateAttr(None)
def propagate_client(self, client, record_base_url):
self._client = client
self._data_url = f"{record_base_url}/native_files/{self.name}/data"
def fetch_all(self):
self._fetch_raw_data()
def _fetch_raw_data(self):
if self.data_ is not None:
return
if self._client is None:
raise RuntimeError("No client to fetch native file data from")
cdata, ctype = self._client.make_request(
"get",
self._data_url,
Tuple[bytes, CompressionEnum],
)
assert self.compression_type == ctype
self.data_ = cdata
@property
def data(self) -> Any:
self._fetch_raw_data()
return decompress(self.data_, self.compression_type)
def save_file(
self, directory: str, new_name: Optional[str] = None, keep_compressed: bool = False, overwrite: bool = False
):
"""
Saves the file to the given directory
"""
if new_name is None:
name = self.name
else:
name = new_name
if keep_compressed:
name += get_compressed_ext(self.compression_type)
full_path = os.path.join(directory, name)
if os.path.exists(full_path) and not overwrite:
raise RuntimeError(f"File {full_path} already exists. Not overwriting")
if keep_compressed:
with open(full_path, "wb") as f:
f.write(self.data)
else:
d = self.data
# TODO - streaming decompression?
if isinstance(d, str):
with open(full_path, "wt") as f:
f.write(self.data)
elif isinstance(d, bytes):
with open(full_path, "wb") as f:
f.write(self.data)
else:
raise RuntimeError(f"Cannot write data of type {type(d)} to a file")
class RecordInfoBackup(BaseModel):
class Config:
extra = Extra.forbid
old_status: RecordStatusEnum
old_tag: Optional[str]
old_priority: Optional[PriorityEnum]
modified_on: datetime
class RecordComment(BaseModel):
class Config:
extra = Extra.forbid
id: int
record_id: int
username: Optional[str]
timestamp: datetime
comment: str
[docs]
class RecordTask(BaseModel):
class Config:
extra = Extra.forbid
id: int
record_id: int
function: Optional[str]
function_kwargs_compressed: Optional[bytes]
tag: str
priority: PriorityEnum
required_programs: List[str]
@property
def function_kwargs(self) -> Optional[Dict[str, Any]]:
if self.function_kwargs_compressed is None:
return None
else:
return decompress(self.function_kwargs_compressed, CompressionEnum.zstd)
class ServiceDependency(BaseModel):
class Config:
extra = Extra.forbid
record_id: int
extras: Dict[str, Any]
[docs]
class RecordService(BaseModel):
class Config:
extra = Extra.forbid
id: int
record_id: int
tag: str
priority: PriorityEnum
find_existing: bool
service_state: Optional[Dict[str, Any]] = None
dependencies: List[ServiceDependency]
[docs]
class BaseRecord(BaseModel):
class Config:
extra = Extra.forbid
allow_mutation = True
validate_assignment = True
id: int
record_type: str
is_service: bool
properties: Optional[Dict[str, Any]]
extras: Optional[Dict[str, Any]]
status: RecordStatusEnum
manager_name: Optional[str]
created_on: datetime
modified_on: datetime
owner_user: Optional[str]
owner_group: Optional[str]
######################################################
# Fields not always included when fetching the record
######################################################
compute_history_: Optional[List[ComputeHistory]] = Field(None, alias="compute_history")
task_: Optional[RecordTask] = Field(None, alias="task")
service_: Optional[RecordService] = Field(None, alias="service")
comments_: Optional[List[RecordComment]] = Field(None, alias="comments")
native_files_: Optional[Dict[str, NativeFile]] = Field(None, alias="native_files")
# Private non-pydantic fields
_client: Any = PrivateAttr(None)
_base_url: str = PrivateAttr(None)
# A dictionary of all subclasses (calculation types) to actual class type
_all_subclasses: ClassVar[Dict[str, Type[BaseRecord]]] = {}
# Local record cache we can use for child records
# This record may also be part of the cache
_record_cache: Optional[RecordCache] = PrivateAttr(None)
_cache_dirty: bool = PrivateAttr(False)
def __init__(self, client=None, **kwargs):
BaseModel.__init__(self, **kwargs)
# Calls derived class propagate_client
# which should filter down to the ones in this (BaseRecord) class
self.propagate_client(client)
assert self._client is client, "Client not set in base record class?"
def __init_subclass__(cls):
"""
Register derived classes for later use
"""
# Get the record type. This is kind of ugly, but works.
# We could use ClassVar, but in my tests it doesn't work for
# disambiguating (ie, via parse_obj_as)
record_type = cls.__fields__["record_type"].default
cls._all_subclasses[record_type] = cls
def __del__(self):
# Sometimes this won't exist if there is an exception during construction
# TODO - we check sys.meta_path. Pydantic attempts an import of something, which is
# not good if the interpreter is shutting down. This is a hack to avoid that.
# Pydantic v2 may fix this
if (
hasattr(self, "_record_cache")
and self._record_cache is not None
and not self._record_cache.read_only
and self._cache_dirty
and sys.meta_path is not None
):
self.sync_to_cache(True) # Don't really *have* to detach, but why not
s = super()
if hasattr(s, "__del__"):
s.__del__(self)
[docs]
@classmethod
def get_subclass(cls, record_type: str) -> Type[BaseRecord]:
"""
Obtain a subclass of this class given its record_type
"""
subcls = cls._all_subclasses.get(record_type)
if subcls is None:
raise RuntimeError(f"Cannot find subclass for record type {record_type}")
return subcls
@classmethod
def _fetch_children_multi(
cls, client, record_cache, records: Iterable[BaseRecord], include: Iterable[str], force_fetch: bool = False
):
"""
Fetches all children of the given records recursively
This tries to work efficiently, fetching larger batches of children
that can span multiple records
Meant to be overridden by derived classes
"""
pass
[docs]
@classmethod
def fetch_children_multi(
cls, records: Iterable[Optional[BaseRecord]], include: Optional[Iterable[str]] = None, force_fetch: bool = False
):
"""
Fetches all children of the given records
This tries to work efficiently, fetching larger batches of children
that can span multiple records
"""
# Remove any None records
# can happen if missing_ok=True in some function calls
records = [r for r in records if r is not None]
if not records:
return
# Get the first record (for the client and other info)
template_record = next(iter(records))
if not all(isinstance(r, type(template_record)) for r in records):
raise RuntimeError("Fetching children of records with different types is not supported.")
if not all(r._client is template_record._client for r in records):
raise RuntimeError("Fetching children of records with different clients is not supported.")
if not all(r._record_cache is template_record._record_cache for r in records):
raise RuntimeError("Fetching children of records with different record caches is not supported.")
# Call the derived class function
if include is None:
include = []
cls._fetch_children_multi(
template_record._client, template_record._record_cache, records, include=include, force_fetch=force_fetch
)
[docs]
def fetch_children(self, include: Optional[Iterable[str]] = None, force_fetch: bool = False):
"""
Fetches all children of this record recursively
"""
self.fetch_children_multi([self], include, force_fetch)
[docs]
def sync_to_cache(self, detach: bool = False):
"""
Syncs this record to the cache
If `detach` is True, then the record will be removed from the cache
"""
if self._record_cache is None:
return
if self._record_cache.read_only:
return
self._record_cache.writeback_record(self)
self._cache_dirty = False
if detach:
self._record_cache = None
def __str__(self) -> str:
return f"<{self.__class__.__name__} id={self.id} status={self.status}>"
[docs]
def propagate_client(self, client):
"""
Propagates a client and related information to this record to any fields within this record that need it
This is expected to be called from derived class propagate_client functions as well
"""
self._client = client
self._base_url = f"api/v1/records/{self.record_type}/{self.id}"
if self.compute_history_ is not None:
for ch in self.compute_history_:
ch.propagate_client(self._client, self._base_url)
if self.native_files_ is not None:
for nf in self.native_files_.values():
nf.propagate_client(self._client, self._base_url)
def _get_child_records(
self,
child_record_ids: Sequence[int],
child_record_type: Type[_Record_T],
include: Optional[Iterable[str]] = None,
) -> List[_Record_T]:
"""
Helper function for obtaining child records either from the cache or from the server
The records are returned in the same order as the `record_ids` parameter.
If `include` is specified, additional fields will be fetched from the server. However, if the records are in the
cache already, they may be missing those fields.
"""
return get_records_with_cache(
self._client, self._record_cache, child_record_type, child_record_ids, include, force_fetch=False
)
def _assert_online(self):
"""Raises an exception if this record does not have an associated client"""
if self.offline:
raise RuntimeError("Record is not connected to a client")
def _fetch_compute_history(self):
self._assert_online()
self.compute_history_ = self._client.make_request(
"get", f"{self._base_url}/compute_history", List[ComputeHistory]
)
self.propagate_client(self._client)
def _fetch_task(self):
if self.is_service:
self.task_ = None
else:
self.task_ = self._client.make_request("get", f"{self._base_url}/task", Optional[RecordTask])
def _fetch_service(self):
if not self.is_service:
self.service_ = None
else:
self.service_ = self._client.make_request("get", f"{self._base_url}/service", Optional[RecordService])
def _fetch_comments(self):
self._assert_online()
self.comments_ = self._client.make_request("get", f"{self._base_url}/comments", List[RecordComment])
def _fetch_native_files(self):
self.native_files_ = self._client.make_request("get", f"{self._base_url}/native_files", Dict[str, NativeFile])
self.propagate_client(self._client)
def _get_output(self, output_type: OutputTypeEnum) -> Optional[Union[str, Dict[str, Any]]]:
history = self.compute_history
if not history:
return None
return history[-1].get_output(output_type)
@property
def offline(self) -> bool:
return self._client is None
@property
def children_status(self) -> Dict[RecordStatusEnum, int]:
"""Returns a dictionary of the status of all children of this record"""
self._assert_online()
return self._client.make_request(
"get",
f"{self._base_url}/children_status",
Dict[RecordStatusEnum, int],
)
@property
def children_errors(self) -> List[BaseRecord]:
"""Returns errored child records"""
self._assert_online()
error_ids = self._client.make_request(
"get",
f"{self._base_url}/children_errors",
List[int],
)
return self._client._get_records_by_type(None, error_ids)
@property
def compute_history(self) -> List[ComputeHistory]:
if self.compute_history_ is None:
self._fetch_compute_history()
return self.compute_history_
@property
def task(self) -> Optional[RecordTask]:
# task_ may be None because it either hasn't been fetched or it doesn't exist
# fetch only if it has been set at some point
if self.task_ is None and "task_" not in self.__fields_set__:
self._fetch_task()
return self.task_
@property
def service(self) -> Optional[RecordService]:
# service_ may be None because it either hasn't been fetched or it doesn't exist
# fetch only if it has been set at some point
if self.service_ is None and "service_" not in self.__fields_set__:
self._fetch_service()
return self.service_
[docs]
def get_waiting_reason(self) -> Dict[str, Any]:
return self._client.make_request("get", f"api/v1/records/{self.id}/waiting_reason", Dict[str, Any])
@property
def comments(self) -> Optional[List[RecordComment]]:
if self.comments_ is None:
self._fetch_comments()
return self.comments_
@property
def native_files(self) -> Optional[Dict[str, NativeFile]]:
if self.native_files_ is None:
self._fetch_native_files()
return self.native_files_
@property
def stdout(self) -> Optional[str]:
return self._get_output(OutputTypeEnum.stdout)
@property
def stderr(self) -> Optional[str]:
return self._get_output(OutputTypeEnum.stderr)
@property
def error(self) -> Optional[Dict[str, Any]]:
return self._get_output(OutputTypeEnum.error)
@property
def provenance(self) -> Optional[Provenance]:
history = self.compute_history
if not history:
return None
return history[-1].provenance
ServiceDependency.update_forward_refs()
_Record_T = TypeVar("_Record_T", bound=BaseRecord)
class RecordAddBodyBase(RestModelBase):
tag: constr(to_lower=True)
priority: PriorityEnum
owner_group: Optional[str]
find_existing: bool = True
class RecordModifyBody(RestModelBase):
record_ids: List[int]
status: Optional[RecordStatusEnum] = None
priority: Optional[PriorityEnum] = None
tag: Optional[str] = None
comment: Optional[str] = None
class RecordDeleteBody(RestModelBase):
record_ids: List[int]
soft_delete: bool
delete_children: bool
class RecordRevertBody(RestModelBase):
revert_status: RecordStatusEnum
record_ids: List[int]
class RecordQueryFilters(QueryModelBase):
record_id: Optional[List[int]] = None
record_type: Optional[List[str]] = None
manager_name: Optional[List[str]] = None
status: Optional[List[RecordStatusEnum]] = None
dataset_id: Optional[List[int]] = None
parent_id: Optional[List[int]] = None
child_id: Optional[List[int]] = None
created_before: Optional[datetime] = None
created_after: Optional[datetime] = None
modified_before: Optional[datetime] = None
modified_after: Optional[datetime] = None
owner_user: Optional[List[Union[int, str]]] = None
owner_group: Optional[List[Union[int, str]]] = None
@validator("created_before", "created_after", "modified_before", "modified_after", pre=True)
def parse_dates(cls, v):
if isinstance(v, str):
return date_parser(v)
return v
[docs]
class RecordQueryIterator(QueryIteratorBase[_Record_T]):
"""
Iterator for all types of record queries
This iterator transparently handles batching and pagination over the results
of a record query, and works with all kinds of records.
"""
def __init__(
self,
client,
query_filters: RecordQueryFilters,
record_type: Type[_Record_T],
include: Optional[Iterable[str]] = None,
):
"""
Construct an iterator
Parameters
----------
client
QCPortal client object used to contact/retrieve data from the server
query_filters
The actual query information to send to the server
record_type
What type of record we are querying for
"""
batch_limit = client.api_limits["get_records"] // 4
self.record_type = record_type
self.include = include
QueryIteratorBase.__init__(self, client, query_filters, batch_limit)
def _request(self) -> List[_Record_T]:
if self.record_type is None:
record_ids = self._client.make_request(
"post",
f"api/v1/records/query",
List[int],
body=self._query_filters,
)
else:
# Get the record type string. This is kind of ugly, but works.
record_type_str = self.record_type.__fields__["record_type"].default
record_ids = self._client.make_request(
"post",
f"api/v1/records/{record_type_str}/query",
List[int],
body=self._query_filters,
)
return self._client._get_records_by_type(self.record_type, record_ids, include=self.include)
def record_from_dict(data: Dict[str, Any], client: Any = None) -> BaseRecord:
"""
Create a record object from a dictionary containing the record information
This determines the appropriate record class (deriving from BaseRecord)
and creates an instance of that class.
"""
record_type = data["record_type"]
cls = BaseRecord.get_subclass(record_type)
return cls(**data, client=client)
def records_from_dicts(
data: Sequence[Optional[Dict[str, Any]]],
client: Any = None,
) -> List[Optional[BaseRecord]]:
"""
Create a list of record objects from a sequence of datamodels
This determines the appropriate record class (deriving from BaseRecord)
and creates an instance of that class.
"""
ret: List[Optional[BaseRecord]] = []
for rd in data:
if rd is None:
ret.append(None)
else:
ret.append(record_from_dict(rd, client))
return ret
