Skip to content

Commit 6e1170e

Browse files
committed
getCCDBTimeMachineTimestamp.py: add CCDB query and upload functionality
1 parent b820139 commit 6e1170e

1 file changed

Lines changed: 159 additions & 7 deletions

File tree

‎GRID/utils/getCCDBTimeMachineTimestamp.py‎

Lines changed: 159 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,40 @@
1313
# started 25.08.2026; Sandro Wenzel
1414

1515
import argparse
16+
import dataclasses
17+
from dataclasses import dataclass, asdict, fields
1618
import json
1719
import os
20+
import requests
1821
import statistics
1922
import subprocess
2023
import sys
2124
import urllib.request
2225
from datetime import datetime, timezone
2326

27+
# hack to find the script with helper functions
28+
o2dpg_root = os.environ.get("O2DPG_ROOT")
29+
if o2dpg_root is None:
30+
raise EnvironmentError("O2DPG_ROOT is not set in the environment.")
31+
mc_prodinfo_path = os.path.abspath(os.path.join(o2dpg_root, "MC", "prodinfo"))
32+
sys.path.append(mc_prodinfo_path)
33+
# FIXME: these should probably be moved to a common script
34+
from mcprodinfo_ccdb_upload import extract_metadata_blocks_from_CCDB
35+
2436
CCDB_HOST = "http://alice-ccdb.cern.ch"
2537

2638

39+
@dataclass(frozen=True)
40+
class CCDBTimeMachineTS:
41+
"""
42+
struct for CCDB time machine timestamp
43+
"""
44+
LPMProductionTag: str
45+
Period: str
46+
Pass: str
47+
NotAfter: int
48+
49+
2750
def alien_json(*args):
2851
"""Run alien.py with -json and return the parsed 'results' list."""
2952
cmd = ["alien.py", "-json"] + list(args)
@@ -111,7 +134,26 @@ def describe_object(o):
111134
f"comment=\"{o.get('comment', '')}\"")
112135

113136

114-
def get_CCDBTimeMachineTS(year, period, run, passname, samples=9, ccdb_check_path="TPC/Calib/CorrectionMapV2", ccdb_check=True):
137+
def make_cert_paths(cert_dir="/tmp"):
138+
key_path = os.environ.get("JALIEN_TOKEN_KEY")
139+
cert_path = os.environ.get("JALIEN_TOKEN_CERT")
140+
if key_path == None and cert_path == None:
141+
uid = os.getuid()
142+
cert_path = os.path.join(cert_dir, f"tokencert_{uid}.pem")
143+
key_path = os.path.join(cert_dir, f"tokenkey_{uid}.pem")
144+
145+
return cert_path, key_path
146+
147+
148+
def get_NotAfter(year, period, run, passname, samples=9, ccdb_check_path="TPC/Calib/CorrectionMapV2", ccdb_check=True):
149+
"""
150+
Estimates the CCDBTimeMachine NotAfter timestamp from the output of the production of a given period/run/pass.
151+
Cross-checks the estimated NotAfter timestamp vs the creation time of a CCDB object which is currently
152+
valid for the given run.
153+
154+
Returns:
155+
int: CCDBTimeMachineTS NotAfter timestamp
156+
"""
115157
pass_dir = f"/alice/data/{year}/{period}/{run}/{passname}/"
116158
print(f"[getCCDBTimeMachineTimestamp] Looking for TF output under {pass_dir}", file=sys.stderr)
117159

@@ -137,7 +179,7 @@ def get_CCDBTimeMachineTS(year, period, run, passname, samples=9, ccdb_check_pat
137179
spread_h = (max(ctimes) - min(ctimes)) / 3.6e6
138180
print(f"[getCCDBTimeMachineTimestamp] {len(ctimes)} samples out of {len(matches)} TF outputs found; "
139181
f"span {human(min(ctimes))} .. {human(max(ctimes))} ({spread_h:.1f} h)", file=sys.stderr)
140-
print(f"[getCCDBTimeMachineTimestamp] Recommended not-after: {not_after} ({human(not_after)})",
182+
print(f"[getCCDBTimeMachineTimestamp] Recommended NotAfter: {not_after} ({human(not_after)})",
141183
file=sys.stderr)
142184

143185
if ccdb_check:
@@ -161,10 +203,111 @@ def get_CCDBTimeMachineTS(year, period, run, passname, samples=9, ccdb_check_pat
161203
return not_after
162204

163205

206+
def query_CCDBTimeMachineTS(base_url, user, period, passname):
207+
"""
208+
Queries CCDBTimeMachineTS from CCDB for a given period and pass.
209+
210+
Returns:
211+
obtained NotAfter timestamp (int) or None
212+
"""
213+
# Construct user ID-specific cert and key paths
214+
cert_path, key_path = make_cert_paths()
215+
216+
# Build full URL
217+
query = f"1/Period={period}/Pass={passname}"
218+
user_path = 'Users/' + user[0] + '/' + user
219+
url = f"{base_url}/browse/{user_path}/CCDBTimeMachineTS/{query}"
220+
221+
print(f"[getCCDBTimeMachineTimestamp] Checking if CCDBTimeMachineTS already exists: {url}")
222+
223+
response = requests.get(url, cert=(cert_path, key_path), verify=False)
224+
if response.status_code != 404:
225+
meta = extract_metadata_blocks_from_CCDB(response.content.decode('utf-8'))
226+
if (len(meta) > 0):
227+
def filter_known_fields(cls, data: dict) -> dict:
228+
valid_keys = {f.name for f in fields(cls)}
229+
return {k: v for k, v in data.items() if k in valid_keys}
230+
231+
clean_meta = filter_known_fields(CCDBTimeMachineTS, meta[0])
232+
print(f"[getCCDBTimeMachineTimestamp] Found object for {clean_meta['Period']}, Pass: {clean_meta['Pass']}, "
233+
f"NotAfter: {clean_meta['NotAfter']} ({human(clean_meta['NotAfter'])}), uploaded by LPMProductionTag {clean_meta['LPMProductionTag']}")
234+
return clean_meta['NotAfter']
235+
236+
return None
237+
238+
239+
def upload_CCDBTimeMachineTS(base_url, user, keys):
240+
"""
241+
Uploads an empty .dat file as CCDBTimeMachineTS with the corresponding meta data provided by `keys`
242+
243+
Returns:
244+
Response object from the POST request
245+
"""
246+
# Create an empty file
247+
empty_file = "empty.dat"
248+
with open(empty_file, "w") as f:
249+
f.write("0")
250+
251+
# Construct user ID-specific cert and key paths
252+
cert_path, key_path = make_cert_paths()
253+
254+
# Build full URL
255+
query = "/".join(f"{k}={v}" for k, v in keys.items())
256+
user_path = 'Users/' + user[0] + '/' + user
257+
start = 1
258+
stop = start + 1
259+
url = f"{base_url}/{user_path}/CCDBTimeMachineTS/{start}/{stop}/{query}"
260+
261+
print(f"[getCCDBTimeMachineTimestamp] Uploading CCDBTimeMachineTS: {url}")
262+
263+
# Prepare request
264+
with open(empty_file, 'rb') as f:
265+
files = {'blob': f}
266+
response = requests.post(url, files=files, cert=(cert_path, key_path), verify=False)
267+
268+
os.remove(empty_file)
269+
270+
return response
271+
272+
273+
def get_CCDBTimeMachineTS(production_tag, year, period, run, passname,
274+
samples=9, ccdb_check_path="TPC/Calib/CorrectionMapV2", ccdb_check=True,
275+
ccdb_url="https://alice-ccdb.cern.ch", username="aliprod",
276+
force_overwrite=False, publish_to_ccdb=True, skip_ccdb_query=False):
277+
"""
278+
Makes a query to CCDB for a CCDBTimeMachineTS for given period and pass.
279+
If an object is found, the NotAfter timestamp is returned.
280+
If no object is found, the NotAfter timestamp is estimated from the reco output of the given period/run/pass.
281+
In case an object is found, a new upload can be forced by setting the argument `force_overwrite=True`
282+
283+
Returns:
284+
int: Either the queried or the estimated NotAfter timestamp
285+
"""
286+
if not skip_ccdb_query:
287+
not_after = query_CCDBTimeMachineTS(ccdb_url, username, period, passname)
288+
else:
289+
not_after = None
290+
291+
if not_after is not None and force_overwrite == True:
292+
print("[getCCDBTimeMachineTimestamp] Existing object will be overwritten")
293+
294+
if not_after == None or force_overwrite == True:
295+
print("[getCCDBTimeMachineTimestamp] Estimating CCDBTimeMachineTS for "
296+
f"LPMProductionTag: {production_tag}, Period: {period}, Pass: {passname}")
297+
not_after = get_NotAfter(year, period, run, passname, samples, ccdb_check_path, ccdb_check)
298+
time_machine_ts = CCDBTimeMachineTS(production_tag, period, passname, not_after)
299+
if publish_to_ccdb:
300+
upload_CCDBTimeMachineTS(ccdb_url, username, dataclasses.asdict(time_machine_ts))
301+
302+
return not_after
303+
304+
164305
def main():
165306
parser = argparse.ArgumentParser(
166307
description="Estimate a CCDB TimeMachine not-after timestamp for a run+pass, "
167308
"from AliEn catalog timestamps of that pass's own reco output.")
309+
parser.add_argument("--production-tag", default=os.environ.get("ALIEN_JDL_LPMPRODUCTIONTAG"),
310+
help="production tag to mark which production uploads an object (default: $ALIEN_JDL_LPMPRODUCTIONTAG)")
168311
parser.add_argument("--year", default=os.environ.get("ALIEN_JDL_LPMANCHORYEAR"),
169312
help="data-taking year, e.g. 2023 (default: $ALIEN_JDL_LPMANCHORYEAR)")
170313
parser.add_argument("--period", default=os.environ.get("ALIEN_JDL_LPMANCHORPRODUCTION"),
@@ -178,17 +321,26 @@ def main():
178321
parser.add_argument("--ccdb-check-path", default="TPC/Calib/CorrectionMapV2",
179322
help="CCDB path to cross-check the estimate against (default: TPC/Calib/CorrectionMapV2)")
180323
parser.add_argument("--no-ccdb-check", action="store_true",
181-
help="skip the CCDB cross-check")
324+
help="skip the CCDB cross-check after calculating the NotAfter")
325+
parser.add_argument("--publish-to-ccdb", action="store_true",
326+
help="upload the CCDBTimeMachineTS object to the CCDB")
327+
parser.add_argument("--skip-ccdb-query", action="store_true",
328+
help="skip CCDB query for existing objects while publishing")
329+
parser.add_argument("--force-overwrite", action="store_true",
330+
help="overwrite the current CCDBTimeMachineTS object in the CCDB if one exists")
331+
parser.add_argument("--ccdb-url", default="https://alice-ccdb.cern.ch",
332+
help="ccdb host for the upload")
333+
parser.add_argument("--username", default="aliprod",
334+
help="username to be used for the upload to CCDB")
182335
args = parser.parse_args()
183336

184337
missing = [n for n in ("year", "period", "run", "passname") if not getattr(args, n)]
185338
if missing:
186339
parser.error("missing required value(s): " + ", ".join(missing))
187340

188-
not_after = get_CCDBTimeMachineTS(args.year, args.period, args.run, args.passname, args.samples,
189-
args.ccdb_check_path, not args.no_ccdb_check)
190-
191-
print(not_after)
341+
get_CCDBTimeMachineTS(args.production_tag, args.year, args.period, args.run, args.passname,
342+
args.samples, args.ccdb_check_path, not args.no_ccdb_check, args.ccdb_url, args.username,
343+
args.force_overwrite, args.publish_to_ccdb, args.skip_ccdb_query)
192344

193345

194346
if __name__ == "__main__":

0 commit comments

Comments
 (0)