labtasker 0.2.11__tar.gz → 0.2.12__tar.gz

{labtasker-0.2.11 → labtasker-0.2.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labtasker
-Version: 0.2.11
+Version: 0.2.12
 Summary: A task queue system for lab experiments
 Author-email: Your Name <your.email@example.com>
 License: Apache License 2.0

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.11"
+__version__ = "0.2.12"
 
 from labtasker.client.client_api import *
 from labtasker.client.core.config import get_client_config

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/client/cli/worker.py

@@ -102,6 +102,12 @@ def ls(
         "--name",
         help="Filter by worker name.",
     ),
+    status: Optional[str] = typer.Option(
+        None,
+        "--status",
+        "-s",
+        help="Filter by worker status. One of `active`, `suspended`, `crashed`.",
+    ),
     extra_filter: Optional[str] = typer.Option(
         None,
         "--extra-filter",
@@ -164,6 +170,7 @@ def ls(
             ls_workers,
             worker_id=worker_id,
             worker_name=worker_name,
+            status=status,
             extra_filter=extra_filter,
         ),
         offset=offset,

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/client/core/api.py

@@ -64,11 +64,15 @@ __all__ = [
 ]
 
 
+def _is_network_transient_error(exception):
+    return isinstance(exception, (httpx.TransportError, ConnectionError, TimeoutError))
+
+
 def _network_err_retry(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         return stamina.retry(
-            on=httpx.TransportError,
+            on=_is_network_transient_error,
             attempts=10,
             timeout=100.0,
             wait_initial=0.5,
@@ -278,12 +282,15 @@ def report_task_status(
 @_network_err_retry
 def refresh_task_heartbeat(
     task_id: str,
+    worker_id: Optional[str] = None,
     client: Optional[httpx.Client] = None,
 ) -> None:
     """Refresh the heartbeat of a task."""
     if client is None:
         client = get_httpx_client()
-    response = client.post(f"/api/v1/queues/me/tasks/{task_id}/heartbeat")
+    response = client.post(
+        f"/api/v1/queues/me/tasks/{task_id}/heartbeat", params={"worker_id": worker_id}
+    )
     raise_for_status(response)
 
 

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/client/core/heartbeat.py

@@ -18,8 +18,9 @@ __all__ = [
 
 class Heartbeat:
 
-    def __init__(self, task_id, heartbeat_interval):
+    def __init__(self, task_id, worker_id, heartbeat_interval):
         self.task_id = task_id
+        self.worker_id = worker_id
         self.heartbeat_interval = heartbeat_interval
 
         self._thread = None
@@ -39,6 +40,7 @@ class Heartbeat:
         self._thread.start()
 
     def delay(self, interval: float) -> bool:
+        """Returns False if it should exit."""
         slice_t = 0.05  # check for stop event
         start_time = time.perf_counter()
         while True:
@@ -67,11 +69,12 @@ class Heartbeat:
         """Refresh heartbeat periodically"""
         while True:
             try:
-                refresh_task_heartbeat(task_id=self.task_id)
+                refresh_task_heartbeat(task_id=self.task_id, worker_id=self.worker_id)
             except Exception as e:
-                logger.error(f"Heartbeat failed for task {self.task_id}: {e}")
+                logger.error(f"Failed to refresh heartbeat: {str(e)}")
+                raise
 
-            # check if heartbeat should stop
+            # Check if heartbeat should stop
             if not self.delay(self.heartbeat_interval):
                 break
 
@@ -85,6 +88,9 @@ class Heartbeat:
             except FileNotFoundError:
                 pass
 
+    def is_alive(self):
+        return self._thread and self._thread.is_alive()
+
 
 _current_heartbeat: ContextVar[Optional[Heartbeat]] = ContextVar(
     "heartbeat", default=None
@@ -92,7 +98,10 @@ _current_heartbeat: ContextVar[Optional[Heartbeat]] = ContextVar(
 
 
 def start_heartbeat(
-    task_id, heartbeat_interval: Optional[float] = None, raise_error=True
+    task_id,
+    worker_id: Optional[str] = None,
+    heartbeat_interval: Optional[float] = None,
+    raise_error=True,
 ):
     logger.debug("Try starting heartbeat.")
     if _current_heartbeat.get() is not None:
@@ -102,12 +111,14 @@ def start_heartbeat(
 
     heartbeat_manager = Heartbeat(
         task_id=task_id,
+        worker_id=worker_id,
         heartbeat_interval=heartbeat_interval
         or get_client_config().task.heartbeat_interval,
     )
     heartbeat_manager.start()
     _current_heartbeat.set(heartbeat_manager)
     logger.debug("Heartbeat started.")
+    return heartbeat_manager
 
 
 def end_heartbeat(raise_error=True):

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/client/core/job_runner.py

@@ -220,7 +220,9 @@ def loop_run(
                     dump_task_info()
 
                     with log_to_file(file_path=get_labtasker_log_dir() / "run.log"):
-                        start_heartbeat(task_id=current_task_id())
+                        start_heartbeat(
+                            task_id=current_task_id(), worker_id=current_worker_id()
+                        )
                         success_flag = False
                         try:
                             func_args = (task.args, *args) if pass_args_dict else args

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/server/database.py

@@ -674,22 +674,52 @@ class DBService:
     @retry_on_transient
     @validate_arg
     def refresh_task_heartbeat(
-        self,
-        queue_id: str,
-        task_id: str,
-    ) -> bool:
+        self, queue_id: str, task_id: str, worker_id: Optional[str] = None
+    ):
         """Update task heartbeat timestamp."""
+        query = {"_id": task_id, "queue_id": queue_id, "status": "running"}
+
         with self._client.start_session() as session:
             with session.start_transaction():
-                return (
-                    self._tasks.update_one(
-                        {"_id": task_id, "queue_id": queue_id},
-                        {"$set": {"last_heartbeat": get_current_time()}},
-                        session=session,
-                    ).modified_count
-                    > 0
+                # Find the task in a single query
+                task = self._tasks.find_one(query)
+                if not task:
+                    raise HTTPException(
+                        status_code=HTTP_404_NOT_FOUND,
+                        detail=f"Task '{task_id}' not found in queue '{queue_id}' or not in 'running' state",
+                    )
+
+                # Validate worker if provided
+                if worker_id:
+                    if task["worker_id"] != worker_id:
+                        raise HTTPException(
+                            status_code=HTTP_403_FORBIDDEN,
+                            detail=f"Task '{task_id}' is assigned to worker '{task['worker_id']}', not '{worker_id}'",
+                        )
+
+                    # Check worker status in a single query
+                    worker = self._workers.find_one(
+                        {"_id": worker_id, "status": WorkerState.ACTIVE}
+                    )
+                    if not worker:
+                        raise HTTPException(
+                            status_code=HTTP_404_NOT_FOUND,
+                            detail=f"Worker '{worker_id}' not found or not active",
+                        )
+
+                # Update the task heartbeat
+                result = self._tasks.update_one(
+                    query,
+                    {"$set": {"last_heartbeat": get_current_time()}},
+                    session=session,
                 )
 
+                if result.modified_count == 0:
+                    raise HTTPException(
+                        status_code=HTTP_404_NOT_FOUND,
+                        detail=f"Failed to update heartbeat for task '{task_id}' - it may have changed state during the operation",
+                    )
+
     @retry_on_transient
     @validate_arg
     def worker_report_task_status(

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/server/endpoints.py

@@ -1,9 +1,9 @@
 import asyncio
 import uuid
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
-from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi import Depends, FastAPI, HTTPException, Query, Request
 from sse_starlette.sse import EventSourceResponse
 from starlette.status import (
     HTTP_201_CREATED,
@@ -309,16 +309,14 @@ def report_task_status(
 )
 def refresh_task_heartbeat(
     task_id: str,
+    worker_id: Optional[str] = Query(None),  # use query param
     queue: Dict[str, Any] = Depends(get_verified_queue_dependency),
     db: DBService = Depends(get_db),
 ):
     """Update task heartbeat timestamp."""
-    done = db.refresh_task_heartbeat(
-        queue_id=queue["_id"],
-        task_id=task_id,
+    db.refresh_task_heartbeat(
+        queue_id=queue["_id"], task_id=task_id, worker_id=worker_id
     )
-    if not done:
-        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Task not found.")
 
 
 @app.get(

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker/server/fsm.py

@@ -43,7 +43,12 @@ class StateTransitionEventHandle:
             self.commit()
 
     def commit(self):
-        event_data = StateTransitionEvent(
+        event_data = self._create_event_data()
+        self._publish_event(event_data)
+        self._entity_data = None
+
+    def _create_event_data(self):
+        return StateTransitionEvent(
             entity_type=self.entity_type,
             queue_id=self.queue_id,
             entity_id=self.entity_id,
@@ -54,9 +59,17 @@ class StateTransitionEventHandle:
             entity_data=self._entity_data,
         )
 
+    def _publish_event(self, event_data):
         # Use fully synchronous event publishing
         event_manager.publish_event(self.queue_id, event_data)
-        self._entity_data = None
+
+
+class NullEventHandle(StateTransitionEventHandle):
+    """A placeholder that does nothing. (Used for cases where triggering event publishing is undesired)"""
+
+    def _publish_event(self, event_data):
+        # Override to do nothing
+        pass
 
 
 class State(str, Enum):
@@ -140,6 +153,18 @@ class BaseFSM:
     def state(self):
         return self._state
 
+    def null_transition(self) -> NullEventHandle:
+        """Perform a null transition and return a handle"""
+        return NullEventHandle(
+            entity_type=self.ENTITY_TYPE,
+            entity_id=self.entity_id,
+            queue_id=self.queue_id,
+            old_state=str(self._state),
+            new_state=str(self._state),
+            transition_time=get_current_time(),
+            metadata=self.metadata,
+        )
+
     def transition_to(self, new_state: State) -> StateTransitionEventHandle:
         """Perform state transition and return a handle"""
         old_state = self._state
@@ -190,6 +215,7 @@ class TaskFSM(BaseFSM):
         TaskState.FAILED: {
             TaskState.PENDING,
             TaskState.CANCELLED,
+            TaskState.FAILED,  # null transition (for more tolerance)
         },  # Can be reset and requeued
         TaskState.CANCELLED: {
             TaskState.PENDING,
@@ -280,11 +306,15 @@ class TaskFSM(BaseFSM):
         Transitions:
         - RUNNING -> PENDING (if retries < max_retries)
         - RUNNING -> FAILED (if retries >= max_retries)
+        - FAILED -> FAILED (null transition, does nothing)
         - Others -> InvalidStateTransition (invalid)
 
         Note: FAILED state can transition back to PENDING for retries
         until max_retries is reached.
         """
+        if self.state == TaskState.FAILED:
+            return self.null_transition()
+
         if self.state != TaskState.RUNNING:
             raise InvalidStateTransition(f"Cannot fail task in {self.state} state")
 
@@ -305,7 +335,10 @@ class WorkerFSM(BaseFSM):
             WorkerState.CRASHED,
         },
         WorkerState.SUSPENDED: {WorkerState.ACTIVE},  # Manual transition
-        WorkerState.CRASHED: {WorkerState.ACTIVE},  # Manual transition
+        WorkerState.CRASHED: {
+            WorkerState.ACTIVE,  # Manual transition
+            WorkerState.CRASHED,  # null transition (for more tolerance)
+        },
     }
 
     def __init__(
@@ -366,7 +399,11 @@ class WorkerFSM(BaseFSM):
         Transitions:
         - ACTIVE -> ACTIVE
         - ACTIVE -> CRASHED (retries >= max_retries)
+        - CRASHED -> CRASHED (null transition, does nothing)
         """
+        if self.state == WorkerState.CRASHED:
+            return self.null_transition()
+
         if self.state != WorkerState.ACTIVE:
             raise InvalidStateTransition(f"Cannot fail worker in {self.state} state")
 

{labtasker-0.2.11 → labtasker-0.2.12}/labtasker.egg-info/SOURCES.txt

@@ -154,7 +154,8 @@ tests/test_client/test_cli/test_task.py
 tests/test_client/test_cli/test_worker.py
 tests/test_client/test_core/__init__.py
 tests/test_client/test_core/test_cli_utils.py
-tests/test_client/test_core/test_heartbeat.py
+tests/test_client/test_core/test_heartbeat_e2e.py
+tests/test_client/test_core/test_heartbeat_unit.py
 tests/test_client/test_core/test_job_runner.py
 tests/test_client/test_core/test_logging.py
 tests/test_client/test_core/test_loop_internal_error_handler.py

{labtasker-0.2.11 → labtasker-0.2.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "labtasker"
-version = "0.2.11"
+version = "0.2.12"
 description = "A task queue system for lab experiments"
 authors = [{ name = "Your Name", email = "your.email@example.com" }]
 license = { text = "Apache License 2.0" }
@@ -133,7 +133,6 @@ disable_error_code = [
     "no-redef",
     "import-untyped"
 ]
-python_version = "0.2.11"
 warn_unused_configs = true
 ignore_missing_imports = true
 show_error_codes = true

labtasker-0.2.12/tests/test_client/test_core/test_heartbeat_e2e.py

@@ -0,0 +1,73 @@
+import time
+
+import pytest
+
+from labtasker.client.core.api import (
+    create_queue,
+    create_worker,
+    fetch_task,
+    submit_task,
+)
+from labtasker.client.core.heartbeat import start_heartbeat
+from labtasker.client.core.paths import set_labtasker_log_dir
+from tests.fixtures.logging import silence_logger
+
+pytestmark = [
+    pytest.mark.e2e,
+    pytest.mark.usefixtures(
+        "silence_logger"
+    ),  # silence logger in testcases of this module
+]
+
+
+@pytest.fixture(autouse=True)
+def setup_queue(client_config):
+    return create_queue(
+        queue_name=client_config.queue.queue_name,
+        password=client_config.queue.password.get_secret_value(),
+        metadata={"tag": "test"},
+    )
+
+
+@pytest.fixture(autouse=True)
+def setup_workers():
+    id1 = create_worker(worker_name="worker_1")
+    id2 = create_worker(worker_name="worker_2")
+
+    return id1, id2
+
+
+@pytest.fixture(autouse=True)
+def setup_task(db_fixture):
+    # relies on db_fixture to clear db after each test case
+
+    fetch_task()
+
+
+def test_heartbeat_worker_mismatch():
+    # Test if heartbeat fails when worker_id does not match
+
+    # setup
+    # 1. create tasks
+    submit_task(task_name="test_task", args={"arg1": 0})
+    # 2. create workers
+    worker_id1 = create_worker(worker_name="worker_1")
+    worker_id2 = create_worker(worker_name="worker_2")
+    # 3. fetch task
+    task_resp = fetch_task(worker_id=worker_id1, heartbeat_timeout=99999.0)
+
+    # 4. set local dir
+    task_id = task_resp.task.task_id
+    set_labtasker_log_dir(
+        task_id=task_id,
+        task_name=task_resp.task.task_name,
+        set_env=True,
+        overwrite=True,
+    )
+
+    # use a different worker_id to start heartbeat
+    heartbeat_manager = start_heartbeat(
+        task_id=task_id, worker_id=worker_id2, heartbeat_interval=0.2
+    )
+    time.sleep(1.0)
+    assert not heartbeat_manager.is_alive()

{labtasker-0.2.11 → labtasker-0.2.12}/tests/test_server/test_fsm.py

@@ -174,7 +174,7 @@ class TestWorkerFSM:
 
     def test_fail_from_invalid_state(self, worker_db_entry):
         """Test failing worker from invalid state."""
-        invalid_states = [WorkerState.SUSPENDED, WorkerState.CRASHED]
+        invalid_states = [WorkerState.SUSPENDED]
         for state in invalid_states:
             worker_db_entry["status"] = state
             fsm = WorkerFSM.from_db_entry(worker_db_entry)