From e5bc48ef80e196aa137fdf1ba981daa33f72de5c Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Tue, 20 Oct 2026 12:15:05 +0200 Subject: [PATCH] PyAPI crash limit --- CHANGELOG.md | 6 +++++ crates/pyhq/python/hyperqueue/ffi/protocol.py | 7 +++-- crates/pyhq/python/hyperqueue/job.py | 12 +++++---- .../hyperqueue/task/function/__init__.py | 4 +-- crates/pyhq/python/hyperqueue/task/program.py | 4 +-- crates/pyhq/python/hyperqueue/task/task.py | 4 +-- crates/pyhq/src/client/job.rs | 26 ++++++++++++++----- 7 files changed, 44 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a0a761ec..abf521d92 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/crates/pyhq/python/hyperqueue/ffi/protocol.py b/crates/pyhq/python/hyperqueue/ffi/protocol.py index a0b7f1b98..72ee58eae 100644 --- a/crates/pyhq/python/hyperqueue/ffi/protocol.py +++ b/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 diff --git a/crates/pyhq/python/hyperqueue/job.py b/crates/pyhq/python/hyperqueue/job.py index 4f8d832d4..68a7ac261 100644 --- a/crates/pyhq/python/hyperqueue/job.py +++ b/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), diff --git a/crates/pyhq/python/hyperqueue/task/function/__init__.py b/crates/pyhq/python/hyperqueue/task/function/__init__.py index ad89d6929..b2baf3742 100644 --- a/crates/pyhq/python/hyperqueue/task/function/__init__.py +++ b/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__( diff --git a/crates/pyhq/python/hyperqueue/task/program.py b/crates/pyhq/python/hyperqueue/task/program.py index 9819fc4eb..60d5a6ee1 100644 --- a/crates/pyhq/python/hyperqueue/task/program.py +++ b/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, diff --git a/crates/pyhq/python/hyperqueue/task/task.py b/crates/pyhq/python/hyperqueue/task/task.py index 49f391d61..38d8d1f40 100644 --- a/crates/pyhq/python/hyperqueue/task/task.py +++ b/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 diff --git a/crates/pyhq/src/client/job.rs b/crates/pyhq/src/client/job.rs index 3f9f0ee83..9468f232c 100644 --- a/crates/pyhq/src/client/job.rs +++ b/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, - crash_limit: Option, + crash_limit: Option, } #[derive(Debug, FromPyObject)] @@ -256,10 +264,16 @@ fn build_task_desc(desc: TaskDescription, submit_dir: &Path) -> anyhow::Result 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)?, + }, }) }