labtasker 0.2.7__tar.gz → 0.2.8__tar.gz

{labtasker-0.2.7 → labtasker-0.2.8}/PKG-INFO

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

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/__init__.py

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

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/api_models.py

@@ -126,6 +126,7 @@ class TaskFetchRequest(BaseRequestModel):
     start_heartbeat: bool = True
     required_fields: Optional[List[str]] = None
     extra_filter: Optional[Dict[str, Any]] = None
+    cmd: Optional[Union[str, List[str]]] = None
 
 
 class Task(
@@ -149,7 +150,7 @@ class Task(
     priority: int
     metadata: Dict
     args: Dict
-    cmd: Union[str, List[str]]
+    cmd: Optional[Union[str, List[str]]]
     summary: Dict
     worker_id: Optional[str]
 

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/client/client_api.py

@@ -1,7 +1,12 @@
 from typing import Any, Dict, List, Optional, Union
 
 from labtasker.client.core.api import *  # noqa: F403
-from labtasker.client.core.context import current_task_id, current_worker_id, task_info
+from labtasker.client.core.context import (
+    current_task_id,
+    current_worker_id,
+    is_enabled,
+    task_info,
+)
 from labtasker.client.core.events import EventListener, connect_events
 from labtasker.client.core.exceptions import LabtaskerTypeError, LabtaskerValueError
 from labtasker.client.core.job_runner import (
@@ -10,6 +15,7 @@ from labtasker.client.core.job_runner import (
     set_loop_internal_error_handler,
     set_prompt_on_task_failure,
 )
+from labtasker.client.core.paths import get_labtasker_log_dir
 from labtasker.client.core.resolver import (
     Required,
     get_params_from_function,
@@ -29,6 +35,8 @@ __all__ = [
     "task_info",
     "current_task_id",
     "current_worker_id",
+    "is_enabled",
+    "get_labtasker_log_dir",
     # event api
     "connect_events",
     "EventListener",

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/client/core/api.py

@@ -195,6 +195,7 @@ def fetch_task(
     required_fields: Optional[List[str]] = None,
     extra_filter: Optional[Dict[str, Any]] = None,
     client: Optional[httpx.Client] = None,
+    cmd: Optional[Union[str, List[str]]] = None,
 ) -> TaskFetchResponse:
     """Fetch the next available task from the queue."""
     if client is None:
@@ -212,6 +213,7 @@ def fetch_task(
         start_heartbeat=start_heartbeat,
         required_fields=required_fields,
         extra_filter=extra_filter,
+        cmd=cmd,
     ).model_dump()
     response = client.post("/api/v1/queues/me/tasks/next", json=payload)
     if response.status_code == HTTP_403_FORBIDDEN:

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/client/core/context.py

@@ -52,3 +52,8 @@ def set_current_task_id(task_id: str):
 def set_current_worker_id(worker_id: Optional[str]):
     os.environ["LABTASKER_WORKER_ID"] = worker_id if worker_id else ""
     _current_worker_id.set(worker_id)
+
+
+def is_enabled() -> bool:
+    """Whether current script is executed under Labtasker context."""
+    return current_worker_id() is not None

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/client/core/job_runner.py

@@ -183,6 +183,7 @@ def loop_run(
                         start_heartbeat=True,
                         required_fields=required_fields,
                         extra_filter=extra_filter,
+                        cmd=cmd,
                     )
                     if not resp.found:  # task run complete
                         logger.info(
@@ -190,11 +191,7 @@ def loop_run(
                         )
                         break
 
-                    # update the "cmd" field of the current task, and get the updated task info
-                    task = update_tasks(
-                        [TaskUpdateRequest(task_id=resp.task.task_id, cmd=cmd)],  # noqa
-                        reset_pending=False,
-                    ).content[0]
+                    task = resp.task
 
                     logger.info(
                         f"Prepared to run task {task.task_id} with args {task.args}."

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/client/core/logging.py

@@ -18,12 +18,14 @@ stderr_console = Console(markup=True, stderr=True)
 LOGGER_FORMAT = (
     "<green>[{time:YYYY-MM-DD HH:mm:ss.SSS}]</green>"
     "[<level>{level: <8}</level>]"
+    "[<cyan>{name}</cyan>]"
     " - <level>{message}</level>"
 )
 
 DEBUG_LOGGER_FORMAT = (
     "<green>[{time:YYYY-MM-DD HH:mm:ss.SSS}]</green>"
     "[<level>{level: <8}</level>]"
+    "[<cyan>{name}</cyan>]"
     "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
 )
 

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/client/core/query_transpiler.py

@@ -726,7 +726,7 @@ class QueryTranspiler(ast.NodeVisitor):
                 operand_info = ""
                 try:
                     if isinstance(node.slice.value.operand, ast.Constant):  # type: ignore[attr-defined]
-                        operand_info = f" (value: -{node.slice.value.operand.value})"  # type: ignore[attr-defined]
+                        operand_info = f" (value: -{node.slice.value.operand.value!s})"  # type: ignore[attr-defined]
                     elif isinstance(node.slice.value.operand, ast.Name):  # type: ignore[attr-defined]
                         operand_info = f" (variable: -{node.slice.value.operand.id})"  # type: ignore[attr-defined]
                 except AttributeError:
@@ -769,7 +769,7 @@ class QueryTranspiler(ast.NodeVisitor):
             operand_info = ""
             try:
                 if isinstance(node.slice.operand, ast.Constant):
-                    operand_info = f" (value: -{node.slice.operand.value})"
+                    operand_info = f" (value: -{node.slice.operand.value})"  # type: ignore[str-bytes-safe]
                 elif isinstance(node.slice.operand, ast.Name):
                     operand_info = f" (variable: -{node.slice.operand.id})"
             except AttributeError:

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/server/database.py

@@ -521,6 +521,7 @@ class DBService:
         start_heartbeat: bool = True,
         required_fields: Optional[List[str]] = None,
         extra_filter: Optional[Dict[str, Any]] = None,
+        cmd: Optional[Union[str, List[str]]] = None,
     ) -> Optional[Mapping[str, Any]]:
         """
         Fetch next available task from queue.
@@ -538,6 +539,7 @@ class DBService:
             start_heartbeat (bool): Whether to start heartbeat.
             required_fields (list, optional): Which fields are required. If None, no constraint is put on which fields should exist in args dict.
             extra_filter (Dict[str, Any], optional): Additional filter criteria for the task.
+            cmd (Optional[Union[str, List[str]]]): The command that runs the job.
         """
         task_timeout = parse_timeout(eta_max) if eta_max else None
 
@@ -611,6 +613,7 @@ class DBService:
                         "last_heartbeat": now if start_heartbeat else None,
                         "last_modified": now,
                         "worker_id": worker_id,
+                        "cmd": cmd,
                     }
                 }
 
@@ -741,7 +744,7 @@ class DBService:
                 )
 
         for event_handle in event_handles:
-            event_handle.update_fsm_event(task, commit=True)
+            event_handle.commit()
 
         return True
 
@@ -774,7 +777,7 @@ class DBService:
                 )
 
         for event_handle in event_handles:
-            event_handle.update_fsm_event(task, commit=True)
+            event_handle.commit()
         return True
 
     def _report_task_status(

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker/server/endpoints.py

@@ -263,6 +263,7 @@ def fetch_task(
         start_heartbeat=task_request.start_heartbeat,
         required_fields=task_request.required_fields,
         extra_filter=task_request.extra_filter,
+        cmd=task_request.cmd,
     )
 
     if not task:

{labtasker-0.2.7 → labtasker-0.2.8}/labtasker.egg-info/SOURCES.txt

@@ -165,8 +165,8 @@ tests/test_client/test_core/test_server_notification_and_client_version.py
 tests/test_client/test_core/test_version_checker.py
 tests/test_client/test_core/test_event/__init__.py
 tests/test_client/test_core/test_event/test_concurrency_job_flow_event.py
-tests/test_client/test_core/test_event/test_event_integration.py
 tests/test_client/test_core/test_event/test_event_listener_basic.py
+tests/test_client/test_core/test_event/test_event_listener_entity_data.py
 tests/test_client/test_core/test_event/test_various_actions.py
 tests/test_client/test_core/test_event/utils.py
 tests/test_client/test_core/test_query_transpiler/__init__.py

{labtasker-0.2.7 → labtasker-0.2.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "labtasker"
-version = "0.2.7"
+version = "0.2.8"
 description = "A task queue system for lab experiments"
 authors = [{ name = "Your Name", email = "your.email@example.com" }]
 license = { text = "Apache License 2.0" }

labtasker-0.2.8/tests/test_client/test_core/test_event/test_event_listener_entity_data.py

@@ -0,0 +1,108 @@
+"""
+End-to-end tests for the EventListener class that checks entity data.
+"""
+
+import threading
+import time
+
+import pytest
+
+from labtasker import Required, create_queue, loop, submit_task
+from labtasker.api_models import EventResponse
+from labtasker.client.core.events import connect_events
+from tests.fixtures.logging import silence_logger
+from tests.test_client.test_core.test_event.utils import dump_events
+
+pytestmark = [
+    pytest.mark.e2e,
+    pytest.mark.usefixtures("silence_logger"),
+]
+
+
+@pytest.fixture(autouse=True)
+def setup_queue(client_config, db_fixture):
+    # relies on db_fixture so that DB is cleaned up after each test
+    return create_queue(
+        queue_name=client_config.queue.queue_name,
+        password=client_config.queue.password.get_secret_value(),
+        metadata={"tag": "test"},
+    )
+
+
+def test_event_listener_entity_data():
+    """Test the basic flow of events when tasks are submitted and processed."""
+    job_finish_event = threading.Event()
+    listener = connect_events(timeout=5)
+
+    def jobflow_thread():
+        try:
+            # Submit tasks to generate events
+            task_id = submit_task(
+                task_name="test_task",
+                args={"foo": "bar"},
+                max_retries=1,  # only 1 attempt
+            ).task_id
+
+            @loop()
+            def dummy(foo=Required()):
+                time.sleep(0.5)
+                raise RuntimeError  # crash task
+
+            dummy()  # fetch and run tasks
+        except Exception as e:
+            pytest.fail(f"Error in jobflow thread: {e}")
+        finally:
+            job_finish_event.set()
+
+    jobflow_thread = threading.Thread(target=jobflow_thread, daemon=True)
+
+    time.sleep(1)  # wait for the listener to start
+    jobflow_thread.start()
+
+    # Wait for job to finish with timeout
+    job_finish_event.wait(timeout=10)
+    # Give some time for all events to be processed
+    time.sleep(2)
+    listener.stop()
+
+    # Check that we received the expected events
+    received_events = dump_events(listener)
+
+    expected_transition_sequence = [
+        # 1 job creation events
+        ("created", "pending"),
+        # worker creation event
+        ("created", "active"),
+        # 1 running and failure events
+        ("pending", "running"),
+        # worker fail() event result in either active -> active or active -> crashed
+        ("active", "active"),
+        # task fail() event
+        ("running", "failed"),
+    ]
+
+    assert len(received_events) == len(
+        expected_transition_sequence
+    ), f"Not enough events, got {len(received_events)} / {len(expected_transition_sequence)}"
+
+    # Verify event transitions match expected sequence
+    for i, event_resp in enumerate(received_events):
+        event_resp: EventResponse
+        from_state, to_state = expected_transition_sequence[i]
+        assert (
+            event_resp.event.old_state == from_state
+        ), f"Event {i} has wrong from_state: {event_resp.event.old_state}, expected {from_state}"
+        assert (
+            event_resp.event.new_state == to_state
+        ), f"Event {i} has wrong to_state: {event_resp.event.new_state}, expected {to_state}"
+
+        print(f"Event data: {event_resp.event.entity_data}")
+
+    # assert cmd in pending -> running event
+    assert received_events[2].event.entity_data["cmd"], "Cmd is missing"
+
+    # assert summary in running -> failed event
+    assert received_events[-1].event.entity_data["summary"], "Summary is missing"
+
+    # Join threads to clean up
+    jobflow_thread.join(timeout=3)