def test_cookiecutter(cookies, monkeypatch):
result = cookies.bake(template=str(pytest.cookie_template))
monkeypatch.syspath_prepend(str(result.project))
from CookieCutter import CookieCutter
assert CookieCutter.SBATCH_DEFAULTS == ""
assert CookieCutter.CLUSTER_NAME == ""
assert CookieCutter.CLUSTER_CONFIG == ""
assert CookieCutter.get_cluster_option() == ""
assert CookieCutter.get_advanced_argument_conversion() is False
sys.modules.pop("CookieCutter")
def test_cookiecutter_extra_context(cookies, monkeypatch):
result = cookies.bake(template=str(pytest.cookie_template),
extra_context={
"sbatch_defaults": "account=foo",
"cluster_name": "dusk",
"cluster_config": "slurm.yaml",
"advanced_argument_conversion": "yes"
})
monkeypatch.syspath_prepend(str(result.project))
from CookieCutter import CookieCutter
assert CookieCutter.SBATCH_DEFAULTS == "account=foo"
assert CookieCutter.CLUSTER_NAME == "dusk"
assert CookieCutter.CLUSTER_CONFIG == "slurm.yaml"
assert CookieCutter.get_cluster_option() == "--cluster=dusk"
assert CookieCutter.get_advanced_argument_conversion() is True
sys.modules.pop("CookieCutter")
def _get_default_partition():
"""Retrieve default partition for cluster"""
cluster = CookieCutter.get_cluster_option()
cmd = f"sinfo -O partition {cluster}"
res = sp.check_output(cmd.split())
m = re.search(r"(?P<partition>\S+)\*", res.decode(), re.M)
partition = m.group("partition")
return partition
def _get_cluster_configuration(partition, constraints=None, memory=0):
"""Retrieve cluster configuration.
Retrieve cluster configuration for a partition filtered by
constraints, memory and cpus
"""
try:
import pandas as pd
except ImportError:
print(
"Error: currently advanced argument conversion "
"depends on 'pandas'.",
file=sys.stderr)
sys.exit(1)
if constraints:
constraint_set = set(constraints.split(","))
cluster = CookieCutter.get_cluster_option()
cmd = f"sinfo -e -o %all -p {partition} {cluster}".split()
try:
output = sp.Popen(" ".join(cmd), shell=True,
stdout=sp.PIPE).communicate()
except Exception as e:
print(e)
raise
data = re.sub(" \\|", "|", output[0].decode())
df = pd.read_csv(StringIO(data), sep="|")
try:
df["TIMELIMIT_MINUTES"] = df["TIMELIMIT"].apply(time_to_minutes)
df["MEMORY_PER_CPU"] = df["MEMORY"] / df["CPUS"]
df["FEATURE_SET"] = df["AVAIL_FEATURES"].str.split(",").apply(set)
except Exception as e:
print(e)
raise
if constraints:
constraint_set = set(constraints.split(","))
i = df["FEATURE_SET"].apply(
lambda x: len(x.intersection(constraint_set)) > 0)
df = df.loc[i]
memory = min(_convert_units_to_mb(memory), max(df["MEMORY"]))
df = df.loc[df["MEMORY"] >= memory]
return df
def tail(path: str, num_lines: int = 10) -> List[bytes]:
if not Path(path).exists():
# allow for filesystem latency
time.sleep(CookieCutter.get_latency_wait())
if not Path(path).exists():
raise FileNotFoundError("{} does not exist.".format(path))
process = subprocess.Popen(
["tail", "-n", str(num_lines), path],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
exit_code = process.wait()
if exit_code != 0:
raise TailError(
"Failed to execute the tail command on the file {} due to the "
"following error:\n{}".format(path,
process.stderr.read().decode()))
return process.stdout.readlines()
),
file=sys.stderr,
)
try:
status = self._query_status_using_log()
except UnknownStatusLine as error:
print("UnknownStatusLine: {}".format(error), file=sys.stderr)
status = self.FAILED
return status
if __name__ == "__main__":
jobid = int(sys.argv[1])
outlog = sys.argv[2]
if CookieCutter.get_unknwn_behaviour().lower() == "wait":
kill_unknown = False
elif CookieCutter.get_unknwn_behaviour().lower() == "kill":
kill_unknown = True
else:
raise ValueError(
"Unknown value for {}_behaviour: {}".format(
UNKNOWN, CookieCutter.get_unknwn_behaviour()
)
)
if CookieCutter.get_zombi_behaviour().lower() == "ignore":
kill_zombie = False
elif CookieCutter.get_zombi_behaviour().lower() == "kill":
kill_zombie = True
else:
raise ValueError(
def mem_mb(self) -> Memory:
mem_value = self.resources.get(
"mem_mb",
self.cluster.get("mem_mb", CookieCutter.get_default_mem_mb()))
return Memory(mem_value, unit=Unit.MEGA)
def threads(self) -> int:
return self.job_properties.get("threads",
CookieCutter.get_default_threads())
def queue(self) -> str:
return self.cluster.get("queue", CookieCutter.get_default_queue())
def logdir(self) -> Path:
project_logdir = Path(
self.cluster.get("logdir", CookieCutter.get_log_dir()))
return project_logdir / self.rule_name / self.wildcards_str
try:
external_job_id = self._submit_cmd_and_get_external_job_id()
parameters_to_status_script = self._get_parameters_to_status_script(
external_job_id)
OSLayer.print(parameters_to_status_script)
except subprocess.CalledProcessError as error:
raise BsubInvocationError(error)
except AttributeError as error:
raise JobidNotFoundError(error)
if __name__ == "__main__":
workdir = Path().resolve()
config_file = workdir / "lsf.yaml"
if config_file.exists():
with config_file.open() as stream:
lsf_config = Config.from_stream(stream)
else:
lsf_config = Config()
jobscript = sys.argv[-1]
cluster_cmds = sys.argv[1:-1]
memory_units = Unit.from_suffix(CookieCutter.get_lsf_unit_for_limits())
lsf_submit = Submitter(
jobscript=jobscript,
memory_units=memory_units,
lsf_config=lsf_config,
cluster_cmds=cluster_cmds,
)
lsf_submit.submit()
#!/usr/bin/env python3
"""
Snakemake SLURM submit script.
"""
from snakemake.utils import read_job_properties
import slurm_utils
from CookieCutter import CookieCutter
# cookiecutter arguments
SBATCH_DEFAULTS = CookieCutter.SBATCH_DEFAULTS
CLUSTER = CookieCutter.get_cluster_option()
#CLUSTER_CONFIG = CookieCutter.CLUSTER_CONFIG
CLUSTER_CONFIG = 'cluster_config.yml'
ADVANCED_ARGUMENT_CONVERSION = CookieCutter.get_advanced_argument_conversion()
RESOURCE_MAPPING = {
"time": ("time", "runtime", "walltime"),
"mem": ("mem", "mem_mb", "ram", "memory"),
"mem-per-cpu": ("mem-per-cpu", "mem_per_cpu", "mem_per_thread"),
"nodes": ("nodes", "nnodes"),
}
# parse job
jobscript = slurm_utils.parse_jobscript()
job_properties = read_job_properties(jobscript)
sbatch_options = {}
cluster_config = slurm_utils.load_cluster_config(CLUSTER_CONFIG)
# 1) sbatch default arguments and cluster
#!/usr/bin/env python3
import re
import subprocess as sp
import shlex
import sys
import time
import logging
from CookieCutter import CookieCutter
logger = logging.getLogger("__name__")
STATUS_ATTEMPTS = 20
jobid = sys.argv[1]
cluster = CookieCutter.get_cluster_option()
for i in range(STATUS_ATTEMPTS):
try:
sacct_res = sp.check_output(
shlex.split(f"sacct {cluster} -P -b -j {jobid} -n"))
res = {
x.split("|")[0]: x.split("|")[1]
for x in sacct_res.decode().strip().split("\n")
}
break
except sp.CalledProcessError as e:
logger.error("sacct process error")
logger.error(e)
except IndexError as e:
logger.error(e)