Slurm normal <> preemptable association hierarchy

coact-facility-overage-daemon.sh

@@ -1,10 +1,11 @@
 #!/bin/sh
 
 export PATH=$PATH:/opt/slurm/slurm-curr/bin
+#export SDF_COACT_URI=coact-dev.slac.stanford.edu:443/graphql-service-dev
 export SDF_COACT_URI=coact.slac.stanford.edu:443/graphql-service
 
 while [ 1 ]; do
     date
-    ./venv/bin/python3  ./sdf_click.py coact overage --password-file ./etc/.secrets/password --grouper-password-file ./etc/.secrets/grouper_password --windows 5 --windows 15 --windows 60 --windows 180 --windows 1440 --verbose --influxdb-url=https://influxdb.slac.stanford.edu:443
+    old_venv/bin/python  ./sdf_click.py coact overage --password-file ./etc/.secrets/password --windows 5 --windows 15 --windows 60 --windows 180 --windows 1440 --verbose --influxdb-url=https://influxdb.slac.stanford.edu:443
     sleep 300
 done

coact-reporegistration-daemon.sh

@@ -1,6 +1,6 @@
-#!/bin/sh
+#!/bin/bash
 
 while [ 1 ]; do
-    SDF_COACT_URI=coact.slac.stanford.edu/graphql-service ./venv/bin/python3 ./sdf_click.py coactd reporegistration --username=sdf-bot --password-file=etc/.secrets/password --grouper-password-file ./etc/.secrets/grouper_password -vv
+    SDF_COACT_URI=coact.slac.stanford.edu/graphql-service old_venv/bin/python ./sdf_click.py coactd reporegistration --username=sdf-bot --password-file=etc/.secrets/password --grouper-password-file /sdf/home/r/ryanw/code/sdf-cli/etc/.secrets/grouper_password -vv
     sleep 1
 done

coact-userregistration-daemon.sh

@@ -1,6 +1,6 @@
 #!/bin/sh
 
 while [ 1 ]; do
-    SDF_COACT_URI=coact.slac.stanford.edu:443/graphql-service ./venv/bin/python3 ./sdf_click.py coactd userregistration --username sdf-bot --password-file ./etc/.secrets/password --grouper-password-file ./etc/.secrets/grouper_password -vv
+    SDF_COACT_URI=coact.slac.stanford.edu:443/graphql-service ./old_venv/bin/python3 ./sdf_click.py coactd userregistration --username sdf-bot --password-file ./etc/.secrets/password -vv
     sleep 5
 done

import-jobs.sh

@@ -3,6 +3,8 @@
 export PATH=$PATH:/opt/slurm/slurm-curr/bin
 export SDF_COACT_URI=coact.slac.stanford.edu:443/graphql-service
 
+export PYTHON_BIN=./old_venv/bin/python3
+
 PASSWORD_FILE=./etc/.secrets/password
 
 if [ ! -z $1 ]; then
@@ -22,11 +24,11 @@ fi
 echo ">" $DATE" ("$(date)")"
 
 # full
-./venv/bin/python3 ./sdf_click.py coact slurmdump --date $DATE \
+$PYTHON_BIN ./sdf_click.py coact slurmdump --date $DATE \
     | tee ../slurm-job-history/$DATE \
-    | ./venv/bin/python3 ./sdf_click.py coact slurmremap \
+    | $PYTHON_BIN ./sdf_click.py coact slurmremap \
     | tee ../slurm-job-remapped/$DATE \
-    | ./venv/bin/python3 ./sdf_click.py coact slurmimport --password-file $PASSWORD_FILE --output=upload >/dev/null
+    | $PYTHON_BIN ./sdf_click.py coact slurmimport --password-file $PASSWORD_FILE --output=upload >/dev/null
 
 # just for 2023 imports
 #cat ../slurm-job-remapped/$DATE | ./sdf.py coact slurmimport --password-file $PASSWORD_FILE --output=upload >/dev/null
@@ -37,4 +39,4 @@ echo ">" $DATE" ("$(date)")"
 ###
 # recalculate summaries
 ###
-./venv/bin/python3  ./sdf_click.py coact slurmrecalculate --password-file=$PASSWORD_FILE --date=$DATE
+$PYTHON_BIN  ./sdf_click.py coact slurmrecalculate --password-file=$PASSWORD_FILE --date=$DATE

modules/coact.py

@@ -75,6 +75,14 @@ def datetime_converter(o: Any) -> Optional[str]:
     return None
 
 
+def parse_account(account: str, default_facility: str = "shared", default_repo: str = "default") -> tuple:
+    """Split a slurm account of the form <facility>:<repo>[@<partition>][^<qos>] into (facility, repo)."""
+    facility, sep, repo = account.partition(":")
+    if not sep or not repo:
+        return default_facility, default_repo
+    return facility, re.split(r"[@^]", repo, maxsplit=1)[0]
+
+
 def time_function(level="INFO"):
     """Decorator to time function execution and log the duration."""
     def decorator(func):
@@ -221,8 +229,11 @@ def remap_job(self, d) -> Optional[dict]:
             a = d["Partition"].split(",")[0]
             d["Partition"] = a
 
-        if "@" in d["Account"]:
-            d["Account"], _ = d["Account"].split("@")
+        if "^preemptable" in d["Account"]:
+            d["QOS"] = "preemptable"
+
+        if "@" in d["Account"] or "^" in d["Account"]:
+            d["Account"] = re.split(r"[@^]", d["Account"])[0]
 
         if d["QOS"] in ("Unknown",):
             d["QOS"] = "normal"
@@ -770,11 +781,8 @@ def calc_resource_hours(startTs, endTs, tres: str, cluster: dict, alloc_nodes: O
             return resource_time, elapsed_secs
 
         d = {field: parts[idx] for field, idx in index.items()}
-        facility = default_facility
-        repo = default_repo
-        try:
-            facility, repo = d["Account"].split(":")
-        except Exception:
+        facility, repo = parse_account(d["Account"], default_facility, default_repo)
+        if (facility, repo) == (default_facility, default_repo) and ":" not in d["Account"]:
             logger.warning(f"could not determine facility and repo from {d['Account']}")
 
         startTs = parse_datetime(int(d["Start"]), force_tz=True)
@@ -804,14 +812,21 @@ def calc_resource_hours(startTs, endTs, tres: str, cluster: dict, alloc_nodes: O
                 sys.exit(1)
             return None
 
-        qos = d["QOS"]
+        raw_qos = d.get("QOS", "")
+        clean_qos = raw_qos
         try:
-            a = qos.split("^")
+            a = raw_qos.split("^")
             b = a[1].split("@")
-            qos = b[0]
-        except:
-            pass
-        if qos not in ("scavenger", "preemptable", "normal"):
+            clean_qos = b[0]
+        except Exception:
+            clean_qos = raw_qos.split("@")[0].split("^")[0]
+
+        if "^preemptable" in d["Account"] or repo == "default" or clean_qos == "preemptable":
+            qos = "preemptable"
+        elif clean_qos in ("preemptable", "normal"):
+            qos = clean_qos
+        else:
+            qos = "normal"
             logger.warning(f"could not determine appropriate qos '{d['QOS']}': line {d}")
 
         out = {

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
 test = [
     "pytest>=9.0.3",
     "pytest-asyncio>=1.3.0",
+    "ansible-lint>=26.8.0",
 ]
 
 [tool.setuptools]

tests/test_slurm_account_parsing.py

@@ -0,0 +1,97 @@
+import pytest
+from modules.coact import parse_account, SlurmRemapper, SlurmImporter
+
+
+@pytest.mark.parametrize(
+    "account,expected",
+    [
+        ("lcls:foo", ("lcls", "foo")),
+        ("lcls:foo@milano", ("lcls", "foo")),
+        ("lcls:foo@milano^preemptable", ("lcls", "foo")),
+        ("lcls:foo^preemptable@milano", ("lcls", "foo")),
+        ("lcls:default@milano", ("lcls", "default")),
+        ("lcls:_regular_@milano", ("lcls", "_regular_")),
+        ("lcls:_preemptable_@milano", ("lcls", "_preemptable_")),
+    ],
+)
+def test_parse_account(account, expected):
+    assert parse_account(account) == expected
+
+
+@pytest.mark.parametrize("account", ["root", "", "nocolon"])
+def test_parse_account_falls_back(account):
+    assert parse_account(account, "shared", "default") == ("shared", "default")
+
+
+@pytest.mark.parametrize(
+    "account_in,qos_in,account_out,qos_out",
+    [
+        # Legacy formats
+        ("lcls:foo", "normal", "lcls:foo", "normal"),
+        ("lcls:foo", "preemptable", "lcls:foo", "preemptable"),
+        ("lcls:foo@milano", "normal", "lcls:foo", "normal"),
+        ("lcls:foo@milano", "preemptable", "lcls:foo", "preemptable"),
+        # New hierarchy formats
+        ("lcls:foo@milano^preemptable", "normal", "lcls:foo", "preemptable"),
+        ("lcls:foo@milano^preemptable", "preemptable", "lcls:foo", "preemptable"),
+        ("lcls:default@milano", "normal", "lcls:default", "normal"),
+    ],
+)
+def test_slurm_remapper(account_in, qos_in, account_out, qos_out):
+    remapper = SlurmRemapper()
+    job_dict = {
+        "JobID": "12345",
+        "User": "testuser",
+        "Account": account_in,
+        "Partition": "milano",
+        "QOS": qos_in,
+    }
+    result = remapper.remap_job(job_dict)
+    assert result is not None
+    assert result["Account"] == account_out
+    assert result["QOS"] == qos_out
+
+
+def test_slurm_importer_convert_dual_hierarchy(monkeypatch):
+    from modules.coact import parse_datetime
+    importer = SlurmImporter(username="test", password_file="dummy")
+    importer._clusters = {"milano": {"cpu": 64, "gpu": 0, "mem": 256 * 1073741824}}
+
+    # Mock allocid lookup
+    start_dt = parse_datetime(1767000000)
+    end_dt = parse_datetime(1768000000)
+    importer._allocid = {("lcls", "foo", "milano"): {(start_dt, end_dt): "alloc_123"}}
+
+    index = {
+        "JobID": 0,
+        "User": 1,
+        "Account": 2,
+        "Partition": 3,
+        "QOS": 4,
+        "Start": 5,
+        "End": 6,
+        "AllocNodes": 7,
+        "NCPUS": 8,
+        "AllocTRES": 9,
+    }
+
+    # Test 1: New preemptable account format
+    parts_new_preempt = ["1001", "user1", "lcls:foo@milano^preemptable", "milano", "normal", "1767225600", "1767229200", "1", "4", "cpu=4,mem=16G"]
+    converted = importer.convert(index, parts_new_preempt)
+    assert converted is not None
+    assert converted["allocationId"] == "alloc_123"
+    assert converted["qos"] == "preemptable"
+
+    # Test 2: New normal account format
+    parts_new_normal = ["1002", "user1", "lcls:foo@milano", "milano", "normal", "1767225600", "1767229200", "1", "4", "cpu=4,mem=16G"]
+    converted = importer.convert(index, parts_new_normal)
+    assert converted is not None
+    assert converted["allocationId"] == "alloc_123"
+    assert converted["qos"] == "normal"
+
+    # Test 3: Legacy preemptable QOS format
+    parts_legacy_preempt = ["1003", "user1", "lcls:foo@milano", "milano", "preemptable", "1767225600", "1767229200", "1", "4", "cpu=4,mem=16G"]
+    converted = importer.convert(index, parts_legacy_preempt)
+    assert converted is not None
+    assert converted["allocationId"] == "alloc_123"
+    assert converted["qos"] == "preemptable"