-
Notifications
You must be signed in to change notification settings - Fork 47
/
dataset_models.py
1289 lines (1015 loc) · 42.6 KB
/
dataset_models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
from __future__ import annotations
from datetime import datetime
from typing import Optional, Dict, Any, List, Iterable, Type, Tuple, Union, Callable, ClassVar
import pandas as pd
import pydantic
from pydantic import BaseModel, Extra, validator, PrivateAttr, Field
from qcelemental.models.types import Array
from tabulate import tabulate
from qcportal.base_models import RestModelBase, validate_list_to_single, CommonBulkGetBody
from qcportal.dataset_view import DatasetViewWrapper
from qcportal.metadata_models import DeleteMetadata
from qcportal.record_models import PriorityEnum, RecordStatusEnum, BaseRecord
from qcportal.utils import make_list, chunk_iterable
class Citation(BaseModel):
"""A literature citation."""
class Config:
extra = Extra.forbid
allow_mutation = False
acs_citation: Optional[str] = None # hand-formatted citation in ACS style
bibtex: Optional[str] = None # bibtex blob for later use with bibtex-renderer
doi: Optional[str] = None
url: Optional[str] = None
def to_acs(self) -> str:
"""Returns an ACS-formatted citation"""
return self.acs_citation
class ContributedValues(BaseModel):
class Config:
extra = Extra.forbid
allow_mutation = False
name: str
values: Any
index: Array[str]
values_structure: Dict[str, Any] = {}
theory_level: Union[str, Dict[str, str]]
units: str
theory_level_details: Optional[Union[str, Dict[str, Optional[str]]]] = None
citations: Optional[List[Citation]] = None
external_url: Optional[str] = None
doi: Optional[str] = None
comments: Optional[str] = None
class BaseDataset(BaseModel):
class Config:
extra = Extra.forbid
allow_mutation = True
validate_assignment = True
id: int
dataset_type: str
name: str
description: str
tagline: str
tags: List[str]
group: str
visibility: bool
provenance: Dict[str, Any]
default_tag: str
default_priority: PriorityEnum
owner_user: Optional[str]
owner_group: Optional[str]
metadata: Dict[str, Any]
extras: Dict[str, Any]
########################################
# Caches of information
########################################
entry_names_: List[str] = Field([], alias="entry_names")
# To be overridden by the derived class with more specific types
specifications_: Dict[str, Any] = {}
entries_: Dict[str, Any] = {}
record_map_: Dict[Tuple[str, str], Any] = {}
# Values computed outside QCA
contributed_values_: Optional[Dict[str, ContributedValues]] = None
#############################
# Private non-pydantic fields
#############################
_client: Any = PrivateAttr(None)
_view_data: Optional[DatasetViewWrapper] = PrivateAttr(None)
# To be overridden by the derived classes
_entry_type: ClassVar[Optional[Type]] = None
_specification_type: ClassVar[Optional[Type]] = None
_record_item_type: ClassVar[Optional[Type]] = None
_record_type: ClassVar[Optional[Type]] = None
# A dictionary of all subclasses (dataset types) to the actual class type
_all_subclasses: ClassVar[Dict[str, Type[BaseDataset]]] = {}
# Some dataset options
auto_fetch_missing: bool = True # Automatically fetch missing records from the server
def __init__(self, client=None, view_data=None, **kwargs):
BaseModel.__init__(self, **kwargs)
# Calls derived class propagate_client
# which should filter down to the ones in this (BaseDataset) class
self.propagate_client(client)
assert self._client is client, "Client not set in base dataset class?"
def __init_subclass__(cls):
"""
Register derived classes for later use
"""
# Get the dataset 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)
dataset_type = cls.__fields__["dataset_type"].default
cls._all_subclasses[dataset_type] = cls
@classmethod
def get_subclass(cls, dataset_type: str):
subcls = cls._all_subclasses.get(dataset_type)
if subcls is None:
raise RuntimeError(f"Cannot find subclass for record type {dataset_type}")
return subcls
def propagate_client(self, client):
"""
Propagates a client to this record to any fields within this record that need it
This may also be called from derived class propagate_client functions as well
"""
self._client = client
for record in self.record_map_.values():
record.populate_client(client)
def _post_add_entries(self, entry_names) -> None:
"""
Perform some actions after entries have been added to the remote server
Parameters
----------
entry_names
Names of the new entries that have been added to the remote server
"""
# If entry names have been fetched, add the new entry names
# This should still be ok if there are no entries - they will be fetched if the list is empty
if self.entry_names_:
self.entry_names_.extend(x for x in entry_names if x not in self.entry_names_)
def _post_add_specification(self, specification_name) -> None:
"""
Perform some actions after specifications have been added to the remote server
Parameters
----------
specification_name
Name of the new specification that has been added to the remote server
"""
# Ignoring the function argument for now... Just get all specs
# TODO - very inefficient for lots of specs
self.fetch_specifications()
def _update_metadata(self, **kwargs):
self.assert_online()
new_body = {
"name": self.name,
"description": self.description,
"tagline": self.tagline,
"tags": self.tags,
"group": self.group,
"visibility": self.visibility,
"provenance": self.provenance,
"default_tag": self.default_tag,
"default_priority": self.default_priority,
"metadata": self.metadata,
}
new_body.update(**kwargs)
body = DatasetModifyMetadata(**new_body)
self._client.make_request("patch", f"api/v1/datasets/{self.dataset_type}/{self.id}", None, body=body)
self.name = body.name
self.description = body.description
self.tagline = body.tagline
self.tags = body.tags
self.group = body.group
self.visibility = body.visibility
self.provenance = body.provenance
self.default_tag = body.default_tag
self.default_priority = body.default_priority
self.metadata = body.metadata
def submit(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
specification_names: Optional[Union[str, Iterable[str]]] = None,
tag: Optional[str] = None,
priority: PriorityEnum = None,
find_existing: bool = True,
):
self.assert_is_not_view()
self.assert_online()
body_data = DatasetSubmitBody(
entry_names=make_list(entry_names),
specification_names=make_list(specification_names),
tag=tag,
priority=priority,
find_existing=find_existing,
)
return self._client.make_request(
"post", f"api/v1/datasets/{self.dataset_type}/{self.id}/submit", Any, body=body_data
)
#########################################
# Various properties and getters/setters
#########################################
def status(self) -> Dict[str, Any]:
self.assert_online()
return self._client.make_request(
"get",
f"api/v1/datasets/{self.dataset_type}/{self.id}/status",
Dict[str, Dict[RecordStatusEnum, int]],
)
def status_table(self) -> str:
"""
Returns the status of the dataset's computations as a table (in a string)
"""
ds_status = self.status()
all_status = {x for y in ds_status.values() for x in y}
ordered_status = RecordStatusEnum.make_ordered_status(all_status)
headers = ["specification"] + [x.value for x in ordered_status]
table = []
for spec, spec_statuses in sorted(ds_status.items()):
row = [spec]
row.extend(spec_statuses.get(s, "") for s in ordered_status)
table.append(row)
return tabulate(table, headers=headers, stralign="right")
def print_status(self) -> None:
print(self.status_table())
def detailed_status(self) -> List[Tuple[str, str, RecordStatusEnum]]:
self.assert_online()
return self._client.make_request(
"get",
f"api/v1/datasets/{self.dataset_type}/{self.id}/detailed_status",
List[Tuple[str, str, RecordStatusEnum]],
)
@property
def offline(self) -> bool:
return self._client is None
def assert_online(self):
if self.offline:
raise RuntimeError("Dataset does not connected to a QCFractal server")
@property
def is_view(self) -> bool:
return self._view_data is not None
def assert_is_not_view(self):
if self.is_view:
raise RuntimeError("Dataset loaded from an offline view")
def set_name(self, new_name: str):
self._update_metadata(name=new_name)
def set_description(self, new_description: str):
self._update_metadata(description=new_description)
def set_visibility(self, new_visibility: bool):
self._update_metadata(visibility=new_visibility)
def set_group(self, new_group: str):
self._update_metadata(group=new_group)
def set_tags(self, new_tags: List[str]):
self._update_metadata(tags=new_tags)
def set_tagline(self, new_tagline: str):
self._update_metadata(tagline=new_tagline)
def set_provenance(self, new_provenance: Dict[str, Any]):
self._update_metadata(provenance=new_provenance)
def set_metadata(self, new_metadata: Dict[str, Any]):
self._update_metadata(metadata=new_metadata)
def set_default_tag(self, new_default_tag: str):
self._update_metadata(default_tag=new_default_tag)
def set_default_priority(self, new_default_priority: PriorityEnum):
self._update_metadata(default_priority=new_default_priority)
###################################
# Specifications
###################################
@property
def specifications(self) -> Dict[str, Any]:
if self.is_view:
return self._view_data.get_specifications(self._specification_type)
else:
if not self.specifications_:
self.fetch_specifications()
return self.specifications_
@property
def specification_names(self) -> List[str]:
return list(self.specifications.keys())
def fetch_specifications(self) -> None:
"""
Fetch all specifications from the remote server
These are fetched and then stored internally, and not returned.
"""
self.assert_is_not_view()
self.assert_online()
self.specifications_ = self._client.make_request(
"get",
f"api/v1/datasets/{self.dataset_type}/{self.id}/specifications",
Dict[str, self._specification_type],
)
def rename_specification(self, old_name: str, new_name: str):
self.assert_is_not_view()
self.assert_online()
name_map = {old_name: new_name}
self._client.make_request(
"patch", f"api/v1/datasets/{self.dataset_type}/{self.id}/specifications", None, body=name_map
)
self.specifications_ = {name_map.get(k, k): v for k, v in self.specifications_.items()}
# Renames the specifications in the record map
self.record_map_ = {(e, name_map.get(s, s)): r for (e, s), r in self.record_map_.items()}
def delete_specification(self, name: str, delete_records: bool = False) -> DeleteMetadata:
self.assert_is_not_view()
self.assert_online()
body = DatasetDeleteStrBody(names=[name], delete_records=delete_records)
ret = self._client.make_request(
"post",
f"api/v1/datasets/{self.dataset_type}/{self.id}/specifications/bulkDelete",
DeleteMetadata,
body=body,
)
# Delete locally-cached stuff
self.specifications_.pop(name, None)
self.record_map_ = {(e, s): r for (e, s), r in self.record_map_.items() if s != name}
return ret
###################################
# Entries
###################################
def fetch_entry_names(self) -> None:
"""
Fetch all entry names from the remote server
These are fetched and then stored internally, and not returned.
"""
self.assert_is_not_view()
self.assert_online()
self.entry_names_ = self._client.make_request(
"get",
f"api/v1/datasets/{self.dataset_type}/{self.id}/entry_names",
List[str],
)
def _internal_fetch_entries(
self,
entry_names: Iterable[str],
) -> None:
"""
Fetches entry information from the remote server, storing it internally
This does not do any checking for existing entries, but is used to actually
request the data from the server.
Note: This function does not do any batching w.r.t. server API limits. It is expected that is done
before this function is called.
Parameters
----------
entry_names
Names of the entries to fetch
"""
if not entry_names:
return
body = DatasetFetchEntryBody(names=entry_names)
fetched_entries = self._client.make_request(
"post",
f"api/v1/datasets/{self.dataset_type}/{self.id}/entries/bulkFetch",
Dict[str, self._entry_type],
body=body,
)
self.entries_.update(fetched_entries)
def fetch_entries(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
force_refetch: bool = False,
) -> None:
"""
Fetches entry information from the remote server, storing it internally
By default, already-fetched entries will not be fetched again, unless
`force_refetch` is True.
Parameters
----------
entry_names
Names of entries to fetch. If None, fetch all entries
force_refetch
If true, fetch data from the server even if it already exists locally
"""
self.assert_is_not_view()
self.assert_online()
# Reload entry names if we are forcing refetching
if force_refetch:
self.fetch_entry_names()
# if not specified, do all entries
if entry_names is None:
entry_names = self.entry_names
else:
entry_names = make_list(entry_names)
# Strip out existing entries if we aren't forcing refetching
if not force_refetch:
entry_names = [x for x in entry_names if x not in self.entries_]
batch_size: int = self._client.api_limits["get_dataset_entries"] // 4
for entry_names_batch in chunk_iterable(entry_names, batch_size):
self._internal_fetch_entries(entry_names_batch)
def get_entry(
self,
entry_name: str,
force_refetch: bool = False,
) -> Optional[Any]:
"""
Obtain entry information
The entry will be automatically fetched from the remote server if needed.
"""
if self.is_view:
entry_dict = self._view_data.get_entries(self._entry_type, [entry_name])
return entry_dict.get(entry_name)
else:
self.fetch_entries(entry_name, force_refetch=force_refetch)
return self.entries_.get(entry_name, None)
def iterate_entries(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
force_refetch: bool = False,
):
"""
Iterate over all entries
This is used as a generator, and automatically fetches entries as needed
Parameters
----------
entry_names
Names of entries to iterate over. If None, iterate over all entries
force_refetch
If true, fetch data from the server even if it already exists locally
"""
#########################################################
# We duplicate a little bit of fetch_entries here, since
# we want to yield in the middle
#########################################################
# Reload entry names if we are forcing refetching
# Nothing to fetch if this is a view
if force_refetch and not self.is_view:
self.fetch_entry_names()
# if not specified, do all entries
if entry_names is None:
entry_names = self.entry_names
else:
entry_names = make_list(entry_names)
if self.is_view:
# Go one at a time. No need to "fetch"
for entry_name in entry_names:
entry = self.get_entry(entry_name)
if entry is not None:
yield entry
else:
# Fetch from server
batch_size: int = self._client.api_limits["get_dataset_entries"] // 4
if self.entries_ is None:
self.entries_ = {}
for entry_names_batch in chunk_iterable(entry_names, batch_size):
# If forcing refetching, then use the whole batch. Otherwise, strip out
# any existing entries
if force_refetch:
names_tofetch = entry_names_batch
else:
names_tofetch = [x for x in entry_names_batch if x not in self.entries_]
self._internal_fetch_entries(names_tofetch)
# Loop over the whole batch (not just what we fetched)
for entry_name in entry_names_batch:
entry = self.entries_.get(entry_name, None)
if entry is not None:
yield entry
@property
def entry_names(self) -> List[str]:
if self.is_view:
return self._view_data.get_entry_names()
if not self.entry_names_:
self.fetch_entry_names()
return self.entry_names_
def rename_entries(self, name_map: Dict[str, str]):
self.assert_is_not_view()
self.assert_online()
self._client.make_request(
"patch", f"api/v1/datasets/{self.dataset_type}/{self.id}/entries", None, body=name_map
)
# rename locally cached entries and stuff
self.entry_names_ = [name_map.get(x, x) for x in self.entry_names_]
self.entries_ = {name_map.get(k, k): v for k, v in self.entries_.items()}
# Renames the entries in the record map
self.record_map_ = {(name_map.get(e, e), s): r for (e, s), r in self.record_map_.items()}
def delete_entries(self, names: Union[str, Iterable[str]], delete_records: bool = False) -> DeleteMetadata:
self.assert_is_not_view()
self.assert_online()
names = make_list(names)
body = DatasetDeleteStrBody(names=names, delete_records=delete_records)
ret = self._client.make_request(
"post",
f"api/v1/datasets/{self.dataset_type}/{self.id}/entries/bulkDelete",
DeleteMetadata,
body=body,
)
# Delete locally-cached stuff
self.entry_names_ = [x for x in self.entry_names_ if x not in names]
self.entries_ = {x: y for x, y in self.entries_.items() if x not in names}
self.record_map_ = {(e, s): r for (e, s), r in self.record_map_.items() if e not in names}
return ret
###########################
# Records
###########################
def _lookup_record(self, entry_name: str, specification_name: str):
return self.record_map_.get((entry_name, specification_name), None)
def _internal_fetch_records(
self,
entry_names: Iterable[str],
specification_names: Iterable[str],
status: Optional[Iterable[RecordStatusEnum]],
include: Optional[Iterable[str]],
) -> None:
"""
Fetches record information from the remote server, storing it internally
This does not do any checking for existing records, but is used to actually
request the data from the server.
Note: This function does not do any batching w.r.t. server API limits. It is expected that is done
before this function is called.
Parameters
----------
entry_names
Names of the entries whose records to fetch. If None, fetch all entries
specification_names
Names of the specifications whose records to fetch. If None, fetch all specifications
status
Fetch only records with these statuses
include
Additional fields/data to include when fetch the entry
"""
if not (entry_names and specification_names):
return
body = DatasetFetchRecordsBody(entry_names=entry_names, specification_names=specification_names, status=status)
record_info = self._client.make_request(
"post",
f"api/v1/datasets/{self.dataset_type}/{self.id}/records/bulkFetch",
List[Tuple[str, str, int]], # (entry_name, spec_name, record_id)
body=body,
)
record_ids = [x[2] for x in record_info]
records = self._client.get_records(record_ids, include=include)
# Update the locally-stored records
for rec_item, rec in zip(record_info, records):
assert rec_item[2] == rec.id
self.record_map_[(rec_item[0], rec_item[1])] = rec
def _internal_update_records(
self,
entry_names: Iterable[str],
specification_names: Iterable[str],
status: Optional[Iterable[RecordStatusEnum]],
include: Optional[Iterable[str]],
):
"""
Update local record information if it has been modified on the server
Note: This function does not do any batching w.r.t. server API limits. It is expected that is done
before this function is called.
Parameters
----------
entry_names
Names of the entries whose records to update. If None, fetch all entries
specification_names
Names of the specifications whose records to update. If None, fetch all specifications
status
Update records that have this status on the server. If None, update records with any status on the server
include
Additional fields/data to include when fetch the entry
"""
if not (entry_names and specification_names):
return
# Get all the record ids that we store that correspond to the entries/specs
existing_record_info = [
(e, s, r) for (e, s), r in self.record_map_.items() if e in entry_names and s in specification_names
]
# Subset that we should check for updated records on the server
# (completed and invalid rarely change)
updateable_record_ids = [
r.id
for e, s, r in existing_record_info
if r.status not in (RecordStatusEnum.complete, RecordStatusEnum.invalid)
]
# Do a raw call to the records/bulkGet endpoint. This allows us to only get
# the 'modified_on' and 'status' fields
batch_size = self._client.api_limits["get_records"] // 4
minfo_dict: Dict[int, datetime] = {} # record_id -> modified time
for record_id_batch in chunk_iterable(updateable_record_ids, batch_size):
body = CommonBulkGetBody(ids=record_id_batch, include=["modified_on", "status"])
modified_info = self._client.make_request(
"post",
f"api/v1/records/bulkGet",
List[Dict[str, Any]],
body=body,
)
# Too lazy to look up how pydantic stores datetime, so use pydantic to parse it
for m in modified_info:
# Only store if the status on the server matches what the caller wants
if status is None or m["status"] in status:
minfo_dict[m["id"]] = pydantic.parse_obj_as(datetime, m["modified_on"])
# Which ones need to be fully updated
need_updating = []
for entry_name, spec_name, existing_record in existing_record_info:
server_mtime = minfo_dict.get(existing_record.id, None)
# Perhaps a status mismatch (status on server isn't one we want)
if server_mtime is None:
continue
if existing_record.modified_on < server_mtime:
need_updating.append((entry_name, spec_name))
# Go via one spec at a time
for spec_name in specification_names:
entries_to_update = [x[0] for x in need_updating if x[1] == spec_name]
for entries_to_update_batch in chunk_iterable(entries_to_update, batch_size):
self._internal_fetch_records(entries_to_update_batch, [spec_name], None, include)
def fetch_records(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
specification_names: Optional[Union[str, Iterable[str]]] = None,
status: Optional[Union[RecordStatusEnum, Iterable[RecordStatusEnum]]] = None,
include: Optional[Iterable[str]] = None,
fetch_updated: bool = True,
force_refetch: bool = False,
):
"""
Fetches record information from the remote server, storing it internally
By default, this function will only fetch records that have not been fetch previously.
If `force_refetch` is True, then this will always fetch the records.
Parameters
----------
entry_names
Names of the entries whose records to fetch. If None, fetch all entries
specification_names
Names of the specifications whose records to fetch. If None, fetch all specifications
status
Fetch only records with these statuses
include
Additional fields to include in the returned record
fetch_updated
Fetch any records that exist locally but have been updated on the server
force_refetch
If true, fetch data from the server even if it already exists locally
"""
self.assert_is_not_view()
self.assert_online()
# Reload entry names if we are forcing refetching
if force_refetch:
self.fetch_entry_names()
self.fetch_specifications()
status = make_list(status)
# if not specified, do all entries and specs
if entry_names is None:
entry_names = self.entry_names
else:
entry_names = make_list(entry_names)
if specification_names is None:
specification_names = self.specification_names
else:
specification_names = make_list(specification_names)
# Determine the number of entries in each batch
# Assume there are many more entries than specifications, and that
# everything has been submitted
# Divide by 4 to go easy on the server
batch_size: int = self._client.api_limits["get_records"] // 4
# Do all entries for one spec. This simplifies things, especially with handling
# existing or update-able records
for spec_name in specification_names:
for entry_names_batch in chunk_iterable(entry_names, batch_size):
# Handle existing records that need to be updated
if fetch_updated and not force_refetch:
self._internal_update_records(entry_names_batch, [spec_name], status, include)
# Prune records that already exist, and then fetch them
if not force_refetch:
entry_names_batch = [x for x in entry_names_batch if (x, spec_name) not in self.record_map_]
if entry_names_batch:
self._internal_fetch_records(entry_names_batch, [spec_name], status, include)
def get_record(
self,
entry_name: str,
specification_name: str,
include: Optional[Iterable[str]] = None,
force_refetch: bool = False,
) -> Optional[BaseRecord]:
"""
Obtain a calculation record related to this dataset
The record will be automatically fetched from the remote server if needed.
"""
if self.is_view:
record_item = self._view_data.get_record_item(self._record_item_type, entry_name, specification_name)
if record_item is None:
return None
else:
return record_item.record
else:
self.fetch_records(entry_name, specification_name, include=include, force_refetch=force_refetch)
return self._lookup_record(entry_name, specification_name)
def iterate_records(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
specification_names: Optional[Union[str, Iterable[str]]] = None,
status: Optional[Union[RecordStatusEnum, Iterable[RecordStatusEnum]]] = None,
include: Optional[Iterable[str]] = None,
fetch_updated: bool = True,
force_refetch: bool = False,
):
#########################################################
# We duplicate a little bit of fetch_records here, since
# we want to yield in the middle
#########################################################
# Get an up-to-date list of entry names and specifications
# Nothing to fetch if this is a view
if force_refetch and not self.is_view:
self.fetch_entry_names()
self.fetch_specifications()
status = make_list(status)
# if not specified, do all entries and specs
if entry_names is None:
entry_names = self.entry_names
else:
entry_names = make_list(entry_names)
if specification_names is None:
specification_names = self.specification_names
else:
specification_names = make_list(specification_names)
if self.is_view:
for spec_name in specification_names:
for entry_name in entry_names:
rec = self.get_record(entry_name, spec_name)
if rec is None:
continue
if status is None or rec.status in status:
yield entry_name, spec_name, rec
else:
# Smaller fetch limit for iteration (than in fetch_records)
fetch_limit: int = self._client.api_limits["get_records"] // 10
n_entries = len(entry_names)
for spec_name in specification_names:
for start_idx in range(0, n_entries, fetch_limit):
entries_batch = entry_names[start_idx : start_idx + fetch_limit]
# Handle existing records that need to be updated
if fetch_updated and not force_refetch:
existing_batch = [x for x in entries_batch if (x, spec_name) in self.record_map_]
self._internal_update_records(existing_batch, [spec_name], status, include)
if force_refetch:
batch_tofetch = entries_batch
else:
# Filter if they already exist
batch_tofetch = [x for x in entries_batch if (x, spec_name) not in self.record_map_]
self._internal_fetch_records(batch_tofetch, [spec_name], status, include)
# Now lookup the just-fetched records and yield them
for entry_name in entries_batch:
rec = self._lookup_record(entry_name, spec_name)
if rec is None:
continue
if status is None or rec.status in status:
yield entry_name, spec_name, rec
def remove_records(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
specification_names: Optional[Union[str, Iterable[str]]] = None,
delete_records: bool = False,
) -> DeleteMetadata:
self.assert_is_not_view()
self.assert_online()
body = DatasetRemoveRecordsBody(
entry_names=make_list(entry_names),
specification_names=make_list(specification_names),
delete_records=delete_records,
)
ret = self._client.make_request(
"post",
f"api/v1/datasets/{self.dataset_type}/{self.id}/records/bulkDelete",
None,
body=body,
)
# Delete locally-cached stuff
self.record_map_ = {
(e, s): r for (e, s), r in self.record_map_.items() if e not in entry_names and s not in specification_names
}
return ret
def modify_records(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
specification_names: Optional[Union[str, Iterable[str]]] = None,
new_tag: Optional[str] = None,
new_priority: Optional[PriorityEnum] = None,
new_comment: Optional[str] = None,
*,
refetch_records: bool = False,
):
self.assert_is_not_view()
self.assert_online()
body = DatasetRecordModifyBody(
entry_names=make_list(entry_names),
specification_names=make_list(specification_names),
tag=new_tag,
priority=new_priority,
comment=new_comment,
)
ret = self._client.make_request(
"patch",
f"api/v1/datasets/{self.dataset_type}/{self.id}/records",
None,
body=body,
)
if refetch_records:
self.fetch_records(entry_names, specification_names)
return ret
def reset_records(
self,
entry_names: Optional[Union[str, Iterable[str]]] = None,
specification_names: Optional[Union[str, Iterable[str]]] = None,
*,
refetch_records: bool = False,
):
self.assert_is_not_view()
self.assert_online()
body = DatasetRecordModifyBody(
entry_names=make_list(entry_names),
specification_names=make_list(specification_names),
status=RecordStatusEnum.waiting,
)
ret = self._client.make_request(
"patch",
f"api/v1/datasets/{self.dataset_type}/{self.id}/records",
None,
body=body,
)