labtasker 0.1.0__py3-none-any.whl

labtasker/__init__.py

@@ -0,0 +1,7 @@
+from labtasker.client.client_api import *
+from labtasker.client.core.exceptions import *
+from labtasker.filtering import install_traceback_filter
+
+__version__ = "0.1.0"
+
+install_traceback_filter()

labtasker/__main__.py

@@ -0,0 +1,6 @@
+"""Entrypoint for the cli app"""
+
+from labtasker.client.cli import app
+
+if __name__ == "__main__":
+    app()

labtasker/api_models.py

@@ -0,0 +1,205 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, SecretStr
+
+from labtasker.constants import Priority
+
+
+class BaseApiModel(BaseModel):
+    """
+    Base API model for all API models.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class HealthCheckResponse(BaseApiModel):
+    status: str = Field(..., pattern=r"^(healthy|unhealthy)$")
+    database: str
+
+
+class QueueCreateRequest(BaseApiModel):
+    queue_name: str = Field(
+        ..., pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
+    )
+    password: SecretStr = Field(..., min_length=1, max_length=100)
+    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+
+    def to_request_dict(self):
+        """
+        Used to form a quest, since password must be revealed
+        """
+        result = self.model_dump()
+        result.update({"password": self.password.get_secret_value()})
+        return result
+
+
+class QueueCreateResponse(BaseApiModel):
+    queue_id: str
+
+
+class QueueGetResponse(BaseApiModel):
+    queue_id: str = Field(alias="_id")
+    queue_name: str
+    created_at: datetime
+    last_modified: datetime
+    metadata: Dict[str, Any]
+
+
+class TaskSubmitRequest(BaseApiModel):
+    """Task submission request."""
+
+    task_name: Optional[str] = Field(
+        None, pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
+    )
+    args: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    cmd: Optional[Union[str, List[str]]] = None
+    heartbeat_timeout: Optional[float] = None
+    task_timeout: Optional[int] = None
+    max_retries: int = 3
+    priority: int = Priority.MEDIUM
+
+
+class TaskFetchRequest(BaseApiModel):
+    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
+    extra_filter: Optional[Dict[str, Any]] = None
+
+
+class Task(BaseApiModel):
+    task_id: str = Field(alias="_id")  # Accepts "_id" as an input field
+    queue_id: str
+    status: str
+    task_name: Optional[str]
+    created_at: datetime
+    start_time: Optional[datetime]
+    last_heartbeat: Optional[datetime]
+    last_modified: datetime
+    heartbeat_timeout: Optional[float]
+    task_timeout: Optional[int]
+    max_retries: int
+    retries: int
+    priority: int
+    metadata: Dict
+    args: Dict
+    cmd: Union[str, List[str]]
+    summary: Dict
+    worker_id: Optional[str]
+
+
+class TaskUpdateRequest(BaseApiModel):
+    """This should be consistent with Task.
+    Fields that disallow manual update are commented out.
+    """
+
+    task_id: str = Field(alias="_id")  # Accepts "_id" as an input field
+    # queue_id: str
+    status: Optional[str] = None
+    task_name: Optional[str] = Field(
+        None, pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
+    )
+    # created_at: datetime
+    # start_time: Optional[datetime]
+    # last_heartbeat: Optional[datetime]
+    # last_modified: datetime
+    heartbeat_timeout: Optional[float] = None
+    task_timeout: Optional[int] = None
+    max_retries: Optional[int] = None
+    retries: Optional[int] = None
+    priority: Optional[int] = None
+    metadata: Optional[Dict] = None
+    args: Optional[Dict] = None
+    cmd: Optional[Union[str, List[str]]] = None
+    summary: Optional[Dict] = None
+    # worker_id: Optional[str]
+
+
+class TaskFetchResponse(BaseApiModel):
+    found: bool = False
+    task: Optional[Task] = None
+
+
+class TaskLsRequest(BaseApiModel):
+    offset: int = Field(0, ge=0)
+    limit: int = Field(100, ge=0, le=1000)
+    task_id: Optional[str] = None
+    task_name: Optional[str] = None
+    extra_filter: Optional[Dict[str, Any]] = None
+
+
+class TaskLsResponse(BaseApiModel):
+    found: bool = False
+    content: List[Task] = Field(default_factory=list)
+
+
+class TaskSubmitResponse(BaseApiModel):
+    task_id: str
+
+
+class TaskStatusUpdateRequest(BaseApiModel):
+    status: str = Field(..., pattern=r"^(success|failed|cancelled)$")
+    worker_id: Optional[str] = None
+    summary: Optional[Dict[str, Any]] = Field(default_factory=dict)
+
+
+class WorkerCreateRequest(BaseApiModel):
+    worker_name: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    max_retries: Optional[int] = 3
+
+
+class WorkerCreateResponse(BaseApiModel):
+    worker_id: str
+
+
+class WorkerStatusUpdateRequest(BaseApiModel):
+    status: str = Field(..., pattern=r"^(active|suspended|failed)$")
+
+
+class WorkerLsRequest(BaseApiModel):
+    offset: int = Field(0, ge=0)
+    limit: int = Field(100, ge=0, le=1000)
+    worker_id: Optional[str] = None
+    worker_name: Optional[str] = None
+    extra_filter: Optional[Dict[str, Any]] = None
+
+
+class Worker(BaseApiModel):
+    worker_id: str = Field(alias="_id")
+    queue_id: str
+    status: str
+    worker_name: Optional[str] = Field(
+        None, pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
+    )
+    metadata: Dict
+    retries: int
+    max_retries: int
+    created_at: datetime
+    last_modified: datetime
+
+
+class WorkerLsResponse(BaseApiModel):
+    found: bool = False
+    content: List[Worker] = Field(default_factory=list)
+
+
+class QueueUpdateRequest(BaseApiModel):
+    new_queue_name: Optional[str] = Field(
+        None, pattern=r"^[a-zA-Z0-9_-]+$", min_length=1, max_length=100
+    )
+    new_password: Optional[SecretStr] = Field(None, min_length=1, max_length=100)
+    metadata_update: Optional[Dict[str, Any]] = Field(default_factory=dict)
+
+    def to_request_dict(self):
+        """
+        Used to form a quest, since password must be revealed
+        """
+        result = self.model_dump()
+        if self.new_password:
+            result.update({"new_password": self.new_password.get_secret_value()})
+        return result

labtasker/client/cli/__init__.py

@@ -0,0 +1,26 @@
+# 1 level commands
+import labtasker.client.cli.config
+import labtasker.client.cli.loop
+
+# multi level commands
+import labtasker.client.cli.queue as queue
+import labtasker.client.cli.task as task
+import labtasker.client.cli.worker as worker
+from labtasker.client.cli.cli import app
+from labtasker.client.core.config import get_client_config
+from labtasker.client.core.logging import stderr_console
+from labtasker.client.core.paths import get_labtasker_client_config_path
+from labtasker.client.core.plugin_utils import load_plugins
+
+app.add_typer(queue.app, name="queue", help=queue.__doc__)
+app.add_typer(task.app, name="task", help=task.__doc__)
+app.add_typer(worker.app, name="worker", help=task.__doc__)
+
+if get_labtasker_client_config_path().exists():
+    load_plugins(group="labtasker.client.cli", config=get_client_config().cli_plugins)
+else:
+    stderr_console.print(
+        f"[bold orange1]Warning:[/bold orange1] config file not found at {get_labtasker_client_config_path()}. "
+        f"[orange1]Skipped plugin loading.[/orange1] "
+        f"It is recommended to initialize config via running [orange1]`labtasker config`[/orange1] first."
+    )

labtasker/client/cli/cli.py

@@ -0,0 +1,43 @@
+"""
+Implements top level cli (mainly callbacks and setup)
+"""
+
+from typing import Optional
+
+import httpx
+import typer
+from typing_extensions import Annotated
+
+from labtasker import __version__
+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
+
+app = typer.Typer(pretty_exceptions_show_locals=False)
+
+
+def version_callback(value: bool):
+    if value:
+        stdout_console.print(f"Labtasker Version: {__version__}")
+        raise typer.Exit()
+
+
+@app.callback()
+def callback(
+    version: Annotated[
+        Optional[bool],
+        typer.Option(
+            ..., "--version", callback=version_callback, help="Print Labtasker version."
+        ),
+    ] = None,
+): ...
+
+
+@app.command()
+@requires_client_config
+def health():
+    """Check server connection and server health."""
+    try:
+        stdout_console.print(health_check())
+    except (httpx.ConnectError, httpx.HTTPStatusError) as e:
+        stderr_console.print(e)

labtasker/client/cli/config.py

@@ -0,0 +1,80 @@
+"""
+Implements `labtasker config`
+"""
+
+import tempfile
+from typing import IO, Optional
+
+import click
+import pydantic
+import tomlkit
+import tomlkit.exceptions
+import typer
+
+from labtasker.client.cli.cli import app
+from labtasker.client.core.config import ClientConfig, init_labtasker_root
+from labtasker.client.core.logging import stderr_console, stdout_console
+from labtasker.client.core.paths import (
+    get_labtasker_client_config_path,
+    get_labtasker_root,
+)
+
+
+@app.command()
+def config(
+    editor: Optional[str] = typer.Option(
+        None,
+        help="Editor to use.",
+    ),
+):
+    """Configure local client. Run `labtasker config` which opens the configuration file using system configured editor"""
+    # 0. Check if labtasker root exists, if not, init
+    if not get_labtasker_root().exists():
+        typer.confirm(
+            "Labtasker root directory not found. Initializing with default template?",
+            abort=True,
+        )
+        init_labtasker_root()
+
+    # 1. Open editor and edit configuration in a temp file
+    with tempfile.NamedTemporaryFile(
+        "w+b",
+        prefix="labtasker.tmp.",
+        suffix=".toml",
+    ) as f:  # type: IO[bytes]
+
+        # 1.1 Copy existing config if exists
+        if get_labtasker_client_config_path().exists():
+            with open(get_labtasker_client_config_path(), "rb") as f_existing:
+                f.write(f_existing.read())
+
+        # 1.2 Edit
+        while True:
+            try:
+                # a. Edit
+                f.seek(0)
+                click.edit(filename=f.name, editor=editor)
+
+                # b. Reload and validate
+                f.seek(0)
+                ClientConfig.model_validate(tomlkit.load(f))
+
+                f.seek(0)
+                updated_content = f.read()
+
+                break
+            except (
+                tomlkit.exceptions.ParseError,
+                pydantic.ValidationError,
+            ) as e:
+                stderr_console.print(
+                    "[bold red]Error:[/bold red] error when parsing config.\n"
+                    f"Detail: {str(e)}"
+                )
+                typer.confirm("Continue to edit?", abort=True)
+
+    # 2. Save
+    with open(get_labtasker_client_config_path(), "wb") as f_existing:
+        f_existing.write(updated_content)
+
+    stdout_console.print("[bold green]Configuration updated successfully.[/bold green]")

labtasker/client/cli/loop.py

@@ -0,0 +1,131 @@
+"""
+Implements `labtasker loop xxx`
+"""
+
+import shlex
+import subprocess
+from collections import defaultdict
+from typing import Optional
+
+import typer
+
+import labtasker
+import labtasker.client.core.context
+from labtasker.client.cli.cli import app
+from labtasker.client.core.cli_utils import (
+    cli_utils_decorator,
+    eta_max_validation,
+    parse_metadata,
+)
+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.logging import logger, stderr_console, stdout_console
+from labtasker.utils import keys_to_query_dict
+
+
+class InfiniteDefaultDict(defaultdict):
+
+    def __getitem__(self, key):
+        if key not in self:
+            self[key] = InfiniteDefaultDict()
+        return super().__getitem__(key)
+
+    def get(self, key, default=None):
+        if key not in self:
+            self[key] = InfiniteDefaultDict()
+        return super().get(key, default)
+
+
+@app.command()
+@cli_utils_decorator
+def loop(
+    cmd: str = typer.Option(
+        ...,
+        "--cmd",
+        "-c",
+        help="Command to run. Support argument auto interpolation, formatted like %(arg1).",
+    ),
+    extra_filter: Optional[str] = typer.Option(
+        None,
+        help='Optional mongodb filter as a dict string (e.g., \'{"key": "value"}\').',
+    ),
+    worker_id: Optional[str] = typer.Option(
+        None,
+        help="Worker ID to run the command under.",
+    ),
+    eta_max: Optional[str] = typer.Option(
+        None,
+        callback=eta_max_validation,
+        help="Maximum ETA for the task. (e.g. '1h', '1h30m', '50s')",
+    ),
+    heartbeat_timeout: Optional[float] = typer.Option(
+        None,
+        help="Heartbeat timeout for the task in seconds.",
+    ),
+):
+    """Run the wrapped job command in loop.
+    Job command follows a template string syntax: e.g. `python main.py --arg1 %(arg1) --arg2 %(arg2)`.
+    The argument inside %(...) will be autofilled by the task args fetched from task queue.
+    """
+    extra_filter = parse_metadata(extra_filter)
+
+    if heartbeat_timeout is None:
+        heartbeat_timeout = get_client_config().task.heartbeat_interval * 3
+
+    # Generate required fields dict
+    dummy_variable_table = InfiniteDefaultDict()
+    try:
+        _, queried_keys = cmd_interpolate(cmd, dummy_variable_table)
+    except (CmdParserError, KeyError, TypeError) as e:
+        raise typer.BadParameter(f"Command error with exception {e}")
+
+    required_fields = keys_to_query_dict(list(queried_keys))
+
+    logger.info(f"Got command: {cmd}")
+
+    @loop_run(
+        required_fields=required_fields,
+        extra_filter=extra_filter,
+        worker_id=worker_id,
+        eta_max=eta_max,
+        heartbeat_timeout=heartbeat_timeout,
+        pass_args_dict=True,
+    )
+    def run_cmd(args):
+        # Interpolate command
+        interpolated_cmd, _ = cmd_interpolate(cmd, args)
+        logger.info(f"Prepared to run interpolated command: {interpolated_cmd}")
+
+        with subprocess.Popen(
+            shlex.split(interpolated_cmd),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        ) as process:
+            while True:
+                output = process.stdout.readline()
+                error = process.stderr.readline()
+
+                if output:
+                    stdout_console.print(output.strip())
+                if error:
+                    stderr_console.print(error.strip())
+
+                # Break loop when process completes and streams are empty
+                if process.poll() is not None and not output and not error:
+                    break
+
+            process.wait()
+            if process.returncode != 0:
+                finish("failed")
+            else:
+                finish("success")
+
+        logger.info(f"Task {labtasker.client.core.context.task_info().task_id} ended.")
+
+    run_cmd()
+
+    logger.info("Loop finished.")

labtasker/client/cli/queue.py

@@ -0,0 +1,148 @@
+"""
+Task queue related CRUD operations.
+"""
+
+from typing import Optional
+
+import typer
+from typing_extensions import Annotated
+
+from labtasker.client.core.api import (
+    create_queue,
+    delete_queue,
+    get_queue,
+    update_queue,
+)
+from labtasker.client.core.cli_utils import cli_utils_decorator, parse_metadata
+from labtasker.client.core.config import get_client_config
+from labtasker.client.core.logging import stdout_console
+
+app = typer.Typer()
+
+
+@app.command()
+@cli_utils_decorator
+def create(
+    queue_name: Annotated[
+        str,
+        typer.Option(
+            prompt=True,
+            envvar="QUEUE_NAME",
+            help="Queue name for current experiment.",
+        ),
+    ],
+    password: Annotated[
+        str,
+        typer.Option(
+            prompt=True,
+            confirmation_prompt=True,
+            hide_input=True,
+            envvar="PASSWORD",
+            help="Password for current queue.",
+        ),
+    ],
+    metadata: Optional[str] = typer.Option(
+        None,
+        help='Optional metadata as a python dict string (e.g., \'{"key": "value"}\').',
+    ),
+):
+    """
+    Create a queue.
+    """
+    metadata = parse_metadata(metadata)
+    stdout_console.print(
+        create_queue(
+            queue_name=queue_name,
+            password=password,
+            metadata=metadata,
+        )
+    )
+
+
+@app.command()
+@cli_utils_decorator
+def create_from_config(
+    metadata: Optional[str] = typer.Option(
+        None,
+        help='Optional metadata as a python dict string (e.g., \'{"key": "value"}\').',
+    )
+):
+    """
+    Create a queue from config in `.labtasker/client.toml`.
+    """
+    metadata = parse_metadata(metadata)
+    config = get_client_config()
+    stdout_console.print(
+        create_queue(
+            queue_name=config.queue.queue_name,
+            password=config.queue.password.get_secret_value(),
+            metadata=metadata,
+        )
+    )
+
+
+@app.command()
+@cli_utils_decorator
+def get():
+    """Get current queue info."""
+    stdout_console.print(get_queue())
+
+
+@app.command()
+@cli_utils_decorator
+def update(
+    new_queue_name: Optional[str] = typer.Option(
+        None,
+        help="New name for the queue.",
+    ),
+    new_password: Optional[str] = typer.Option(
+        None,
+        prompt=True,
+        confirmation_prompt=True,
+        hide_input=True,
+        prompt_required=False,  # only trigger interactive prompt when `--new-password` is provided
+        help="New password for the queue.",
+    ),
+    metadata: Optional[str] = typer.Option(
+        None,
+        help='Optional metadata update as a python dict string (e.g., \'{"key": "value"}\').',
+    ),
+):
+    """
+    Update the current queue.
+    If you do not wish to expose password in command (e.g. `labtasker queue update --new-password my-pass --new-queue-name my-name`),
+    omit the content of `--new-password` and an interactive prompt will show up (i.e. labtasker queue update --new-password  --new-queue-name my-name).
+    """
+    # Parse metadata
+    parsed_metadata = parse_metadata(metadata)
+
+    # Proceed with the update logic
+    stdout_console.print(
+        f"Updating queue with:\n"
+        f"  New Queue Name: {new_queue_name or 'No change'}\n"
+        f"  New Password: {'******' if new_password else 'No change'}\n"
+        f"  Metadata: {parsed_metadata or 'No change'}"
+    )
+
+    updated_queue = update_queue(
+        new_queue_name=new_queue_name,
+        new_password=new_password,
+        metadata_update=parsed_metadata,
+    )
+    stdout_console.print(updated_queue)
+
+
+@app.command()
+@cli_utils_decorator
+def delete(
+    cascade: bool = typer.Option(False, help="Delete all tasks in the queue."),
+    yes: bool = typer.Option(False, "--yes", "-y", help="Confirm the operation."),
+):
+    """Delete current queue."""
+    if not yes:
+        typer.confirm(
+            f"Are you sure you want to delete current queue '{get_queue().queue_name}' with cascade={cascade}?",
+            abort=True,
+        )
+    delete_queue(cascade_delete=cascade)
+    stdout_console.print("Queue deleted.")