Skip to content
Snippets Groups Projects
Commit ecad5416 authored by Daan Bijl's avatar Daan Bijl
Browse files

draft implementation sqdl-writer

parent a9e8fb9f
No related branches found
No related tags found
No related merge requests found
......@@ -44,6 +44,7 @@ class MetadataFormatter:
'variables_measured': self.validate(desc['vars']),
'dimensions': self.validate(desc['dims']),
# TODO more metadata ?
"fridge": "test-parameter", # todo:
}
if "project" in desc:
metadata["project"] = desc["project"]
......
......@@ -4,13 +4,22 @@ import re
import time
from collections.abc import Mapping
from datetime import datetime
from typing import Dict
from .dataset_scanner import DatasetScanner, FileInfo
from .exceptions import InvalidNameError, NoScopeError, DatasetError
from .metadata_formatter import MetadataFormatter
from core_tools.data.sqdl.model.task_queue import TaskQueueOperations
from core_tools.data.sqdl.model.upload import UploadOperations
from core_tools.data.sqdl.model.log import LogOperations
import psutil
import core_tools as ct
from core_tools.startup.config import get_configuration
from sqdl_client.api.v1.dataset import Dataset
from sqdl_client.api.v1.file import File
from sqdl_client.api.v1.scope import Scope
from sqdl_client.client import QDLClient
from sqdl_client.exceptions import (
ObjectNotFoundException,
......@@ -18,22 +27,14 @@ from sqdl_client.exceptions import (
UniqueConstraintViolationException
)
from requests.exceptions import ConnectionError, ReadTimeout
from .dataset_scanner import DatasetScanner, FileInfo
from .exceptions import InvalidNameError, NoScopeError, DatasetError
from .metadata_formatter import MetadataFormatter
from .uploader_db import UploaderDb
from .uploader_task_queue import UploaderTaskQueue
from .uploader_registry import UploadRegistry
from .upload_logger import UploadLogger
from psycopg2._psycopg import connection as Connection
logger = logging.getLogger(__name__)
class SqdlUploader:
def __init__(self, cfg, client=None):
def __init__(self, cfg: Dict, conn: Connection, client=None):
self.cfg = cfg
if client is None:
self.client = QDLClient()
......@@ -44,16 +45,18 @@ class SqdlUploader:
if api_key:
self.client.use_api_key(api_key)
# db_path = self.cfg.get('uploader.database', '~/.sqdl_uploader/uploader.db')
self.db = UploaderDb(self.cfg)
self.task_queue = UploaderTaskQueue(self.db)
self.upload_registry = UploadRegistry(self.db)
self.logger = UploadLogger(self.db)
self.connection = conn
self.task_queue = TaskQueueOperations()
self.upload_registry = UploadOperations()
self.logger = LogOperations()
self.metadata_formatter = MetadataFormatter()
# load scopes to fix them when not set during export
self.scopes = cfg.get('export.scopes', {})
if cfg.get('uploader.retry_failed', False):
self.task_queue.retry_all_failed()
with self.connection:
c = self.connection.cursor()
self.task_queue.retry_all_failed(c)
self.pid = os.getpid()
self.cleanup_abandoned_tasks()
logger.info(f"Started uploader, pid:{self.pid}")
......@@ -62,26 +65,33 @@ class SqdlUploader:
def process_task(self) -> bool:
start = time.perf_counter()
task = self.task_queue.get_oldest_task(self.pid)
if task is None:
task = self.task_queue.get_newest_retry_task(self.pid)
# task = self.task_queue.get_oldest_task(self.pid)
with self.connection:
c = self.connection.cursor()
task = self.task_queue.claim_oldest_task(c, self.pid)
if task is None:
task = self.task_queue.claim_newest_retry_task(c, self.pid)
if task is None:
return False
try:
duration = time.perf_counter() - start
logger.info(f'Uploading {task.uid} (query: {duration * 1000:.1f} ms) {task.ds_path}')
ds_scanner = DatasetScanner(task.uid, task.ds_path)
logger.info(f'Uploading {task.coretools_uid} (query: {duration * 1000:.1f} ms) {task.dataset_path}')
ds_scanner = DatasetScanner(task.coretools_uid, task.dataset_path)
desc = ds_scanner.get_description()
if not task.scope:
task.scope = self.get_scope(desc)
self.log(task, f"Resolved scope using project {desc['project']}")
logger.debug(f'Get/create {task.uid}')
logger.debug(f'Get/create {task.coretools_uid}')
# get / create sQDL dataset
sqdl_ds = self.get_create_sqdl_dataset(task.scope, desc)
if task.update_dataset or task.set_raw_final:
ds_upload_id = self.upload_registry.get_create_dataset(sqdl_ds.uuid, task.scope, task.uid)
if task.update_dataset or task.is_ready:
with self.connection:
c = self.connection.cursor()
ds_upload_id = self.upload_registry.create_dataset(c, task.scope, task.coretools_uid, sqdl_ds.uuid)
files = self.sort_files(ds_scanner.get_files(), desc)
self.upload_files(task, ds_upload_id, files, sqdl_ds)
......@@ -91,34 +101,43 @@ class SqdlUploader:
if task.update_rating:
new_rating = 1 if desc['starred'] else 0
if new_rating != sqdl_ds.rating:
logger.info(f"update rating {task.uid} {sqdl_ds.rating} -> {new_rating}")
logger.info(f"update rating {task.coretools_uid} {sqdl_ds.rating} -> {new_rating}")
sqdl_ds.update_rating(new_rating)
deleted = self.task_queue.delete_task(task)
if not deleted:
logger.debug(f'Task {task.uid} has been modified during upload. Release task')
self.task_queue.release_task(task)
with self.connection:
c = self.connection.cursor()
deleted = self.task_queue.delete_task(c, task)
if not deleted:
logger.debug(f'Task {task.coretools_uid} has been modified during upload. Release task')
self.task_queue.release_task(c, task)
# log success
self.log(task, 'Uploaded')
duration = time.perf_counter() - start
logger.info(f'Uploaded {task.uid} in {duration * 1000:5.1f} ms')
logger.info(f'Uploaded {task.coretools_uid} in {duration * 1000:5.1f} ms')
except DatasetError as ex:
logger.error(f"Exception processing {task.uid} '{ex}' {task.ds_path}")
self.task_queue.set_failed(task)
logger.error(f"Exception processing {task.coretools_uid} '{ex}' {task.dataset_path}")
with self.connection:
c = self.connection.cursor()
self.task_queue.set_failed(c, task)
self.log(task, f'{type(ex)}: {ex}')
except (ConnectionError, ReadTimeout) as ex:
# server cannot be reached
logger.error(f"Exception processing {task.uid} '{ex}'", exc_info=True)
self.task_queue.release_task(task)
logger.error(f"Exception processing {task.coretools_uid} '{ex}'", exc_info=True)
with self.connection:
c = self.connection.cursor()
self.task_queue.release_task(c, task)
time.sleep(1.0)
except RequestException as ex:
logger.error(f'Exception processing {task.uid} {task.ds_path}. Response:{ex.response}', exc_info=True)
logger.error(f'Exception processing {task.coretools_uid} {task.dataset_path}. Response:{ex.response}', exc_info=True)
if ex.response is not None:
logger.info(f'Response: {ex.response.url}; {ex.response.headers}')
self.task_queue.set_failed(task)
with self.connection:
c = self.connection.cursor()
self.task_queue.set_failed(c, task)
self.log(task, f'{type(ex)}: {ex}')
time.sleep(0.5)
......@@ -126,14 +145,18 @@ class SqdlUploader:
# TODO: Catch all should be split in dataset related errors and connection errors @@@
# database connection failures, sQDL connecton failures should be given a retry.
# dataset errors should mark the task as failed.
logger.error(f'Exception processing {task.uid} {task.ds_path}', exc_info=True)
self.task_queue.set_failed(task)
logger.error(f'Exception processing {task.coretools_uid} {task.dataset_path}', exc_info=True)
with self.connection:
c = self.connection.cursor()
self.task_queue.set_failed(c, task)
self.log(task, f'{type(ex)}: {ex}')
time.sleep(0.5)
return True
def log(self, task, message):
self.logger.log(task.scope, task.uid, message)
with self.connection:
c = self.connection.cursor()
self.logger.log(c, task.scope, task.coretools_uid, message)
def get_scope(self, desc):
# is it in the json file?
......@@ -150,10 +173,17 @@ class SqdlUploader:
self._validate_dataset_name(desc['name'])
metadata = self.metadata_formatter.format(desc)
sqdl_api = self.client.api
logger.info(sqdl_api.version)
logger.info(sqdl_api.get_user_info())
try:
logger.info("retrieving scope by name: '{}'".format(scope_name))
scope = sqdl_api.scope.retrieve_from_name(scope_name)
except ObjectNotFoundException:
logger.warning("no scope of corresponding name")
scope = None
if scope is None:
scope = sqdl_api.scope.create(
name=scope_name,
......@@ -213,10 +243,12 @@ class SqdlUploader:
def upload_files(self, task, ds_upload_id, file_entries: list[FileInfo], sqdl_ds: Dataset):
uploaded_files_list = self.upload_registry.get_files(ds_upload_id)
logger.debug(f"{len(uploaded_files_list)} uploaded files registry")
uploaded_files = {uf.filename: uf for uf in self.upload_registry.get_files(ds_upload_id)}
logger.debug(f"{len(uploaded_files)} different files in registry for {task.uid}: {[uploaded_files.keys()]}")
with self.connection:
c = self.connection.cursor()
uploaded_files_list = self.upload_registry.get_files_for_dataset(c, parent_idx=ds_upload_id)
uploaded_files = {uf.filename: uf for uf in uploaded_files_list}
logger.debug(f"{len(uploaded_files)} different files in registry for {task.coretools_uid}: {[uploaded_files.keys()]}")
# get file list of sQDL dataset
sqdl_file_list = sqdl_ds.files
......@@ -238,17 +270,35 @@ class SqdlUploader:
if sqdl_file is None:
sqdl_file = sqdl_ds.create_new_file(fi.name, fi.file_type, fi.mimetype,
sequence_number=fi.seq_number)
elif not sqdl_file.is_mutable:
logger.error(f"Cannot upload modified file '{fi.name}'. It's immutable. {task}; {fi}; {uploaded_files.get(fi.name)}")
continue
make_immutable = fi.file_type == 'raw' and task.set_raw_final
logger.info("FILE DATA: file url '{}'".format(sqdl_file.url))
logger.info("FILE DATA: file name '{}'".format(sqdl_file.name))
logger.info("FILE DATA: presigned url '{}'".format(sqdl_file.presigned_url))
logger.info("FILE DATA: presigned upload data:")
for k, v in sqdl_file.presigned_upload_data.items():
logger.info("FILE DATA: key = '{}', value = '{}'".format(k, v))
make_immutable = fi.file_type == 'raw' and task.is_ready
self.upload_file(fi, sqdl_file, make_immutable)
self.upload_registry.add_update_file(ds_upload_id, sqdl_file.uuid, fi.name, fi.st_mtime_us)
with self.connection:
c = self.connection.cursor()
self.upload_registry.create_or_update_file(
c,
parent_idx=ds_upload_id,
sqdl_uuid=sqdl_file.uuid,
filename=fi.name,
last_modified=fi.st_mtime_us
)
elif fi.name not in sqdl_files:
# Strange: it is registered as uploaded, but not there?
raise Exception(f'File {fi.name} of {task.scope}:{task.uid} missing in sQDL')
raise Exception(f'File {fi.name} of {task.scope}:{task.coretools_uid} missing in sQDL')
def upload_file(self, fi: FileInfo, sqdl_file: File, make_immutable):
def upload_file(self, fi: FileInfo, sqdl_file: File, make_immutable: bool):
tries = 2
while tries:
try:
......@@ -265,13 +315,21 @@ class SqdlUploader:
def cleanup_abandoned_tasks(self):
pids = psutil.pids()
for task in self.task_queue.get_claimed_tasks():
alive = task.claimed_by in pids
logger.info(f"task (uid:{task.uid}) is claimed by {task.claimed_by}, alive: {alive}")
with self.connection:
c = self.connection.cursor()
tasks = self.task_queue.get_claimed_tasks(c)
for task in tasks:
alive = task.is_claimed_by in pids
logger.info(f"task (uid:{task.coretools_uid}) is claimed by {task.is_claimed_by}, alive: {alive}")
if not alive:
self.task_queue.release_task(task)
# todo: map more efficiently - collect all dead tasks, then release them all in one transaction, instead of doing a transaction per loop
with self.connection:
c = self.connection.cursor()
self.task_queue.release_task(c, task)
def poll(self) -> None:
def poll(self, conn: Connection) -> None:
self.connection = conn
# NOTE: KeyboardInterrupt and SystemExit will not be caught.
try:
work_done = self.process_task()
......
import os
import logging
from sqlalchemy import create_engine, URL
from sqlalchemy.orm import sessionmaker, Session
from core_tools.data.sqdl.model import create_database
logger = logging.getLogger(__name__)
class UploaderDb:
echo_sql = False # enable for debugging
def __init__(self, config: dict = None, db_file=None):
# if db_file is None:
# db_file = ':memory:'
# else:
# db_file = os.path.expanduser(db_file)
# print(db_file)
# os.makedirs(os.path.dirname(db_file), exist_ok=True)
#
# self.engine = create_engine("sqlite+pysqlite:///" + db_file,
# echo=UploaderDb.echo_sql,
# )
if db_file is not None:
logger.warning("SQDL Uploader no longer uses SQLite: db_file parameter depricated")
if config is None:
config = {}
connection_url = URL.create(
drivername="postgresql+psycopg2",
host=config.get("host", default="localhost"),
port=config.get("port", default=5432),
username=config.get("username", default="dbijl"),
database=config.get("database", default="core-tools")
)
self.engine = create_engine(
connection_url,
echo=UploaderDb.echo_sql,
)
create_database(self.engine)
self.sessionmaker = sessionmaker(bind=self.engine)
def session(self) -> Session:
return self.sessionmaker()
from sqlalchemy import select, insert, delete, update, func
from sqlalchemy.exc import IntegrityError
from sqlalchemy.dialects.postgresql import insert as psql_insert
from core_tools.data.sqdl.model import UploadedDataset, UploadedFile
from core_tools.data.sqdl.uploader_db import UploaderDb
# %%
class UploadRegistry:
def __init__(self, db: UploaderDb):
self.db = db
def get_create_dataset(self, sqdl_uuid, scope, uid) -> int:
try:
with self.db.session() as session:
stmt = insert(UploadedDataset).values(sqdl_uuid=sqdl_uuid, scope=scope, uid=uid)
result = session.scalars(stmt.returning(UploadedDataset.id)).first()
session.commit()
return result
except IntegrityError:
with self.db.session() as session:
stmt = select(UploadedDataset.id).where(UploadedDataset.sqdl_uuid == sqdl_uuid)
return session.scalars(stmt).first()
def get_files(self, dataset_id) -> list[UploadedFile]:
with self.db.session() as session:
stmt = select(UploadedFile).where(UploadedFile.dataset_id == dataset_id)
return session.scalars(stmt).all()
def add_update_file(self, dataset_id, file_sqdl_uuid, filename, st_mtime_ns) -> None:
with self.db.session() as session:
# stmt = sqlite_upsert(UploadedFile).values(
stmt = psql_insert(UploadedFile).values(
dataset_id=dataset_id,
sqdl_uuid=file_sqdl_uuid,
filename=filename,
st_mtime_us=st_mtime_ns // 1000)
stmt = stmt.on_conflict_do_update(
constraint="uploaded_file_sqdl_uuid_key",
set_=dict(sqdl_uuid=file_sqdl_uuid, st_mtime_us=st_mtime_ns // 1000))
session.execute(stmt)
session.commit()
def get_counts(self) -> dict[str, int]:
counts = {}
with self.db.session() as session:
stmt = select(func.count()).select_from(UploadedDataset)
counts["uploaded datasets"] = session.scalar(stmt)
stmt = select(func.count()).select_from(UploadedFile)
counts["uploaded files"] = session.scalar(stmt)
return counts
import logging
from dataclasses import dataclass
from typing import List
from sqlalchemy import select, delete, update, func
from sqlalchemy.dialects.postgresql import insert as psql_insert
from core_tools.data.sqdl.model import UploadTask
from core_tools.data.sqdl.uploader_db import UploaderDb
logger = logging.getLogger(__name__)
@dataclass
class DatasetLocator:
scope: str
uid: int
path: str
class UploaderTaskQueue:
def __init__(self, db: UploaderDb):
self.db = db
def get_tasks(self) -> list[UploadTask]:
with self.db.session() as session:
stmt = select(UploadTask)
result = session.scalars(stmt).all()
return result
def get_oldest_task(self, pid: int) -> UploadTask:
with self.db.session() as session:
stmt = select(UploadTask).where(~UploadTask.failed & UploadTask.claimed_by.is_(None))
stmt = stmt.order_by(UploadTask.id)
stmt = stmt.limit(1)
task = session.scalars(stmt).first()
if task is None:
return None
if not self.claim_task(task, pid):
return None
return task
def get_newest_retry_task(self, pid: int) -> UploadTask:
with self.db.session() as session:
stmt = select(UploadTask).where(UploadTask.retry & UploadTask.claimed_by.is_(None))
stmt = stmt.order_by(UploadTask.id.desc())
stmt = stmt.limit(1)
task = session.scalars(stmt).first()
if task is None:
return None
if not self.claim_task(task, pid):
return None
return task
def claim_task(self, task: UploadTask, pid: int) -> bool:
with self.db.session() as session:
stmt = update(UploadTask).values(
version_id=UploadTask.version_id + 1,
claimed_by=pid,
).where((UploadTask.id == task.id) & (UploadTask.version_id == task.version_id))
update_result = session.execute(stmt).rowcount
session.commit()
if update_result != 1:
logger.info(f"Failed to claim task uid {task.uid}")
return False
task.version_id += 1
return True
def get_claimed_tasks(self) -> List[UploadTask]:
with self.db.session() as session:
stmt = select(UploadTask).where(UploadTask.claimed_by.is_not(None))
stmt = stmt.order_by(UploadTask.id)
return session.scalars(stmt).all()
def release_task(self, task):
with self.db.session() as session:
stmt = update(UploadTask).values(
claimed_by=None,
).where(UploadTask.id == task.id)
update_result = session.execute(stmt).rowcount
session.commit()
if update_result != 1:
logger.warning(f"Failed to release task uid {task.uid}")
def list_all_failed(self) -> UploadTask:
with self.db.session() as session:
stmt = select(UploadTask).where(UploadTask.failed)
stmt = stmt.order_by(UploadTask.id)
return session.scalars(stmt).all()
def add_dataset(self, ds_locator: DatasetLocator, final=False) -> None:
self._insert_or_update_task(ds_locator, update_dataset=True, set_raw_final=final)
def update_dataset(self, ds_locator: DatasetLocator, final=False) -> None:
self._insert_or_update_task(ds_locator, update_dataset=True, set_raw_final=final,
failed=False)
def reload_dataset(self, ds_locator: DatasetLocator, final=False) -> None:
self._insert_or_update_task(ds_locator, update_dataset=True, set_raw_final=final, retry=True)
def update_name(self, ds_locator: DatasetLocator) -> None:
self._insert_or_update_task(ds_locator, update_name=True)
def update_rating(self, ds_locator: DatasetLocator):
self._insert_or_update_task(ds_locator, update_rating=True)
def delete_task(self, task) -> bool:
with self.db.session() as session:
stmt = delete(UploadTask).where(
(UploadTask.id == task.id) & (UploadTask.version_id == task.version_id))
delete_result = session.execute(stmt)
deleted = delete_result.rowcount == 1
session.commit()
return deleted
def set_failed(self, task) -> None:
with self.db.session() as session:
stmt = update(UploadTask).values(
version_id=UploadTask.version_id + 1,
failed=True,
retry=False,
claimed_by=None,
).where(
(UploadTask.id == task.id) & (UploadTask.version_id == task.version_id))
update_result = session.execute(stmt).rowcount
session.commit()
if not update_result:
self.release_task(task)
def retry_all_failed(self) -> None:
with self.db.session() as session:
stmt = update(UploadTask).values(
version_id=UploadTask.version_id + 1,
retry=True,
).where(UploadTask.failed)
session.execute(stmt)
session.commit()
def _insert_or_update_task(self, ds_locator: DatasetLocator, **kwargs) -> None:
scope = ds_locator.scope
uid = ds_locator.uid
path = ds_locator.path
with self.db.session() as session:
stmt = psql_insert(UploadTask).values(version_id=1, scope=scope, uid=uid, ds_path=path, **kwargs)
stmt = stmt.on_conflict_do_update(
constraint="upload_task_uid_key",
set_=dict(version_id=UploadTask.version_id + 1, **kwargs))
session.execute(stmt)
session.commit()
def get_counts(self) -> dict[str, int]:
counts = {}
with self.db.session() as session:
stmt = select(UploadTask.failed, UploadTask.retry, func.count()).select_from(UploadTask)
stmt = stmt.group_by(UploadTask.failed, UploadTask.retry)
results = {}
for failed, retry, count in session.execute(stmt).all():
results[(failed, retry)] = count
counts["new"] = results.get((False, False), 0)
counts["reload"] = results.get((False, True), 0)
counts["failed"] = results.get((True, False), 0)
counts["retry"] = results.get((True, True), 0)
return counts
......@@ -88,6 +88,11 @@ def _configure_logging(cfg, app_name):
file_handler.setFormatter(logging.Formatter(file_format))
root_logger.addHandler(file_handler)
stream_handler = logging.StreamHandler(stream=sys.stderr)
stream_handler.setLevel("INFO")
stream_handler.setFormatter(logging.Formatter(file_format))
root_logger.addHandler(stream_handler)
logger.info(f'Start {app_name} logging')
for name in ['matplotlib', 'h5py', 'qcodes']:
......
from core_tools.startup.app_launcher import launch_app
module_name = 'core_tools.startup.sqdl_sync'
def launch_sqdl_sync(kill=False, close_at_exit=False):
launch_app('SQDL Sync', module_name,
kill=kill, close_at_exit=close_at_exit)
......@@ -5,7 +5,7 @@ from core_tools.data.sqdl.sqdl_writer import SQDLWriter
def sync_init():
# to-do: validate if setting sample info is still required
# todo: validate if setting sample info is still required
set_sample_info('Any', 'Any', 'Any')
connect_local_db()
print('Starting SQDL Sync')
......
from core_tools.data.sqdl import init_sqdl, sqdl_query, load_by_uuid
def main():
init_sqdl("Test")
records = sqdl_query()
for r in records:
ds = load_by_uuid(r.uuid)
print("--- new dataset ---")
print("dataset name: {}".format(ds.exp_name))
print("dataset uuid: {}".format(ds.exp_uuid))
print("dataset timestamp: {}".format(ds.run_timestamp))
if __name__ == "__main__":
main()
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment