From 3b2df6a178698c59df8ef2ca63b42f8d9e8a81cf Mon Sep 17 00:00:00 2001 From: John Zhuge Date: Wed, 5 Jun 2024 23:12:03 -0700 Subject: [PATCH] Set env var BD_USER --- .../launchers/impl/TitusAgentLauncherImpl.java | 15 ++++++++++----- .../impl/TitusAgentLauncherImplSpec.groovy | 7 ++++--- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/genie-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java b/genie-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java index 7784e78dbd..62a8fdd7ae 100644 --- a/genie-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java +++ b/genie-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java @@ -81,6 +81,7 @@ public class TitusAgentLauncherImpl implements AgentLauncher { static final int MEGABYTE_TO_MEGABIT = 8; + private static final String BD_USER_ENV_VAR = "BD_USER"; private static final String GENIE_USER_ATTR = "genie.user"; private static final String GENIE_SOURCE_HOST_ATTR = "genie.sourceHost"; private static final String GENIE_ENDPOINT_ATTR = "genie.endpoint"; @@ -246,6 +247,8 @@ public Health health() { private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) throws AgentLaunchException { final String jobId = resolvedJob.getJobSpecification().getJob().getId(); + final String jobUser = resolvedJob.getJobMetadata().getUser(); + // Map placeholders in entry point template to their values final Map placeholdersMap = Map.of( TitusAgentLauncherProperties.JOB_ID_PLACEHOLDER, @@ -267,7 +270,7 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) thr ); final Duration runtimeLimit = this.titusAgentLauncherProperties.getRuntimeLimit(); - final Map jobAttributes = this.createJobAttributes(jobId, resolvedJob); + final Map jobAttributes = this.createJobAttributes(jobId, jobUser); final TitusBatchJobRequest.TitusBatchJobRequestBuilder requestBuilder = TitusBatchJobRequest.builder() .owner( @@ -299,7 +302,7 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) thr .image(this.getTitusImage(resolvedJob)) .entryPoint(entryPoint) .command(command) - .env(this.createJobEnvironment()) + .env(this.createJobEnvironment(jobUser)) .attributes( this.binder .bind( @@ -408,9 +411,9 @@ private DataSize getDataSizeProperty(final String propertyKey, final DataSize de } } - private Map createJobAttributes(final String jobId, final ResolvedJob resolvedJob) { + private Map createJobAttributes(final String jobId, final String jobUser) { final Map jobAttributes = new HashMap<>(); - jobAttributes.put(GENIE_USER_ATTR, resolvedJob.getJobMetadata().getUser()); + jobAttributes.put(GENIE_USER_ATTR, jobUser); jobAttributes.put(GENIE_SOURCE_HOST_ATTR, this.genieHostInfo.getHostname()); jobAttributes.put(GENIE_ENDPOINT_ATTR, this.titusAgentLauncherProperties.getGenieServerHost()); jobAttributes.put(GENIE_JOB_ID_ATTR, jobId); @@ -425,7 +428,7 @@ private Map createJobAttributes(final String jobId, final Resolv return jobAttributes; } - private Map createJobEnvironment() { + private Map createJobEnvironment(final String jobUser) { final Map jobEnvironment = this.binder .bind( TitusAgentLauncherProperties.ADDITIONAL_ENVIRONMENT_PROPERTY, @@ -438,6 +441,8 @@ private Map createJobEnvironment() { jobEnvironment.putAll(this.tracePropagator.injectForAgent(currentSpan.context())); } + jobEnvironment.putIfAbsent(BD_USER_ENV_VAR, jobUser); + return jobEnvironment; } diff --git a/genie-web/src/test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy b/genie-web/src/test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy index 9edf89d5fc..15c0109936 100644 --- a/genie-web/src/test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy +++ b/genie-web/src/test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy @@ -91,6 +91,7 @@ class TitusAgentLauncherImplSpec extends Specification { this.requestedMemory = 1024L this.launcherProperties = new TitusAgentLauncherProperties() this.launcherProperties.setEndpoint(URI.create(TITUS_ENDPOINT_PREFIX)) + this.launcherProperties.setAdditionalEnvironment(["BD_USER": USER]) this.job = Mock(JobSpecification.ExecutionResource) { getId() >> JOB_ID @@ -508,7 +509,7 @@ class TitusAgentLauncherImplSpec extends Specification { 1 * this.cache.put(JOB_ID, TITUS_JOB_ID) launcherExt.isPresent() requestCapture != null - requestCapture.getContainer().getEnv().isEmpty() + requestCapture.getContainer().getEnv() == ["BD_USER": USER] when: def prop1Key = "${UUID.randomUUID()}.${UUID.randomUUID()}.${UUID.randomUUID()}".toString() @@ -529,7 +530,7 @@ class TitusAgentLauncherImplSpec extends Specification { 1 * this.cache.put(JOB_ID, TITUS_JOB_ID) launcherExt.isPresent() requestCapture != null - requestCapture.getContainer().getEnv().size() == 1 + requestCapture.getContainer().getEnv().size() == 2 requestCapture.getContainer().getEnv().get(prop1Key) == prop1Value when: @@ -551,7 +552,7 @@ class TitusAgentLauncherImplSpec extends Specification { 1 * this.cache.put(JOB_ID, TITUS_JOB_ID) launcherExt.isPresent() requestCapture != null - requestCapture.getContainer().getEnv().size() == 2 + requestCapture.getContainer().getEnv().size() == 3 requestCapture.getContainer().getEnv().get(prop1Key) == prop1Value requestCapture.getContainer().getEnv().get(prop2Key) == prop2Value }