circle_ci_artifact.py

circle_ci_artifact.py

# Documentation: https://circleci.com/docs/api/v2/

import os
import sys
import time
import re
import shutil
import tempfile

import requests

CIRCLECI_TOKEN = '<YOUR_CIRCLECI_TOKEN>'
CIRCLECI_BASE_URL = 'https://circleci.com/api/v2/'

TARGET_CIRCLECI_BUILD_NAME = 'pytorch_windows_vs2019_py36_cuda10.1_build'

CIRCLECI_PROJECT_SLUG = 'gh/pytorch/pytorch'
TARGET_BRANCH = 'refs/heads/master'
TARGET_COMMIT = '173577566275cc3a54e535eb8e83e73ebee998fc'

if TARGET_BRANCH.startswith('refs/heads'):
      match = re.search('refs/heads/(.+)', TARGET_BRANCH)
      branch_name = match.group(1)
      TARGET_BRANCH_NAME = branch_name
elif TARGET_BRANCH.startswith('refs/pull'):
      match = re.search('refs/pull/(\d+)/head', TARGET_BRANCH)
      pr_num = match.group(1)
      TARGET_BRANCH_NAME = 'pull/' + pr_num

s = requests.Session()
s.headers.update({'Circle-Token': CIRCLECI_TOKEN})

def get_job_detail(job_number):
      job_url = CIRCLECI_BASE_URL + 'project/' + CIRCLECI_PROJECT_SLUG + '/job/' + str(job_number)
      job_raw = s.get(job_url)
      job_json = job_raw.json()
      return job_json

def get_job_artifact(job_number):
      job_artifacts_url = CIRCLECI_BASE_URL + 'project/' + CIRCLECI_PROJECT_SLUG + '/' + str(job_number) + '/artifacts'
      job_artifacts_raw = s.get(job_artifacts_url)
      job_artifacts_json = job_artifacts_raw.json()
      return job_artifacts_json

print("Fetching pipeline for branch: " + TARGET_BRANCH_NAME)
pipelines_url = CIRCLECI_BASE_URL + 'project/' + CIRCLECI_PROJECT_SLUG + '/pipeline'
next_page_token = None
while True:
      if next_page_token is None:
            params = {'branch': TARGET_BRANCH_NAME}
      else:
            params = {'branch': TARGET_BRANCH_NAME, "page-token": next_page_token }
      pipelines_raw = s.get(pipelines_url, params=params)
      pipelines_json = pipelines_raw.json()

      pipelines = pipelines_json['items']
      next_page_token = pipelines_json['next_page_token']
      if next_page_token is None:
            break
      vcs_pipeline = [x for x in pipelines if x['trigger']['type'] == 'webhook']

      if len(TARGET_COMMIT) > 0:
            vcs_pipeline = [x for x in vcs_pipeline if x['vcs']['revision'] == TARGET_COMMIT]

      if len(vcs_pipeline) != 0:
            break

if len(vcs_pipeline) == 0:
      print("Pipelines not found on CircleCI. Please check if the branch/ref name is correct. Aborting.")
      sys.exit(-1)

first_pipeline = vcs_pipeline[0]

pipeline_id = first_pipeline['id']

print("Fetching workflow for pipeline: " + pipeline_id)
workflow_url = CIRCLECI_BASE_URL + 'pipeline/' + pipeline_id + '/workflow'
workflow_raw = s.get(workflow_url)
workflow_json = workflow_raw.json()
workflows = workflow_json['items']
first_workflow = workflows[0]

workflow_id = first_workflow['id']

print("Fetching job for workflow: " + workflow_id)
job_url = CIRCLECI_BASE_URL + 'workflow/' + workflow_id + '/job'
job_raw = s.get(job_url)
job_json = job_raw.json()
jobs = job_json['items']

for job in jobs:
      print('\t' + job['id'] + ': ' + job['name'])

sys.stdout.flush()

win_cuda_build_jobs = [x for x in jobs if x['name'] == TARGET_CIRCLECI_BUILD_NAME]

if len(win_cuda_build_jobs) == 0:
      print("PyTorch Windows Job not found on CircleCI. Aborting.")
      sys.exit(-1)

win_cuda_build_job = win_cuda_build_jobs[0]
job_number = win_cuda_build_job['job_number']
print('Job Number: ' + str(job_number))
print('Job Status: ' + win_cuda_build_job['status'])
job_json = get_job_detail(job_number)
print('Job URL: ' + job_json['web_url'])
job_status = job_json['status']

while job_status != 'success':
      sys.stdout.flush()
      print('Waiting for job: ' + str(job_number))
      time.sleep(60)
      job_json = get_job_detail(job_number)
      job_status = job_json['status']

print("Fetching artifact for job: " + str(job_number))

job_artifacts_json = get_job_artifact(job_number)
all_artifacts = job_artifacts_json['items']
target_artifact = None
for a in all_artifacts:
      url = a['url']
      if url.endswith("7z"):
            target_artifact = a

if target_artifact is None:
      print("Artifact not found. Aborting.")
      print(f"Available artifacts: {all_artifacts}")
      sys.exit(-1)


print('Artifact URL: ' + target_artifact['url'])

print('Download artifact...')
download_url = target_artifact['url']
local_filename = download_url.split('/')[-1]

dest_file_path = os.path.join(tempfile.gettempdir(), local_filename)

if os.path.exists(dest_file_path):
    os.unlink(dest_file_path)

with requests.get(download_url, stream=True) as r:
    with open(local_filename, 'wb') as f:
        shutil.copyfileobj(r.raw, f)

print('Moving artifact to tempdir: ' + tempfile.gettempdir())
shutil.move(local_filename, tempfile.gettempdir())

print('Download OK')