import sys
from time import sleep
import json
NUM_RUNS = 70
import requests
from funcx.sdk.client import FuncXClient
pyhf_endpoint = 'a727e996-7836-4bec-9fa2-44ebf7ca5302'
fxc = FuncXClient()
fxc.max_requests = 200
def prepare_workspace(data):
import pyhf
w = pyhf.Workspace(data)
return w
prepare_func = fxc.register_function(prepare_workspace)
def infer_hypotest(w, metadata, doc):
import pyhf
import time
tick = time.time()
m = w.model(patches=[doc],
modifier_settings={
"normsys": {
"interpcode": "code4"
def main(args):
if args.config_file is not None:
with open(args.config_file, "r") as infile:
config = json.load(infile)
backend = args.backend
pallet_path = Path(config["input_prefix"]).joinpath(config["pallet_name"])
# locally get pyhf pallet for analysis
if not pallet_path.exists():
download(config["pallet_url"], pallet_path)
analysis_name = config["analysis_name"]
analysis_prefix_str = "" if analysis_name is None else f"{analysis_name}_"
if config["analysis_dir"] is not None:
pallet_path = pallet_path.joinpath(config["analysis_dir"])
with open(pallet_path.joinpath(
f"{analysis_prefix_str}BkgOnly.json")) as bkgonly_json:
bkgonly_workspace = json.load(bkgonly_json)
# Initialize funcX client
fxc = FuncXClient()
fxc.max_requests = 200
with open("endpoint_id.txt") as endpoint_file:
pyhf_endpoint = str(endpoint_file.read().rstrip())
# register functions
prepare_func = fxc.register_function(prepare_workspace)
infer_func = fxc.register_function(infer_hypotest)
# execute background only workspace
prepare_task = fxc.run(bkgonly_workspace,
backend,
endpoint_id=pyhf_endpoint,
function_id=prepare_func)
# Read patchset in while background only workspace running
with open(pallet_path.joinpath(
f"{analysis_prefix_str}patchset.json")) as patchset_json:
patchset = pyhf.PatchSet(json.load(patchset_json))
workspace = None
while not workspace:
try:
workspace = fxc.get_result(prepare_task)
except Exception as excep:
print(f"prepare: {excep}")
sleep(10)
print("--------------------")
print(workspace)
# execute patch fits across workers and retrieve them when done
n_patches = len(patchset.patches)
tasks = {}
for patch_idx in range(n_patches):
patch = patchset.patches[patch_idx]
task_id = fxc.run(
workspace,
patch.metadata,
[patch.patch],
backend,
endpoint_id=pyhf_endpoint,
function_id=infer_func,
)
tasks[patch.name] = {"id": task_id, "result": None}
while count_complete(tasks.values()) < n_patches:
for task in tasks.keys():
if not tasks[task]["result"]:
try:
result = fxc.get_result(tasks[task]["id"])
print(
f"Task {task} complete, there are {count_complete(tasks.values())+1} results now"
)
tasks[task]["result"] = result
except Exception as excep:
print(f"inference: {excep}")
sleep(15)
print("--------------------")
print(tasks.values())
