Skip to content

Commit 73e8f73

Browse files
committed
sidecar: classical resources should not be dispatched on fail/cancel
Signed-off-by: vsoch <vsoch@users.noreply.github.com>
1 parent e242622 commit 73e8f73

3 files changed

Lines changed: 80 additions & 14 deletions

File tree

‎python/fluence/providers/base.py‎

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,26 @@
3131
TAG_KEY = "fluence-pod-uid"
3232

3333

34+
def ungate_position(default=1):
35+
"""Queue position we ungate at. One means next in line."""
36+
try:
37+
return int(os.environ.get("FLUENCE_UNGATE_POSITION", default))
38+
except ValueError:
39+
return default
40+
41+
42+
def position_at_most(pos, threshold):
43+
"""True when the queue position is at or under threshold.
44+
45+
Braket reports anything over 2000 as the string >2000, so a position that
46+
is not a number counts as far away.
47+
"""
48+
try:
49+
return int(pos) <= int(threshold)
50+
except (TypeError, ValueError):
51+
return False
52+
53+
3454
def log(msg: str) -> None:
3555
ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
3656
print(f"[fluence] {ts} {msg}", flush=True)
@@ -79,11 +99,21 @@ def find_my_task(self, pod_uid: str, backend: str, timeout: int) -> "Task | None
7999
found or timeout. Returns an opaque Task or None."""
80100
raise NotImplementedError
81101

82-
def is_ready_to_ungate(self, task: "Task") -> bool:
83-
"""True when the gang should be ungated — queue position == 1 or the task
84-
is already RUNNING/terminal. Always implementable."""
102+
def is_ready_to_ungate(self, task: "Task", position=None) -> bool:
103+
"""True when the gang should be ungated.
104+
105+
Either the queue position is at or under position, or the task already
106+
left the queue and can still return a result. Ungating earlier gives
107+
the gang time to start up, so raise position when it is slow.
108+
"""
85109
raise NotImplementedError
86110

111+
def task_failed(self, task: "Task") -> bool:
112+
"""True when the task ended with no result, so there is nothing to
113+
fetch. Defaults to False for vendors that cannot tell.
114+
"""
115+
return False
116+
87117
def queue_position(self, task: "Task") -> "int | None":
88118
"""Optional richer telemetry: integer queue position (1 == next), or None
89119
if the vendor does not expose one. Not required for the ungate decision."""

‎python/fluence/providers/braket.py‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,15 @@
1616
import os
1717
import time
1818

19-
from fluence.providers.base import Provider, Task, TAG_KEY, log, register
19+
from fluence.providers.base import (
20+
Provider,
21+
Task,
22+
TAG_KEY,
23+
log,
24+
position_at_most,
25+
register,
26+
ungate_position,
27+
)
2028

2129

2230
class BraketTask(Task):
@@ -210,12 +218,22 @@ def _aws_task(self, task: BraketTask):
210218
from braket.aws import AwsQuantumTask
211219
return AwsQuantumTask(arn=task.arn)
212220

213-
def is_ready_to_ungate(self, task: BraketTask) -> bool:
221+
def is_ready_to_ungate(self, task: BraketTask, position=None) -> bool:
214222
t = self._aws_task(task)
215-
if t.state() in ("RUNNING", "COMPLETED", "FAILED", "CANCELLED"):
223+
# a task that left the queue reports no position, so this is how we
224+
# notice we missed the window
225+
if t.state() in ("RUNNING", "COMPLETED"):
216226
return True
227+
if position is None:
228+
position = ungate_position()
217229
try:
218-
return str(t.queue_position().queue_position) == "1"
230+
return position_at_most(t.queue_position().queue_position, position)
231+
except Exception:
232+
return False
233+
234+
def task_failed(self, task: BraketTask) -> bool:
235+
try:
236+
return self._aws_task(task).state() in ("FAILED", "CANCELLED")
219237
except Exception:
220238
return False
221239

‎python/fluence/sidecar.py‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
FLUXION_BACKEND / FLUXION_VENDOR scheduler-chosen backend / vendor
1919
FLUENCE_TASK_DISCOVERY_TIMEOUT seconds to wait for discovery (default 300)
2020
FLUENCE_POLL_INTERVAL seconds between polls (default 30)
21+
FLUENCE_UNGATE_POSITION ungate at this queue position or closer
22+
(default 1, next in line)
2123
"""
2224

2325
from __future__ import annotations
@@ -27,20 +29,24 @@
2729
import time
2830

2931
from fluence.providers import resolve_from_env
30-
from fluence.providers.base import log
32+
from fluence.providers.base import log, ungate_position
3133
from fluence.ungate import ungate_pods, gated_pods_from_env, namespace_from_env, wait_for_gated_pods
3234

3335

3436

35-
def _poll(provider, task, poll_interval, ungate):
37+
def _poll(provider, task, poll_interval, ungate, position=1):
38+
"""Poll until the task is ready or has failed. True when ready."""
3639
mode = "gang" if ungate else "observe-only"
37-
log(f"{mode} mode: polling queue position")
40+
log(f"{mode} mode: polling queue position, ungating at {position} or closer")
3841
last = object()
3942
while True:
4043
try:
41-
if provider.is_ready_to_ungate(task):
44+
if provider.task_failed(task):
45+
log("ERROR: task reached a terminal state with no result")
46+
return False
47+
if provider.is_ready_to_ungate(task, position):
4248
log(f"task ready (position={provider.queue_position(task)})")
43-
return
49+
return True
4450
pos = provider.queue_position(task)
4551
if pos != last:
4652
log(f"queue position: {pos}")
@@ -64,6 +70,7 @@ def main():
6470
discovery_timeout = int(os.environ.get("FLUENCE_TASK_DISCOVERY_TIMEOUT", 300))
6571
poll_interval = int(os.environ.get("FLUENCE_POLL_INTERVAL", 30))
6672
ungate_timeout = int(os.environ.get("FLUENCE_UNGATE_TIMEOUT", 120))
73+
ungate_at = ungate_position()
6774

6875
namespace = namespace_from_env()
6976

@@ -90,11 +97,22 @@ def main():
9097
job_id = provider.job_id(task)
9198
log(f"discovered task, job_id={job_id}")
9299

93-
_poll(provider, task, poll_interval, ungate=not observe)
100+
ready = _poll(provider, task, poll_interval, ungate=not observe,
101+
position=ungate_at)
94102

95103
if observe:
96104
log("observe-only run complete")
97-
return
105+
sys.exit(0 if ready else 1)
106+
107+
if not ready:
108+
# fail open like a discovery failure does, so the gang is not stranded.
109+
# The pods find no result and exit, and we exit non zero so it shows
110+
log("ERROR: ungating anyway so the gang is not stranded, but the task "
111+
"produced no result")
112+
ungate_pods(gated_pods_from_env() or wait_for_gated_pods(
113+
namespace, gang_group, exclude=pod_name, timeout=ungate_timeout),
114+
job_id, namespace)
115+
sys.exit(1)
98116

99117
# Ungate the gang: discover the gated pods in the gang group and remove their
100118
# gate, stamping the job-id so each can fetch results by id. The gang pods are

0 commit comments

Comments
 (0)