Writing scripts for the OET

The Observation Execution Tool (OET) can run observing scripts in a headless non-interactive manner. For efficiency, OET script execution is split into two phases: an initialisation phase and an execution phase. Scripts that are expected to be run by the OET should be structured to have two entry points corresponding to these two phases, as the template below:

Observing script template

def init(subarray: int, *args, **kwargs):
    # Called by the OET when the script is loaded and initialised by someone
    # calling 'oet prepare'. Add your script initialisation code here. Note that
    # the target subarray is supplied to this function as the first argument.
    pass

def main(*args, **kwargs):
    # Called by the OET when the prepared script is told to run by someone
    # calling 'oet start'. Add the main body of your script to this function.
    pass

The initialisation phase occurs when the script is loaded and the script’s init function is called (if defined) to perform any preparation and/or initialisation. Expensive and slow operations that can be performed ahead of the main body of script execution can be run in the initialisation phase. Typical actions performed in init are I/O intensive operations, e.g., cloning a git repository, creating multiple Tango device proxies, subscribing to Tango events, etc. When run by the Observation Execution Tool (OET), the init function is passed an integer subarray ID declaring which subarray the control script is intended to control.

Subsequently, at some point a user may call oet start, requesting that the initialised script begin the main body of its execution. When this occurs, the OET calls the script’s main function, which should performs the main function of the script. For an observing script, this would involve the configuration and control of a subarray.

below is the real example script in the scripts folder of this project.

SKA : Allocate Resources and Perform Observation

Allocating resources and performing scans requires communication with TMC CentralNode and TMC SubarrayNode, and targets a specific subarray. This script’s init function pre-applies the subarray ID argument to the main function. Note that this script does not perform any Tango calls directly, but uses ska_oso_scripting.api functions to perform all the required Tango interactions (command invocation; event subscriptions; event monitoring).

Resource allocation and perform observation script for an SKA MID/LOW subarray

"""
Standard observing script for SBDefinition execution.

Creates an AssignResources command from the SBDefinition and sends to TMC CentralNode, then
creates ConfigureRequests and Scan commands to send to TMC SubarrayNode.

At the end of the observation, all resources are released.

This script and the transform functions to telescope commands should support
all functionality currently available in the ODT (e.g. PST, 5 point scans, etc)
"""
import functools
import logging
import os
from typing import Any

from ska_oso_pdm import SBDefinition, TelescopeType

from ska_oso_scripting import api
from ska_oso_scripting.api import execution_block
from ska_oso_scripting.api.objects import SubArray
from ska_oso_scripting.core.execution import ValueTransitionError
from ska_oso_scripting.pdm_transforms import (
    create_cdm_assign_resources_request_from_scheduling_block,
    create_cdm_configure_requests_from_scheduling_block,
)

from ska_oso_scripting.topics import user_topics
from ska_oso_scripting.engineering.low import workarounds as LOW_WORKAROUNDS
from ska_oso_scripting.engineering.scripting_workarounds import TEMP_WORKAROUNDS

LOG = logging.getLogger(__name__)
FORMAT = "%(asctime)-15s %(message)s"

logging.basicConfig(level=logging.INFO, format=FORMAT)


def init(subarray_id: int | str, **init_args):
    """
    Initialise the script, binding the script runtime args to the script.
    """
    try:
        subarray_id = int(subarray_id)
    except ValueError as err:
        raise TypeError("subarray_id must be an integer") from err

    LOG.debug(f"Initializing script {__name__} with subarray_id={subarray_id}")

    for arg,value in init_args.items():
        LOG.debug(f"Initializing script {__name__} with argument(s) {arg}={value}")

    global main

    main = functools.partial(_main, subarray_id=subarray_id)
    LOG.info(f"Script bound to sub-array {subarray_id}")


def assign_resources(subarray: SubArray, sbi: SBDefinition, apply_low_workarounds: bool =False):
    """
    assign resources to a target sub-array using a Scheduling Block (SB).
    :param subarray: subarray ID
    :param sbi: ska_oso_pdm.SBDefinition
    :return:
    """
    LOG.info(
        f"Running assign_resources(subarray={subarray.id}, sbi.sbd_id={sbi.sbd_id})"
    )

    cdm_allocation = create_cdm_assign_resources_request_from_scheduling_block(
            subarray.id, sbi
        )

    if apply_low_workarounds:
        cdm_allocation = LOW_WORKAROUNDS.update_assign_resources_request(cdm_allocation)
        LOW_WORKAROUNDS.pre_obs_checks(
            cdm_allocation.model_dump(mode="json", exclude_none=True, by_alias=True)
        )

    response = api.assign_resources_from_cdm(subarray.id, cdm_allocation)
    LOG.info(f"Resources Allocated: {response}")

    LOG.info("Allocation complete")


def observe(*, subarray: SubArray, sbi: SBDefinition, runtime_args: dict[str,Any], apply_low_workarounds:bool =False):
    """
    Observe using a Scheduling Block (SB) and template CDM file.

    :param subarray:  SubArray instance containing subarray ID
    :param sbi: Instance of a SBDefinition
    :param runtime_args: Any script runtime arguments to be considered.
    :return:
    """

    LOG.info(
        f"Starting observing for Scheduling Block: {sbi.sbd_id}, subarray_id={subarray.id})"
    )

    if runtime_args:
        for arg,value in runtime_args.items():
            LOG.info(f"runtime argument(s) {arg}={value}")

    scan_sequence = sbi.dish_allocations.scan_sequence if sbi.telescope == TelescopeType.SKA_MID else sbi.mccs_allocation.subarray_beams[0].scan_sequence

    if not scan_sequence:
        LOG.info(f"No scans defined in Scheduling Block {sbi.sbd_id}. No observation performed.")
        return

    cdm_configure_requests = (
        create_cdm_configure_requests_from_scheduling_block(sbi,runtime_args)
    )

    if apply_low_workarounds:
        LOW_WORKAROUNDS.update_configure_requests(cdm_configure_requests)

    for scan_definition_idx in range(len(scan_sequence)):
        cdm_configs = cdm_configure_requests[scan_definition_idx]
        for index, cdm_config in enumerate(cdm_configs):
            scan_id_string = (f"{scan_definition_idx} "
                              f"({str(index + 1)}/{str(len(cdm_configs))})" if len(cdm_configs) > 1 else '')
            try:
                # With the CDM modified, we can now issue the Configure instruction...
                LOG.info(f"Configuring subarray {subarray.id} for scan: {scan_id_string}")
                api.send_message(
                    user_topics.script.announce,
                    msg=f"Configuring subarray {subarray.id} for scan: {scan_id_string}"
                )
                api.configure_from_cdm(subarray.id, cdm_config)
            except ValueTransitionError as err:
                LOG.error(f"Error configuring subarray: {err}")
                api.send_message(
                    user_topics.script.announce,
                    msg=f"Error configuring subarray for scan {scan_id_string}"
                )
                raise err
            else:
                LOG.info(f"Configuration for scan {scan_id_string} complete")
                api.send_message(
                    user_topics.script.announce,
                    msg=f"Configuration for scan {scan_id_string} complete"
                )
            try:
                # with configuration complete, we can begin the scan.
                LOG.info(f"Starting scan: {scan_id_string}")
                api.send_message(
                    user_topics.script.announce,
                    msg=f"Starting scan: {scan_id_string}"
                )
                api.scan(subarray.id)
            except ValueTransitionError as err:
                LOG.error(f"Error when executing scan: {scan_id_string}: {err}")
                api.send_message(
                    user_topics.script.announce,
                    msg=f"Error when executing scan: {scan_id_string}"
                )
                raise err
            else:
                LOG.info(f"Scan {scan_id_string} complete")
                api.send_message(
                    user_topics.script.announce,
                    msg=f"Scan {scan_id_string} complete"
                )

    # All scans are complete. Observations are concluded with an 'end'
    # command.
    LOG.info(f"End scheduling block: {sbi.sbd_id}")
    api.end(subarray.id)

    LOG.info("Observation script complete")


def _resolve_context_bool(
    context: dict | None,
    *,
    key: str,
    fallback: bool,
) -> bool:
    """
    Resolve a boolean context value, falling back when not provided.

    :param context: Optional script context.
    :param key: Context key to resolve.
    :param fallback: Value to use when key is absent.
    :raises TypeError: If the provided context value is not boolean.
    """
    if not context or key not in context:
        return fallback

    value = context[key]
    if isinstance(value, bool):
        return value

    raise TypeError(
        f"context['{key}'] must be a boolean, got {value!r} "
        f"(type={type(value).__name__})"
    )


def _main(subarray_id: int, sb_json: str, sbi_id: str, context: dict | None = None, **runtime_args):
    LOG.info(f"Running OS process {os.getpid()}")
    LOG.info(f"Called with main(subarray_id={subarray_id}, sbi_id={sbi_id})")
    LOG.debug(f"main() sb_json={sb_json}")
    sbd: SBDefinition = api.load_sbd(sb_json)

    qa_enabled = _resolve_context_bool(
        context,
        key='wait_for_qa_ready',
        fallback=TEMP_WORKAROUNDS.enable_adr111,
    )
    mccs_early_scan = _resolve_context_bool(
        context,
        key='mccs_early_scan',
        fallback=TEMP_WORKAROUNDS.mccs_early_scan,
    )

    api.init_qa_wait(qa_enabled, mccs_early_scan=mccs_early_scan)

    apply_low_workarounds = apply_low_workarounds_in_args(runtime_args)

    if apply_low_workarounds:
        sbd = LOW_WORKAROUNDS.update_scheduling_block_definition(sbd)

    eb_id = execution_block.create_eb(telescope=sbd.telescope, sbi_ref=sbi_id)
    try:
        LOG.info(f"Created Execution Block {eb_id}")
        subarray = SubArray(subarray_id)
        assign_resources(subarray, sbd, apply_low_workarounds)
        api.configure_subarray_quality_monitor(subarray.id, sbd)
        if apply_low_workarounds:
            LOW_WORKAROUNDS.check_mccs_subarray_leaf_node_obsstate(subarray_id)
        observe(subarray=subarray, sbi=sbd, apply_low_workarounds=apply_low_workarounds, runtime_args=runtime_args)
        api.release_all_resources(subarray_id)
        execution_block.mark_execution_block_observed(eb_id)
    except Exception as err:
        LOG.error(f"Scheduling block execution failed: {err}", exc_info=True)
        if apply_low_workarounds:
            LOW_WORKAROUNDS.back_to_empty(subarray_id)
        execution_block.mark_execution_block_failed(eb_id=eb_id, error_message=str(err))
        raise


def apply_low_workarounds_in_args(runtime_args: dict) -> bool:
    if (workarounds := runtime_args.get("workarounds")) is None:
        return False

    if "low" in str(workarounds).lower():
        return True

    return False