SpeechToTextWorkflowSchedulerQuartz.java

/*
 * Licensed to The Apereo Foundation under one or more contributor license
 * agreements. See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership.
 *
 *
 * The Apereo Foundation licenses this file to you under the Educational
 * Community License, Version 2.0 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of the License
 * at:
 *
 *   http://opensource.org/licenses/ecl2.txt
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
 * License for the specific language governing permissions and limitations under
 * the License.
 *
 */
package org.opencastproject.speechtotext.async.impl;

import static org.opencastproject.speechtotext.async.api.SpeechToTextAsyncTracker.JOBS_WORKFLOW_CONFIGURATION;

import org.opencastproject.assetmanager.api.AssetManager;
import org.opencastproject.assetmanager.util.Workflows;
import org.opencastproject.job.api.Job;
import org.opencastproject.job.jpa.JpaJob;
import org.opencastproject.kernel.scanner.AbstractScanner;
import org.opencastproject.security.api.Organization;
import org.opencastproject.security.api.OrganizationDirectoryService;
import org.opencastproject.security.api.SecurityService;
import org.opencastproject.serviceregistry.api.ServiceRegistry;
import org.opencastproject.speechtotext.async.api.SpeechToTextAsyncException;
import org.opencastproject.speechtotext.async.persistence.SpeechToTextControl;
import org.opencastproject.speechtotext.async.persistence.SpeechToTextDatabase;
import org.opencastproject.util.NeedleEye;
import org.opencastproject.util.NotFoundException;
import org.opencastproject.util.OsgiUtil;
import org.opencastproject.workflow.api.ConfiguredWorkflow;
import org.opencastproject.workflow.api.WorkflowDatabaseException;
import org.opencastproject.workflow.api.WorkflowDefinition;
import org.opencastproject.workflow.api.WorkflowInstance;
import org.opencastproject.workflow.api.WorkflowService;

import org.apache.commons.lang3.StringUtils;
import org.osgi.service.cm.ConfigurationException;
import org.osgi.service.cm.ManagedService;
import org.osgi.service.component.ComponentContext;
import org.osgi.service.component.annotations.Activate;
import org.osgi.service.component.annotations.Component;
import org.osgi.service.component.annotations.Reference;
import org.quartz.CronExpression;
import org.quartz.JobDetail;
import org.quartz.JobExecutionContext;
import org.quartz.impl.StdSchedulerFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Instant;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.Dictionary;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * Quartz job to start workflows to asynchronously attach captions generated by the SpeechToText service.
 */
@Component(
    immediate = true,
    service = {
        ManagedService.class
    },
    property = {
        "service.description=Speech to Text Workflow Scheduler"
    })
public class SpeechToTextWorkflowSchedulerQuartz extends AbstractScanner implements ManagedService {
  private static final Logger logger = LoggerFactory.getLogger(SpeechToTextWorkflowSchedulerQuartz.class);

  public static final String JOB_GROUP = "org-opencast-stt-wf-scheduler-group";
  public static final String JOB_NAME = "org-opencast-stt-wf-scheduler-job";
  public static final String SCANNER_NAME = "STT Worflow Scheduler Quartz";
  public static final String TRIGGER_GROUP = "org-opencast-stt-wf-scheduler-trigger-group";
  public static final String TRIGGER_NAME = "org-opencast-stt-wf-scheduler-trigger";

  // === Configuration options
  /* Number of seconds to abandon retrying attaching subtitles after they are generated. */
  public static final String ABANDON_AFTER_SECS = "abandon-after-secs";
  /* Workflow definition to use to attach the subtitles */
  public static final String WORKFLOW = "workflow";
  /* Workflow definition to use to retry transcription when error occurs */
  public static final String WORKFLOW_RETRY = "retry-workflow";
  /* Max tries if an error happens */
  public static final String MAX_TRIES = "max-tries";

  // === Default values
  /* Default number of seconds to abandon retrying attaching subtitles. */
  private static final long DEFAULT_ABANDON_AFTER_SECS = 24 * 60 * 60;

  /* Workflow definition to start to attach the subtitles */
  private String attachWorkflowDef;
  /* If STT jobs not finished after this interval, they won't be tracked anymore. */
  private long abandonAfterMs = DEFAULT_ABANDON_AFTER_SECS * 1000;
  /*
   * If STT job fails, maximum number of retries. A retry will start the workflow specified in retry-workflow. Set to 0
   * or do not specify retry-workflow to disable retries.
   */
  private int maxTries = 0;

  private String retryWorkflowDef = null;

  private AssetManager assetManager;
  private WorkflowService workflowService;
  private SpeechToTextDatabase database;

  // Only used by unit tests!
  private Workflows wfUtil;

  @Override
  @Activate
  protected void activate(ComponentContext cc) {
    logger.info("Activating!");
    super.activate(cc);
  }

  public SpeechToTextWorkflowSchedulerQuartz() {
    try {
      quartz = new StdSchedulerFactory().getScheduler();
      quartz.start();
      // create and set the job. To actually run it call schedule(..)
      final JobDetail job = new JobDetail(getJobName(), getJobGroup(), Runner.class);
      job.setDurability(false);
      job.setVolatility(true);
      job.getJobDataMap().put(JOB_PARAM_PARENT, this);
      quartz.addJob(job, true);
    } catch (org.quartz.SchedulerException e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public void updated(Dictionary<String, ?> properties) throws ConfigurationException {
    logger.info("Updating!");
    unschedule();

    if (properties != null) {
      logger.debug("Updating configuration...");

      setEnabled(properties.get(PARAM_KEY_ENABLED) != null
          && Boolean.valueOf(properties.get(PARAM_KEY_ENABLED).toString()));
      if (!isEnabled()) {
        logger.info("Speech to text workflow scheduler quartz job is disabled");
        return;
      }

      String cronExpression = (String) properties.get(PARAM_KEY_CRON_EXPR);
      if (StringUtils.isBlank(cronExpression) || !CronExpression.isValidExpression(cronExpression)) {
        throw new ConfigurationException(PARAM_KEY_CRON_EXPR, "Cron expression must be valid");
      }
      setCronExpression(cronExpression);
      logger.info("{} : {}", PARAM_KEY_CRON_EXPR, cronExpression);

      attachWorkflowDef = (String) properties.get(WORKFLOW);
      if (StringUtils.isBlank(attachWorkflowDef)) {
        throw new ConfigurationException(WORKFLOW, "Attach workflow definition is missing");
      }
      logger.info("Attach workflow definition is {}", attachWorkflowDef);

      abandonAfterMs = DEFAULT_ABANDON_AFTER_SECS * 1000;
      if (properties.get(ABANDON_AFTER_SECS) != null) {
        try {
          abandonAfterMs = Integer.parseInt((String) properties.get(ABANDON_AFTER_SECS)) * 1000;
        } catch (Exception e) {
          // Log and use default
          logger.warn("Invalid configuration for {}: '{}', defaulting to '{}'", ABANDON_AFTER_SECS,
                  properties.get(ABANDON_AFTER_SECS), DEFAULT_ABANDON_AFTER_SECS, e);
        }
      }
      logger.info("Abandon attempts to start attach workflow after {} ms", abandonAfterMs);

      // Reset values before update
      retryWorkflowDef = null;
      maxTries = 0;
      Optional<String> retryWfOpt = OsgiUtil.getOptCfg(properties, WORKFLOW_RETRY);
      if (retryWfOpt.isPresent()) {
        retryWorkflowDef = retryWfOpt.get();
        logger.info("Retry workflow definition: {}", retryWorkflowDef);
        Optional<Integer> maxTriesOpt = OsgiUtil.getOptCfgAsInt(properties, MAX_TRIES);
        if (maxTriesOpt.isPresent()) {
          maxTries = maxTriesOpt.get();
          logger.info("If a transcription error occurs, it will be retried for {} times total.", maxTries);
        } else {
          logger.info("If a transcription error occurs, it will NOT be retried");
        }
      }
    }
    schedule();
  }

  @Override
  public String getJobGroup() {
    return JOB_GROUP;
  }

  @Override
  public String getJobName() {
    return JOB_NAME;
  }

  @Override
  public String getTriggerGroupName() {
    return TRIGGER_GROUP;
  }

  @Override
  public String getTriggerName() {
    return TRIGGER_NAME;
  }

  /* Used by unit tests */
  public void setWfUtil(Workflows wfUtil) {
    this.wfUtil = wfUtil;
  }

  @Override
  public void scan() {
    logger.debug("Waking up...");

    try {
      handleTranscriptionInProgress();

      expireOldTranscriptionNotDone();

      handleTranscriptionFinished();

    } catch (Exception e) {
      logger.warn("Could not read/update speech to text database.", e);
    }
  }

  /**
   * Handle speech to text controls that have status 'in progress'. Checks the associated job status, which may still be
   * in progress, failed, or completed.
   *
   * @throws SpeechToTextAsyncException
   */
  void handleTranscriptionInProgress() throws SpeechToTextAsyncException {
    // Get all STT with transcription in progress status
    List<SpeechToTextControl> stts = database.findByStatus(SpeechToTextControl.Status.InProgress);
    for (SpeechToTextControl stt : stts) {
      try {
        // Make sure we get the updated job from db. Note that if job was deleted, the stt control should also have been
        // deleted because of the cascade delete policy.
        Job job = getServiceRegistry().getJob(stt.getJob().getId());
        // If job failed, set stt control status to error
        if (job.getStatus() == Job.Status.FAILED) {
          database.updateStatusByJob(SpeechToTextControl.Status.TranscriptionError, JpaJob.from(job));
        } else if (job.getStatus() == Job.Status.FINISHED) {
          // If job finished ok, update state accordingly
          database.updateStatusByJob(SpeechToTextControl.Status.TranscriptionDone, JpaJob.from(job));
        }
      } catch (Exception e) {
        logger.warn("Exception when handling in progress speech to text for media package: {}, job: {}",
                stt.getMediaPackageId(), stt.getJob(), e);
      }
    }
  }

  /**
   * Transition tracked STTs after the configured interval: if still in progress, they are changed to error; if in
   * 'workflow started' state, they are changed to done.
   */
  void expireOldTranscriptionNotDone() {
    Date nowMinusInterval = Date.from(Instant.now().minusMillis(abandonAfterMs));

    try {
      // STTs that are stuck in progress for a long time will be transitioned to error so that they will be retried.
      database.transitionStatusByDate(SpeechToTextControl.Status.TranscriptionError, nowMinusInterval,
              SpeechToTextControl.Status.InProgress);

      // STTs that had a workflow to attach subtitles started are transitioned to done because afaik everything worked
      // as supposed to
      database.transitionStatusByDate(SpeechToTextControl.Status.Done, nowMinusInterval,
              SpeechToTextControl.Status.WorkflowInProgress);

    } catch (Exception e) {
      logger.warn("Exception when expiring old speech to text controls", e);
    }
  }

  /**
   * Apply next action after subtitle generation finishes if Whisper is running asynchronously. If all jobs completed
   * successfully, start a workflow to attach subtitles for all STTs that have finished (one per workflow id that
   * created them). If one of the jobs had an error, start a workflow to retry subtitle generation if configured for it.
   *
   * @throws SpeechToTextAsyncException
   */
  void handleTranscriptionFinished() throws SpeechToTextAsyncException {
    if (attachWorkflowDef == null) {
      logger.info("Workflow to attach subtitles not configured. Skipping.");
      return;
    }

    // Get list of all workflows that have STT jobs finished (ok and in error)
    List<Long> wfIds = database.findDistinctWorkflowIdByStatus(SpeechToTextControl.Status.TranscriptionDone,
            SpeechToTextControl.Status.TranscriptionError);
    String mpId = "unknown";
    for (long wfId : wfIds) {
      // From those, double-check if ALL jobs were finished with no errors.
      try {
        List<SpeechToTextControl> stts = database.findByWorkflowId(wfId);
        if (stts.isEmpty()) {
          continue;
        }
        mpId = stts.get(0).getMediaPackageId();

        boolean allFinished = stts.stream().noneMatch(stt -> SpeechToTextControl.Status.InProgress == stt.getStatus());
        if (!allFinished) {
          logger.debug(
                  "Will not process media package {} workflow {} this time because not all STT jobs have finished.",
                  mpId, wfId);
          continue;
        }

        // Get list of all jobs pointed by the stt controls (started by the current workflow)
        List<JpaJob> jobs = stts.stream().map(sst -> sst.getJob()).collect(Collectors.toList());

        // Check if media package still exists
        if (!assetManager.snapshotExists(mpId)) {
          // Are there any workflows for it?
          if (!workflowService.getWorkflowInstancesByMediaPackage(mpId).isEmpty()) {
            // Media package exists, but probably has not been archived yet; will try next time
            logger.info("Media package {} has not been archived yet.", mpId);
            continue;
          }
          logger.info("Media package {} does not exist anymore so marking STTs as canceled.", mpId);
          database.updateStatusByJob(SpeechToTextControl.Status.Canceled, jobs.toArray(new JpaJob[0]));
          continue;
        }

        // All successful?
        boolean doAttach = stts.stream()
                .allMatch(stt -> SpeechToTextControl.Status.TranscriptionDone == stt.getStatus());

        if (doAttach) {
          // Start workflow to attach subtitles to that media package
          WorkflowInstance wfInstance = startWorkflow(mpId, attachWorkflowDef, jobs);
          if (wfInstance == null) {
            logger.warn("Could not start workflow to attach subtitles for mp {}", mpId);
            // Will try again in the next run; there may be another workflow in progress for that media package
            continue;
          }
          logger.info("Workflow to attach subtitles started for mp: {}, wf instance id: {}", mpId, wfInstance.getId());
          // If success, update status in database
          database.updateStatusByJob(SpeechToTextControl.Status.WorkflowInProgress, jobs.toArray(new JpaJob[0]));
        } else {
          // All jobs finished, but at least one in error

          // Is retry configured?
          if (retryWorkflowDef != null && maxTries > 1) {
            // Get all workflows that started subtitle generation for this media package
            List<Long> mpWfs = database.findDistinctWorkflowIdByMediaPackageId(mpId);
            // Has maximum already be reached? Check the number of workflows that started subtitle generation
            if (mpWfs.size() < maxTries) {
              // Start workflow to retry transcript generation

              WorkflowInstance wfInstance = startWorkflow(mpId, retryWorkflowDef, new ArrayList<JpaJob>());
              if (wfInstance == null) {
                logger.warn("Could not start workflow to retry subtitle generation for mp {}", mpId);
                // Will try again in the next run; there may be another workflow in progress for that media package
                continue;
              }
              logger.info("Workflow to retry subtitle generation started for mp: {}, wf instance id: {}", mpId,
                      wfInstance.getId());
            } else {
              logger.info(
                      "Media package {} already has {} workflows that started subtitle generation and "
                              + "max tries is {} so we won't retry automatically anymore.",
                      mpId, mpWfs.size(), maxTries);
            }
          }

          // Update status in database to Canceled so that next time we don't check the same ones
          database.updateStatusByJob(SpeechToTextControl.Status.Canceled, jobs.toArray(new JpaJob[0]));
        }
      } catch (Exception e) {
        logger.warn("Error when starting workflow to attach subtitles to media package: {}", mpId, e);
      }
    }
  }

  @Override
  public String getScannerName() {
    return SCANNER_NAME;
  }

  /**
   * Starts a workflow on the latest version of the media package.
   *
   * @param mpId
   *          The media package id
   * @param wfDefinitionId
   *          The workflow definition id
   * @param jobs
   *          The list of jobs associated to the STT controls
   * @return Workflow instance started
   */
  WorkflowInstance startWorkflow(String mpId, String wfDefinitionId, List<JpaJob> jobs) {
    try {
      WorkflowDefinition wfDef = workflowService.getWorkflowDefinitionById(wfDefinitionId);
      Workflows workflows = wfUtil != null ? wfUtil : new Workflows(assetManager, workflowService);
      Set<String> mpIds = Collections.singleton(mpId);
      Map<String, String> wfConfig = new HashMap<String, String>();
      if (jobs.size() > 0) {
        wfConfig.put(JOBS_WORKFLOW_CONFIGURATION,
                jobs.stream().map(j -> String.valueOf(j.getId())).collect(Collectors.joining(",")));
      }
      List<WorkflowInstance> wfList = workflows
              .applyWorkflowToLatestVersion(mpIds, ConfiguredWorkflow.workflow(wfDef, wfConfig));
      WorkflowInstance wf = wfList.size() > 0 ? wfList.get(0) : null;
      if (wf != null) {
        logger.info("Workflow {} started for mp {}: {}", attachWorkflowDef, mpId, wf.getId());
      } else {
        logger.info("Workflow {} NOT started for mp {}", attachWorkflowDef, mpId);
      }
      return wf;
    } catch (NotFoundException | WorkflowDatabaseException e) {
      logger.warn("Could not get workflow definition: {}", wfDefinitionId);
    }

    return null;
  }

  /** Quartz job for automatically starting workflows */
  public static class Runner extends TypedQuartzJob<AbstractScanner> {
    private static final NeedleEye eye = new NeedleEye();

    public Runner() {
      super(Optional.of(eye));
    }

    @Override
    protected void execute(final AbstractScanner parameters, JobExecutionContext ctx) {
      logger.debug("Starting " + parameters.getScannerName() + " job.");

      // iterate all organizations
      for (final Organization org : parameters.getOrganizationDirectoryService().getOrganizations()) {
        // set the organization on the current thread
        parameters.getAdminContextFor(org.getId()).runInContext(parameters::scan);
      }

      logger.debug("Finished " + parameters.getScannerName() + " job.");
    }
  }

  @Reference
  public void setDatabase(SpeechToTextDatabase database) {
    this.database = database;
  }

  @Reference
  public void setAssetManager(AssetManager service) {
    this.assetManager = service;
  }

  @Reference
  public void setWorkflowService(WorkflowService workflowService) {
    logger.debug("Setting workflow service!");
    this.workflowService = workflowService;
  }

  @Override
  @Reference
  public void bindSecurityService(SecurityService service) {
    super.bindSecurityService(service);
  }

  @Override
  @Reference
  public void bindOrganizationDirectoryService(OrganizationDirectoryService service) {
    super.bindOrganizationDirectoryService(service);
  }

  @Override
  @Reference
  public void bindServiceRegistry(ServiceRegistry service) {
    super.bindServiceRegistry(service);
  }
}