support AWS resource cleanup when walltime hits

Author: lahirujayathilake

airavata-api/src/main/java/org/apache/airavata/helix/impl/participant/GlobalParticipant.java

@@ -52,6 +52,7 @@ public class GlobalParticipant extends HelixParticipant<AbstractTask> {
         "org.apache.airavata.helix.impl.task.aws.CreateEC2InstanceTask",
         "org.apache.airavata.helix.impl.task.aws.NoOperationTask",
         "org.apache.airavata.helix.impl.task.aws.AWSJobSubmissionTask",
+        "org.apache.airavata.helix.impl.task.aws.AWSCompletingTask",
     };
 
     @SuppressWarnings("WeakerAccess")

airavata-api/src/main/java/org/apache/airavata/helix/impl/task/AWSTaskFactory.java

@@ -19,10 +19,10 @@
 */
 package org.apache.airavata.helix.impl.task;
 
+import org.apache.airavata.helix.impl.task.aws.AWSCompletingTask;
 import org.apache.airavata.helix.impl.task.aws.AWSJobSubmissionTask;
 import org.apache.airavata.helix.impl.task.aws.CreateEC2InstanceTask;
 import org.apache.airavata.helix.impl.task.aws.NoOperationTask;
-import org.apache.airavata.helix.impl.task.completing.CompletingTask;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -63,7 +63,7 @@ public AiravataTask createJobVerificationTask(String processId) {
 
     @Override
     public AiravataTask createCompletingTask(String processId) {
-        return new CompletingTask();
+        return new AWSCompletingTask();
     }
 
     @Override

airavata-api/src/main/java/org/apache/airavata/helix/impl/task/aws/AWSCompletingTask.java

@@ -0,0 +1,49 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache 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
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.apache.airavata.helix.impl.task.aws;
+
+import org.apache.airavata.helix.impl.task.AiravataTask;
+import org.apache.airavata.helix.impl.task.TaskContext;
+import org.apache.airavata.helix.impl.task.aws.utils.AWSTaskUtil;
+import org.apache.airavata.helix.task.api.TaskHelper;
+import org.apache.airavata.helix.task.api.annotation.TaskDef;
+import org.apache.airavata.model.status.ProcessState;
+import org.apache.helix.task.TaskResult;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@TaskDef(name = "AWS_COMPLETING_TASK")
+public class AWSCompletingTask extends AiravataTask {
+
+    private static final Logger logger = LoggerFactory.getLogger(AWSCompletingTask.class);
+
+    @Override
+    public TaskResult onRun(TaskHelper helper, TaskContext taskContext) {
+        logger.info("Starting completing task for task {}, experiment id {}", getTaskId(), getExperimentId());
+        logger.info("Process {} successfully completed", getProcessId());
+        saveAndPublishProcessStatus(ProcessState.COMPLETED);
+        cleanup();
+        AWSTaskUtil.terminateEC2Instance(getTaskContext(), getGatewayId());
+        return onSuccess("Process " + getProcessId() + " successfully completed");
+    }
+
+    @Override
+    public void onCancel(TaskContext taskContext) {
+    }
+}

airavata-api/src/main/java/org/apache/airavata/helix/impl/task/submission/config/GroovyMapBuilder.java

@@ -80,6 +80,7 @@ public GroovyMapData build() throws Exception {
         mapData.setWorkingDirectory(taskContext.getWorkingDir());
         mapData.setTaskId(taskContext.getTaskId());
         mapData.setExperimentDataDir(taskContext.getProcessModel().getExperimentDataDir());
+        mapData.setExperimentId(taskContext.getExperimentId());
 
         SimpleDateFormat gmtDateFormat = new SimpleDateFormat("yyyy-MM-dd+HH:mmZ");
         gmtDateFormat.setTimeZone(TimeZone.getTimeZone("EST"));
@@ -121,6 +122,7 @@ public GroovyMapData build() throws Exception {
                         ((JobSubmissionTaskModel) taskContext.getSubTaskModel());
                 if (jobSubmissionTaskModel.getWallTime() > 0) {
                     mapData.setMaxWallTime(maxWallTimeCalculator(jobSubmissionTaskModel.getWallTime()));
+                    mapData.setWallTimeInSeconds(jobSubmissionTaskModel.getWallTime() * 60);
                     // TODO fix this
                     /*if (resourceJobManager != null) {
                         if (resourceJobManager.getResourceJobManagerType().equals(ResourceJobManagerType.LSF)) {
@@ -161,6 +163,7 @@ public GroovyMapData build() throws Exception {
             // if so we ignore scheduling configuration.
             if (scheduling.getWallTimeLimit() > 0 && mapData.getMaxWallTime() == null) {
                 mapData.setMaxWallTime(maxWallTimeCalculator(scheduling.getWallTimeLimit()));
+                mapData.setWallTimeInSeconds(scheduling.getWallTimeLimit() * 60);
 
                 // TODO fix this
                 /*

airavata-api/src/main/java/org/apache/airavata/helix/impl/task/submission/config/GroovyMapData.java

@@ -107,6 +107,9 @@ public class GroovyMapData {
     @ScriptTag(name = "maxWallTime")
     private String maxWallTime;
 
+    @ScriptTag(name = "wallTimeInSeconds")
+    private Integer wallTimeInSeconds;
+
     @ScriptTag(name = "qualityOfService")
     private String qualityOfService;
 
@@ -152,6 +155,9 @@ public class GroovyMapData {
     @ScriptTag(name = "experimentDataDir")
     private String experimentDataDir;
 
+    @ScriptTag(name = "experimentId")
+    private String experimentId;
+
     @ScriptTag(name = "computeHostName")
     private String computeHostName;
 
@@ -363,6 +369,14 @@ public GroovyMapData setMaxWallTime(String maxWallTime) {
         return this;
     }
 
+    public Integer getWallTimeInSeconds() {
+        return wallTimeInSeconds;
+    }
+
+    public void setWallTimeInSeconds(Integer wallTimeInSeconds) {
+        this.wallTimeInSeconds = wallTimeInSeconds;
+    }
+
     public String getQualityOfService() {
         return qualityOfService;
     }
@@ -496,6 +510,14 @@ public void setExperimentDataDir(String experimentDataDir) {
         this.experimentDataDir = experimentDataDir;
     }
 
+    public String getExperimentId() {
+        return experimentId;
+    }
+
+    public void setExperimentId(String experimentId) {
+        this.experimentId = experimentId;
+    }
+
     public String getCurrentTime() {
         return currentTime;
     }

airavata-api/src/main/resources/templates/CLOUD_Groovy.template

@@ -1,6 +1,9 @@
 #!${shellName}
 
 # Cloud execution script generated by Apache Airavata
+# User: ${gatewayUserName}
+# Experiment ID: ${experimentId}
+# Walltime (seconds): ${wallTimeInSeconds}
 <%
     if (exports != null) for(com in exports) out.print 'export ' + com +'\n'
     if (moduleCommands != null) for(mc in moduleCommands) out.print mc +'\n'
@@ -9,6 +12,22 @@
     if (jobSubmitterCommand != null && jobSubmitterCommand != "") out.print jobSubmitterCommand + ' '
     if (executablePath != null && executablePath != "") out.print  executablePath + ' '
     if (inputs != null) for(input in inputs) out.print input + ' '
-    out.print '\n'
-    if (postJobCommands != null) for(pjc in postJobCommands) out.print pjc +'\n'
-%>
+    out.print '&\n'
+    out.print 'MAIN_JOB_PID=\\$!\n'
+%>
+
+(
+  sleep ${wallTimeInSeconds}
+
+  if ps -p \$MAIN_JOB_PID > /dev/null; then
+    echo "Walltime of ${wallTimeInSeconds} seconds exceeded. Terminating job PID \$MAIN_JOB_PID." >&2
+    pkill -P \$MAIN_JOB_PID
+    kill -9 \$MAIN_JOB_PID
+  fi
+) &
+
+WATCHDOG_PID=\$!
+wait \$MAIN_JOB_PID
+kill \$WATCHDOG_PID 2>/dev/null
+
+<% if (postJobCommands != null) for(pjc in postJobCommands) out.print pjc +'\n' %>