labtasker 0.1.2__tar.gz → 0.1.3__tar.gz

labtasker-0.1.3/.gitmodules

@@ -0,0 +1,4 @@
+[submodule "plugins/labtasker_plugin_task_count"]
+	path = plugins/labtasker_plugin_task_count
+	url = https://github.com/fkcptlst/labtasker-plugin-task-count.git
+	branch = main

{labtasker-0.1.2 → labtasker-0.1.3}/MANIFEST.in

@@ -11,6 +11,8 @@ prune .venv
 prune .tox
 prune .benchmarks
 
+prune plugins
+
 exclude Makefile
 exclude coverage.xml
 

{labtasker-0.1.2 → labtasker-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: labtasker
-Version: 0.1.2
+Version: 0.1.3
 Summary: A task queue system for lab experiments
 Author-email: Your Name <your.email@example.com>
 License: MIT
@@ -54,6 +54,8 @@ Requires-Dist: pytest-dependency<0.7.0,>=0.6.0; extra == "dev"
 Provides-Extra: doc
 Requires-Dist: mkdocs-material<9.7.0,>=9.6.5; extra == "doc"
 Requires-Dist: mike<2.2.0,>=2.1.3; extra == "doc"
+Provides-Extra: plugins
+Requires-Dist: labtasker-plugin-task-count; extra == "plugins"
 
 # Labtasker
 

labtasker-0.1.3/demo/advanced/custom_resolver/submit_job.py

@@ -0,0 +1,11 @@
+import labtasker
+
+if __name__ == "__main__":
+    for i in range(5):
+        print(f"Submitting i={i}")
+        labtasker.submit_task(
+            args={
+                "args_a": {"a": i, "b": "boy"},
+                "args_b": {"foo": 2 * i, "bar": "baz"},
+            }
+        )

labtasker-0.1.3/demo/advanced/custom_resolver/w.py

@@ -0,0 +1,38 @@
+import time
+from dataclasses import dataclass
+from typing import Any, Dict
+
+from typing_extensions import Annotated
+
+import labtasker
+from labtasker import Required
+
+
+@dataclass
+class ArgsGroupA:
+    a: int
+    b: str
+
+
+@dataclass
+class ArgsGroupB:
+    foo: int
+    bar: str
+
+
+@labtasker.loop()
+def main(
+    # use type annotation/default values to automatically resolve the required_fields
+    # use the self-defined resolver to convert the task args into custom types
+    args_a: Annotated[
+        Dict[str, Any], Required(resolver=lambda a: ArgsGroupA(**a))
+    ],  # option1. use Annotated
+    args_b=Required(resolver=lambda b: ArgsGroupB(**b)),  # option2. use default kwarg
+):
+    print(f"got args_a: {args_a}")
+    print(f"got args_b: {args_b}")
+    time.sleep(0.5)
+
+
+if __name__ == "__main__":
+    main()

labtasker-0.1.3/demo/advanced/custom_resolver/wo.py

@@ -0,0 +1,29 @@
+import time
+from dataclasses import dataclass
+
+import labtasker
+
+
+@dataclass
+class ArgsGroupA:
+    a: int
+    b: str
+
+
+@dataclass
+class ArgsGroupB:
+    foo: int
+    bar: str
+
+
+@labtasker.loop(required_fields=["args_a", "args_b"], pass_args_dict=True)
+def main(args):
+    args_a = ArgsGroupA(**args["args_a"])
+    args_b = ArgsGroupB(**args["args_b"])
+    print(f"got args_a: {args_a}")
+    print(f"got args_b: {args_b}")
+    time.sleep(0.5)
+
+
+if __name__ == "__main__":
+    main()

labtasker-0.1.3/demo/basic/python_demo/run_job.py

@@ -0,0 +1,25 @@
+import time
+
+import labtasker
+from labtasker import Required
+
+
+def job(arg1: int, arg2: int):
+    """Simulate a long-running job"""
+    time.sleep(3)  # simulate a long-running job
+    return arg1 + arg2
+
+
+@labtasker.loop()
+def main(arg1: int = Required(), arg2: int = Required()):
+    # The labtasker autofills the parameter specified by Required()
+    # or Annotated[Any, Required()]
+    # Alternatively, you can fill in the required fields
+    # in loop(required_fields=["arg1", "arg2"]
+    # and access it via labtasker.task_info().args
+    result = job(arg1, arg2)
+    print(f"The result is {result}")
+
+
+if __name__ == "__main__":
+    main()

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/__init__.py

@@ -1,12 +1,14 @@
+__version__ = "0.1.3"
+
 from labtasker.client.client_api import *
 from labtasker.client.core.config import get_client_config
 from labtasker.client.core.exceptions import *
 from labtasker.client.core.paths import get_labtasker_client_config_path
+from labtasker.client.core.version_checker import check_pypi_status
 from labtasker.filtering import install_traceback_filter, set_traceback_filter_hook
 
-__version__ = "0.1.2"
-
 install_traceback_filter()
+check_pypi_status()
 
 # by default, traceback filter is enabled.
 # you may disable it via client config

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/api_models.py

@@ -1,9 +1,12 @@
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, Field, SecretStr
+from packaging.version import Version
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
 
+from labtasker import __version__
 from labtasker.constants import Priority
+from labtasker.utils import validate_dict_keys
 
 
 class BaseApiModel(BaseModel):
@@ -14,12 +17,61 @@ class BaseApiModel(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
 
 
-class HealthCheckResponse(BaseApiModel):
+class BaseRequestModel(BaseApiModel):
+    client_version: str = __version__
+
+    @field_validator("client_version")
+    def validate_client_version(cls, v, field):
+        # make sure it is a valid version
+        Version(v)
+        return v
+
+
+class Notification(BaseModel):
+    """Server notification such as compatibility warning etc."""
+
+    type: str = Field(..., pattern=r"^(info|warning|error)$")
+    level: str = Field(..., pattern=r"^(low|medium|high)$")
+    details: Union[str, Dict[str, Any]]
+
+
+class BaseResponseModel(BaseApiModel):
+    notification: Optional[List[Notification]] = None
+
+
+class MetadataKeyValidateMixin:
+
+    @field_validator("metadata")
+    def validate_keys(cls, v, field):
+        if v:
+            validate_dict_keys(v)
+        return v
+
+
+class ArgsKeyValidateMixin:
+
+    @field_validator("args")
+    def validate_keys(cls, v, field):
+        if v:
+            validate_dict_keys(v)
+        return v
+
+
+class SummaryKeyValidateMixin:
+
+    @field_validator("summary")
+    def validate_keys(cls, v, field):
+        if v:
+            validate_dict_keys(v)
+        return v
+
+
+class HealthCheckResponse(BaseResponseModel):
     status: str = Field(..., pattern=r"^(healthy|unhealthy)$")
     database: str
 
 
-class QueueCreateRequest(BaseApiModel):
+class QueueCreateRequest(BaseRequestModel, MetadataKeyValidateMixin):
     queue_name: str = Field(
         ..., pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
     )
@@ -35,11 +87,11 @@ class QueueCreateRequest(BaseApiModel):
         return result
 
 
-class QueueCreateResponse(BaseApiModel):
+class QueueCreateResponse(BaseResponseModel):
     queue_id: str
 
 
-class QueueGetResponse(BaseApiModel):
+class QueueGetResponse(BaseResponseModel, MetadataKeyValidateMixin):
     queue_id: str = Field(alias="_id")
     queue_name: str
     created_at: datetime
@@ -47,7 +99,12 @@ class QueueGetResponse(BaseApiModel):
     metadata: Dict[str, Any]
 
 
-class TaskSubmitRequest(BaseApiModel):
+class TaskSubmitRequest(
+    BaseRequestModel,
+    ArgsKeyValidateMixin,
+    MetadataKeyValidateMixin,
+    SummaryKeyValidateMixin,
+):
     """Task submission request."""
 
     task_name: Optional[str] = Field(
@@ -62,16 +119,21 @@ class TaskSubmitRequest(BaseApiModel):
     priority: int = Priority.MEDIUM
 
 
-class TaskFetchRequest(BaseApiModel):
+class TaskFetchRequest(BaseRequestModel):
     worker_id: Optional[str] = None
     eta_max: Optional[str] = None
     heartbeat_timeout: Optional[float] = None
     start_heartbeat: bool = True
-    required_fields: Optional[Dict[str, Any]] = None
+    required_fields: Optional[List[str]] = None
     extra_filter: Optional[Dict[str, Any]] = None
 
 
-class Task(BaseApiModel):
+class Task(
+    BaseApiModel,
+    ArgsKeyValidateMixin,
+    MetadataKeyValidateMixin,
+    SummaryKeyValidateMixin,
+):
     task_id: str = Field(alias="_id")  # Accepts "_id" as an input field
     queue_id: str
     status: str
@@ -92,7 +154,12 @@ class Task(BaseApiModel):
     worker_id: Optional[str]
 
 
-class TaskUpdateRequest(BaseApiModel):
+class TaskUpdateRequest(
+    BaseRequestModel,
+    ArgsKeyValidateMixin,
+    MetadataKeyValidateMixin,
+    SummaryKeyValidateMixin,
+):
     """This should be consistent with Task.
     Fields that disallow manual update are commented out.
     """
@@ -130,12 +197,12 @@ class TaskUpdateRequest(BaseApiModel):
     # worker_id: Optional[str]
 
 
-class TaskFetchResponse(BaseApiModel):
+class TaskFetchResponse(BaseResponseModel):
     found: bool = False
     task: Optional[Task] = None
 
 
-class TaskLsRequest(BaseApiModel):
+class TaskLsRequest(BaseRequestModel):
     offset: int = Field(0, ge=0)
     limit: int = Field(100, ge=0, le=1000)
     task_id: Optional[str] = None
@@ -143,36 +210,36 @@ class TaskLsRequest(BaseApiModel):
     extra_filter: Optional[Dict[str, Any]] = None
 
 
-class TaskLsResponse(BaseApiModel):
+class TaskLsResponse(BaseResponseModel):
     found: bool = False
     content: List[Task] = Field(default_factory=list)
 
 
-class TaskSubmitResponse(BaseApiModel):
+class TaskSubmitResponse(BaseResponseModel):
     task_id: str
 
 
-class TaskStatusUpdateRequest(BaseApiModel):
+class TaskStatusUpdateRequest(BaseRequestModel):
     status: str = Field(..., pattern=r"^(success|failed|cancelled)$")
     worker_id: Optional[str] = None
     summary: Optional[Dict[str, Any]] = None
 
 
-class WorkerCreateRequest(BaseApiModel):
+class WorkerCreateRequest(BaseRequestModel, MetadataKeyValidateMixin):
     worker_name: Optional[str] = None
     metadata: Optional[Dict[str, Any]] = None
     max_retries: Optional[int] = 3
 
 
-class WorkerCreateResponse(BaseApiModel):
+class WorkerCreateResponse(BaseResponseModel):
     worker_id: str
 
 
-class WorkerStatusUpdateRequest(BaseApiModel):
+class WorkerStatusUpdateRequest(BaseRequestModel):
     status: str = Field(..., pattern=r"^(active|suspended|failed)$")
 
 
-class WorkerLsRequest(BaseApiModel):
+class WorkerLsRequest(BaseRequestModel):
     offset: int = Field(0, ge=0)
     limit: int = Field(100, ge=0, le=1000)
     worker_id: Optional[str] = None
@@ -180,7 +247,7 @@ class WorkerLsRequest(BaseApiModel):
     extra_filter: Optional[Dict[str, Any]] = None
 
 
-class Worker(BaseApiModel):
+class Worker(BaseApiModel, MetadataKeyValidateMixin):
     worker_id: str = Field(alias="_id")
     queue_id: str
     status: str
@@ -194,12 +261,12 @@ class Worker(BaseApiModel):
     last_modified: datetime
 
 
-class WorkerLsResponse(BaseApiModel):
+class WorkerLsResponse(BaseResponseModel):
     found: bool = False
     content: List[Worker] = Field(default_factory=list)
 
 
-class QueueUpdateRequest(BaseApiModel):
+class QueueUpdateRequest(BaseRequestModel):
     new_queue_name: Optional[str] = Field(
         None, pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
     )

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/client/cli/cli.py

@@ -8,7 +8,7 @@ import httpx
 import typer
 from typing_extensions import Annotated
 
-from labtasker import __version__
+from labtasker import __version__, check_pypi_status
 from labtasker.client.core.api import health_check
 from labtasker.client.core.config import requires_client_config
 from labtasker.client.core.logging import stderr_console, stdout_console
@@ -19,6 +19,7 @@ app = typer.Typer(context_settings={"help_option_names": ["-h", "--help"]})
 def version_callback(value: bool):
     if value:
         stdout_console.print(f"Labtasker Version: {__version__}")
+        check_pypi_status(force_check=True, blocking=True)
         raise typer.Exit()
 
 

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/client/cli/loop.py

@@ -22,10 +22,8 @@ from labtasker.client.core.cli_utils import (
 from labtasker.client.core.cmd_parser import cmd_interpolate
 from labtasker.client.core.config import get_client_config
 from labtasker.client.core.exceptions import CmdParserError
-from labtasker.client.core.job_runner import finish
-from labtasker.client.core.job_runner import loop as loop_run
+from labtasker.client.core.job_runner import finish, loop_run
 from labtasker.client.core.logging import logger, stderr_console, stdout_console
-from labtasker.utils import keys_to_query_dict
 
 
 class InfiniteDefaultDict(defaultdict):
@@ -59,6 +57,8 @@ def loop(
     ),
     extra_filter: Optional[str] = typer.Option(
         None,
+        "--extra-filter",
+        "-f",
         help='Optional mongodb filter as a dict string (e.g., \'{"key": "value"}\').',
     ),
     worker_id: Optional[str] = typer.Option(
@@ -102,7 +102,7 @@ def loop(
     except (CmdParserError, KeyError, TypeError) as e:
         raise typer.BadParameter(f"Command error with exception {e}")
 
-    required_fields = keys_to_query_dict(list(queried_keys))
+    required_fields = list(queried_keys)
 
     logger.info(f"Got command: {cmd}")
 

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/client/cli/queue.py

@@ -74,6 +74,8 @@ def create(
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Only show queue ID string, rather than full response. Useful when using in bash scripts.",
     ),
 ):
@@ -99,6 +101,8 @@ def create_from_config(
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Only show queue ID string, rather than full response. Useful when using in bash scripts.",
     ),
 ):
@@ -118,7 +122,9 @@ def create_from_config(
 @app.command()
 @cli_utils_decorator
 def get(
-    quiet: bool = typer.Option(False, help="Only show queue ID string."),
+    quiet: bool = typer.Option(
+        False, "--quiet", "-q", help="Only show queue ID string."
+    ),
 ):
     """Get current queue info."""
     resp = get_queue()
@@ -148,6 +154,8 @@ def update(
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Suppress the output. Execution result is only available via status code.",
     ),
 ):

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/client/cli/task.py

@@ -274,6 +274,8 @@ def ls(
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Only show task IDs that match the query, rather than full entry. "
         "Useful when using in bash scripts.",
     ),
@@ -380,6 +382,8 @@ def update(
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Disable interactive mode and confirmations. Set this to true if you are using this in a bash script.",
     ),
     editor: Optional[str] = typer.Option(

{labtasker-0.1.2 → labtasker-0.1.3}/labtasker/client/cli/worker.py

@@ -59,6 +59,8 @@ def create(
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Only show worker ID string, rather than full response. Useful when using in bash scripts.",
     ),
 ):
@@ -95,10 +97,14 @@ def ls(
     ),
     extra_filter: Optional[str] = typer.Option(
         None,
+        "--extra-filter",
+        "-f",
         help='Optional mongodb filter as a dict string (e.g., \'{"key": "value"}\').',
     ),
     quiet: bool = typer.Option(
         False,
+        "--quiet",
+        "-q",
         help="Only show worker IDs that match the query, rather than full entry. Useful when using in bash scripts.",
     ),
     pager: bool = typer.Option(

labtasker-0.1.3/labtasker/client/client_api.py

@@ -0,0 +1,118 @@
+from typing import Any, Dict, List, Optional, Union
+
+from labtasker.client.core.api import *
+from labtasker.client.core.context import current_task_id, current_worker_id, task_info
+from labtasker.client.core.exceptions import LabtaskerTypeError, LabtaskerValueError
+from labtasker.client.core.job_runner import finish, loop_run
+from labtasker.client.core.resolver import (
+    Required,
+    get_params_from_function,
+    get_required_fields,
+    resolve_args_partial,
+)
+from labtasker.utils import validate_required_fields
+
+__all__ = [
+    # python job runner api
+    "loop",
+    "finish",
+    "Required",
+    # context api
+    "task_info",
+    "current_task_id",
+    "current_worker_id",
+    # http api (you should be careful with these unless you know what you are doing)
+    "close_httpx_client",
+    "health_check",
+    "submit_task",
+    "delete_worker",
+    "create_queue",
+    "create_worker",
+    "delete_queue",
+    "delete_task",
+    "delete_worker",
+    "fetch_task",
+    "get_queue",
+    "health_check",
+    "ls_tasks",
+    "ls_worker",
+    "refresh_task_heartbeat",
+    "report_task_status",
+]
+
+
+def loop(
+    required_fields: List[str] = None,
+    extra_filter: Optional[Dict[str, Any]] = None,
+    cmd: Optional[Union[str, List[str]]] = None,
+    worker_id: Optional[str] = None,
+    create_worker_kwargs: Optional[Dict[str, Any]] = None,
+    eta_max: Optional[str] = None,
+    heartbeat_timeout: Optional[float] = None,
+    pass_args_dict: bool = False,
+):
+    """Run the wrapped job function in 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
+        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
+        eta_max: Maximum ETA for task execution.
+        heartbeat_timeout: Heartbeat timeout in seconds. Default to 3 times the send interval.
+        pass_args_dict: If True, passes task_info().args as first argument
+
+    Returns:
+        The decorated function
+
+    """
+    try:
+        if required_fields is not None:
+            validate_required_fields(required_fields)
+    except ValueError as e:
+        raise LabtaskerValueError(str(e)) from e
+    except TypeError as e:
+        raise LabtaskerTypeError(str(e)) from e
+
+    def decorator(func):
+        """
+        Steps:
+            1. Try get required fields from type annotations.
+            2. Wrap the job function with an args resolver wrapper
+            3. Wrap the resolver-wrapped job function with loop_run
+            4. Return the wrapped function
+
+        Args:
+            func:
+
+        Returns:
+
+        """
+        param_metas = get_params_from_function(func)
+
+        # if required_fields is provided, merge them with the ones specified in param_metas
+        all_required_fields = get_required_fields(
+            param_metas=param_metas, extra_required_fields=required_fields
+        )
+
+        # wrap the job function with args resolver
+        # that is, fill in the positional and keyword arguments for the required fields, type cast them
+        # into custom types (e.g. dataclasses) if the type_caster is provided
+        func = resolve_args_partial(
+            func, param_metas=param_metas, pass_args_dict=pass_args_dict
+        )
+
+        # return the decorated function
+        return loop_run(
+            required_fields=all_required_fields,
+            extra_filter=extra_filter,
+            cmd=cmd,
+            worker_id=worker_id,
+            create_worker_kwargs=create_worker_kwargs,
+            eta_max=eta_max,
+            heartbeat_timeout=heartbeat_timeout,
+            pass_args_dict=True,
+        )(func)
+
+    return decorator