@pradyumn-el/pollycli/src/job.js

pollycli lets users access the functionalities of Polly over a command line interface

import chalk from 'chalk';
import inquirer from 'inquirer';

const fs = require('fs');
const axios = require('axios');
const pollyEnv = require('./env.json');
const pollymsg = require('./message');
const { handleApiError,formatJobStatusReason } = require('./helper-functions');

import {
  organizationDetails
} from './organization';
import {
  getAllDockerRepos,
  getAllDockerCommits
} from './docker';
import { pollyApi } from './api-client';
async function isFileValid(jobFile) {
  try {
    if (fs.existsSync(jobFile)) {
      // extracting file
      let jobData = null;
      try {
        jobData = fs.readFileSync(jobFile)
        jobData = JSON.parse(jobData);
      } catch (error) {
        pollymsg.pollyError(`File ${jobFile} not valid.`)
      }

      // checking if necessary fiels are present
      const minRequiredKeys = ['cpu', 'memory', 'machineType', 'image', 'tag', 'name'];
      const arrayDiff = minRequiredKeys.filter(x => !Object.keys(jobData).includes(x));
      if (arrayDiff.length > 0) {
        if (!((!arrayDiff.includes('cpu') && !arrayDiff.includes('memory')) || !arrayDiff.includes('machineType'))) {
          pollymsg.pollyError("Please mention either cpu and memory or machine type in the .json file")
        }

        if (arrayDiff.includes('cpu')) {
          arrayDiff.splice(arrayDiff.indexOf('cpu'), 1);
        }

        if (arrayDiff.includes('memory')) {
          arrayDiff.splice(arrayDiff.indexOf('memory'), 1);
        }

        if (arrayDiff.includes('machineType')) {
          arrayDiff.splice(arrayDiff.indexOf('machineType'), 1);
        }

        if (arrayDiff.length > 0) {
          pollymsg.pollyError(`Following entry(s) are missing in the discription file: ${arrayDiff.join(",")}`)
        }
      }
      if (Object.keys(jobData).includes('machineType')) {
        const resourceTypes = Object.keys(pollyEnv.machines);
        if (!resourceTypes.includes(jobData.machineType)) {
          pollymsg.pollyErrorNoExit(`Machine type should be one of the above ${resourceTypes.join(", ")}`)
          pollymsg.pollyMessage("Here is a list of machine types we support:")
          let machineTypeTable = [
            ['machine type', 'specifications']
          ]
          for (const macType in pollyEnv.machines) {
            machineTypeTable.push([macType, pollyEnv.machines[macType]])
          }
          pollymsg.pollyTable(machineTypeTable);
          process.exit(1);
        }
      }

      // checking if CPU entry is valid
      if (Object.keys(jobData).includes('cpu')) {
        isCPUValid(jobData);
        if (convertCPU(jobData.cpu) > 2 && !Object.keys(jobData).includes('machineType')) {
          pollymsg.pollyErrorNoExit("Job seems to need more resources. Please use machine type to specify resources. Here is a list of machine types we support:")
          let machineTypeTable = [
            ['machine type', 'specifications']
          ]
          for (const macType in pollyEnv.machines) {
            machineTypeTable.push([macType, pollyEnv.machines[macType]])
          }
          pollymsg.pollyTable(machineTypeTable);
          process.exit(1);
        }
      }

      // checking if the memory entry is valid
      if (Object.keys(jobData).includes('cpu')) {
        isMemoryValid(jobData);
        if (convertMemory(jobData.memory) > 8589934592 && !Object.keys(jobData).includes('machineType')) {
          pollymsg.pollyErrorNoExit("Job seems to need more resources. Please use machine type to specify resourcing. Here is a list of machine types we support")
          let machineTypeTable = [
            ['machine type', 'specifications']
          ]
          for (const macType in pollyEnv.machines) {
            machineTypeTable.push([macType, pollyEnv.machines[macType]])
          }
          pollymsg.pollyTable(machineTypeTable);
          process.exit(1);
        }
      }

      if (Object.keys(jobData).includes('machineType') && Object.keys(jobData).includes('memory') && Object.keys(jobData).includes('cpu')) {
        pollymsg.pollyMessage(`As both machine type and (cpu and memory) are given. Job will be run with machine type ${jobData.machineType} specification.`);
        pollymsg.pollyMessage(`Type ${jobData.machineType} has ${pollyEnv.machines[jobData.machineType]} as resource.`)
        pollymsg.pollyMessage(`If this much resource is not required try running by explicitly specifying memory and cpu without specifying the machine type`)
        if (!pollyUserConfirmation) {
          const questions = [{
            type: 'confirm',
            name: 'runconfirm',
            message: 'Would you like to continue'
          }]
          const answers = await inquirer.prompt(questions);
          if (!answers.runconfirm) {
            process.exit(0)
          }
        }
      }
      // add other validity checks after here
      if (!await isDockerValid(jobData)) {
        process.exit(1);
      }

      return true;
    } else {
      pollymsg.pollyError(`File ${jobFile} does not exist.`)
    }
  } catch (error) {
    if(error.isAxiosError) {
      handleApiError(error, "not able to validate config");
    }
    pollymsg.pollyError(`File path ${jobFile} is not valid.`)
  }
}

function convertCPU(cpu) {
  const cpuRegex = new RegExp("^[0-9]+m$");
  if (cpuRegex.test(cpu)) {
    return parseInt(cpu.slice(0, -1)) * 0.001;
  } else {
    return parseInt(cpu);
  }
}

const convertMemory = (memory) => {
  const memoryUnits = {
    "K": 1000,
    "Ki": 1024,
    "M": 1000 * 1000,
    "Mi": 1024 * 1024,
    "G": 1000 * 1000 * 1000,
    "Gi": 1024 * 1024 * 1024
  };

  const memoryRegex = new RegExp("^[0-9]+[GMK]i$");
  if (memoryRegex.test(memory)) {
    return parseInt(memory.slice(0, -2)) * memoryUnits[memory.slice(-2)];
  } else {
    return parseInt(memory.slice(0, -1)) * memoryUnits[memory.slice(-1)];
  }
}

function isCPUValid(jobData) {
  // Regex from kubernectics
  // https://github.com/kubernetes/kubernetes/issues/24925#issuecomment-229716184
  var cpuRegex = new RegExp('^([+-]?[0-9.]+)([m]*[-+]?[0-9]*)$')
  if (cpuRegex.test(jobData.cpu)) {
    return true
  } else {
    pollymsg.pollyError("CPU given is not valid. Please check the documantation for more information.");
  }
}

function isMemoryValid(jobData) {
  // Regex from kubernectics
  // https://github.com/kubernetes/kubernetes/issues/24925#issuecomment-229716184
  const memoryRegex = new RegExp("^[0-9]+[GMK]i$");
  if (memoryRegex.test(jobData.memory)) {
    return true
  } else {
    pollymsg.pollyError("Memory given is not valid. Please check the documantation for more information.");
  }
}


async function isDockerValid(jobData) {
  if (jobData.image.includes("dkr.ecr")) {
    if (!jobData.secret) {
      pollymsg.pollyError("secret is required to run an image from AWS ECR")
    } else {
      const requestUrl = `https://${jobData.image.split('/')[0]}/v2/${jobData.image.split('/').splice(1).join("/")}/manifests/${jobData.tag}`;
      let buff = Buffer.from(jobData.secret, 'base64');
      let extractedSecret = buff.toString('ascii');
      extractedSecret = JSON.parse(extractedSecret)
      if ('auths' in extractedSecret) {
        if (jobData.image.split('/')[0] in extractedSecret.auths) {
            const awsHeaders = {
              headers: {
                Authorization: `Basic ${extractedSecret.auths[jobData.image.split('/')[0]].auth}`
              }
            };
        
            try {
              await axios.get(requestUrl, awsHeaders);
            } catch (error) {
              if (error.response.status == "404"){
                pollymsg.pollyError(`Not able to find the image ${jobData.image} and the corresponding tag ${jobData.tag}`);
              } else if(error.response.status == "401"){
                pollymsg.pollyError("Unauthorized - Please make sure you have required access to use this image")
              } else {
                pollymsg.pollyError(`Not able to validate image ${jobData.image}:${jobData.tag}\nError code: ${error.response.status}\nError message: ${error.response.statusText}`)                
              }
            }
         
        } else {
          pollymsg.pollyError("secret does not contain the credentials to pull this image");
        }
      } else {
        pollymsg.pollyError("Invalid secret")
        pollymsg.pollyMessage("See miscellaneous options in pollycli. It will help you create secret");
      }
    }
  } else if (jobData.image.includes("gcr.io")) {
    pollymsg.pollyError("We do not support GCP images right now. This feature will come soon.");
  } else if (jobData.image.includes("polly.elucidata.io")) {
    let foundDocker = false;
    const organizationInfo = await organizationDetails();
    const orgId = organizationInfo.org_id;

    const allDockerImages = await getAllDockerRepos(orgId);
    for (const dockers of allDockerImages) {
      if (`${pollyEnv.dockerDomain}/${dockers.attributes.resource_name}` == jobData.image) {
        const dockerCommits = await getAllDockerCommits(orgId, dockers.attributes.resource_id, 'public' == dockers.attributes.resource_access);
        for (const dockerCommit of dockerCommits) {
          if (dockerCommit.attributes.tag_name == jobData.tag) {
            foundDocker = true;
          }
        }
      }
    }
    if (!foundDocker) {
      pollymsg.pollyError(`Not able to find the image ${jobData.image} and the corresponding tag ${jobData.tag}`);
    }
  } else {
    let authHeader = {};
    const requestUrl = `https://hub.docker.com/v2/repositories/${jobData.image}/tags/${jobData.tag}`
    if (jobData.secret) {
      let buff = Buffer.from(jobData.secret, 'base64');
      let extractedSecret = buff.toString('ascii');
      extractedSecret = JSON.parse(extractedSecret)
      if ('auths' in extractedSecret) {
        let idDockerCredPresent = false;
        let dockerCrdKey = null;
        for (let aut of Object.keys(extractedSecret.auths)) {
          if (aut.indexOf('index.docker') > -1) {
            idDockerCredPresent = true;
            dockerCrdKey = aut
          }
        }
        if (!idDockerCredPresent) {
          pollymsg.pollyError("Docker credentials not present")
        }
        extractedSecret.auths[dockerCrdKey].auth
        authHeader = {
          headers: {
            Authorization: `Basic ${extractedSecret.auths[dockerCrdKey].auth}`
          }
        };
      } else {
        pollymsg.pollyError("Invalid secret")
        pollymsg.pollyMessage("See miscellaneous options in pollycli. It will help you create secret");
      }
    }
    try {
      await axios.get(requestUrl, authHeader);
    } catch (error) {
      if (error.response.status == "404"){
        pollymsg.pollyError(`Not able to find the image ${jobData.image} and the corresponding tag ${jobData.tag}`);
      } else {
        pollymsg.pollyError(`Not able to validate image ${jobData.image}:${jobData.tag}\nError code: ${error.response.status}\nError message: ${error.response.statusText}`)
      }
    }
  }
  return true;
}

async function submitJobToPolly(projectId, jobData) {
  try {
    const submitBody = {
      "data": {
        "type": "jobs",
        "attributes": jobData
      }
    }
    const jobUrl = `/projects/${projectId}/jobs`
    const postData = await pollyApi.post(jobUrl, submitBody);
    return postData;
  } catch (error) {
    const err_msg = "Not able to submit job";
    if(error.isAxiosError) {
      handleApiError(error, err_msg);
    }

    pollymsg.pollyError(err_msg);
  }
}

export async function submitJob(projectId, jobFile, internalCalls = false) {
  if (!projectId) { 
    pollymsg.pollyError("No workspace ID given.")
  } else if (!jobFile) {
    pollymsg.pollyError("No job file given.")
  }

   

  if (await isFileValid(jobFile)) {
    let jobData = fs.readFileSync(jobFile)
    jobData = JSON.parse(jobData);

    //Adding polly related creds as environment variables
    const email = pollystore.get('pollyUser').pollyemail;
    const refreshToken = pollystore.get(`${email}`).pollyrefreshToken;
    const idToken = pollystore.get(`${email}`).pollyIdToken;
    const exp = pollystore.get(`${email}`).pollyExp;
    const sub = pollystore.get(`${email}`).pollySub;
    const aud = pollystore.get(`${email}`).pollyAud;

    if (Object.keys(jobData).indexOf('secret_env') < 0) {
      jobData['secret_env'] = {};
    }

    jobData['secret_env']['POLLY_REFRESH_TOKEN'] = refreshToken;
    jobData['secret_env']['POLLY_ID_TOKEN'] = idToken;
    jobData['secret_env']['POLLY_WORKSPACE_ID'] = projectId;
    jobData['secret_env']['POLLY_USER'] = email;
    jobData['secret_env']['POLLY_SUB'] = sub;
    jobData['secret_env']['POLLY_EXP'] = exp;
    jobData['secret_env']['POLLY_AUD'] = aud;

    let submittedProject = await submitJobToPolly(projectId, jobData);
    submittedProject = submittedProject.data

    if (internalCalls) {
      return submittedProject;
    }
    let tableData = [['Workspace ID', 'Job ID']]
    tableData.push([submittedProject.data.project_id, submittedProject.data.job_id]);
    pollymsg.pollySuccess("Submitted the job to run!");
    pollymsg.pollyTable(tableData);
    pollymsg.pollyMessage(chalk.white(`Useful commands:
status : ${chalk.italic('polly jobs status --workspace-id')} ${chalk.italic(submittedProject.data.project_id)} ${chalk.italic('--job-id')} ${chalk.italic(submittedProject.data.job_id)}
logs   : ${chalk.italic('polly jobs logs --workspace-id')} ${chalk.italic(submittedProject.data.project_id)} ${chalk.italic('--job-id')} ${chalk.italic(submittedProject.data.job_id)} ${chalk.italic('--attempt-no 1')}
cancel : ${chalk.italic('polly jobs cancel --workspace-id')} ${chalk.italic(submittedProject.data.project_id)} ${chalk.italic('--job-id')} ${chalk.italic(submittedProject.data.job_id)}`))
    return submittedProject;
  }
}

export async function cancelJob(projectId, jobId) {
  if (!jobId) {
    pollymsg.pollyError("Please specify the job ID of the job to be cancelled.");
  }
  const jobUrl = `${pollyEnv.baseV2Api}/projects/${projectId}/jobs/${jobId}`;
  try {
    await pollyApi.delete(jobUrl);
    pollymsg.pollySuccess("Job cancelled.")
    process.exit(0);
  } catch (error) {
    const err_msg = "Failed to cancel the job.";
    if(error.isAxiosError) {
      handleApiError(error, err_msg);
    }
   
    pollymsg.pollyError(err_msg);
  }
}

export async function jobStatus(projectId, jobId, internalCalls = false){
  /* 
  The function is used for internal calls, basically functions within pollycli calls this function. 
  Returns list of job status data arranged in desc order of created_ts.
  */
  if (!jobId) {
    jobId = ''
  }
  const jobBaseUrl = `${pollyEnv.baseV2Api}`;
  let sortedJobsTemp = null;
  let jobUrl = `${jobBaseUrl}/projects/${projectId}/jobs/${jobId}`;
  const sortedJobs = [];

  try {
    do {
      let postDatas = await pollyApi.get(jobUrl);
      sortedJobsTemp = postDatas.data
      sortedJobs.push(...sortedJobsTemp.data);
      if (sortedJobsTemp.links){
        jobUrl = `${jobBaseUrl}${sortedJobsTemp.links.next}`
      }
    } while (sortedJobsTemp.data.length > 0 && sortedJobsTemp.links && sortedJobsTemp.links.next)

  } catch (error) {
    const err_msg = "Not able to get the status of the Job(s)";
    if(error.isAxiosError) {
      handleApiError(error, err_msg);
    }

    pollymsg.pollyError(err_msg);
  } 

  return sortedJobs;

}

export async function getJobStatusBatch(page_link="", serial_no=1){
  /* 
  The function fetches data from compute lambda with a certain batch size defined.
  Displays users table of job status data arranged in desc order of created_ts.
  */
  const jobBaseUrl = `${pollyEnv.baseV2Api}`;
  let sortedJobsTemp = null;
  let jobUrl = `${jobBaseUrl}${page_link}`;
  const sortedJobs = [];
  
  try {
    let postDatas = await pollyApi.get(jobUrl);
    sortedJobsTemp = postDatas.data;
    sortedJobs.push(...sortedJobsTemp.data);
    if (sortedJobsTemp.links){
      page_link = sortedJobsTemp.links.next
    }
  } catch (error) {
    const err_msg = "Not able to get the status of the Job(s)";
    if(error.isAxiosError) {
      handleApiError(error, err_msg);
    }

    pollymsg.pollyError(err_msg);
  } 

  let dataArray = [
    ['S.No', 'Job ID', 'Attempt No.','Job Name', 'State', 'Reason']
  ]

  let job_attempt_num = 1
  for (let jobData of sortedJobs) {
    if (jobData.attributes.hasOwnProperty("session_info")){
      for(let session_info of jobData.attributes.session_info){
        if (session_info.hasOwnProperty("reason")){
          session_info.reason = await formatJobStatusReason(session_info.reason)
          dataArray.push([serial_no, jobData.id, job_attempt_num, jobData.attributes.job_name, session_info.state, session_info.reason])
        }else{
          dataArray.push([serial_no, jobData.id, job_attempt_num,jobData.attributes.job_name, session_info.state, "NA"])
        }
        job_attempt_num+=1;
        serial_no+=1;
      }
      job_attempt_num = 1;
    }else{
      dataArray.push([serial_no, jobData.id, "1", jobData.attributes.job_name, jobData.attributes.state, "NA"]);
      serial_no+=1;
    }
  }
  pollymsg.pollyTable(dataArray);
  return {"page_link": page_link, "serial_no": serial_no};
}


export async function displayJobStatus(projectId, jobId){
  /* 
  The function is called when job status data is displayed to user.
  Iteration to fetch data is controlled by user(pagination is implemented for job status).  
  Displays table of data with batch size set in compute lambda or displays only one job if job_id is provided by user.
  */
  const inquirer = require('inquirer');
  let serial_no = 1;
  let resp = null;
  if (jobId == ''){
    let page_link = `/projects/${projectId}/jobs/`;
    resp = await getJobStatusBatch(page_link, serial_no);
    page_link = resp["page_link"];
    serial_no = resp["serial_no"];
    if (!page_link){
      process.exit(0);
    }
    await showMenu();
    
    async function showMenu () {
      await inquirer
      .prompt([
          {
          name: 'check',
          type: 'input',
          message: "Press Enter to view more. Press any key to exit",
          default: 'Yes'
      }]
      ).then(async (answers) => {
          if (answers.check === 'Yes') {
            resp = await getJobStatusBatch(page_link, serial_no);
            page_link = resp["page_link"];
            serial_no = resp["serial_no"];
            if (!page_link){
              process.exit(0);
            }
            return showMenu();
          }
          else {
            process.exit(0);
          }
      })
      .catch((err) => {
        pollymsg.pollyError("Please try again something went wrong.");
      });
  }
  }else {
    let page_link = `/projects/${projectId}/jobs/${jobId}`;
    await getJobStatusBatch(page_link, serial_no);
    process.exit(0);
  }
}