labtasker 0.2.9__tar.gz → 0.2.11__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labtasker
-Version: 0.2.9
+Version: 0.2.11
 Summary: A task queue system for lab experiments
 Author-email: Your Name <your.email@example.com>
 License: Apache License 2.0
@@ -32,7 +32,7 @@ Requires-Dist: pyyaml<7.0.0,>=6.0.0
 Requires-Dist: tomlkit<0.14.0,>=0.13.2
 Requires-Dist: importlib-metadata<9.0.0,>=8.5.0
 Requires-Dist: packaging<26.0,>=24.2
-Requires-Dist: sse-starlette<3.0.0,>=2.1.3
+Requires-Dist: sse-starlette<4.0.0,>=2.1.3
 Requires-Dist: httpx-sse<0.5.0,>=0.4.0
 Requires-Dist: stamina<26.0.0,>=25.1.0
 Requires-Dist: noneprompt<0.2.0,>=0.1.9
@@ -57,6 +57,7 @@ Requires-Dist: asgi-lifespan<3.0.0,>=2.1.0; extra == "dev"
 Requires-Dist: tox<4.29.0,>=4.24.0; extra == "dev"
 Requires-Dist: pytest-dependency<0.7.0,>=0.6.0; extra == "dev"
 Requires-Dist: pytest-sugar<2.0.0,>=1.0.0; extra == "dev"
+Requires-Dist: rust-just<2.0.0,>=1.42.4; extra == "dev"
 Provides-Extra: doc
 Requires-Dist: mkdocs-material<9.7.0,>=9.6.5; extra == "doc"
 Requires-Dist: mkdocs-glightbox<0.5.0,>=0.4.0; extra == "doc"

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

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

{labtasker-0.2.9 → labtasker-0.2.11}/labtasker/client/client_api.py

@@ -42,31 +42,35 @@ __all__ = [
     "connect_events",
     "EventListener",
     # http api (you should be careful with these unless you know what you are doing)
+    "get_httpx_client",
     "close_httpx_client",
     "health_check",
-    "submit_task",
-    "delete_worker",
     "create_queue",
-    "create_worker",
+    "get_queue",
     "delete_queue",
-    "delete_task",
-    "delete_worker",
+    "submit_task",
     "fetch_task",
-    "get_queue",
-    "health_check",
-    "ls_tasks",
-    "ls_workers",
-    "refresh_task_heartbeat",
     "report_task_status",
+    "refresh_task_heartbeat",
+    "create_worker",
+    "ls_workers",
+    "report_worker_status",
+    "ls_tasks",
+    "update_tasks",
+    "delete_task",
+    "update_queue",
+    "delete_worker",
     # utilities
     "run_with_pty",
     "run_with_subprocess",
 ]
 
+assert len(set(__all__)) == len(__all__), "Duplicated symbols in __all__"
+
 
 def loop(
     required_fields: Optional[List[str]] = None,
-    extra_filter: Optional[Dict[str, Any]] = None,
+    extra_filter: Optional[Union[str, Dict[str, Any]]] = None,
     cmd: Optional[Union[str, List[str]]] = None,
     worker_id: Optional[str] = None,
     create_worker_kwargs: Optional[Dict[str, Any]] = None,
@@ -78,7 +82,7 @@ def loop(
 
     Args:
         required_fields: Fields (or extra fields other than specified using Required(...)) required for task execution in a dot-separated manner. E.g. ["arg1.arg11", "arg2.arg22"]
-        extra_filter: Additional filtering criteria for tasks
+        extra_filter: Additional filtering criteria for tasks. Dict in MongoDB syntax or string in Python syntax is allowed.
         cmd: Command line arguments that runs current process. Default to sys.argv
         worker_id: Specific worker ID to use
         create_worker_kwargs: Arguments for default worker creation

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

@@ -35,6 +35,7 @@ from labtasker.client.core.utils import (
     cast_http_error,
     display_server_notifications,
     raise_for_status,
+    transpile_query_safe,
 )
 from labtasker.constants import Priority
 from labtasker.security import SecretStr, get_auth_headers
@@ -67,7 +68,13 @@ def _network_err_retry(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         return stamina.retry(
-            on=httpx.TransportError, attempts=5, wait_initial=0.5, wait_max=10.0
+            on=httpx.TransportError,
+            attempts=10,
+            timeout=100.0,
+            wait_initial=0.5,
+            wait_max=16.0,
+            wait_jitter=1.0,
+            wait_exp_base=2.0,
         )(func)(*args, **kwargs)
 
     return wrapper
@@ -193,7 +200,7 @@ def fetch_task(
     heartbeat_timeout: Optional[float] = None,
     start_heartbeat: bool = True,
     required_fields: Optional[List[str]] = None,
-    extra_filter: Optional[Dict[str, Any]] = None,
+    extra_filter: Optional[Union[str, Dict[str, Any]]] = None,
     client: Optional[httpx.Client] = None,
     cmd: Optional[Union[str, List[str]]] = None,
 ) -> TaskFetchResponse:
@@ -206,6 +213,9 @@ def fetch_task(
             "Either eta_max or start_heartbeat must be specified."
         )
 
+    if isinstance(extra_filter, str):  # transpile to mongodb query
+        extra_filter = transpile_query_safe(query_str=extra_filter)
+
     payload = TaskFetchRequest(
         worker_id=worker_id,
         eta_max=eta_max,
@@ -303,7 +313,7 @@ def ls_workers(
     worker_id: Optional[str] = None,
     worker_name: Optional[str] = None,
     status: Optional[str] = None,
-    extra_filter: Optional[Dict[str, Any]] = None,
+    extra_filter: Optional[Union[str, Dict[str, Any]]] = None,
     limit: int = 100,
     offset: int = 0,
     sort: Optional[List[Tuple[str, int]]] = None,
@@ -312,6 +322,10 @@ def ls_workers(
     """List workers."""
     if client is None:
         client = get_httpx_client()
+
+    if isinstance(extra_filter, str):  # transpile to mongodb query
+        extra_filter = transpile_query_safe(query_str=extra_filter)
+
     payload = WorkerLsRequest(
         worker_id=worker_id,
         worker_name=worker_name,
@@ -364,7 +378,7 @@ def ls_tasks(
     task_id: Optional[str] = None,
     task_name: Optional[str] = None,
     status: Optional[str] = None,
-    extra_filter: Optional[Dict[str, Any]] = None,
+    extra_filter: Optional[Union[str, Dict[str, Any]]] = None,
     limit: int = 100,
     offset: int = 0,
     sort: Optional[List[Tuple[str, int]]] = None,
@@ -373,6 +387,10 @@ def ls_tasks(
     """List tasks in a queue."""
     if client is None:
         client = get_httpx_client()
+
+    if isinstance(extra_filter, str):  # transpile to mongodb query
+        extra_filter = transpile_query_safe(query_str=extra_filter)
+
     payload = TaskLsRequest(
         task_id=task_id,
         task_name=task_name,

{labtasker-0.2.9 → labtasker-0.2.11}/labtasker/client/core/cli_utils.py

@@ -34,7 +34,7 @@ from labtasker.client.core.exceptions import (
     QueryTranspilerError,
 )
 from labtasker.client.core.logging import stderr_console
-from labtasker.client.core.query_transpiler import transpile_query
+from labtasker.client.core.utils import transpile_query_safe
 from labtasker.utils import parse_time_interval, unflatten_dict
 
 DT = TypeVar("DT")
@@ -84,7 +84,7 @@ def parse_filter(filter_str: Optional[str]) -> Optional[Dict[str, Any]]:
         return parse_dict(d_str=filter_str)
     except typer.BadParameter:
         try:
-            return transpile_query(query_str=filter_str)  # type: ignore[arg-type]
+            return transpile_query_safe(query_str=filter_str)  # type: ignore[arg-type]
         except QueryTranspilerError as e:
             raise typer.BadParameter(f"Invalid filter str: {e}") from e
 

{labtasker-0.2.9 → labtasker-0.2.11}/labtasker/client/core/events.py

@@ -176,9 +176,9 @@ class EventListener:
             self._retry_context_iter = stamina.retry_context(
                 on=httpx.TransportError,
                 attempts=10,
-                timeout=60,
+                timeout=100.0,
                 wait_initial=0.5,
-                wait_max=8.0,
+                wait_max=16.0,
                 wait_jitter=1.0,
                 wait_exp_base=2.0,
             ).__iter__()

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

@@ -12,6 +12,7 @@ import labtasker
 from labtasker.api_models import TaskUpdateRequest
 from labtasker.client.core.api import (
     create_worker,
+    delete_worker,
     fetch_task,
     get_queue,
     report_task_status,
@@ -37,6 +38,7 @@ from labtasker.client.core.exceptions import (
 from labtasker.client.core.heartbeat import end_heartbeat, start_heartbeat
 from labtasker.client.core.logging import log_to_file, logger, stderr_console
 from labtasker.client.core.paths import get_labtasker_log_dir, set_labtasker_log_dir
+from labtasker.client.core.utils import transpile_query_safe
 from labtasker.utils import parse_time_interval
 
 __all__ = [
@@ -105,7 +107,7 @@ def dump_task_info():
 
 def loop_run(
     required_fields: List[str],
-    extra_filter: Optional[Dict[str, Any]] = None,
+    extra_filter: Optional[Union[str, Dict[str, Any]]] = None,
     cmd: Optional[Union[str, List[str]]] = None,
     worker_id: Optional[str] = None,
     create_worker_kwargs: Optional[Dict[str, Any]] = None,
@@ -141,6 +143,9 @@ def loop_run(
                 f"Invalid eta_max {eta_max}. ETA max must be a valid duration string (e.g. '1h', '1h30m', '50s')"
             )
 
+    if isinstance(extra_filter, str):  # transpile to mongodb query
+        extra_filter = transpile_query_safe(query_str=extra_filter)
+
     # Check connection and authentication
     try:
         get_queue()
@@ -155,9 +160,12 @@ def loop_run(
         raise e
 
     # Create worker if not exists
+    auto_create_worker = False
+
     if current_worker_id() is None:
-        new_worker_id = worker_id or create_worker(**(create_worker_kwargs or {}))
-        set_current_worker_id(new_worker_id)
+        auto_create_worker = worker_id is None
+        worker_id = worker_id or create_worker(**(create_worker_kwargs or {}))
+        set_current_worker_id(worker_id)
 
     if cmd is None:
         cmd = sys.argv
@@ -270,12 +278,9 @@ def loop_run(
                                 resp = update_tasks(
                                     task_updates=[
                                         TaskUpdateRequest(
-                                            **{
-                                                "task_id": current_task_id(),
-                                                "status": "pending",  # running -> pending
-                                                "retries": 0,
-                                                "worker_id": None,
-                                            }
+                                            task_id=current_task_id(),  # noqa
+                                            status="pending",  # running -> pending
+                                            retries=0,
                                         )
                                     ]
                                 )
@@ -336,6 +341,10 @@ def loop_run(
                                 finish(status="success")
                             end_heartbeat()
                 except _LabtaskerLoopExit:
+                    # clean up the worker
+                    if auto_create_worker:  # worker is managed automatically
+                        delete_worker(worker_id=current_worker_id())
+
                     logger.info("Exiting task loop.")
                     break
                 except WorkerSuspended:

{labtasker-0.2.9 → labtasker-0.2.11}/labtasker/client/core/query_transpiler.py

@@ -1,6 +1,7 @@
 import ast
+import difflib
 from datetime import timezone
-from typing import Any, Dict, List, NoReturn, Type
+from typing import Any, Dict, List, NoReturn, Optional, Type
 
 import dateparser
 from rich.console import Console
@@ -154,9 +155,10 @@ class QueryTranspiler(ast.NodeVisitor):
         ast.Or: "$or",  # a or b -> {$or: [a, b]}
     }
 
-    def __init__(self, query_str: str):
+    def __init__(self, query_str: str, allowed_fields: Optional[List[str]] = None):
         super().__init__()
         self.query_str = query_str
+        self.allowed_fields = allowed_fields
 
     def _report_error(
         self, node: ast.AST, msg: str, exception: Type[QueryTranspilerError]
@@ -701,6 +703,20 @@ class QueryTranspiler(ast.NodeVisitor):
         Python: field_name
         MongoDB: "field_name" (as a field reference)
         """
+        if self.allowed_fields and node.id not in self.allowed_fields:
+            suggestions = difflib.get_close_matches(
+                node.id, self.allowed_fields, n=1, cutoff=0.6
+            )
+            suggestion_msg = (
+                f" Maybe you meant '{suggestions[0]}'?" if suggestions else ""
+            )
+
+            self._report_error(
+                node=node,
+                msg=f"Field '{node.id}' is unknown or not allowed.{suggestion_msg}"
+                f"\nAllowed fields: {', '.join(sorted(self.allowed_fields))}",
+                exception=QueryTranspilerValueError,
+            )
         return node.id
 
     def visit_Attribute(self, node: ast.Attribute) -> str:
@@ -960,7 +976,9 @@ class QueryTranspiler(ast.NodeVisitor):
         )
 
 
-def transpile_query(query_str: str) -> Dict[str, Any]:
+def transpile_query(
+    query_str: str, allowed_fields: Optional[List[str]] = None
+) -> Dict[str, Any]:
     """
     Transpile a Python-like query string and convert it to a MongoDB query object.
 
@@ -969,6 +987,7 @@ def transpile_query(query_str: str) -> Dict[str, Any]:
 
     Args:
         query_str: A string containing a Python-like expression to be converted
+        allowed_fields: A list of allowed fields (e.g. ast.Name.id). If None, allowed_fields will not be checked.
 
     Returns:
         A dictionary representing the equivalent MongoDB query
@@ -979,7 +998,7 @@ def transpile_query(query_str: str) -> Dict[str, Any]:
     try:
         query_str = query_str.strip()
         parsed_ast = ast.parse(query_str)
-        visitor = QueryTranspiler(query_str=query_str)
+        visitor = QueryTranspiler(query_str=query_str, allowed_fields=allowed_fields)
         result = visitor.visit(parsed_ast)
 
         if isinstance(result, bool) or isinstance(result, int):

{labtasker-0.2.9 → labtasker-0.2.11}/labtasker/client/core/utils.py

@@ -3,7 +3,7 @@ import os
 import subprocess
 import sys
 import threading
-from functools import wraps
+from functools import partial, wraps
 from typing import Any, Callable, Optional
 
 import httpx
@@ -20,6 +20,7 @@ from labtasker.client.core.exceptions import (
 )
 from labtasker.client.core.logging import stderr_console, stdout_console
 from labtasker.client.core.paths import get_labtasker_client_config_path
+from labtasker.client.core.query_transpiler import transpile_query
 
 server_notification_prefix = {
     "info": "[bold dodger_blue1]INFO(notification):[/bold dodger_blue1] ",
@@ -33,6 +34,31 @@ server_notification_level = {
     "high": 2,
 }
 
+transpile_query_safe = partial(
+    transpile_query,
+    allowed_fields=[
+        "task_id",
+        "queue_id",
+        "status",
+        "task_name",
+        "created_at",
+        "start_time",
+        "last_heartbeat",
+        "last_modified",
+        "heartbeat_timeout",
+        "task_timeout",
+        "max_retries",
+        "retries",
+        "priority",
+        "metadata",
+        "args",
+        "cmd",
+        "summary",
+        "worker_id",
+        "worker_name",
+    ],
+)
+
 
 def json_serializer(obj: Any, **kwargs) -> str:
     return json.dumps(to_jsonable_python(obj), **kwargs)

{labtasker-0.2.9 → labtasker-0.2.11}/labtasker/constants.py

@@ -7,5 +7,4 @@ class Priority(int, Enum):
     HIGH = 20
 
 
-KEY_PATTERN = r"^[a-zA-Z0-9_-]+$"
 DOT_SEPARATED_KEY_PATTERN = r"^[a-zA-Z0-9_-]+(\.[a-zA-Z0-9_-]+)*$"

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

@@ -614,6 +614,7 @@ class DBService:
                         "last_modified": now,
                         "worker_id": worker_id,
                         "cmd": cmd,
+                        "summary": {},  # clear previous summary before fetched
                     }
                 }
 

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "labtasker"
-version = "0.2.9"
+version = "0.2.11"
 description = "A task queue system for lab experiments"
 authors = [{ name = "Your Name", email = "your.email@example.com" }]
 license = { text = "Apache License 2.0" }
@@ -36,7 +36,7 @@ dependencies = [
     "tomlkit (>=0.13.2,<0.14.0)",
     "importlib-metadata (>=8.5.0,<9.0.0)",
     "packaging (>=24.2,<26.0)",
-    "sse-starlette (>=2.1.3,<3.0.0)",
+    "sse-starlette (>=2.1.3,<4.0.0)",
     "httpx-sse (>=0.4.0,<0.5.0)",
     "stamina (>=25.1.0,<26.0.0)",
     "noneprompt (>=0.1.9,<0.2.0)",
@@ -64,6 +64,7 @@ dev = [
     "tox (>=4.24.0,<4.29.0)",
     "pytest-dependency (>=0.6.0,<0.7.0)",
     "pytest-sugar (>=1.0.0,<2.0.0)",
+    "rust-just (>=1.42.4,<2.0.0)",
 ]
 doc = [
     "mkdocs-material (>=9.6.5,<9.7.0)",
@@ -132,7 +133,7 @@ disable_error_code = [
     "no-redef",
     "import-untyped"
 ]
-python_version = "3.10"
+python_version = "0.2.11"
 warn_unused_configs = true
 ignore_missing_imports = true
 show_error_codes = true

{labtasker-0.2.9 → labtasker-0.2.11}/tests/test_client/test_core/test_job_runner.py

@@ -90,6 +90,42 @@ def test_job_success(setup_tasks):
         assert task.status == "success"
 
 
+def test_job_str_filter(setup_tasks):
+    """Test if Pythonic query str as extra_filter works"""
+    tasks = ls_tasks()
+    assert tasks.found
+    assert len(tasks.content) == TOTAL_TASKS
+
+    idx = -1
+
+    @loop_run(
+        required_fields=["arg1", "arg2"],
+        eta_max="1h",
+        pass_args_dict=True,
+        extra_filter="args.arg2.arg3 < 2",  # only run the first two tasks
+    )
+    def job(args):
+        nonlocal idx
+        idx += 1
+        task_name = task_info().task_name
+        assert task_name == f"test_task_{idx}"
+        assert args["arg1"] == idx
+        assert args["arg2"]["arg3"] == idx
+
+        time.sleep(0.5)  # a tiny delay to ensure the tasks api request are processed
+
+        finish("success")
+
+    job()
+
+    assert idx + 1 == 2, idx
+
+    tasks = ls_tasks(extra_filter="args.arg2.arg3 < 2")
+    assert tasks.found
+    for task in tasks.content:
+        assert task.status == "success"
+
+
 def test_job_manual_failure(setup_tasks):
     cnt = 0