Arvados

# Copyright (C) The Arvados Authors. All rights reserved.

#

# SPDX-License-Identifier: Apache-2.0

import itertools

from typing import (

    Container,

    Iterator,

    List,

    Optional,

    Union,

)

import arvados.util

Filter = List[Union[None, bool, float, str, 'Filter']]

def lookup(

        client: 'arvados.api_resources.ArvadosAPIClient',

        container: Union[str, 'arvados.api_resources.Container', 'arvados.api_resources.ContainerRequest'],

) -> Optional['arvados.api_resources.Container']:

    """Retrieve and return an Arvados container

    `container` can be an Arvados UUID, an Arvados container request object, or

    an Arvados container object.

    If `container` is a UUID for one of the other types, get that object

    using `client`.

    Then if the object is a container, return it.

    Otherwise if the object is a container request, if the request has no

    `container_uuid` (indicating it has not run), return None. Otherwise,

    get the corresponding container object using `client`, and return that.

    Raises `ValueError` if `container` is not a valid input for any reason

    (malformed UUID, Arvados object missing fields, etc.).

    """

    if isinstance(container, str):

        if not arvados.util.uuid_pattern.fullmatch(container):

            raise ValueError(f"{container!r} is not a valid Arvados UUID")

        _, uuid_kind, _ = container.split('-')

        if uuid_kind == 'dz642':

            resource = client.containers()

        elif uuid_kind == 'xvhdp':

            resource = client.container_requests()

        else:

            raise ValueError(f"{container!r} is not a container or request UUID")

        container = resource.get(uuid=container).execute()

    try:

        match = arvados.util.uuid_pattern.fullmatch(container['uuid'])

    except (KeyError, TypeError):

        match = None

    if match is None:

        raise ValueError("object does not have a valid Arvados UUID")

    _, uuid_kind, _ = container['uuid'].split('-')

    if uuid_kind == 'dz642':

        return container

    elif uuid_kind != 'xvhdp':

        raise ValueError("object does not have a container or request UUID")

    else:

        try:

            container_uuid = container['container_uuid']

        except KeyError:

            raise ValueError("container request object missing container_uuid field") from None

        if container_uuid is None:

            return None

        try:

            match = arvados.util.container_uuid_pattern.fullmatch(container_uuid)

        except TypeError:

            match = None

        if match is None:

            raise ValueError(f"container request object has invalid container_uuid {container_uuid!r}")

        else:

            return client.containers().get(uuid=container_uuid).execute()

def container_started(

        client: 'arvados.api_resources.ArvadosAPIClient',

        container: Union[str, 'arvados.api_resources.Container', 'arvados.api_resources.ContainerRequest'],

) -> bool:

    container_obj = lookup(client, container)

    if container_obj is None:

        return False

    else:

        return container_obj['status'] not in {'Queued', 'Locked'}

def container_finished(

        client: 'arvados.api_resources.ArvadosAPIClient',

        container: Union[str, 'arvados.api_resources.Container', 'arvados.api_resources.ContainerRequest'],

) -> bool:

    container_obj = lookup(client, container)

    if container_obj is None:

        return False

    else:

        return container_obj['status'] in {'Cancelled', 'Complete'}

def container_succeeded(

        client: 'arvados.api_resources.ArvadosAPIClient',

        container: Union[str, 'arvados.api_resources.Container', 'arvados.api_resources.ContainerRequest'],

        success: Container[int]=frozenset([0]),

) -> bool:

    container_obj = lookup(client, container)

    return (

        container_obj is not None

        and container_obj['status'] == 'Complete'

        and container_obj['exit_code'] in success

    )

def child_requests(

        client: 'arvados.api_resources.ArvadosAPIClient',

        container: Union[str, 'arvados.api_resources.Container', 'arvados.api_resources.ContainerRequest'],

        filters: List[Filter]=[],

        select: Optional[List[str]]=None,

        depth: Optional[int]=None,

) -> Iterator['arvados.api_resources.ContainerRequest']:

    """Iterate child requests of an Arvados container or request

    Given an Arvados containter, request, or UUID, iterate the child requests

    of the corresponding container.

    `filters` and `select` correspond to the Arvados API `list` call. `filters`

    sets additional search criteria for requests to iterate. `select` chooses

    which fields are available in each request object.

    `depth` limits how many times the search recurses. For example `depth=1`

    will only iterate the immediate child requests of `container`, and not any

    of their children.

    """

    if select is not None and 'container_uuid' not in select:

        select = ['container_uuid', *select]

    if depth is None:

        depth_iter = itertools.repeat(None)

    else:

        depth_iter = range(depth - 1, -1, -1)

    container = lookup(client, container)

    if container is None:

        return

    container_uuids = {container['uuid']}

    for level in depth_iter:

        child_filters = [

            ['requesting_container_uuid', 'in', list(container_uuids)],

        ]

        container_uuids = set()

        for child in arvados.util.keyset_list_all(

                client.container_requests().list,

                filters=filters + child_filters,

                select=select,

        ):

            yield child

            container_uuids.add(child['container_uuid'])

        container_uuids.discard(None)

        if not container_uuids:

            break

def child_containers(

        client: 'arvados.api_resources.ArvadosAPIClient',

        container: Union[str, 'arvados.api_resources.Container', 'arvados.api_resources.ContainerRequest'],

        request_filters: List[Filter]=[],

        container_filters: List[Filter]=[],

        select: Optional[List[str]]=None,

        depth: Optional[int]=None,

) -> Iterator['arvados.api_resources.Container']:

    """Iterate child containers of an Arvados container or request

    Given an Arvados containter, request, or UUID, iterate the child containers

    of the corresponding container.

    `request_filters`, `container_filters`, and `select` correspond to the

    Arvados API `list` call. `request_filters` and `container_filters` set

    additional search criteria for requests to search and containers to

    iterate, respectively. `select` chooses which fields are available in

    each container object.

    `depth` limits how many times the search recurses. For example `depth=1`

    will only iterate the immediate child containers of `container`, and not any

    of their children.

    """

    if select is not None and 'uuid' not in select:

        select = ['uuid', *select]

    container_uuids = {req['container_uuid'] for req in child_requests(

            client, container,

            filters=request_filters + [['container_uuid', '!=', None]],

            select=['container_uuid'],

            depth=depth,

    )}

    container_uuids.discard(None)

    while container_uuids:

        next_uuids = list(itertools.islice(container_uuids, 100))

        for container in arvados.util.keyset_list_all(

                client.containers().list,

                filters=container_filters + [['uuid', 'in', next_uuids]],

                select=select,

        ):

            yield container

            container_uuids.remove(container['uuid'])

if __name__ == '__main__':

    import arvados, argparse, csv, functools, re, sys

    client = arvados.api('v1')

    parser = argparse.ArgumentParser(

        description="Write a CSV report of all child containers of an Arvados container",

        epilog="This tool demonstrates functions that walk an Arvados container hierarchy.",

    )

    parser.add_argument(

        '--fields', '-f',

        type=re.compile(r'\W+', re.ASCII).split,

        default=[

            'uuid',

            'state',

            'started_at',

            'finished_at',

            'exit_code',

            'output',

            'log',

        ],

        help="""container fields to include in the report,

        separated by commas or whitespace""",

    )

    parser.add_argument(

        'containers',

        nargs=argparse.REMAINDER,

        metavar='UUID',

        type=functools.partial(lookup, client),

        help="UUID(s) of container or request object(s) to report",

    )

    args = parser.parse_args()

    out_csv = csv.DictWriter(

        sys.stdout,

        extrasaction='ignore',

        fieldnames=args.fields,

    )

    out_csv.writeheader()

    for container in args.containers:

        out_csv.writerows(child_containers(client, container, select=args.fields))