Acceptance Criteria

• Pass any required attributes from credentials_file, region, s3_bucket_name, profile, execution_role, and log_group_name to AWSExecutor super class
• Re-structure executor in order to accord with abstract class methods as defined in RemoteExecutor, namely _upload_task, submit_task, get_status, _poll_task, _query_result, and cancel when appropriate
• Update unit tests for executor
• Update doc strings for executor

• The batch executor should assume that infrastructure already exists. Infrastructure provisioning does not need to happen in this executor.
• Test the existing executor and figure out what enhancements, fixes, tests, and docs need to be done for release.

Once the execution of a function finishes using this executor, there doesn't exist a mechanism to delete all the intermediate files created on remote and locally. We need a cleanup mechanism to take care of that similar to some of the other executors like SSH, AWSLambda, etc.

The purpose of this issue is to ensure that the Batch executor

• returns any stdout and stderr printed by tasks using the mechanism defined in #1380. To do this, the run() implementation needs to
  • retrieves the stdout and stderr from the executor backend -- in the case of AWS Executors, by parsing Cloudwatch logs (see how Braket does this)
  • printing those strings to self.task_stdout and self.task_stderr.
• distinguish task runtime exceptions -- those raised in the executor backend -- from other exceptions originating from the interaction of Covalent and the executor backend. This is explained in more detail in #1390. When an runtime exception occurs, the run() implementation should:
  1. Retrieve whatever stdout and stderr messages have already been printed by the task
  2. Ensure that the exception traceback is appended to the task's stderr.
  3. Print stdout and stderr to self.task_stdout and self.task_stderr, respectively.
  4. Raise a TaskRuntimeException.
     For examples, see how the dask executor now deals with task runtime exceptions.

Note: These changes should be implemented in a backward-compatible manner -- so that the new AWSExecutors should work with Covalent 0.202.0post1 (the current stable AWSExecutors work with the latest develop).

Acceptance criteria:

• Any stdout and stderr printed by a task before raising an unhandled exception is retrieved and printed to self.task_stdout and self.task_stderr respectively, where self is the executor plugin instance.
• If a task raises an exception:
  • The traceback is included in the task’s stderr.
  • The run() method raises a TaskRuntimeError.
• The executor plugin remains compatible with Covalent Core `0.202.0post1.

For the most part, the executor seems to be running fine but there are some cases where the workflow fails. One of these cases is:

import numpy as np

import covalent as ct
from  covalent_awsbatch_plugin.awsbatch import AWSBatchExecutor

N = 1
STALL = 1e5
X = 20

executor = AWSBatchExecutor()

@ct.electron(executor=executor)
def matrix_workload(mat_1, mat_2, stall=STALL, x=X):

    i = 0
    while i < stall:
        x * x
        i += 1

    return mat_1

@ct.lattice
def inflate(mat_1, dim=3):
    mat_2 = np.random.default_rng().integers(3, size=(dim, dim))
    return matrix_workload(mat_1, mat_2)

mat_1 = np.random.default_rng().integers(10, size=(3, 3))
dispatch_id = ct.dispatch(inflate)(mat_1)

Not sure why it fails but the status returned in the logs is FAILED.

There should be a suite of unit tests and a functional test for the executor. See the tests in the custom executor template for guidance (https://github.com/AgnostiqHQ/covalent-executor-template)

Partially fixes https://github.com/AgnostiqHQ/covalent-staging/issues/413

Overview

• The idea is to update the Readme to contain more information about required/optional config values, and provide information about each cloud resource that needs to be provisioned prior to running a live example
• A more complete example (not just containing one task) but also a workflow

Overview

The changes required to the plugin repo are minimal. We assets/infra folder where the terraform files are stored. For example, the repo structure would look like this:

- .github
- <PLUGIN_FOLDER>
    - assets
        - **infra**
            - main.tf
            - ...

Also we would need to add pydantic classes for validation of infra and executor arguments, for example:

class ExecutorPluginDefaults(BaseModel):
    """
    Default configuration values for the executor
    """

    credentials: str = ""
    profile: str = ""
    region: str = "us-east-1"
    ...
    retry_attempts: int = 3
    time_limit: int = 300
    poll_freq: int = 10

class ExecutorInfraDefaults(BaseModel):
    """
    Configuration values for provisioning AWS Batch cloud  infrastructure
    """
    prefix: str
    aws_region: str = "us-east-1"
    ...
    credentials: Optional[str] = ""
    profile: Optional[str] = ""
    retry_attempts: Optional[int] = 3
    time_limit: Optional[int] = 300
    poll_freq: Optional[int] = 10


_EXECUTOR_PLUGIN_DEFAULTS = ExecutorPluginDefaults().dict()

EXECUTOR_PLUGIN_NAME = "<PLUGIN_NAME>"

Acceptance Criteria

• Move infra to repo
• Create pydantic validation models

• Move the relevant terraform scripts to the plugin
• Update setup.py to include the infra scripts as part of the plugin sdist

Tracking progress and any discussions when performing QA for the AWSBatchExecutor in this issue.

Description

As a part of a recent change for self-hosted support, the base AWSExecutor class is solely responsible for resolving default values for profile, credentials file, and region instead of having individual aws executors determine defaults. As a result the region argument can be left unspecified as the AWSExecutor will resolve this via boto3's default behavior (to check AWS_DEFAULT_REGION environment variable).

However there is a need to explicitly obtain the region for log configuration, even if the region attribute may be empty in the Executor & AWSExecutor. We then need a robust way to obtain the region without duplicating the logic boto3 may use to resolve the region.

Proposed solution: we must use the region directly supplied by a boto3 session

• Update the CI so that tests are run using the most recent Covalent prerelease, rather than the stable release
• Re-add Python 3.10 to the list of automated test environments

Description

Covalent now supports dispatch cancellation. Executor plugins can opt in to this functionality by registering a job identifier for each task (in this case the job id or arn) and implementing a cancel method. This method will be invoked by Covalent when the user requests that the task be cancelled.

Methods provided by Covalent

All executor plugins automatically implement the following methods to be invoked in the plugin's run() method.

async def set_job_handle(self, job_handle)

saves a job identifier (job_handle) in Covalent's database. The job_handle can be any JSONable type, typically a string or integer, and should contain whatever information is needed to cancel the job later. For instance, AWS Batch's cancel_job method expects the job's jobId.

async def get_cancel_requested(self) -> bool

queries Covalent if task cancellation has been requested. This can be called at various junctures in the run() method if desired.

The run() method should raise a TaskCancelledError exception upon ascertaining that the backend job has been cancelled.

Methods to be implemented by each plugin:

Each plugin defines the following abstract method to be overridden with backend-specific logic:

async def cancel(self, task_metadata: Dict, job_handle: str) -> bool`

Upon receiving a cancellation request, the Covalent server will invoke this with the following inputs:

• task_metadata: currently contains the keys dispatch_id and node_id as for the run() method.
• job_handle: the task's job identifier previously saved using set_job_handle().

In addition to querying Covalent for cancellation requests, the run() method may need to query the compute backend to determine whether the job is in fact cancelled and raise TaskCancelledError if that is the case.

The code below sketches how to use the above methods:

from covalent._shared_files.exceptions import TaskCancelledError

...
async def proceed_if_task_not_cancelled(self):
  if await self.get_cancel_requested():
     self._debug_log(f"Task Cancelled")
     raise TaskCancelledError(f"Batch job {batch_job_name} cancelled")

async def run(self, function: Callable, args: List, kwargs: Dict, task_metadata: Dict) -> Any:
    ...
    await self.proceed_if_task_not_cancelled()
    # upload pickled assets
    ...
    await self.proceed_if_task_not_cancelled()
    # invoke job/task
    await self.set_job_handle(handle=job_handle)
    ...
    # await self.poll_job_status(task_metadata, job_handle)

async def poll_job_status(self, task_metadata, job_id):
    # Boto3 client invocations to check the job status
    # while timeout_not_exceeded:
    #    job_state = client.describe_job(job_id)    
    #    if job_state == "CANCELLED":
    #        raise TaskCancelledError
    #    ...
    #    await asyncio.sleep(poll_freq)
        
async def cancel(self, task_metadata: Dict, job_handle: str) -> None:
  """
    Cancel the batch job
    
    Arg(s)
      task_metadata: Dictionary with the task's dispatch_id and node id
      job_handle: Unique job handle assigned to the task by Batch
  
    Return(s)
      True/False indicating if the job was cancelled
  """
   # boto client invocations to cancel the task
   ...
    if job_cancelled:
        return True
    else:
        # Job cannot be cancelled for one reason or another
        return False

Acceptance Criteria

• In the run() method:
  • Save the job handle for the task once that is determined.
  • Check whether the job has been cancelled at various junctures
• Implement cancel method
• Ensure that the a workflow is tested with cancellation to ensure cancel functionality correctly integrated

Relevant PR: #88

Issue

This issue concerns the output of covalent deploy up awsbatch --help

Consider also, changes to Covalent OS:

• distinguish/omit user-settable vs not parameters in/from help table; some param defaults a uniquely auto-generated based on the prefix (see relevant .tf files)

See here for notes.

• The following infra defaults are not used for resource creation:

vcpus
memory
num_gpus
retry_attempts
time_limit
poll_freq
container_image_uri

• Empty defaults in the help menu should be replaced with

credentials: str = str(Path.home() / ".aws/credentials")
profile: str = "default"
...
vpc_id: str = "default"
subnet_id: str = "default"

• The parameter named vcpu here is named vcpus in GCP Batch. Neither should be an infra parameter, else names should be resolved.

• The following names conflict across the plugin and infra defaults

Acceptance Criteria

• Create a file covalent_googlebatch_plugin/gcpbatch.py
• Create the executor_plugin_name variable according to covalent_gcpbatch_plugin
• The plugin GCPBatchExecutor should be importable with from covalent_gcpbatch_plugin.gcpbatch import GCPBatchExecutor
• Add _EXECUTOR_PLUGIN_DEFAULTS with the values provided below
• Implement the executor by using the google-cloud-batch & google-cloud-storage GCP python client libraries. The following guide can be used as a reference for the necessary objects required for the implementation
• The executor methods outlined below should be implemented, as used a a reference (function names can change based on issue owner's discretion)

Note: Only implementation is in scope for this issue, unit/functional tests will come later.

Expected executor methods:

• __init__()
• run (async)
• _create_job (async)
• _upload_to_bucket(async)
• _create_job_sync
• _pickle_func(async)
• _pickle_func_sync
• _submit_job (async)
• get_status (async)
• _poll_task (async)
• cancel (async)
• _query_result (async)
• _get_service_client

Expected executor plugin defaults

_EXECUTOR_PLUGIN_DEFAULTS = {
	"storage_bucket_name": "",
	"image_uri": "",
	"service_account_email": "",
	"retries": 3,
	"time_limit": 300,
	"cache_dir": "/tmp/covalent",
	"poll_freq": 5,
	"vcpus": 1,
	"memory": 256
}

When QA testing the ECS executor in https://app.zenhub.com/workspaces/executors-62e0403830817d00118788a4/issues/agnostiqhq/covalent-staging/594, certain errors in electron execution were detected in CloudWatch's event logs, but these logs were not redirected to Covalent's logs. The CloudWatch logs were useful in debugging those errors and adding those logs to Covalent would be super-helpful in debugging on the client side.

Once separation of workflow and electron statuses is done, the electron level statuses need to be updated to accommodate executor dependent statuses. In this case the following status definitions will be updated:

• REGISTERING - Uploading task and registering the Batch job definition

• PROVISIONING - Corresponds to the Batch jobs state Submitted

• PENDING_BACKEND - Corresponds to Batch job states Pending and Runnable

• STARTING - Corresponds to Batch job state Starting

• RUNNING - Corresponds to Batch job state Running

• DEPROVISIONING - Corresponds to the Batch job states Deactivating, Stopping, and Deprovisioning

• COMPLETING - Batch job is in the Stopped state, results files being retrieved, temporary files are being deleted

For the end state, there is a separate class called EndState and contains all possible end state statuses like COMPLETED, FAILED, etc. and thus does not need to be added to the executor dependent status.

Acceptance Criteria:

• Above mentioned statuses need to be updated inside the local executor
• Tests need to be added to verify if those definitions are as expected

Spend no more than 30 minutes on this task

Write an explicit list of instructions for how this executor should be validated for QA. Assume the QA engineer has not used the executor before (but otherwise is familiar with Covalent as an Agnostiq engineer).

The list of instructions should be posted as a comment on this issue. Mention @wjcunningham7 in the comment.

If such instructions already exist in a README or RTD, you can instead link to those instructions.

Acceptance Criteria

• An explicit list of instructions, or a link to instructions, for how to validate the executor for QA

Description

As the executors are intended to be async-aware they are expected to use asynchronous calls for i/o bound operations, currently the official AWS SDK client boto3 only performs synchronous / blocking calls when interacting with Amazon services.

As this executor utilizes boto3 for making requests to AWS services (such as uploading function files, ect.) it currently blocks the main thread affecting parallelization and increasing latency to the covalent server (also viewing the UI). Hence we will need to leverage other client libs that support asyncio like aioboto3 and extend them if necessary. Since these don't seem to be mature at yet, we should use the more conservative approach of offloading all boto3 calls, such as this one, to a thread pool.

Disk I/O (including serializing inputs) can also block and therefore should also be performed in a separate thread.

Acceptance criteria

• All boto3 calls are delegated to the default threadpool executor using asyncio().get_running_loop().run_in_executor(None, ...).
• All disk I/O is delegated to the default threadpool executor

Consider as an example the Docker script in the covalent-awsbatch-plugin repo.

There is a similar (or identical) script in the Braket, Lambda, and ECS plugins.

This script is used to build a Docker image locally and upload it to ECR to execute the task on the cloud.

We want to standardize and optimize this procedure. Instead of building the Docker image locally for each plugin, we will host a common base image on ECR that the plugins can pull.

Covalent's braket extras requires >=1.28.5 which is incompatible with the strict pin here.

Acceptance criteria:

• Ensure that the instructions in the Readme are correct.
• Double check that the permissions instructions are are correct.
• Check that the links go to the correct place.