Repository: radiasoft/sirepo

Issue:
Check for NERSC account fails incorrectly
In `job_agent._sbatch_script()` this call is made:
```python
o = subprocess.check_output(["hpssquota"], text=True)
assert re.search(
    r"{}\s+\d+\.".format(p), o
), f"sbatchProject={p} is invalid. hpssquota={o}"
```
The `hpssquota` command on NERSC yields e.g.:
```bash
+-----------------------------------------+------------+-------------+----------------+------------+-------------+----------------+
|                             File system | Space used | Space quota | Space used (%) | Inode used | Inode quota | Inode used (%) |
+-----------------------------------------+------------+-------------+----------------+------------+-------------+----------------+
|                                    home |  822.75MiB |    40.00GiB |           2.0% |     30.84K |       1.00M |           3.1% |
|                                cscratch |    3.68MiB |    20.00TiB |           0.0% |     124.00 |      10.00M |           0.0% |
| mkeilman usage on HPSS charged to m4091 |      0.00B |     7.20TiB |           0.0% |          - |           - |              - |
+-----------------------------------------+------------+-------------+----------------+------------+-------------+----------------+
```
which does contain the valid account m4091 but is not found by the regex. The fancy table is new; `re.search(r"{}\s+\|?\s+\d+\.".format(p), o)` catching the pipe (optionally) works, or running `hpssquota --space` to output the old format and keep the regex the same. There are also options to output JSON (`hpssquota -J`) and other formats if that is more robust



Discussion:
robnagler: @mkeilman @moellep I did not mean to approve and I don't see how to revert.
mkeilman: Changes in place. I used some other examples in sirepo as a model for the monkeypatching - lmk if that's the right approach
mkeilman: A couple of quick notes: first, as per @robnagler I will change this to be "pessimistic" and look for a good result and not any specific error text. Also, it turns out that you can call hpssquota with any valid project code, even one you are not a part of, and it will return successfully (though the sbatch script will fail). If you leave it out you get only those projects assigned to you, so that is what I will do, plus I'll move to json as discussed earlier
moellep: When I call "hpssquote -N -J" with no project, it gives me one entry, which has "fs": "moellep usage on HPSS charged to m4091", but I don't see how that gets parsed.
mkeilman: > When I call "hpssquote -N -J" with no project, it gives me one entry, which has "fs": "moellep usage on HPSS charged to m4091", but I don't see how that gets parsed.

That's an error - in the previous commit I was passing the project, in which case there is an entry with "fs": <project> - but we can't do only that because hpssquota will happily show info for projects you don't belong to. So I will have to add a check for the project. I'll also update the testing monkeypatch to more closely hew to the actual output.

moellep: @robnagler could you take another look at this PR. I don't know if some users have multiple project codes on NERSC or how this would handle it.

Patch:
diff --git a/sirepo/nersc.py b/sirepo/nersc.py
new file mode 100644
index 000000000..f833419e6
--- /dev/null
+++ b/sirepo/nersc.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+"""Methods related to running on NERSC
+
+:copyright: Copyright (c) 2023 RadiaSoft LLC.  All Rights Reserved.
+:license: http://www.apache.org/licenses/LICENSE-2.0.html
+"""
+import re
+import subprocess
+import sirepo.util
+
+
+def sbatch_project_option(project):
+    from pykern.pkjson import load_any
+
+    if not project:
+        return ""
+    try:
+        if not list(filter(lambda x: project in x.fs, load_any(_hpssquota(project)))):
+            raise sirepo.util.UserAlert(f"Project {project} not found on NERSC")
+    except Exception as e:
+        raise sirepo.util.UserAlert(f"Cannot determine quota for {project}; error={e}")
+    return f"#SBATCH --account={project}"
+
+
+# Note: project is only used for unit tests
+def _hpssquota(project=None):
+    # -N excludes home and scratch file systems; -J outputs json
+    return subprocess.run(
+        ("hpssquota", "-N", "-J"),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
diff --git a/sirepo/pkcli/job_agent.py b/sirepo/pkcli/job_agent.py
index 89c472778..72e55340a 100644
--- a/sirepo/pkcli/job_agent.py
+++ b/sirepo/pkcli/job_agent.py
@@ -18,6 +18,7 @@ import re
 import shutil
 import signal
 import sirepo.modules
+import sirepo.nersc
 import sirepo.tornado
 import socket
 import subprocess
@@ -751,17 +752,6 @@ class _SbatchRun(_SbatchCmd):
         await c._await_exit()
 
     def _sbatch_script(self):
-        def _assert_project():
-            p = self.msg.sbatchProject
-            if not p:
-                return ""
-            o = subprocess.check_output(["hpssquota"], text=True)
-            assert re.search(r"^[-\w]+$", p), f"invalid NERSC project={p}"
-            assert re.search(
-                r"{}\s+\d+\.".format(p), o
-            ), f"sbatchProject={p} is invalid. hpssquota={o}"
-            return f"#SBATCH --account={p}"
-
         def _processor():
             if self.msg.sbatchQueue == "debug" and pkconfig.channel_in("dev"):
                 return "knl"
@@ -775,7 +765,7 @@ class _SbatchRun(_SbatchCmd):
 #SBATCH --constraint={_processor()}
 #SBATCH --qos={self.msg.sbatchQueue}
 #SBATCH --tasks-per-node={self.msg.tasksPerNode}
-{_assert_project()}"""
+{sirepo.nersc.sbatch_project_option(self.msg.sbatchProject)}"""
             s = "--cpu-bind=cores shifter --entrypoint"
         m = "--mpi=pmi2" if pkconfig.channel_in("dev") else ""
         f = self.run_dir.join(self.jid + ".sbatch")


Test patch:
diff --git a/tests/nersc_test.py b/tests/nersc_test.py
new file mode 100644
index 000000000..4954cb7cb
--- /dev/null
+++ b/tests/nersc_test.py
@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+"""NERSC-related tests.
+
+:copyright: Copyright (c) 2023 RadiaSoft LLC.  All Rights Reserved.
+:license: http://www.apache.org/licenses/LICENSE-2.0.html
+"""
+import sirepo.util
+from pykern.pkcollections import PKDict
+
+
+def test_nersc_project(monkeypatch):
+    from sirepo import nersc
+    from pykern.pkunit import pkeq, pkok, pkre, pkfail, pkexcept
+
+    _ERROR = "ERROR"
+    _VALID_PROJECT = "VALID_TEST_PROJECT"
+    _NO_SUCH_PROJECT = "NOT_" + _VALID_PROJECT
+
+    mock_output = PKDict()
+    mock_output[
+        _VALID_PROJECT
+    ] = f'[{{"fs": "user usage on HPSS charged to {_VALID_PROJECT}"}}]'
+    mock_output[_NO_SUCH_PROJECT] = "[]"
+    mock_output[_ERROR] = "arbitrary hpssquota error"
+
+    monkeypatch.setattr(nersc, "_hpssquota", lambda p: mock_output[p])
+
+    with pkexcept(sirepo.util.UserAlert):
+        nersc.sbatch_project_option(_ERROR)
+
+    with pkexcept(sirepo.util.UserAlert):
+        nersc.sbatch_project_option(_NO_SUCH_PROJECT)
+
+    pkeq("", nersc.sbatch_project_option(""))
+    pkre(_VALID_PROJECT, nersc.sbatch_project_option(_VALID_PROJECT))