Relax canceled code model to accept arbitrary ints

Fixes #268

CHANGELOG.md

@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Fixed
+
+- Relaxed `CanceledCode` schema to accept arbitrary integers.
+
 ## [v1.26.3](https://github.com/allenai/beaker-py/releases/tag/v1.26.3) - 2024-03-20
 
 ### Fixed

beaker/data_model/job.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from pydantic import Field
 
@@ -64,7 +64,7 @@ class JobStatus(BaseModel):
     finalized: Optional[datetime] = None
     canceled: Optional[datetime] = None
     canceled_for: Optional[str] = None
-    canceled_code: Optional[CanceledCode] = None
+    canceled_code: Optional[Union[CanceledCode, int]] = None
     idle_since: Optional[datetime] = None
     exit_code: Optional[int] = None
     message: Optional[str] = None
@@ -245,7 +245,7 @@ class JobStatusUpdate(BaseModel):
     finalized: Optional[bool] = None
     canceled: Optional[bool] = None
     canceled_for: Optional[str] = None
-    canceled_code: Optional[CanceledCode] = None
+    canceled_code: Optional[Union[CanceledCode, int]] = None
     idle: Optional[bool] = None
     message: Optional[str] = None
 

tests/data_model_test.py

@@ -175,3 +175,13 @@ def test_constraints_behave_like_dictionaries():
 def test_constraints_extra_fields():
     c = Constraints(cluster=["ai2/general-cirrascale"], gpus=["A100"])  # type: ignore
     assert hasattr(c, "gpus")
+
+
+def test_job_status_with_canceled_code():
+    from datetime import datetime
+
+    status = JobStatus(created=datetime.utcnow(), canceled_code=0)
+    assert status.canceled_code == CanceledCode.not_set
+
+    status = JobStatus(created=datetime.utcnow(), canceled_code=6)
+    assert status.canceled_code == 6