diff --git a/bin/slurm_calibrate.sh b/bin/slurm_calibrate.sh
index 05a04cba4e88fec483f37b293e8aba569b7f5a1a..9e5088757b161bb0206da82517a2f1fd6210184a 100755
--- a/bin/slurm_calibrate.sh
+++ b/bin/slurm_calibrate.sh
@@ -10,6 +10,7 @@ caltype=$6
final=$7
finalize=$8
cluster_cores=$9
+notebook_python_path=$10
echo "Running with the following parameters:"
echo "Notebook path: $nb_path"
@@ -43,8 +44,8 @@ fi
echo "Running notebook"
-${python_path} -m princess ${nb_path} --save --on-error-resume-next
-${python_path} -m nbconvert --to rst --TemplateExporter.exclude_input=True ${nb_path}
+${notebook_python_path} -m princess ${nb_path} --save --on-error-resume-next
+${notebook_python_path} -m nbconvert --to rst --TemplateExporter.exclude_input=True ${nb_path}
# stop the cluster if requested
if [ "${ipcluster_profile}" != "NO_CLUSTER" ]
diff --git a/notebooks/test/test-cli.ipynb b/notebooks/test/test-cli.ipynb
index ef866040aa10b5ebfb5df881f3c493c4f3724652..2874fa87593b5d841e964d74920156a1896e1758 100644
--- a/notebooks/test/test-cli.ipynb
+++ b/notebooks/test/test-cli.ipynb
@@ -19,8 +19,11 @@
"metadata": {},
"outputs": [],
"source": [
- "in_folder = \"/dev/null\" # input folder\n",
- "out_folder = \"/dev/null\" # output folder\n",
+ "root = \"root/path/for/nb\" # Variables included in the user notebook path must\n",
+ "# be included in the notebook for reasons\n",
+ "\n",
+ "in_folder = \"./\" # input folder\n",
+ "out_folder = \"./\" # output folder\n",
"list_normal = [10] # parameterized list, range allowed\n",
"list_intellilist = [2345] # parameterized list with ranges, range allowed\n",
"concurrency_parameter = [1] # concurrency parameter, range allowed\n",
@@ -56,15 +59,9 @@
}
],
"metadata": {
- "kernelspec": {
- "display_name": "Python 3.6.8 64-bit",
- "name": "python3"
- },
- "language_info": {
- "name": "python",
- "version": ""
- }
+ "language_info": {},
+ "orig_nbformat": 3
},
"nbformat": 4,
"nbformat_minor": 0
-}
+}
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 7a65bbfa5b7dc6d52ebce011678d601677284a85..51e7fa68530c330db255bf5d2db6b01781ce8c39 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..4068a07b66bfc13128dfcf01dc7d57a58a740ef8 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", {}).get("notebook") 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 = 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,22 @@ 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
+ 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),
+ user_venv + "/bin/python" if user_venv else python_path # used for nb execution
+ ]
output = check_output(srun_base).decode('utf8')
jobid = None
@@ -924,6 +974,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 +989,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 +1000,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 a98aeef536ec49d0a51cccb2abe789203c1c935a..b9ccbab6c2c09da5f60cdedcdc3e1175468373a0 100644
--- a/src/xfel_calibrate/notebooks.py
+++ b/src/xfel_calibrate/notebooks.py
@@ -239,6 +239,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/sw/software/exfel_environments/sqs-remi-preview"
+ },
+ "concurrency": {
+ "parameter": None,
+ "use function": None,
+ "default concurrency": None,
+ "cluster cores": 1
+ },
+ },
+ },
"TEST": {
"TEST-CLI": {
"notebook": "notebooks/test/test-cli.ipynb",
@@ -248,5 +263,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",
+ }
},
}
diff --git a/tests/test_xfel_calibrate/conftest.py b/tests/test_xfel_calibrate/conftest.py
index c00cfda97c78621930ce5fe2797f94b66ed9068e..61fb9c6f1fb24e61c59bc8fcfe6c66ad321c36ee 100644
--- a/tests/test_xfel_calibrate/conftest.py
+++ b/tests/test_xfel_calibrate/conftest.py
@@ -99,6 +99,11 @@ class MockProposal:
self.path = tmp_path / instrument / cycle / proposal
self.path_raw = self.path / "raw"
self.path_proc = self.path / "proc"
+ self.path_usr = self.path / "usr"
+ self.path_scratch = self.path / "scratch"
+
+ self.path_usr.mkdir(parents=True)
+ self.path_scratch.mkdir()
self.runs = self.create_runs(runs)
diff --git a/tests/test_xfel_calibrate/test_cli.py b/tests/test_xfel_calibrate/test_cli.py
index f0ffeb4af21ededff97abe66e22f669266517e6d..4a75d6242c97644a249379cce32a6ec93b38c622 100644
--- a/tests/test_xfel_calibrate/test_cli.py
+++ b/tests/test_xfel_calibrate/test_cli.py
@@ -55,6 +55,23 @@ class TestBasicCalls:
assert err == ""
+ @mock.patch.object(sys, "argv", ["xfel-calibrate", "TEST", "-h"])
+ def test_help_user_notebook(self, capsys):
+ with pytest.raises(SystemExit):
+ calibrate.run()
+
+ out, err = capsys.readouterr()
+
+ assert "TEST-USER-NB" in out
+ assert "/{root}/test-cli.ipynb" in out
+
+ assert err == ""
+
+ @mock.patch.object(sys, "argv", ["xfel-calibrate", "TEST-RAISES-ERRORS", "--help"])
+ def test_help_bad_config(self):
+ with pytest.raises(KeyError):
+ calibrate.run()
+
@mock.patch.object(sys, "argv", ["xfel-calibrate", "NotADetector", "beep", "-h"])
def test_unknown_detector(self, capsys):
with pytest.raises(SystemExit) as exit_exception:
diff --git a/tests/test_xfel_calibrate/test_user_configs.py b/tests/test_xfel_calibrate/test_user_configs.py
new file mode 100644
index 0000000000000000000000000000000000000000..e064ca7311e0f1895169c8166bbf7845c8a9585a
--- /dev/null
+++ b/tests/test_xfel_calibrate/test_user_configs.py
@@ -0,0 +1,144 @@
+# pylint: disable=missing-class-docstring, missing-function-docstring, no-self-use
+""" Tests for user configurable notebooks
+
+These tests cover functionality enabled by setting `"notebook": None` and
+instead providing path templates under the `user` key in the `notebooks.py`
+configuration file.
+"""
+import shlex
+import shutil
+from pathlib import Path
+
+import nbformat
+import pytest
+from nbparameterise import extract_parameters
+
+from tests.test_xfel_calibrate.conftest import (
+ CalibrateCall,
+ FakeProcessCalibrate,
+ MockProposal,
+)
+
+TEST_CLI_PATH = Path("./notebooks/test/test-cli.ipynb")
+
+
+class TestUserNotebooks:
+ user_nb = None # Gets set in `self.calibrate_call`
+
+ @pytest.fixture(scope="class", autouse=True)
+ def fake_process_calibrate(self):
+ with FakeProcessCalibrate() as fake_process:
+ yield fake_process
+
+ @pytest.fixture(scope="class")
+ def mock_proposal(self, tmpdir_factory):
+ return MockProposal(
+ tmp_path=Path(tmpdir_factory.mktemp("exp")),
+ instrument="AGIPD",
+ cycle="202031",
+ proposal="p900113",
+ runs=1,
+ sequences=1,
+ )
+
+ @pytest.fixture(scope="function")
+ def calibrate_call(self, mock_proposal: MockProposal, capsys, tmp_path: Path):
+ # The user notebook path template for `TEST`/`TEST-USER-NB` is:
+ # `"/{instrument}/{cycle}/p{proposal}/test-cli.ipynb"`
+ #
+ # So by setting `instrument` to tmp_path/TEST we set the user notebook
+ # location to one in our tmp dir
+
+ self.user_nb = mock_proposal.path_usr / "test-cli.ipynb"
+
+ shutil.copy(TEST_CLI_PATH.absolute(), str(self.user_nb))
+
+ return CalibrateCall(
+ tmp_path,
+ capsys,
+ in_folder=mock_proposal.path_raw,
+ out_folder=mock_proposal.path_proc,
+ detector="TEST",
+ cal_type="TEST-USER-NB",
+ extra_args=[
+ "--root",
+ str(mock_proposal.path_usr),
+ "--number",
+ "10",
+ ],
+ )
+
+ def test_user_notebook_is_test_notebook(self, calibrate_call: CalibrateCall):
+ assert TEST_CLI_PATH.read_bytes() == self.user_nb.read_bytes()
+
+ def test_output_ipynb(self, calibrate_call: CalibrateCall):
+ notebook_path = calibrate_call.paths.notebooks
+ assert len(notebook_path) == 1
+
+ with notebook_path[0].open() as file:
+ notebook = nbformat.read(file, as_version=4)
+
+ parameters = {p.name: p.value for p in extract_parameters(notebook)}
+
+ assert parameters["number"] == 10
+
+
+class TestUserVenv:
+ @pytest.fixture(scope="class", autouse=True)
+ def fake_process_calibrate(self):
+ with FakeProcessCalibrate() as fake_process:
+ yield fake_process
+
+ @pytest.fixture(scope="class")
+ def mock_proposal(self, tmpdir_factory):
+ return MockProposal(
+ tmp_path=Path(tmpdir_factory.mktemp("exp")),
+ instrument="AGIPD",
+ cycle="202031",
+ proposal="p900113",
+ runs=1,
+ sequences=1,
+ )
+
+ @pytest.fixture(scope="function")
+ def calibrate_call(self, mock_proposal: MockProposal, capsys, tmp_path: Path):
+ self.user_nb = mock_proposal.path_usr / "test-cli.ipynb"
+
+ shutil.copy(TEST_CLI_PATH.absolute(), str(self.user_nb))
+
+ return CalibrateCall(
+ tmp_path,
+ capsys,
+ in_folder=mock_proposal.path_raw,
+ out_folder=mock_proposal.path_proc,
+ detector="TEST",
+ cal_type="TEST-USER-NB-VENV",
+ extra_args=[
+ "--root",
+ str(mock_proposal.path_usr),
+ "--number",
+ "10",
+ ],
+ )
+
+ def test_call(self, calibrate_call: CalibrateCall):
+ assert "Running job in user venv at" in calibrate_call.out
+
+ def test_expected_processes_called(
+ self,
+ mock_proposal: MockProposal,
+ fake_process_calibrate: FakeProcessCalibrate,
+ ):
+ process_calls = [
+ list(shlex.shlex(p, posix=True, punctuation_chars=True))
+ if isinstance(p, str)
+ else p
+ for p in fake_process_calibrate.calls
+ ]
+
+ processes_called = [p[0] for p in process_calls] # List of the process names
+
+ assert "sbatch" in processes_called
+
+ sbatch_cmd = process_calls[processes_called.index("sbatch")]
+ assert any(str(mock_proposal.path_usr / ".venv") in arg for arg in sbatch_cmd)