diff --git a/setup.py b/setup.py
index 3ad381ac0a111e758849cbc05e2ad3a9819e4b8a..16eb60ef0c25370658d8bc8270db0a40c8618789 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ from subprocess import check_output
import numpy
from Cython.Build import cythonize
from Cython.Distutils import build_ext
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
from src.xfel_calibrate.notebooks import notebooks
@@ -36,10 +36,12 @@ class PreInstallCommand(build):
data_files = []
for ctypes in notebooks.values():
for nb in ctypes.values():
- data_files.append(nb["notebook"])
+ data_files.append(nb.get("notebook"))
data_files += nb.get("dep_notebooks", [])
data_files += nb.get("pre_notebooks", [])
+data_files = list(filter(None, data_files)) # Get rid of `None` entries
+
setup(
name="European XFEL Offline Calibration",
version="1.0",
diff --git a/src/xfel_calibrate/calibrate.py b/src/xfel_calibrate/calibrate.py
index 951dba8ea94b711cff5ce1796b7486633696f107..79ecf612087d5cdc16231d9d1e9e2a375c9f38f9 100755
--- a/src/xfel_calibrate/calibrate.py
+++ b/src/xfel_calibrate/calibrate.py
@@ -1,6 +1,7 @@
#!/usr/bin/env python
import argparse
+import ast
import inspect
import math
import os
@@ -13,7 +14,7 @@ import warnings
from datetime import datetime
from pathlib import Path
from subprocess import DEVNULL, check_output
-from typing import List, Union
+from typing import List, Optional, Union
import nbformat
import numpy as np
@@ -359,9 +360,25 @@ def make_extended_parser() -> argparse.ArgumentParser:
# The information is extracted from the first markdown cell of
# the notebook.
for caltype, notebook in det_notebooks.items():
- nbpath = os.path.join(PKG_DIR, notebook["notebook"])
- nb = nbformat.read(nbpath, as_version=4)
- msg += make_epilog(nb, caltype=caltype)
+ if notebook.get("notebook") is None:
+ if notebook.get("user") is None:
+ raise KeyError(
+ f"`{detector}` does not have a notebook path, for "
+ "notebooks that are stored in pycalibration set the "
+ "`notebook` key to a relative path or set the "
+ "`['user']['notebook']` key to an absolute path/path "
+ "pattern. Notebook configuration dictionary contains "
+ f"only: `{notebook}`"
+ ""
+ )
+ # Everything should be indented by 17 spaces
+ msg += caltype.ljust(17) + "User defined notebook, arguments may vary\n"
+ msg += " "*17 + "User notebook expected to be at path:\n"
+ msg += " "*17 + notebook["user"]["notebook"] + "\n"
+ else:
+ nbpath = os.path.join(PKG_DIR, notebook["notebook"])
+ nb = nbformat.read(nbpath, as_version=4)
+ msg += make_epilog(nb, caltype=caltype)
return make_initial_parser(epilog=msg)
elif len(sys.argv) <= 3:
@@ -376,12 +393,37 @@ def make_extended_parser() -> argparse.ArgumentParser:
print("Not one of the known calibrations or detectors")
sys.exit(1)
- notebook = os.path.join(PKG_DIR, nb_info["notebook"])
- cvar = nb_info.get("concurrency", {}).get("parameter", None)
+ if nb_info["notebook"]:
+ notebook = os.path.join(PKG_DIR, nb_info["notebook"])
+ else:
+ # If `"notebook"` entry is None, then set it to the user provided
+ # notebook TODO: This is a very hacky workaround, better implementation
+ # is not really possible with the current state of this module
+ user_notebook_path = nb_info["user"]["notebook"]
+ # Pulls out the variables in the templated path string, so that they
+ # can be added to the argument parser
+ user_notebook_variables = [
+ k.value.id
+ for k
+ in ast.walk(ast.parse(f"f'{user_notebook_path}'"))
+ if isinstance(k, ast.FormattedValue)
+ ]
+
+ user_notebook_parser = argparse.ArgumentParser()
+
+ for var in user_notebook_variables:
+ user_notebook_parser.add_argument(f"--{var}")
+
+ user_notebook_args, _ = user_notebook_parser.parse_known_args(
+ args=list(filter(lambda x: x != "-h", _)) # Drop help from args
+ )
- nb = nbformat.read(notebook, as_version=4)
+ nb_info["notebook"] = nb_info["user"]["notebook"].format(**vars(user_notebook_args))
+ notebook = str(nb_info["notebook"])
+ cvar = nb_info.get("concurrency", {}).get("parameter", None)
+ nb = nbformat.read(notebook, as_version=4)
# extend parameters if needed
ext_func = nb_info.get("extend parms", None)
@@ -663,10 +705,12 @@ def remove_duplications(l):
return unique_l
-def concurrent_run(temp_path, nb, nbname, args, cparm=None, cval=None,
- final_job=False, job_list=[], fmt_args={}, cluster_cores=8,
- sequential=False, dep_jids=[],
- show_title=True):
+def concurrent_run(
+ temp_path, nb, nbname, args, cparm=None, cval=None,
+ final_job=False, job_list=[], fmt_args={}, cluster_cores=8,
+ sequential=False, dep_jids=[],
+ show_title=True, user_venv: Optional[Path] = None,
+):
""" Launch a concurrent job on the cluster via SLURM
"""
@@ -707,16 +751,21 @@ def concurrent_run(temp_path, nb, nbname, args, cparm=None, cval=None,
srun_base = get_launcher_command(args, temp_path, dep_jids)
print(" ".join(srun_base))
- srun_base += [os.path.join(PKG_DIR, "bin", "slurm_calibrate.sh"), # path to helper sh
- os.path.abspath(nbpath), # path to notebook
- python_path, # path to python
- cluster_profile,
- '"{}"'.format(base_name.upper()),
- '"{}"'.format(args["detector"].upper()),
- '"{}"'.format(args["type"].upper()),
- "FINAL" if final_job else "NONFINAL",
- "{}/finalize.py".format(os.path.abspath(temp_path)),
- str(cluster_cores)]
+ if user_venv:
+ print(f"Running job in user venv at {user_venv}\n")
+
+ srun_base += [
+ os.path.join(PKG_DIR, "bin", "slurm_calibrate.sh"), # path to helper sh
+ os.path.abspath(nbpath), # path to notebook
+ user_venv + "/bin/python" if user_venv else python_path, # path to python
+ cluster_profile,
+ '"{}"'.format(base_name.upper()),
+ '"{}"'.format(args["detector"].upper()),
+ '"{}"'.format(args["type"].upper()),
+ "FINAL" if final_job else "NONFINAL",
+ "{}/finalize.py".format(os.path.abspath(temp_path)),
+ str(cluster_cores)
+ ]
output = check_output(srun_base).decode('utf8')
jobid = None
@@ -924,6 +973,10 @@ def run():
'submission_time': submission_time
}
+ user_venv = None
+ if nb_info.get("user", {}).get("venv"):
+ user_venv = nb_info["user"]["venv"].format(**args)
+
joblist = []
cluster_cores = concurrency.get("cluster cores", 8)
# Check if there are pre-notebooks
@@ -935,7 +988,7 @@ def run():
args,
job_list=joblist, fmt_args=fmt_args,
cluster_cores=cluster_cores,
- sequential=sequential,
+ sequential=sequential, user_venv=user_venv
)
joblist.append(jobid)
@@ -946,7 +999,7 @@ def run():
fmt_args=fmt_args,
cluster_cores=cluster_cores,
sequential=sequential,
- dep_jids=joblist,
+ dep_jids=joblist, user_venv=user_venv
)
joblist.append(jobid)
else:
diff --git a/src/xfel_calibrate/notebooks.py b/src/xfel_calibrate/notebooks.py
index 1b4e4cd965172286162c83590c30175bf17041a8..d21d2d5cf5a4fe877e0d8e83298fdc7dc53e96dd 100644
--- a/src/xfel_calibrate/notebooks.py
+++ b/src/xfel_calibrate/notebooks.py
@@ -233,6 +233,21 @@ notebooks = {
"cluster cores": 16},
},
},
+ "REMI": {
+ "CORRECT": {
+ "notebook": None,
+ "user": {
+ "notebook": "/gpfs/exfel/exp/{instrument}/{cycle}/p{proposal}/usr/calibration/notebooks/correct.ipynb",
+ "venv":"/gpfs/exfel/exp/{instrument}/{cycle}/p{proposal}/usr/calibration/.venv"
+ },
+ "concurrency": {
+ "parameter": None,
+ "use function": None,
+ "default concurrency": None,
+ "cluster cores": 1
+ },
+ },
+ },
"TEST": {
"TEST-CLI": {
"notebook": "notebooks/test/test-cli.ipynb",
@@ -242,5 +257,36 @@ notebooks = {
"cluster cores": 1,
},
},
+ "TEST-USER-NB": {
+ "notebook": None,
+ "user": {
+ "notebook": "/{root}/test-cli.ipynb",
+ "venv": None, # default, pycalibration environment
+ },
+ "concurrency": {
+ "parameter": None,
+ "use function": None,
+ "default concurrency": None,
+ "cluster cores": 1
+ },
+ },
+ "TEST-USER-NB-VENV": {
+ "notebook": None,
+ "user": {
+ "notebook": "/{root}/test-cli.ipynb",
+ "venv": "/{root}/.venv",
+ },
+ "concurrency": {
+ "parameter": None,
+ "use function": None,
+ "default concurrency": None,
+ "cluster cores": 1
+ },
+ },
+ },
+ "TEST-RAISES-ERRORS": {
+ "TEST-BAD-KEY": {
+ "noteboke": "a typo",
+ }
},
}