Skip to content

Commit

Permalink
Add JobTimeLimit class
Browse files Browse the repository at this point in the history
Added a representation of batch job time limits, similar to `JobSize`,
along with corresponding documentation and unit tests.
  • Loading branch information
TimothyWillard committed Nov 5, 2024
1 parent 51fc0bd commit 74b4ccd
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 2 deletions.
58 changes: 56 additions & 2 deletions flepimop/gempyor_pkg/src/gempyor/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,12 @@
metadata and job size calculations for example.
"""

__all__ = ["JobSize", "write_manifest"]
__all__ = ["JobSize", "JobTimeLimit", "write_manifest"]


import click
from dataclasses import dataclass
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
import json
import math
from pathlib import Path
Expand Down Expand Up @@ -149,6 +149,60 @@ def size_from_jobs_sims_blocks(
return cls(jobs=jobs, simulations=simulations, blocks=blocks)


@dataclass(frozen=True, slots=True)
class JobTimeLimit:
"""
A batch submission job time limit.
Attributes:
time_limit: The time limit of the batch job.
Raises:
ValueError: If the `time_limit` attribute is not positive.
"""

time_limit: timedelta

def __post_init__(self) -> None:
if (total_seconds := self.time_limit.total_seconds()) <= 0.0:
raise ValueError(
f"The `time_limit` attribute has {math.floor(total_seconds):,} "
"seconds, which is less than or equal to 0."
)

def __str__(self) -> str:
return self.format()

def format(self, batch_system: Literal["aws", "local", "slurm"] | None = None) -> str:
"""
Format the job time limit as a string appropriate for a given batch system.
Args:
batch_system: The batch system the format should be formatted for.
Returns:
The time limit formatted for the batch system.
Examples:
>>> from datetime import timedelta
>>> job_time_limit = JobTimeLimit(
... time_limit=timedelta(days=1, hours=2, minutes=34, seconds=5)
... )
>>> job_time_limit.format()
'1595'
>>> job_time_limit.format(batch_system="slurm")
'26:34:05'
"""
if batch_system == "slurm":
total_seconds = self.time_limit.total_seconds()
hours = math.floor(total_seconds / (60.0 * 60.0))
minutes = math.floor((total_seconds - (60.0 * 60.0 * hours)) / 60.0)
seconds = math.ceil(total_seconds - (60.0 * minutes) - (60.0 * 60.0 * hours))
return f"{hours}:{minutes:02d}:{seconds:02d}"
limit_in_mins = math.ceil(self.time_limit.total_seconds() / 60.0)
return str(limit_in_mins)


def write_manifest(
job_name: str,
flepi_path: Path,
Expand Down
63 changes: 63 additions & 0 deletions flepimop/gempyor_pkg/tests/batch/test_job_time_limit_class.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
from datetime import timedelta
import re
from typing import Literal

import pytest

from gempyor.batch import JobTimeLimit


@pytest.mark.parametrize(
"time_limit", (timedelta(), timedelta(hours=-1.0), timedelta(days=-3.0))
)
def test_time_limit_non_positive_value_error(time_limit: timedelta) -> None:
with pytest.raises(
ValueError,
match=(
r"^The \`time\_limit\` attribute has [0-9\,\-]+ seconds\, "
r"which is less than or equal to 0\.$"
),
):
JobTimeLimit(time_limit=time_limit)


@pytest.mark.parametrize(
"time_limit",
(
timedelta(hours=1),
timedelta(hours=2, minutes=34, seconds=56),
timedelta(days=1, seconds=3),
timedelta(minutes=12345),
),
)
@pytest.mark.parametrize("batch_system", (None, "aws", "local", "slurm"))
def test_format_output_validation(
time_limit: timedelta, batch_system: Literal["aws", "local", "slurm"] | None
) -> None:
job_time_limit = JobTimeLimit(time_limit=time_limit)
formatted_time_limit = job_time_limit.format(batch_system=batch_system)
assert isinstance(formatted_time_limit, str)
if batch_system == "slurm":
assert re.match(r"^[0-9]+\:[0-9]{2}\:[0-9]{2}$", formatted_time_limit)
else:
assert formatted_time_limit.isdigit()


@pytest.mark.parametrize(
("time_limit", "batch_system", "expected"),
(
(timedelta(hours=1), None, "60"),
(timedelta(seconds=20), None, "1"),
(timedelta(days=2, hours=3, minutes=45), None, "3105"),
(timedelta(hours=1), "slurm", "1:00:00"),
(timedelta(seconds=20), "slurm", "0:00:20"),
(timedelta(days=1, hours=2, minutes=34, seconds=5), "slurm", "26:34:05"),
),
)
def test_format_exact_results(
time_limit: timedelta,
batch_system: Literal["aws", "local", "slurm"] | None,
expected: str,
) -> None:
job_time_limit = JobTimeLimit(time_limit=time_limit)
assert job_time_limit.format(batch_system=batch_system) == expected

0 comments on commit 74b4ccd

Please sign in to comment.