Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes compare_xml method for env_batch #4706

Merged
merged 3 commits into from
Nov 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 104 additions & 26 deletions CIME/XML/env_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from collections import OrderedDict
import stat, re, math
import pathlib
from itertools import zip_longest

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -1412,41 +1413,118 @@ def cancel_job(self, jobid):
else:
return True

def zip(self, other, name):
for self_pnode in self.get_children(name):
try:
other_pnode = other.get_children(name, attributes=self_pnode.attrib)[0]
except (TypeError, IndexError):
other_pnode = None

for node1 in self.get_children(root=self_pnode):
for node2 in other.scan_children(
node1.name, attributes=node1.attrib, root=other_pnode
):
yield node1, node2

def _compare_arg(self, index, arg1, arg2):
try:
flag1 = arg1.attrib["flag"]
name1 = arg1.attrib.get("name", "")
except AttributeError:
flag2, name2 = arg2.attrib["flag"], arg2.attrib["name"]

return {f"arg{index}": ["", f"{flag2} {name2}"]}

try:
flag2 = arg2.attrib["flag"]
name2 = arg2.attrib.get("name", "")
except AttributeError:
return {f"arg{index}": [f"{flag1} {name1}", ""]}

if flag1 != flag2 or name1 != name2:
return {f"arg{index}": [f"{flag1} {name1}", f"{flag2} {name2}"]}

return {}

def _compare_argument(self, index, arg1, arg2):
if arg1.text != arg2.text:
return {f"argument{index}": [arg1.text, arg2.text]}

return {}

def compare_xml(self, other):
xmldiffs = {}
f1batchnodes = self.get_children("batch_system")
for bnode in f1batchnodes:
f2bnodes = other.get_children("batch_system", attributes=self.attrib(bnode))
f2bnode = None
if len(f2bnodes):
f2bnode = f2bnodes[0]
f1batchnodes = self.get_children(root=bnode)
for node in f1batchnodes:
name = self.name(node)
text1 = self.text(node)
text2 = ""
attribs = self.attrib(node)
f2matches = other.scan_children(name, attributes=attribs, root=f2bnode)
foundmatch = False
for chkmatch in f2matches:
name2 = other.name(chkmatch)
attribs2 = other.attrib(chkmatch)
text2 = other.text(chkmatch)
if name == name2 and attribs == attribs2 and text1 == text2:
foundmatch = True
break
if not foundmatch:
xmldiffs[name] = [text1, text2]

f1groups = self.get_children("group")
for node in f1groups:

for node1, node2 in self.zip(other, "batch_system"):
if node1.name == "submit_args":
self_nodes = self.get_children(root=node1)
other_nodes = other.get_children(root=node2)
for i, (x, y) in enumerate(
zip_longest(self_nodes, other_nodes, fillvalue=None)
):
if (x is not None and x.name == "arg") or (
y is not None and y.name == "arg"
):
xmldiffs.update(self._compare_arg(i, x, y))
elif (x is not None and x.name == "argument") or (
y is not None and y.name == "argument"
):
xmldiffs.update(self._compare_node(x, y, i))
elif node1.name == "directives":
self_nodes = self.get_children(root=node1)
other_nodes = other.get_children(root=node2)
for i, (x, y) in enumerate(
zip_longest(self_nodes, other_nodes, fillvalue=None)
):
xmldiffs.update(self._compare_node(x, y, i))
elif node1.name == "queues":
self_nodes = self.get_children(root=node1)
other_nodes = other.get_children(root=node2)
for i, (x, y) in enumerate(
zip_longest(self_nodes, other_nodes, fillvalue=None)
):
xmldiffs.update(self._compare_node(x, y, i))
else:
xmldiffs.update(self._compare_node(node1, node2))

for node in self.get_children("group"):
group = self.get(node, "id")
f2group = other.get_child("group", attributes={"id": group})
xmldiffs.update(
super(EnvBatch, self).compare_xml(other, root=node, otherroot=f2group)
)
return xmldiffs

def _compare_node(self, x, y, index=None):
"""Compares two XML nodes and returns diff.

Compares the attributes and text of two XML nodes. Handles the case when either node is `None`.

The `index` argument can be used to append the nodes tag. This can be useful when comparing a list
of XML nodes that all have the same tag to differentiate which nodes are different.

Args:
x (:obj:`CIME.XML.generic_xml._Element`): First node.
y (:obj:`CIME.XML.generic_xml._Element`): Second node.
index (int, optional): Index of the nodes.

Returns:
dict: Key is the tag and value is the difference.
"""
diff = {}

if index is None:
index = ""

if x is None:
diff[f"{y.name}{index}"] = ["", y.text]
elif y is None:
diff[f"{x.name}{index}"] = [x.text, ""]
elif x.text != y.text or x.attrib != y.attrib:
diff[f"{x.name}{index}"] = [x.text, y.text]

return diff

def make_all_batch_files(self, case):
machdir = case.get_value("MACHDIR")
logger.info("Creating batch scripts")
Expand Down
18 changes: 18 additions & 0 deletions CIME/XML/generic_xml.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,24 @@ def __hash__(self):
def __deepcopy__(self, _):
return _Element(deepcopy(self.xml_element))

def __str__(self):
return str(self.xml_element)

def __repr__(self):
return repr(self.xml_element)

@property
def name(self):
return self.xml_element.tag

@property
def text(self):
return self.xml_element.text

@property
def attrib(self):
return dict(self.xml_element.attrib)


class GenericXML(object):

Expand Down
167 changes: 165 additions & 2 deletions CIME/tests/test_unit_xml_env_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,178 @@
import os
import unittest
import tempfile
from contextlib import ExitStack
from unittest import mock

from CIME.utils import CIMEError
from CIME.utils import CIMEError, expect
from CIME.XML.env_batch import EnvBatch, get_job_deps

# pylint: disable=unused-argument

XML_BASE = b"""<?xml version="1.0"?>
<file id="env_batch.xml" version="2.0">
<header>
These variables may be changed anytime during a run, they
control arguments to the batch submit command.
</header>
<group id="config_batch">
<entry id="BATCH_SYSTEM" value="slurm">
<type>char</type>
<valid_values>miller_slurm,nersc_slurm,lc_slurm,moab,pbs,lsf,slurm,cobalt,cobalt_theta,none</valid_values>
<desc>The batch system type to use for this machine.</desc>
</entry>
</group>
<group id="job_submission">
<entry id="PROJECT_REQUIRED" value="FALSE">
<type>logical</type>
<valid_values>TRUE,FALSE</valid_values>
<desc>whether the PROJECT value is required on this machine</desc>
</entry>
</group>
<batch_system type="slurm">
<batch_query per_job_arg="-j">squeue</batch_query>
<batch_submit>sbatch</batch_submit>
<batch_cancel>scancel</batch_cancel>
<batch_directive>#SBATCH</batch_directive>
<jobid_pattern>(\\d+)$</jobid_pattern>
<depend_string>--dependency=afterok:jobid</depend_string>
<depend_allow_string>--dependency=afterany:jobid</depend_allow_string>
<depend_separator>:</depend_separator>
<walltime_format>%H:%M:%S</walltime_format>
<batch_mail_flag>--mail-user</batch_mail_flag>
<batch_mail_type_flag>--mail-type</batch_mail_type_flag>
<batch_mail_type>none, all, begin, end, fail</batch_mail_type>
<submit_args>
<arg flag="--time" name="$JOB_WALLCLOCK_TIME"/>
<arg flag="-p" name="$JOB_QUEUE"/>
<arg flag="--account" name="$PROJECT"/>
</submit_args>
<directives>
<directive> --job-name={{ job_id }}</directive>
<directive> --nodes={{ num_nodes }}</directive>
<directive> --output={{ job_id }}.%j </directive>
<directive> --exclusive </directive>
</directives>
</batch_system>
<batch_system MACH="docker" type="slurm">
<submit_args>
<argument>-w docker</argument>
</submit_args>
<queues>
<queue walltimemax="00:15:00" nodemax="1">debug</queue>
<queue walltimemax="24:00:00" nodemax="20" nodemin="5">big</queue>
<queue walltimemax="00:30:00" nodemax="5" default="true">smallfast</queue>
</queues>
</batch_system>
</file>"""

XML_DIFF = b"""<?xml version="1.0"?>
<file id="env_batch.xml" version="2.0">
<header>
These variables may be changed anytime during a run, they
control arguments to the batch submit command.
</header>
<group id="config_batch">
<entry id="BATCH_SYSTEM" value="pbs">
<type>char</type>
<valid_values>miller_slurm,nersc_slurm,lc_slurm,moab,pbs,lsf,slurm,cobalt,cobalt_theta,none</valid_values>
<desc>The batch system type to use for this machine.</desc>
</entry>
</group>
<group id="job_submission">
<entry id="PROJECT_REQUIRED" value="FALSE">
<type>logical</type>
<valid_values>TRUE,FALSE</valid_values>
<desc>whether the PROJECT value is required on this machine</desc>
</entry>
</group>
<batch_system type="slurm">
<batch_query per_job_arg="-j">squeue</batch_query>
<batch_submit>batch</batch_submit>
<batch_cancel>scancel</batch_cancel>
<batch_directive>#SBATCH</batch_directive>
<jobid_pattern>(\\d+)$</jobid_pattern>
<depend_string>--dependency=afterok:jobid</depend_string>
<depend_allow_string>--dependency=afterany:jobid</depend_allow_string>
<depend_separator>:</depend_separator>
<walltime_format>%H:%M:%S</walltime_format>
<batch_mail_flag>--mail-user</batch_mail_flag>
<batch_mail_type_flag>--mail-type</batch_mail_type_flag>
<batch_mail_type>none, all, begin, end, fail</batch_mail_type>
<submit_args>
<arg flag="--time" name="$JOB_WALLCLOCK_TIME"/>
<arg flag="-p" name="pbatch"/>
<arg flag="--account" name="$PROJECT"/>
<arg flag="-m" name="plane"/>
</submit_args>
<directives>
<directive> --job-name={{ job_id }}</directive>
<directive> --nodes=10</directive>
<directive> --output={{ job_id }}.%j </directive>
<directive> --exclusive </directive>
<directive> --qos=high </directive>
</directives>
</batch_system>
<batch_system MACH="docker" type="slurm">
<submit_args>
<argument>-w docker</argument>
</submit_args>
<queues>
<queue walltimemax="00:15:00" nodemax="1">debug</queue>
<queue walltimemax="24:00:00" nodemax="20" nodemin="10">big</queue>
</queues>
</batch_system>
</file>"""


def _open_temp_file(stack, data):
tfile = stack.enter_context(tempfile.NamedTemporaryFile())

tfile.write(data)

tfile.seek(0)

return tfile


class TestXMLEnvBatch(unittest.TestCase):
def test_compare_xml(self):
with ExitStack() as stack:
file1 = _open_temp_file(stack, XML_DIFF)
batch1 = EnvBatch(infile=file1.name)

file2 = _open_temp_file(stack, XML_BASE)
batch2 = EnvBatch(infile=file2.name)

diff = batch1.compare_xml(batch2)
diff2 = batch2.compare_xml(batch1)

expected_diff = {
"BATCH_SYSTEM": ["pbs", "slurm"],
"arg1": ["-p pbatch", "-p $JOB_QUEUE"],
"arg3": ["-m plane", ""],
"batch_submit": ["batch", "sbatch"],
"directive1": [" --nodes=10", " --nodes={{ num_nodes }}"],
"directive4": [" --qos=high ", ""],
"queue1": ["big", "big"],
"queue2": ["", "smallfast"],
}

assert diff == expected_diff

expected_diff2 = {
"BATCH_SYSTEM": ["slurm", "pbs"],
"arg1": ["-p $JOB_QUEUE", "-p pbatch"],
"arg3": ["", "-m plane"],
"batch_submit": ["sbatch", "batch"],
"directive1": [" --nodes={{ num_nodes }}", " --nodes=10"],
"directive4": ["", " --qos=high "],
"queue1": ["big", "big"],
"queue2": ["smallfast", ""],
}

assert diff2 == expected_diff2

@mock.patch("CIME.XML.env_batch.EnvBatch._submit_single_job")
def test_submit_jobs(self, _submit_single_job):
case = mock.MagicMock()
Expand Down Expand Up @@ -273,7 +436,7 @@ def test_get_submit_args(self):
<batch_submit>sbatch</batch_submit>
<batch_cancel>scancel</batch_cancel>
<batch_directive>#SBATCH</batch_directive>
<jobid_pattern>(\d+)$</jobid_pattern>
<jobid_pattern>(\\d+)$</jobid_pattern>
<depend_string>--dependency=afterok:jobid</depend_string>
<depend_allow_string>--dependency=afterany:jobid</depend_allow_string>
<depend_separator>:</depend_separator>
Expand Down
Loading