Issue #3534: Tag compute resources according to the tool's metadata -…

… initial API part
epam · May 22, 2024 · b45f2b9 · b45f2b9
1 parent 6a0ecce
commit b45f2b9
Show file tree

Hide file tree

Showing 20 changed files with 231 additions and 52 deletions.
diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/CloudFacade.java b/api/src/main/java/com/epam/pipeline/manager/cloud/CloudFacade.java
@@ -36,7 +36,8 @@
 import java.util.Optional;
 
 public interface CloudFacade {
-    RunInstance scaleUpNode(Long runId, RunInstance instance, Map<String, String> runtimeParameters);
+    RunInstance scaleUpNode(Long runId, RunInstance instance, Map<String, String> runtimeParameters,
+                            Map<String, String> customTags);
 
     RunInstance scaleUpPoolNode(String nodeId, NodePool node);
 
@@ -48,9 +49,9 @@ public interface CloudFacade {
 
     boolean isNodeExpired(Long runId);
 
-    boolean reassignNode(Long oldId, Long newId);
+    boolean reassignNode(Long oldId, Long newId, Map<String, String> customTags);
 
-    boolean reassignPoolNode(String nodeLabel, Long newId);
+    boolean reassignPoolNode(String nodeLabel, Long newId, Map<String, String> customTags);
 
     /**
      * Fills in provider related data for running instance associated with run,

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/CloudFacadeImpl.java b/api/src/main/java/com/epam/pipeline/manager/cloud/CloudFacadeImpl.java
@@ -87,10 +87,10 @@ public CloudFacadeImpl(final MessageHelper messageHelper,
 
     @Override
     public RunInstance scaleUpNode(final Long runId, final RunInstance instance,
-                                   final Map<String, String> runtimeParameters) {
+                                   final Map<String, String> runtimeParameters, final Map<String, String> customTags) {
         final AbstractCloudRegion region = regionManager.loadOrDefault(instance.getCloudRegionId());
         final RunInstance scaledUpInstance = getInstanceService(region)
-                .scaleUpNode(region, runId, instance, runtimeParameters);
+                .scaleUpNode(region, runId, instance, runtimeParameters, customTags);
         kubernetesManager.createNodeService(scaledUpInstance);
         return scaledUpInstance;
     }
@@ -145,15 +145,15 @@ public boolean isNodeExpired(final Long runId) {
     }
 
     @Override
-    public boolean reassignNode(final Long oldId, final Long newId) {
+    public boolean reassignNode(final Long oldId, final Long newId, final Map<String, String> customTags) {
         final AbstractCloudRegion region = getRegionByRunId(oldId);
-        return getInstanceService(region).reassignNode(region, oldId, newId);
+        return getInstanceService(region).reassignNode(region, oldId, newId, customTags);
     }
 
     @Override
-    public boolean reassignPoolNode(final String nodeLabel, final Long newId) {
+    public boolean reassignPoolNode(final String nodeLabel, final Long newId, final Map<String, String> customTags) {
         final AbstractCloudRegion region = loadRegionFromNodeLabels(nodeLabel);
-        return getInstanceService(region).reassignPoolNode(region, nodeLabel, newId);
+        return getInstanceService(region).reassignPoolNode(region, nodeLabel, newId, customTags);
     }
 
     @Override

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/CloudInstanceService.java b/api/src/main/java/com/epam/pipeline/manager/cloud/CloudInstanceService.java
@@ -59,7 +59,8 @@ default Map<String, String> getPoolLabels(final NodePool pool) {
      * @param instance
      * @return
      */
-    RunInstance scaleUpNode(T region, Long runId, RunInstance instance, Map<String, String> runtimeParameters);
+    RunInstance scaleUpNode(T region, Long runId, RunInstance instance, Map<String, String> runtimeParameters,
+                            Map<String, String> customTags);
 
     RunInstance scaleUpPoolNode(T region, String nodeId, NodePool node);
 
@@ -140,8 +141,8 @@ default Map<String, String> getPoolLabels(final NodePool pool) {
      * @param newId
      * @return {@code true} if operation was successful
      */
-    boolean reassignNode(T region, Long oldId, Long newId);
-    boolean reassignPoolNode(T region, String nodeLabel, Long newId);
+    boolean reassignNode(T region, Long oldId, Long newId, Map<String, String> customTags);
+    boolean reassignPoolNode(T region, String nodeLabel, Long newId, Map<String, String> customTags);
 
     /**
      * Builds environment variables required for running a container in provided region

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/aws/AWSInstanceService.java b/api/src/main/java/com/epam/pipeline/manager/cloud/aws/AWSInstanceService.java
@@ -112,9 +112,10 @@ public AWSInstanceService(final EC2Helper ec2Helper,
     public RunInstance scaleUpNode(final AwsRegion region,
                                    final Long runId,
                                    final RunInstance instance,
-                                   final Map<String, String> runtimeParameters) {
+                                   final Map<String, String> runtimeParameters,
+                                   final Map<String, String> customTags) {
         final String command = buildNodeUpCommand(region, String.valueOf(runId), instance,
-                Collections.emptyMap(), runtimeParameters);
+                Collections.emptyMap(), runtimeParameters, customTags);
         return instanceService.runNodeUpScript(cmdExecutor, runId, instance, command, buildScriptEnvVars(region));
     }
 
@@ -124,7 +125,7 @@ public RunInstance scaleUpPoolNode(final AwsRegion region,
                                        final NodePool node) {
         final RunInstance instance = node.toRunInstance();
         final String command = buildNodeUpCommand(region, nodeIdLabel, instance, getPoolLabels(node),
-                Collections.emptyMap());
+                Collections.emptyMap(), Collections.emptyMap());
         return instanceService.runNodeUpScript(cmdExecutor, null, instance, command, buildScriptEnvVars(region));
     }
 
@@ -141,16 +142,18 @@ public void scaleDownPoolNode(final AwsRegion region, final String nodeLabel) {
     }
 
     @Override
-    public boolean reassignNode(final AwsRegion region, final Long oldId, final Long newId) {
+    public boolean reassignNode(final AwsRegion region, final Long oldId, final Long newId,
+                                final Map<String, String> customTags) {
         final String command = commandService.buildNodeReassignCommand(
-                nodeReassignScript, oldId, newId, getProvider().name());
+                nodeReassignScript, oldId, newId, getProvider().name(), customTags);
         return instanceService.runNodeReassignScript(cmdExecutor, command, oldId, newId, buildScriptEnvVars(region));
     }
 
     @Override
-    public boolean reassignPoolNode(final AwsRegion region, final String nodeLabel, final Long newId) {
+    public boolean reassignPoolNode(final AwsRegion region, final String nodeLabel, final Long newId,
+                                    final Map<String, String> customTags) {
         final String command = commandService.buildNodeReassignCommand(
-                nodeReassignScript, nodeLabel, String.valueOf(newId), getProvider().name());
+                nodeReassignScript, nodeLabel, String.valueOf(newId), getProvider().name(), customTags);
         return instanceService.runNodeReassignScript(cmdExecutor, command, nodeLabel,
                 String.valueOf(newId), buildScriptEnvVars(region));
     }
@@ -363,10 +366,12 @@ private String buildNodeUpCommand(final AwsRegion region,
                                       final String nodeLabel,
                                       final RunInstance instance,
                                       final Map<String, String> labels,
-                                      final Map<String, String> runtimeParameters) {
+                                      final Map<String, String> runtimeParameters,
+                                      final Map<String, String> customTags) {
         final NodeUpCommand.NodeUpCommandBuilder commandBuilder = commandService
                 .buildNodeUpCommand(nodeUpScript, region, nodeLabel, instance, getProviderName(), runtimeParameters)
-                .sshKey(region.getSshKeyName());
+                .sshKey(region.getSshKeyName())
+                .customTags(customTags);
 
         if (StringUtils.isNotBlank(region.getKmsKeyId())) {
             commandBuilder.encryptionKey(region.getKmsKeyId());

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/azure/AzureInstanceService.java b/api/src/main/java/com/epam/pipeline/manager/cloud/azure/AzureInstanceService.java
@@ -103,7 +103,8 @@ public AzureInstanceService(final CommonCloudInstanceService instanceService,
     public RunInstance scaleUpNode(final AzureRegion region,
                                    final Long runId,
                                    final RunInstance instance,
-                                   final Map<String, String> runtimeParameters) {
+                                   final Map<String, String> runtimeParameters,
+                                   final Map<String, String> customTags) {
         final String command = buildNodeUpCommand(region, String.valueOf(runId), instance, Collections.emptyMap(),
                 runtimeParameters);
         return instanceService.runNodeUpScript(cmdExecutor, runId, instance, command, buildScriptAzureEnvVars(region));
@@ -137,17 +138,20 @@ public void scaleDownPoolNode(final AzureRegion region, final String nodeLabel)
     }
 
     @Override
-    public boolean reassignNode(final AzureRegion region, final Long oldId, final Long newId) {
+    public boolean reassignNode(final AzureRegion region, final Long oldId, final Long newId,
+                                final Map<String, String> customTags) {
         final String command = commandService.buildNodeReassignCommand(
-                nodeReassignScript, oldId, newId, getProvider().name());
+                nodeReassignScript, oldId, newId, getProvider().name(), customTags);
         return instanceService.runNodeReassignScript(cmdExecutor, command, oldId, newId,
                 buildScriptAzureEnvVars(region));
     }
 
     @Override
-    public boolean reassignPoolNode(final AzureRegion region, final String nodeLabel, final Long newId) {
+    public boolean reassignPoolNode(final AzureRegion region, final String nodeLabel, final Long newId,
+                                    final Map<String, String> customTags) {
         final String command = commandService.
-            buildNodeReassignCommand(nodeReassignScript, nodeLabel, String.valueOf(newId), getProvider().name());
+            buildNodeReassignCommand(nodeReassignScript, nodeLabel, String.valueOf(newId), getProvider().name(),
+                    customTags);
         return instanceService.runNodeReassignScript(cmdExecutor, command, nodeLabel,
                                                      String.valueOf(newId), buildScriptAzureEnvVars(region));
     }

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/commands/ClusterCommandService.java b/api/src/main/java/com/epam/pipeline/manager/cloud/commands/ClusterCommandService.java
@@ -85,20 +85,24 @@ public String buildNodeDownCommand(final String nodeDownScript,
     public String buildNodeReassignCommand(final String nodeReassignScript,
                                            final Long oldId,
                                            final Long newId,
-                                           final String cloud) {
-        return buildNodeReassignCommand(nodeReassignScript, String.valueOf(oldId), String.valueOf(newId), cloud);
+                                           final String cloud,
+                                           final Map<String, String> customTags) {
+        return buildNodeReassignCommand(nodeReassignScript, String.valueOf(oldId), String.valueOf(newId), cloud,
+                customTags);
     }
 
     public String buildNodeReassignCommand(final String nodeReassignScript,
                                            final String oldId,
                                            final String newId,
-                                           final String cloud) {
+                                           final String cloud,
+                                           final Map<String, String> customTags) {
         return ReassignCommand.builder()
                 .executable(AbstractClusterCommand.EXECUTABLE)
                 .script(nodeReassignScript)
                 .oldRunId(oldId)
                 .newRunId(newId)
                 .cloud(cloud)
+                .customTags(customTags)
                 .build()
                 .getCommand();
     }

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/commands/NodeUpCommand.java b/api/src/main/java/com/epam/pipeline/manager/cloud/commands/NodeUpCommand.java
@@ -51,6 +51,7 @@ public class NodeUpCommand extends AbstractClusterCommand {
     private final Map<String, String> additionalLabels;
     private final Set<String> prePulledImages;
     private final Map<String, String> runtimeParameters;
+    private final Map<String, String> customTags;
 
     @Override
     protected List<String> buildCommandArguments() {
@@ -105,6 +106,10 @@ protected List<String> buildCommandArguments() {
             commands.add(argName);
             commands.add(argValue);
         });
+        MapUtils.emptyIfNull(customTags).forEach((key, value) -> {
+            commands.add("--custom_tags");
+            commands.add(key + "=" + value);
+        });
         return commands;
     }
 }
diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/commands/ReassignCommand.java b/api/src/main/java/com/epam/pipeline/manager/cloud/commands/ReassignCommand.java
@@ -17,9 +17,11 @@
 package com.epam.pipeline.manager.cloud.commands;
 
 import lombok.Builder;
+import org.apache.commons.collections4.MapUtils;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 @Builder
 public class ReassignCommand extends AbstractClusterCommand {
@@ -31,6 +33,7 @@ public class ReassignCommand extends AbstractClusterCommand {
     private final String oldRunId;
     private final String newRunId;
     private final String cloud;
+    private final Map<String, String> customTags;
 
     @Override
     protected List<String> buildCommandArguments() {
@@ -43,6 +46,10 @@ protected List<String> buildCommandArguments() {
         commands.add(newRunId);
         commands.add(CLOUD_PARAMETER);
         commands.add(cloud);
+        MapUtils.emptyIfNull(customTags).forEach((key, value) -> {
+            commands.add("--custom_tags");
+            commands.add(key + "=" + value);
+        });
         return commands;
     }
 }
diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/gcp/GCPInstanceService.java b/api/src/main/java/com/epam/pipeline/manager/cloud/gcp/GCPInstanceService.java
@@ -97,7 +97,7 @@ public GCPInstanceService(final ClusterCommandService commandService,
 
     @Override
     public RunInstance scaleUpNode(final GCPRegion region, final Long runId, final RunInstance instance,
-                                   final Map<String, String> runtimeParameters) {
+                                   final Map<String, String> runtimeParameters, final Map<String, String> customTags) {
         final String command = buildNodeUpCommand(region, String.valueOf(runId), instance, Collections.emptyMap(),
                 runtimeParameters);
         return instanceService.runNodeUpScript(cmdExecutor, runId, instance, command, buildScriptGCPEnvVars(region));
@@ -126,17 +126,20 @@ public void scaleDownPoolNode(final GCPRegion region, final String nodeLabel) {
     }
 
     @Override
-    public boolean reassignNode(final GCPRegion region, final Long oldId, final Long newId) {
+    public boolean reassignNode(final GCPRegion region, final Long oldId, final Long newId,
+                                final Map<String, String> customTags) {
         final String command = commandService.buildNodeReassignCommand(
-                nodeReassignScript, oldId, newId, getProviderName());
+                nodeReassignScript, oldId, newId, getProviderName(), customTags);
         return instanceService.runNodeReassignScript(cmdExecutor, command, oldId, newId,
                 buildScriptGCPEnvVars(region));
     }
 
     @Override
-    public boolean reassignPoolNode(final GCPRegion region, final String nodeLabel, final Long newId) {
+    public boolean reassignPoolNode(final GCPRegion region, final String nodeLabel, final Long newId,
+                                    final Map<String, String> customTags) {
         final String command = commandService
-            .buildNodeReassignCommand(nodeReassignScript, nodeLabel, String.valueOf(newId), getProviderName());
+            .buildNodeReassignCommand(nodeReassignScript, nodeLabel, String.valueOf(newId), getProviderName(),
+                    customTags);
         return instanceService.runNodeReassignScript(cmdExecutor, command, nodeLabel, String.valueOf(newId),
                                                      buildScriptGCPEnvVars(region));
     }

diff --git a/api/src/main/java/com/epam/pipeline/manager/cloud/local/LocalInstanceService.java b/api/src/main/java/com/epam/pipeline/manager/cloud/local/LocalInstanceService.java
@@ -58,7 +58,8 @@ public CloudProvider getProvider() {
     public RunInstance scaleUpNode(final LocalRegion region,
                                    final Long runId,
                                    final RunInstance instance,
-                                   final Map<String, String> runtimeParameters) {
+                                   final Map<String, String> runtimeParameters,
+                                   final Map<String, String> customTags) {
         throw new UnsupportedOperationException(
                 messageHelper.getMessage(MessageConstants.ERROR_SCALING_LOCAL_CLUSTER));
     }
@@ -128,12 +129,14 @@ public RunInstance describeAliveInstance(LocalRegion region, String nodeLabel, R
     }
 
     @Override
-    public boolean reassignNode(final LocalRegion region, final Long oldId, final Long newId) {
+    public boolean reassignNode(final LocalRegion region, final Long oldId, final Long newId,
+                                final Map<String, String> customTags) {
         return reassignKubeNode(String.valueOf(oldId), String.valueOf(newId));
     }
 
     @Override
-    public boolean reassignPoolNode(final LocalRegion region, final String nodeLabel, final Long newId) {
+    public boolean reassignPoolNode(final LocalRegion region, final String nodeLabel, final Long newId,
+                                    final Map<String, String> customTags) {
         return reassignKubeNode(nodeLabel, String.valueOf(newId));
     }
 

diff --git a/api/src/main/java/com/epam/pipeline/manager/cluster/autoscale/AutoscaleManager.java b/api/src/main/java/com/epam/pipeline/manager/cluster/autoscale/AutoscaleManager.java
@@ -520,7 +520,8 @@ private void createNodeForRun(List<CompletableFuture<Void>> tasks, String runId,
                 //save required instance
                 pipelineRunManager.updateRunInstance(longId, requiredInstance.getInstance());
                 RunInstance instance = cloudFacade
-                        .scaleUpNode(longId, requiredInstance.getInstance(), requiredInstance.getRuntimeParameters());
+                        .scaleUpNode(longId, requiredInstance.getInstance(), requiredInstance.getRuntimeParameters(),
+                                requiredInstance.getCustomTags());
                 //save instance ID and IP
                 pipelineRunManager.updateRunInstance(longId, instance);
                 pipelineRunManager.updateRunInstanceStartDate(longId, DateUtils.nowUTC());
@@ -607,6 +608,7 @@ public InstanceRequest getNewRunInstance(final PipelineRun run) throws GitClient
             instanceRequest.setInstance(instance);
             instanceRequest.setRequestedImage(run.getActualDockerImage());
             instanceRequest.setRuntimeParameters(buildRuntimeParameters(run));
+            instanceRequest.setCustomTags(autoscalerService.buildCustomInstanceTags(run));
             return instanceRequest;
         }
 

diff --git a/api/src/main/java/com/epam/pipeline/manager/cluster/autoscale/AutoscalerService.java b/api/src/main/java/com/epam/pipeline/manager/cluster/autoscale/AutoscalerService.java
@@ -20,14 +20,16 @@
 import com.epam.pipeline.entity.cluster.pool.NodePool;
 import com.epam.pipeline.entity.cluster.pool.RunningInstance;
 import com.epam.pipeline.entity.configuration.PipelineConfiguration;
+import com.epam.pipeline.entity.pipeline.PipelineRun;
 import com.epam.pipeline.entity.pipeline.RunInstance;
 import io.fabric8.kubernetes.client.KubernetesClient;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /**
- * Provides helper nethods for {@link AutoscaleManager}
+ * Provides helper methods for {@link AutoscaleManager}
  */
 public interface AutoscalerService {
 
@@ -39,4 +41,5 @@ public interface AutoscalerService {
     void adjustRunPrices(long runId, List<InstanceDisk> disks);
     Optional<NodePool> findPool(String nodeLabel, KubernetesClient client);
     void registerDisks(Long runId, RunInstance instance);
+    Map<String, String> buildCustomInstanceTags(PipelineRun run);
 }