Source code for ewoksid13.mcp.tools.submit

"""Template tool module — copy this file to add a new MCP tool.

Every tool is a plain function with type hints and a docstring: FastMCP turns
those into the tool's input schema and description. Once defined, import the
function in ``ewoksid13/mcp/tools/__init__.py`` and add it to ``TOOLS``.
"""

from __future__ import annotations

import logging
import os
from typing import Any, Dict
from ewokstools.submit import save_and_execute

logger = logging.getLogger(__name__)




def submit_test(queue: str = None) -> Dict[str, Any]:
    """Submit a trivial test job to check that job submission works.

    Use this tool whenever the user wants to *test*, *check* or *verify* the job
    submission pipeline (rather than run a real data-processing job). It submits
    a minimal throwaway ewoks workflow — a single ``SumTask`` that adds two
    integers — and returns the submitted job's details, so the user can confirm
    a worker is reachable and accepting jobs.

    The job can run either on the SLURM cluster or on a Celery worker, chosen
    from the ``queue`` argument:

    - **no queue given** (``queue`` empty/``None``): submit to the **SLURM**
      cluster.
    - **a queue name given**: submit to that **Celery** queue (worker).

    Parameters
    ----------
    queue : str, optional
        Name of the Celery queue (worker) to submit the test job to. Leave
        empty/``None`` to submit to the SLURM cluster instead.

    Returns
    -------
    dict
        The submitted job's details: ``"worker_type"`` (``"slurm"`` or
        ``"celery"``) and ``"job_id"``.
    """
    workflow = {
        "graph": {"id": "test"},
        "nodes": [
            {
                "id": "n",
                "task_type": "class",
                "task_identifier": "ewokscore.tests.examples.tasks.sumtask.SumTask",
            }
        ],
    }
    inputs = [
        {
            "name": "a",
            "value": 10,
            "id": "n",
        },
        {
            "name": "b",
            "value": 3,
            "id": "n",
        },
    ]
    if queue:
        submit_parameters = {"queue": queue}
        destination_filename = None
    else:
        submit_parameters = {}
        destination_filename = f"{os.environ['HOME']}/.ewoksid13/workflows/test.json"

    submitted = save_and_execute(
        python_package="ewoksid13",
        workflow=workflow,
        inputs=inputs,
        destination_filename=destination_filename,
        **submit_parameters,
    )
    return {
        "worker_type": submitted.worker_type,
        "job_id": submitted.job_id,
    }