PyAPI crash limit

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+### New features
+
+* Python API `crash_limit` now accepts `"unlimited"` and `"never-restart"` (and `0` is rejected)
+
 ## v0.27.0
 
 ### New features

crates/pyhq/python/hyperqueue/ffi/protocol.py

@@ -1,9 +1,12 @@
 import dataclasses
 import datetime
-from typing import Dict, List, Optional, Sequence, Union
+from typing import Dict, List, Literal, Optional, Sequence, Union
 
 from ..output import StdioDef
 
+# Either a positive number of crashes, "unlimited" or "never-restart"
+CrashLimit = Union[int, Literal["unlimited", "never-restart"]]
+
 
 class ResourceRequest:
     n_nodes: int = 0
@@ -42,7 +45,7 @@ class TaskDescription:
     task_dir: bool
     priority: int
     resource_request: Sequence[ResourceRequest]
-    crash_limit: Optional[int]
+    crash_limit: Optional[CrashLimit]
 
 
 @dataclasses.dataclass

crates/pyhq/python/hyperqueue/job.py

@@ -4,7 +4,7 @@
 
 from .common import GenericPath
 from .ffi import JobId, TaskId
-from .ffi.protocol import JobDescription, ResourceRequest
+from .ffi.protocol import CrashLimit, JobDescription, ResourceRequest
 from .output import Stdio, default_stderr, default_stdout
 from .task.function import PythonFunction
 from .task.program import ExternalProgram, ProgramArgs
@@ -57,7 +57,7 @@ def program(
         task_dir: bool = False,
         priority: int = 0,
         resources: Optional[Union[ResourceRequest, Sequence[ResourceRequest]]] = None,
-        crash_limit: Optional[int] = None,
+        crash_limit: Optional[CrashLimit] = None,
     ) -> ExternalProgram:
         """
         Creates a new task that will execute the provided command.
@@ -77,7 +77,8 @@ def program(
         :param priority: Priority of the created task.
         :param resources: List of resource requests required by this task.
         :param crash_limit: How many times can a worker that executes this task crash before the
-        task is considered to be failed.
+        task is considered to be failed. Either a positive number,
+        ``"unlimited"`` or ``"never-restart"``.
         """
         task = ExternalProgram(
             len(self.tasks),
@@ -111,7 +112,7 @@ def function(
         name: Optional[str] = None,
         priority: int = 0,
         resources: Optional[Union[ResourceRequest, Sequence[ResourceRequest]]] = None,
-        crash_limit: Optional[int] = None,
+        crash_limit: Optional[CrashLimit] = None,
     ) -> PythonFunction:
         """
         Creates a new task that will execute the provided Python function.
@@ -129,7 +130,8 @@ def function(
         :param priority: Priority of the created task.
         :param resources: List of resource requests required by this task.
         :param crash_limit: How many times can a worker that executes this task crash before the
-        task is considered to be failed.
+        task is considered to be failed. Either a positive number,
+        ``"unlimited"`` or ``"never-restart"``.
         """
         task = PythonFunction(
             len(self.tasks),

crates/pyhq/python/hyperqueue/task/function/__init__.py

@@ -5,7 +5,7 @@
 
 from ...common import GenericPath
 from ...ffi import TaskId
-from ...ffi.protocol import ResourceRequest, TaskDescription
+from ...ffi.protocol import CrashLimit, ResourceRequest, TaskDescription
 from ..task import EnvType, Task, _make_ffi_requests
 from .wrapper import CloudWrapper
 
@@ -103,7 +103,7 @@ def __init__(
         dependencies=(),
         priority: int = 0,
         resources: Optional[Union[ResourceRequest, Sequence[ResourceRequest]]] = None,
-        crash_limit: Optional[int] = None,
+        crash_limit: Optional[CrashLimit] = None,
     ):
         name = generate_task_name(task_id, name, fn)
         super().__init__(

crates/pyhq/python/hyperqueue/task/program.py

@@ -2,7 +2,7 @@
 
 from ..common import GenericPath
 from ..ffi import TaskId
-from ..ffi.protocol import ResourceRequest, TaskDescription
+from ..ffi.protocol import CrashLimit, ResourceRequest, TaskDescription
 from ..output import Output, Stdio, gather_outputs
 from ..validation import ValidationException, validate_args
 from .task import EnvType, Task, _make_ffi_requests
@@ -30,7 +30,7 @@ def __init__(
         task_dir: bool = False,
         priority: int = 0,
         resources: Optional[ResourceRequest],
-        crash_limit: Optional[int] = None,
+        crash_limit: Optional[CrashLimit] = None,
     ):
         super().__init__(
             task_id,

crates/pyhq/python/hyperqueue/task/task.py

@@ -3,7 +3,7 @@
 
 from ..common import GenericPath
 from ..ffi import TaskId
-from ..ffi.protocol import ResourceRequest
+from ..ffi.protocol import CrashLimit, ResourceRequest
 from ..output import Stdio, StdioDef
 from ..validation import ValidationException
 
@@ -33,7 +33,7 @@ def __init__(
         stdout: Optional[Stdio] = None,
         stderr: Optional[Stdio] = None,
         name: Optional[str] = None,
-        crash_limit: Optional[int] = None,
+        crash_limit: Optional[CrashLimit] = None,
     ):
         assert dependencies is not None
         self.task_id = task_id

crates/pyhq/src/client/job.rs

@@ -1,7 +1,9 @@
 use crate::marshal::FromPy;
 use crate::utils::error::ToPyResult;
 use crate::{ClientContextPtr, FromPyObject, PyJobId, PyTaskId, borrow_mut, run_future};
-use hyperqueue::client::commands::submit::command::{DEFAULT_STDERR_PATH, DEFAULT_STDOUT_PATH};
+use hyperqueue::client::commands::submit::command::{
+    DEFAULT_STDERR_PATH, DEFAULT_STDOUT_PATH, parse_crash_limit,
+};
 use hyperqueue::client::commands::submit::resource_rq_map_to_vec;
 use hyperqueue::client::output::resolve_task_paths;
 use hyperqueue::client::resources::parse_allocation_request;
@@ -38,6 +40,12 @@ enum AllocationValue {
     String(String),
 }
 
+#[derive(Debug, FromPyObject)]
+enum CrashLimitValue {
+    Int(u16),
+    String(String),
+}
+
 #[derive(Debug, FromPyObject)]
 pub struct ResourceRequestDescription {
     n_nodes: NumOfNodes,
@@ -64,7 +72,7 @@ pub struct TaskDescription {
     task_dir: bool,
     priority: i32,
     resource_request: Vec<ResourceRequestDescription>,
-    crash_limit: Option<u16>,
+    crash_limit: Option<CrashLimitValue>,
 }
 
 #[derive(Debug, FromPyObject)]
@@ -256,10 +264,16 @@ fn build_task_desc(desc: TaskDescription, submit_dir: &Path) -> anyhow::Result<H
         }),
         priority: UserPriority::new(desc.priority),
         time_limit: None,
-        crash_limit: desc
-            .crash_limit
-            .map(CrashLimit::MaxCrashes)
-            .unwrap_or_default(),
+        crash_limit: match desc.crash_limit {
+            None => CrashLimit::default(),
+            Some(CrashLimitValue::Int(0)) => {
+                anyhow::bail!(
+                    "Crash limit cannot be 0. Use `never-restart` or `unlimited` instead."
+                );
+            }
+            Some(CrashLimitValue::Int(v)) => CrashLimit::MaxCrashes(v),
+            Some(CrashLimitValue::String(v)) => parse_crash_limit(&v)?,
+        },
     })
 }