1313from datetime import UTC , datetime
1414from enum import StrEnum
1515from pathlib import Path
16- from typing import Any
16+ from typing import Any , Protocol
1717
1818import humanize
1919from loguru import logger
3939RUN_FAILED_MESSAGE = "Failed to get status for run with ID '%s'"
4040
4141
42- def validate_due_date ( due_date : str | None ) -> None :
43- """Validate that due_date is in ISO 8601 format and in the future.
42+ def _validate_scheduling_datetime ( value : str | None , field_name : str ) -> None :
43+ """Validate that a scheduling datetime field is in ISO 8601 format and in the future.
4444
4545 Args:
46- due_date (str | None): The datetime string to validate.
46+ value (str | None): The datetime string to validate.
47+ field_name (str): The name of the field (e.g. 'due_date', 'deadline') for error messages.
4748
4849 Raises:
4950 ValueError: If
50- the format is invalid
51- or the due_date is not in the future.
51+ the format is invalid,
52+ the datetime is timezone-naive,
53+ or the datetime is not in the future.
5254 """
53- if due_date is None :
55+ if value is None :
5456 return
5557
5658 # Try parsing with fromisoformat (handles most ISO 8601 formats)
5759 try :
5860 # Handle 'Z' suffix by replacing with '+00:00'
59- normalized = due_date .replace ("Z" , "+00:00" )
61+ normalized = value .replace ("Z" , "+00:00" )
6062 parsed_dt = datetime .fromisoformat (normalized )
6163 except (ValueError , TypeError ) as e :
6264 message = (
63- f"Invalid ISO 8601 format for due_date . "
65+ f"Invalid ISO 8601 format for { field_name }
6466 f"Expected format like '2025-10-19T19:53:00+00:00' or '2025-10-19T19:53:00Z', "
65- f"but got: '{ due_date } { e }
67+ f"but got: '{ value } { e }
6668 )
6769 raise ValueError (message ) from e
6870
6971 # Ensure the datetime is timezone-aware (reject naive datetimes)
7072 if parsed_dt .tzinfo is None :
7173 message = (
72- f"Invalid ISO 8601 format for due_date . "
74+ f"Invalid ISO 8601 format for { field_name }
7375 f"Expected format with timezone like '2025-10-19T19:53:00+00:00' or '2025-10-19T19:53:00Z', "
74- f"but got: '{ due_date }
76+ f"but got: '{ value }
7577 )
7678 raise ValueError (message )
7779
7880 # Check that the datetime is in the future
7981 now = datetime .now (UTC )
8082 if parsed_dt <= now :
8183 message = (
82- f"due_date must be in the future. "
83- f"Got '{ due_date } { parsed_dt .isoformat ()}
84+ f"{ field_name }
85+ f"Got '{ value } { parsed_dt .isoformat ()}
8486 f"but current UTC time is { now .isoformat ()}
8587 )
8688 raise ValueError (message )
8789
8890
91+ def validate_due_date (due_date : str | None ) -> None :
92+ """Validate that due_date is in ISO 8601 format and in the future.
93+
94+ Args:
95+ due_date (str | None): The datetime string to validate.
96+
97+ Raises:
98+ ValueError: If
99+ the format is invalid
100+ or the due_date is not in the future.
101+ """
102+ _validate_scheduling_datetime (due_date , "due_date" )
103+
104+
105+ def validate_deadline (deadline : str | None ) -> None :
106+ """Validate that deadline is in ISO 8601 format and in the future.
107+
108+ Args:
109+ deadline (str | None): The datetime string to validate.
110+
111+ Raises:
112+ ValueError: If
113+ the format is invalid
114+ or the deadline is not in the future.
115+ """
116+ _validate_scheduling_datetime (deadline , "deadline" )
117+
118+
119+ def _parse_scheduling_datetime (value : str ) -> datetime :
120+ """Parse an ISO 8601 scheduling datetime string, handling 'Z' suffix.
121+
122+ Args:
123+ value (str): The datetime string to parse.
124+
125+ Returns:
126+ datetime: The parsed timezone-aware datetime.
127+ """
128+ normalized = value .replace ("Z" , "+00:00" )
129+ return datetime .fromisoformat (normalized )
130+
131+
132+ def validate_scheduling_constraints (due_date : str | None , deadline : str | None ) -> None :
133+ """Validate cross-field scheduling constraints.
134+
135+ When both due_date and deadline are provided, due_date must be before deadline.
136+
137+ Args:
138+ due_date (str | None): The due date string (already individually validated).
139+ deadline (str | None): The deadline string (already individually validated).
140+
141+ Raises:
142+ ValueError: If due_date is not before deadline.
143+ """
144+ if due_date is None or deadline is None :
145+ return
146+
147+ parsed_due_date = _parse_scheduling_datetime (due_date )
148+ parsed_deadline = _parse_scheduling_datetime (deadline )
149+
150+ if parsed_due_date >= parsed_deadline :
151+ message = (
152+ f"due_date must be before deadline. "
153+ f"Got due_date='{ due_date } { parsed_due_date .isoformat ()}
154+ f"and deadline='{ deadline } { parsed_deadline .isoformat ()}
155+ )
156+ raise ValueError (message )
157+
158+
89159def validate_mappings (mappings : list [str ] | None ) -> None :
90160 """Validate mapping format for file metadata amendment.
91161
@@ -124,9 +194,16 @@ def validate_mappings(mappings: list[str] | None) -> None:
124194 raise ValueError (msg ) from e
125195
126196
197+ class _SchedulingLike (Protocol ):
198+ """Protocol for objects with an optional deadline attribute."""
199+
200+ deadline : datetime | None
201+
202+
127203def is_not_terminated_with_deadline_exceeded (
128204 run_state : RunState ,
129- custom_metadata : dict [str , Any ] | None ,
205+ scheduling : _SchedulingLike | None = None ,
206+ custom_metadata : dict [str , Any ] | None = None ,
130207) -> bool | None :
131208 """Check if the run is not terminated and the deadline has been exceeded.
132209
@@ -135,7 +212,11 @@ def is_not_terminated_with_deadline_exceeded(
135212
136213 Args:
137214 run_state (RunState): The current state of the run.
138- custom_metadata (dict[str, Any] | None): The custom metadata containing optional deadline information.
215+ scheduling: The scheduling response object from the API (has .deadline attribute),
216+ or None if no scheduling constraints were set.
217+ custom_metadata (dict[str, Any] | None): Legacy fallback - the custom metadata containing
218+ optional deadline information in sdk.scheduling.deadline. Used only when scheduling
219+ response field is not available.
139220
140221 Returns:
141222 bool | None: True if run is not terminated and deadline exceeded,
@@ -146,16 +227,29 @@ def is_not_terminated_with_deadline_exceeded(
146227 if run_state == RunState .TERMINATED :
147228 return None
148229
149- if not custom_metadata :
150- return None
151-
152- deadline_str = custom_metadata .get ("sdk" , {}).get ("scheduling" , {}).get ("deadline" )
153- if not deadline_str :
230+ # Try the first-class scheduling response field first
231+ deadline_value = None
232+ if scheduling is not None :
233+ deadline_value = getattr (scheduling , "deadline" , None )
234+
235+ # Fallback to custom_metadata for backward compatibility with older runs
236+ if deadline_value is None and custom_metadata :
237+ deadline_str = custom_metadata .get ("sdk" , {}).get ("scheduling" , {}).get ("deadline" )
238+ if deadline_str :
239+ try :
240+ deadline_value = _parse_scheduling_datetime (deadline_str )
241+ except (ValueError , TypeError , AttributeError ):
242+ return None
243+
244+ if deadline_value is None :
154245 return None
155246
156247 try :
157248 now = datetime .now (tz = UTC )
158- deadline_dt = datetime .fromisoformat (deadline_str )
249+ if isinstance (deadline_value , datetime ):
250+ return now > deadline_value
251+ # Handle string values
252+ deadline_dt = _parse_scheduling_datetime (str (deadline_value ))
159253 return now > deadline_dt
160254 except (ValueError , TypeError , AttributeError ):
161255 # Invalid deadline format, return None
0 commit comments