run a set of commands on multiple taskcluster worker types

.taskcluster.yml

provisionerId: aws-provisioner-v1
workerType: github-worker
retries: 5
priority: highest
created: '2019-05-23T11:09:17.216Z'
deadline: '2019-05-24T11:09:17.216Z'
expires: '2019-06-20T11:09:17.216Z'
scopes:
  - 'queue:scheduler-id:taskcluster-github'
  - 'queue:create-task:aws-provisioner-v1/gecko-t-win*'
  - 'queue:create-task:releng-hardware/gecko-t-win*'
  - 'queue:create-task:aws-provisioner-v1/gecko-1-b-win*'
  - 'queue:create-task:aws-provisioner-v1/gecko-2-b-win*'
  - 'queue:create-task:aws-provisioner-v1/gecko-3-b-win*'
  - 'queue:create-task:releng-hardware/gecko-1-b-win*'
  - 'queue:create-task:releng-hardware/gecko-2-b-win*'
  - 'queue:create-task:releng-hardware/gecko-3-b-win*'
  - 'queue:task-group-id:taskcluster-github/*'
  - 'queue:schedule-task:taskcluster-github/*'
payload:
  image: python
  maxRunTime: 1800
  features:
    taskclusterProxy: true
  artifacts:
    public/results.json:
      path: results.json
      type: file
  command:
    - /bin/bash
    - '--login'
    - '-c'
    - >-
      pip install asyncio slugid taskcluster && python <(curl -s
      https://gist.githubusercontent.com/grenade/a2ff8966607583fbc1944fccc256a80c/raw/decision-task.py)
metadata:
  owner: grenade@mozilla.com
  source: 'https://gist.github.com/grenade/a2ff8966607583fbc1944fccc256a80c'
  name: decision-task
  description: a gist based task group
tags: {}
extra: {}

config.json

{
  "task": {
    "name": {
      "prefix": "software-versions-",
      "suffix": ""
    },
    "description": {
      "prefix": "retrieve software versions on ",
      "suffix": ""
    },
    "owner": "grenade@mozilla.com",
    "command": [
      "mkdir .\\public",
      "chcp 65001 > nul && ver > .\\public\\windows-version.txt",
      "chcp 65001 > nul && C:\\mozilla-build\\python\\python.exe --version 2> .\\public\\python2-version.txt",
      "chcp 65001 > nul && C:\\mozilla-build\\python3\\python3.exe --version > .\\public\\python3-version.txt",
      "chcp 65001 > nul && hg --version > .\\public\\hg-version.txt",
      "chcp 65001 > nul && C:\\generic-worker\\generic-worker.exe --version > .\\public\\generic-worker-version.txt",
      "chcp 65001 > nul && reg query HKLM\\SOFTWARE\\Microsoft\\PowerShell\\1\\PowerShellEngine /v PowerShellVersion > .\\public\\powershell-version.txt",
      "chcp 65001 > nul && reg query HKLM\\SOFTWARE\\Mozilla\\OpenCloudConfig\\Source /v Organisation > .\\public\\occ-source-organisation.txt",
      "chcp 65001 > nul && reg query HKLM\\SOFTWARE\\Mozilla\\OpenCloudConfig\\Source /v Repository > .\\public\\occ-source-repository.txt",
      "chcp 65001 > nul && reg query HKLM\\SOFTWARE\\Mozilla\\OpenCloudConfig\\Source /v Revision > .\\public\\occ-source-revision.txt"
    ],
    "artifacts": [
      {
        "type": "file",
        "name": "public/windows-version.txt",
        "path": "public/windows-version.txt"
      },
      {
        "type": "file",
        "name": "public/python2-version.txt",
        "path": "public/python2-version.txt"
      },
      {
        "type": "file",
        "name": "public/python3-version.txt",
        "path": "public/python3-version.txt"
      },
      {
        "type": "file",
        "name": "public/hg-version.txt",
        "path": "public/hg-version.txt",
        "line": 0
      },
      {
        "type": "file",
        "name": "public/generic-worker-version.txt",
        "path": "public/generic-worker-version.txt"
      },
      {
        "type": "file",
        "name": "public/powershell-version.txt",
        "path": "public/powershell-version.txt",
        "line": 1,
        "split": {
          "separator": "    ",
          "index": 2
        }
      },
      {
        "type": "file",
        "name": "public/occ-source-organisation.txt",
        "path": "public/occ-source-organisation.txt",
        "line": 1,
        "split": {
          "separator": "    ",
          "index": 2
        }
      },
      {
        "type": "file",
        "name": "public/occ-source-repository.txt",
        "path": "public/occ-source-repository.txt",
        "line": 1,
        "split": {
          "separator": "    ",
          "index": 2
        }
      },
      {
        "type": "file",
        "name": "public/occ-source-revision.txt",
        "path": "public/occ-source-revision.txt",
        "line": 1,
        "split": {
          "separator": "    ",
          "index": 2
        }
      }
    ],
    "features": {
      "taskclusterProxy": true
    },
    "maxruntime": 600
  },
  "targets": [
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win7-32"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win7-32-beta"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win7-32-gpu"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win7-32-gpu-b"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win10-64"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win10-64-alpha"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win10-64-beta"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win10-64-gpu"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win10-64-gpu-a"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-t-win10-64-gpu-b"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-1-b-win2012"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-1-b-win2012-beta"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-2-b-win2012"
    },
    {
      "provisioner": "aws-provisioner-v1",
      "workertype": "gecko-3-b-win2012"
    },
    {
      "provisioner": "releng-hardware",
      "workertype": "gecko-1-b-win2012-gamma"
    },
    {
      "provisioner": "releng-hardware",
      "workertype": "gecko-2-b-win2012-gamma"
    },
    {
      "provisioner": "releng-hardware",
      "workertype": "gecko-3-b-win2012-gamma"
    }
  ]
}

decision-task.py

import asyncio
import os
import json
import slugid
import taskcluster
import taskcluster.aio
import time
import urllib
import urllib.request
from datetime import datetime, timedelta
from gzip import decompress

GIST_USER = 'grenade'
GIST_SHA = 'a2ff8966607583fbc1944fccc256a80c'


async def create_task(provisioner, workerType, taskGroupId, task):
  taskId = slugid.nice()
  payload = {
    'created': '{}Z'.format(datetime.utcnow().isoformat()[:-3]),
    'deadline': '{}Z'.format((datetime.utcnow() + timedelta(days=3)).isoformat()[:-3]),
    'provisionerId': provisioner,
    'workerType': workerType,
    'taskGroupId': taskGroupId,
    'routes': [],
    'scopes': [],
    'payload': {
      'maxRunTime': task['maxruntime'],
      'command': task['command'],
      'artifacts': list(map(lambda x: {'type': x['type'],'name': x['name'],'path': x['path']}, task['artifacts'])),
      'features': task['features']
    },
    'metadata': {
      'name': '{}{}/{}{}'.format(task['name']['prefix'], provisioner, workerType, task['name']['suffix']),
      'description': '{}{}/{}{}'.format(task['description']['prefix'], provisioner, workerType, task['description']['suffix']),
      'owner': task['owner'],
      'source': 'https://gist.github.com/{}/{}'.format(GIST_USER, GIST_SHA)
    }
  }
  print('creating {}/{} task {} (https://tools.taskcluster.net/groups/{}/tasks/{})'.format(provisioner, workerType, taskId, os.environ.get('TASK_ID'), taskId))
  return queue.createTask(taskId, payload)

  
async def print_task_artifacts(provisioner, workerType, taskGroupId, task):
  taskStatus = await create_task(provisioner, workerType, taskGroupId, task)
  print('{}/{} - {}: {}'.format(provisioner, workerType, taskStatus['status']['taskId'], taskStatus['status']['state']))
  while taskStatus['status']['state'] != 'completed':
    time.sleep(2)
    print('{}/{} - {}: {}'.format(provisioner, workerType, taskStatus['status']['taskId'], taskStatus['status']['state']))
    taskStatus = await asyncQueue.status(taskStatus['status']['taskId'])
  print('{}/{} - {}: {} on run {}'.format(provisioner, workerType, taskStatus['status']['taskId'], taskStatus['status']['state'], taskStatus['status']['runs'][-1]['runId']))  
  for artifactDefinition in task['artifacts']:
    artifactUrl = 'https://taskcluster-artifacts.net/{}/{}/{}'.format(taskStatus['status']['taskId'], taskStatus['status']['runs'][-1]['runId'], artifactDefinition['name'])
    print('{}/{} - {}'.format(provisioner, workerType, artifactUrl))
    if 'line' in artifactDefinition:
      if 'split' in artifactDefinition:
        artifactText = decompress(urllib.request.urlopen(urllib.request.Request(artifactUrl)).read()).decode('utf-8').strip().split('\n', 1)[artifactDefinition['line']].strip().split(artifactDefinition['split']['separator'])[artifactDefinition['split']['index']]
      else:
        artifactText = decompress(urllib.request.urlopen(urllib.request.Request(artifactUrl)).read()).decode('utf-8').strip().split('\n', 1)[artifactDefinition['line']].strip()
    else:
      artifactText = decompress(urllib.request.urlopen(urllib.request.Request(artifactUrl)).read()).decode('utf-8').strip()
    print('{}/{} - {}: {}'.format(provisioner, workerType, artifactDefinition['name'], artifactText))
    if workerType in results:
      results[workerType]['artifacts'].update({
        os.path.splitext(os.path.basename(artifactDefinition['name']))[0]: artifactText
      })
    else:
      results.update({
        workerType: {
          'artifacts': {
            os.path.splitext(os.path.basename(artifactDefinition['name']))[0]: artifactText
          }
        },
      })


config = json.loads(urllib.request.urlopen('https://gist.githubusercontent.com/{}/{}/raw/config.json?{}'.format(GIST_USER, GIST_SHA, slugid.nice())).read())
taskclusterOptions = {
  'rootUrl': os.environ['TASKCLUSTER_PROXY_URL']
}
queue = taskcluster.Queue(taskclusterOptions)

start = time.time()
loop = asyncio.get_event_loop()

session = taskcluster.aio.createSession(loop=loop)
asyncQueue = taskcluster.aio.Queue(taskclusterOptions, session=session)

tasks = []
results = {}
for target in config['targets']:
  tasks.append(asyncio.ensure_future(print_task_artifacts(target['provisioner'], target['workertype'], os.environ.get('TASK_ID'), config['task'])))

loop.run_until_complete(asyncio.wait(tasks, timeout=1200))
loop.close()
print(results)
with open('results.json', 'w') as fp:
    json.dump(results, fp, indent=2, sort_keys=True)
end = time.time()
print("total time: {}".format(end - start))