Merge pull request #166 from isofinly/main

feat: Implement window constraints with new command-line options.

Author: abhidg

cats/CI_api_interface.py

@@ -47,7 +47,7 @@ def ciuk_parse_response_data(response: dict):
     and is set up to cache data from call to call even accross different
     processes within the same half hour window. The returned prediction data
     is in half hour blocks starting from the half hour containing the current
-    time and extending for 48 hours into the future.
+    time and extending up to 48 hours into the future.
 
     :param response:
     :return:

cats/__init__.py

@@ -15,7 +15,10 @@
 from .configure import get_runtime_config
 from .constants import CATS_ASCII_BANNER_COLOUR, CATS_ASCII_BANNER_NO_COLOUR
 from .plotting import plotplan
-from .forecast import CarbonIntensityAverageEstimate, WindowedForecast
+from .forecast import (
+    CarbonIntensityAverageEstimate,
+    WindowedForecast,
+)
 
 __version__ = "1.1.0"
 
@@ -28,6 +31,68 @@ def indent_lines(lines, spaces):
     return "\n".join(" " * spaces + line for line in lines.split("\n"))
 
 
+def parse_time_constraint(
+    time_str: str, timezone_info=None
+) -> Optional[datetime.datetime]:
+    """
+    Parse a time constraint string into a datetime object.
+
+    :param time_str: Time string in various formats (HH:MM, YYYY-MM-DDTHH:MM, etc.)
+    :param timezone_info: Default timezone if not specified in the string
+    :return: Parsed datetime object
+    :raises ValueError: If the time string cannot be parsed
+    """
+    if not time_str:
+        return None
+
+    # If timezone_info is not provided, use system local timezone
+    if timezone_info is None:
+        timezone_info = datetime.datetime.now().astimezone().tzinfo
+
+    # Try to parse as full ISO format first
+    try:
+        if "T" in time_str:
+            # Full datetime string
+            if time_str.endswith("Z"):
+                time_str = time_str[:-1] + "+00:00"
+            elif time_str[-6] not in ["+", "-"] and time_str[-3] != ":":
+                # No timezone info, add default
+                dt = datetime.datetime.fromisoformat(time_str)
+                return dt.replace(tzinfo=timezone_info)
+            return datetime.datetime.fromisoformat(time_str)
+        else:
+            # Time only (HH:MM or HH:MM:SS)
+            time_part = datetime.time.fromisoformat(time_str)
+            today = datetime.datetime.now().date()
+            return datetime.datetime.combine(today, time_part, tzinfo=timezone_info)
+    except ValueError as e:
+        raise ValueError(f"Unable to parse time constraint '{time_str}': {e}")
+
+
+def validate_window_constraints(
+    start_window: Optional[datetime.datetime],
+    end_window: Optional[datetime.datetime],
+    window_minutes: int,
+) -> tuple[Optional[datetime.datetime], Optional[datetime.datetime], int]:
+    """
+    Validate window constraints.
+
+    :param start_window: Start window constraint datetime
+    :param end_window: End window constraint datetime
+    :param window_minutes: Maximum window duration in minutes
+    :return: Tuple of (start_datetime, end_datetime, validated_window_minutes)
+    :raises ValueError: If constraints are invalid
+    """
+    if window_minutes < 1 or window_minutes > 2820:
+        raise ValueError("Window must be between 1 and 2820 minutes (47 hours)")
+
+    if start_window and end_window:
+        if start_window >= end_window:
+            raise ValueError("Start window must be before end window")
+
+    return start_window, end_window, window_minutes
+
+
 def parse_arguments():
     """
     Parse command line arguments
@@ -201,6 +266,27 @@ def positive_integer(string):
         "\"pip install 'climate-aware-task-scheduler[plots]'\"",
         action="store_true",
     )
+    parser.add_argument(
+        "--window",
+        type=positive_integer,
+        help="Maximum time window to search for optimal start time, in minutes. "
+        "Must be between 1 and 2820 (47 hours). Default: 2820 minutes (47 hours).",
+        default=2820,
+    )
+    parser.add_argument(
+        "--start-window",
+        type=parse_time_constraint,
+        help="Earliest time the job is allowed to start, in ISO format (e.g., '2024-01-15T09:00'). "
+        "If only time is provided (e.g., '09:00'), today's date is assumed. "
+        "Timezone info is optional and defaults to system timezone.",
+    )
+    parser.add_argument(
+        "--end-window",
+        type=parse_time_constraint,
+        help="Latest time the job is allowed to start, in ISO format (e.g., '2024-01-15T17:00'). "
+        "If only time is provided (e.g., '17:00'), today's date is assumed. "
+        "Timezone info is optional and defaults to system timezone.",
+    )
 
     return parser
 
@@ -324,9 +410,6 @@ def main(arguments=None) -> int:
     args = parser.parse_args(arguments)
     colour_output = args.no_colour or args.no_color
 
-    # Print CATS ASCII art banner, before any output from printing or logging
-    print_banner(colour_output)
-
     if args.command and not args.scheduler:
         print(
             "cats: To run a command or sbatch script with the -c or --command option, you must\n"
@@ -335,11 +418,27 @@ def main(arguments=None) -> int:
         return 1
 
     CI_API_interface, location, duration, jobinfo, PUE = get_runtime_config(args)
-    if duration > CI_API_interface.max_duration:
-        print(
-            f"""API allows a maximum job duration of {CI_API_interface.max_duration} minutes.
-This is usually due to forecast limitations."""
+
+    # Validate and parse window constraints
+    try:
+        start_constraint, end_constraint, max_window = validate_window_constraints(
+            args.start_window, args.end_window, args.window
         )
+    except ValueError as e:
+        print(f"Error in window constraints: {e}")
+        return 1
+    # Check against both API limit and user-specified window
+    effective_max_duration = min(CI_API_interface.max_duration, max_window)
+    if duration > effective_max_duration:
+        if max_window < CI_API_interface.max_duration:
+            print(
+                f"""Job duration ({duration} minutes) exceeds specified window ({max_window} minutes)."""
+            )
+        else:
+            print(
+                f"""API allows a maximum job duration of {CI_API_interface.max_duration} minutes.
+This is usually due to forecast limitations."""
+            )
         return 1
 
     ########################
@@ -362,8 +461,21 @@ def main(arguments=None) -> int:
 
     # Find best possible average carbon intensity, along
     # with corresponding job start time.
+    search_start = datetime.datetime.now().astimezone()
+
+    # Apply start window constraint if provided
+    if start_constraint:
+        # Ensure start constraint is in the same timezone as search_start
+        if start_constraint.tzinfo != search_start.tzinfo:
+            start_constraint = start_constraint.astimezone(search_start.tzinfo)
+        search_start = max(search_start, start_constraint)
+
     wf = WindowedForecast(
-        CI_forecast, duration, start=datetime.datetime.now().astimezone()
+        CI_forecast,
+        duration,
+        start=search_start,
+        max_window_minutes=max_window,
+        end_constraint=end_constraint,
     )
     now_avg, best_avg = wf[0], min(wf)
     output = CATSOutput(now_avg, best_avg, location, "GBR", colour=not colour_output)
@@ -390,6 +502,8 @@ def main(arguments=None) -> int:
             dateformat = args.dateformat or ""
         print(output.to_json(dateformat, sort_keys=True, indent=2))
     else:
+        # Print CATS ASCII art banner, before any output from printing or logging
+        print_banner(colour_output)
         print(output)
     if args.plot:
         plotplan(CI_forecast, output)

cats/forecast.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, InitVar
+from dataclasses import dataclass
 from typing import Optional
 from datetime import datetime, timedelta
 
@@ -29,8 +29,8 @@ class CarbonIntensityAverageEstimate:
     value: float
     start: datetime  # Start of the time-integration window
     end: datetime  # End of the time-integration window
-    start_value: float # CI point estimate at start time
-    end_value: float # CI point estimate at end time
+    start_value: float  # CI point estimate at start time
+    end_value: float  # CI point estimate at end time
 
 
 class WindowedForecast:
@@ -39,10 +39,23 @@ def __init__(
         data: list[CarbonIntensityPointEstimate],
         duration: int,  # in minutes
         start: datetime,
+        max_window_minutes: Optional[int] = None,
+        end_constraint: Optional[datetime] = None,
     ):
-        self.data_stepsize = data[1].datetime - data[0].datetime
+        self.duration = duration
+        self.max_window_minutes = max_window_minutes
+        self.end_constraint = end_constraint
+
+        # Filter data based on constraints if any are specified
+        if max_window_minutes is not None or end_constraint is not None:
+            filtered_data = self._filter_data_by_constraints(
+                data, start, duration, max_window_minutes or 2820, end_constraint
+            )
+        else:
+            filtered_data = data
+
+        self.data_stepsize = filtered_data[1].datetime - filtered_data[0].datetime
         self.start = start
-        # TODO: Expect duration as a timedelta directly
         self.end = start + timedelta(minutes=duration)
 
         # Restrict data points so that start time falls within the
@@ -57,7 +70,7 @@ def bisect_right(data, t):
         # bisect_right(data, start) returns the index of the first
         # data point with datetime value immediately preceding the job
         # start time
-        self.data = data[bisect_right(data, start) :]
+        self.data = filtered_data[bisect_right(filtered_data, start) :]
 
         # Find number of data points in a window, by finding the index
         # of the closest data point past the job end time. Could be
@@ -74,6 +87,46 @@ def bisect_left(data, t):
 
         self.ndata = bisect_left(self.data, self.end)  # window size
 
+    def _filter_data_by_constraints(
+        self,
+        data: list[CarbonIntensityPointEstimate],
+        start: datetime,
+        duration: int,
+        max_window_minutes: int,
+        end_constraint: Optional[datetime],
+    ) -> list[CarbonIntensityPointEstimate]:
+        """Filter forecast data based on time constraints."""
+
+        # Calculate the maximum time we need data for
+        search_window_end = start + timedelta(minutes=max_window_minutes)
+
+        if end_constraint:
+            # Ensure timezone compatibility
+            if end_constraint.tzinfo != start.tzinfo:
+                end_constraint = end_constraint.astimezone(start.tzinfo)
+            # Jobs must start before end_constraint
+            search_window_end = min(search_window_end, end_constraint)
+
+        # We need data points to cover jobs starting up to search_window_end
+        # plus the duration of those jobs
+        max_data_time = search_window_end + timedelta(minutes=duration)
+
+        # Filter data to respect the constraints
+        filtered_data = []
+        for d in data:
+            if d.datetime <= max_data_time:
+                filtered_data.append(d)
+            else:
+                break
+
+        if len(filtered_data) < 2:
+            raise ValueError(
+                "Insufficient forecast data for the specified time window constraints. "
+                "Try increasing --window or adjusting --end-window."
+            )
+
+        return filtered_data
+
     def __getitem__(self, index: int) -> CarbonIntensityAverageEstimate:
         """Return the average of timeseries data from index over the
         window size.  Data points are integrated using the trapeziodal
@@ -86,6 +139,9 @@ def __getitem__(self, index: int) -> CarbonIntensityAverageEstimate:
         dividing the total integral value by the number of intervals.
         """
 
+        if index >= len(self):
+            raise IndexError("Window index out of range")
+
         # Account for the fact that the start and end of each window
         # might not fall exactly on data points.  The starting
         # intensity is interpolated between the first (index) and
@@ -149,8 +205,36 @@ def interp(
         )
 
     def __iter__(self):
-        for index in range(self.__len__()):
-            yield self.__getitem__(index)
+        for index in range(len(self)):
+            yield self[index]
 
     def __len__(self):
-        return len(self.data) - self.ndata
+        """Return number of valid forecast windows respecting all constraints."""
+        base_length = len(self.data) - self.ndata
+
+        if base_length <= 0:
+            return 0
+
+        max_valid_index = base_length - 1
+
+        # Check max window constraint only if specified
+        if self.max_window_minutes is not None:
+            data_stepsize_minutes = self.data_stepsize.total_seconds() / 60
+            max_index_by_window = int(self.max_window_minutes / data_stepsize_minutes)
+            max_valid_index = min(max_valid_index, max_index_by_window)
+
+        # Check end constraint
+        if self.end_constraint:
+            if self.end_constraint.tzinfo != self.start.tzinfo:
+                end_constraint = self.end_constraint.astimezone(self.start.tzinfo)
+            else:
+                end_constraint = self.end_constraint
+
+            # Find the maximum index where job start time is before end_constraint
+            for i in range(min(base_length, max_valid_index + 1)):
+                window_start = self.start + i * self.data_stepsize
+                if window_start >= end_constraint:
+                    max_valid_index = i - 1
+                    break
+
+        return max(0, max_valid_index + 1)