/encoder.py

#!/usr/bin/env python3
'''Encodes files for Auto-Encode and returns them for insertion'''

import os
import re
import json
import time
import yaml
import logging
import logging.handlers
import pexpect
import argparse
import platform
import subprocess

from threading import Thread, Lock

# AE api client
import capi.capi as capi

# AE utlities
import utils.ldir as ldir 
import utils.utils as utils

class Encoder(object):
    '''Encoder for Auto-Encode'''
    def __init__(self, cfg_path="confs/ae.yaml", log=None):
        '''
        inits Encoder
        
        Args:
            cfg_path (string): path to config file
            log (logging object): logging object
        '''
        # load yaml
        #with open("/Auto-Encode/confs/ae.yaml") as fp:
        with open(cfg_path) as fp:
            self.fconfig = yaml.load(fp, Loader=yaml.FullLoader)
            # extract capi config
            self.config = self.fconfig["encoder"]
       
        # setup logging
        if log is not None:
            self.log = log
        else:
            # ensure log path exists
            parent = os.path.dirname(self.config["log"]["path"])
            os.makedirs(parent, exist_ok=True)

            # setup logging
            self.log = logging.getLogger(platform.node())
            self.log.setLevel(logging.DEBUG)

            # create console handler
            console = logging.StreamHandler()
            console.setLevel(logging.DEBUG)

            # create a file handler
            fhandler = logging.handlers.TimedRotatingFileHandler(\
                            self.config["log"]["path"], \
                            when='D', \
                            interval=1, \
                            backupCount=3)
            fhandler.setLevel(self.config["log"]["level"].upper())

            # create a logging format
            formatter = logging.Formatter('%(asctime)s [%(name)s] [%(processName)s] [%(levelname)s] %(message)s')
            fhandler.setFormatter(formatter)
            console.setFormatter(formatter)

            # add the handlers to the logger
            self.log.addHandler(fhandler)
            # use console_log
            if self.config["log"]["console"] == True:
                self.log.addHandler(console)

        # api client
        self.capi = capi.Capi(cfg_path, self.log)
        
        # ae utilities
        self.utils = utils.Utils(self.config, self.log, self.capi)

        # setup vars
        self.node = platform.node()
        self.threads = []

        # get all wdir data
        self.wdirs = []
        for wdir in self.fconfig["watch_dirs"]:
            self.log.info("Adding job type {}".format(wdir["name"]))
            dir_obj = ldir.Load_dir(wdir["name"], \
                               wdir["load_dir"], \
                               wdir["targ_dir"], \
                               wdir["share_root"], \
                               wdir["presets"], \
                               wdir["type"], \
                               wdir["staging_root"], \
                               wdir["container"])
            self.wdirs.append(dir_obj)
         
        # log spam protection
        self.spam = []

    def _download(self, jdata, src, dest):
        '''
        downloads media to worker assuming you have a share mounted from
        cifs/nfs/samba/gluster/ceph

        Args:
            jdata (dict): job data
            src (string): path to copy file from
            dest (string): path to copy file to
        Returns:
            bool: True if download succeded, False if it didn't
        '''
        # set job status
        self.capi.set_job_status(jdata["job_id"], "downloading")

        # copy file
        return self.utils.copy(src, dest, jdata["job_id"])
    
    def _upload(self, jdata, src, dest):
        '''
        uploads media to share assuming you have a share mounted from
        cifs/nfs/samba/gluster/ceph

        Args:
            jdata (dict): job data
            src (string): path to copy file from
            dest (string): path to copy file to
        '''
        # set job status
        self.capi.set_job_status(jdata["job_id"], "uploading")

        # copy file
        while 1:
            ustatus = self.utils.copy(src, dest, jdata["job_id"])

            # check if upload was successfull
            if ustatus == True:
                # set job status
                self.capi.set_job_status(jdata["job_id"], "uploaded")
                break

            # upload failed
            time.sleep(60)

    def _hdr_scan(self, path):
        '''detects if a video is 4k'''
        out = subprocess.check_output("mediainfo \"" + path + "\"", shell=True)
        if ": BT.2020" in str(out):
            self.log.info("HDR detected injecting HDR preset")
            return '-x265-params "level=5.2:colorprim=bt2020:colormatrix=bt2020nc:transfer=smpte2084"'
        else:
            return ""

    def _extra_preset_scan(self, path):
        '''
        scans and sets extra presets

        Args:
            path (string): path to file to scan
        Returns:
            string: extra presets to set
        '''
        expresets = " "
        # list of scan results
        scanres = []
        # hdr presets
        scanres.append(self._hdr_scan(path))
        # combine all scans into list of extra presets
        return expresets.join(scanres)
    
    def _calc_secs(self, time):
        '''
        gets the seconds from a ffmpeg timestamp
        
        Args:
            time (string): ffmpeg timestamp
        Returns:
            int: seconds from ffmpeg timestamp
        '''
        subtime_re = re.compile("\d\d") 
        t_groups = subtime_re.findall(time)
        total_s = 0
        # count num of seconds
        h_s = int(t_groups[0]) * 60 * 60
        m_s = int(t_groups[1]) * 60
        s = int(t_groups[2])
        total_s = h_s + m_s + s + int(t_groups[3])/100
        return total_s

    def _transcode(self, src, dest, presets, jdata, expresets=True):
        '''
        transcodes codes a file using ffmpeg

        Args:
            src (string): path to file to transcode
            dest (string): destination to place transcoded file
            presets (string): presets to to use
            expresets (bool, optional): whether to scan for extra presets
                defaults to True
        Retuns:
            bool: True if successful, False if not
        '''
        # scan for extra presets to set
        if expresets == True:
            presets = presets + " " + self._extra_preset_scan(src)

        print(presets)

        # build ffmpeg command
        ffcmd = 'ffmpeg -i "{}" {} "{}"'.format(src, \
                                                presets, \
                                                dest)
        self.log.debug(ffcmd)
        
        # start transcode job
        proc = pexpect.spawn(ffcmd)
        status = proc.compile_pattern_list([
            pexpect.EOF,
            "frame=(.)*"])
        self.log.info("Started encoding " + jdata["name"])
        self.capi.set_job_status(jdata["job_id"], "converting")

        # update job progress
        time_re = re.compile("time=(\d)*:(\d)*:(\d)*.(\d)*")
        spam_prot = 60
        while 1:
            i = proc.expect_list(status, timeout=None)
            if i == 0:
                break

            if i == 1:
                # get current transcoded time from ffmpeg
                line = proc.match.group(0)
                line = line.decode("utf-8")
                # log once every 60 attempts to keep log spam down
                if spam_prot == 60:
                    self.log.info("%s", line.rstrip())
                    spam_prot = 0
                else:
                    spam_prot += 1
                t_groups = time_re.search(line)
                # make sure a time was actually pulled
                if t_groups is not None:
                    total_s = self._calc_secs(t_groups[0])
                    # update mongodb
                    self.capi.set_job_transcoded_time(jdata["job_id"], total_s)
                    prog = str((total_s / int(jdata["total_time"])) * 100).split(".")[0]
                    self.capi.set_job_progress(jdata["job_id"], prog)
                continue

        # job ended
        proc.close()

        # update job status if we errored out
        if proc.exitstatus != 0:
            # job errored out
            self.log.error("ffmpeg ran into a problem")
            self.capi.set_job_status(jdata["job_id"], "error")
            return False
        return True

    def worker(self, wdir):
        '''
        Completes one Auto-Encode job

        Args:
            wdir (dict): load directory config
        Returns:
            bool: True if job completed, False if no job
        '''
        # get job
        try:
            jdata = self.capi.get_job(jtype=wdir.name)
            
            # return False if we dont have a job
            if jdata is None:
                return False

            # set jobs worker
            self.capi.set_job_worker(jdata["job_id"], self.node)

            # get name with new extension
            ncname = self.utils.get_ncname(jdata["name"], wdir)
            # build paths
            local_tmp = os.path.join(self.config["local"]["tmp"], jdata["name"])
            local_staging  = os.path.join(self.config["local"]["staging"], \
                                          ncname)
            server_tmp = os.path.join(self.config["share"]["tmp"], \
                                      jdata["path"])
            server_staging  = os.path.join(self.config["share"]["staging"], \
                                           jdata["staging_path"])

            # download media
            if self._download(jdata, server_tmp, local_tmp) == False:
                # failed to download job
                self.capi.set_job_status(jdata["job_id"], "recieved")
                return False
            
            # pause before transcode to let updates_dl finish
            time.sleep(10)
            # transcode file
            if self._transcode(local_tmp, local_staging, wdir.presets, jdata) == False:
                # failed to transcode job
                self.capi.set_job_status(jdata["job_id"], "error")
                return False
            
            # copy transcoded file to share
            self._upload(jdata, local_staging, server_staging)

            # pause before removal
            time.sleep(30)
            # remove local tmp/staging
            os.remove(local_tmp)
            os.remove(local_staging)
            return True
        except KeyboardInterrupt:
            self.log.info("CTRL^C detected")
            self.log.info("Current job being reset to recieved")
            self.capi.set_job_status(jdata["job_id"], "recieved")
            exit(1)


    def start(self):
        '''Start encoder'''
        # enter main work loop
        self.log.info("Encoder started")
        while 1:
            # check if cluster is shutting down
            if self.capi.get_cstatus() == "shutdown":
                self.log.info("Shutting down")
                break

            # start working on jobs
            for wdir in self.wdirs:
                # get one job if one exists and complete it
                if self.worker(wdir):
                    # if we completed a job check for cluster shutdown
                    break

            time.sleep(10)

def argparser():
    '''parses arguments for loader'''
    parser = argparse.ArgumentParser(description=\
                                     "Loads files into Auto-Encode")
    parser.add_argument('--conf', help="path to config file")
    return parser.parse_args()

if __name__ == "__main__":
    args = argparser()
    enc = Encoder(args.conf)
    enc.start()