| /* |
| * 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 |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * 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.beam.runners.core.construction; |
| |
| import com.fasterxml.jackson.core.Base64Variants; |
| import com.fasterxml.jackson.databind.ObjectMapper; |
| import java.io.File; |
| import java.io.FileOutputStream; |
| import java.io.IOException; |
| import java.util.LinkedHashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Optional; |
| import java.util.Set; |
| import java.util.UUID; |
| import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ArtifactInformation; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Components; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.DockerPayload; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Environment; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ExternalPayload; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ProcessPayload; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.StandardArtifacts; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.StandardEnvironments; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.SplittableParDoComponents; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.StandardProtocols; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.options.PortablePipelineOptions; |
| import org.apache.beam.sdk.util.ReleaseInfo; |
| import org.apache.beam.sdk.util.ZipFiles; |
| import org.apache.beam.sdk.util.common.ReflectHelpers; |
| import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString; |
| import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.HashCode; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Files; |
| import org.checkerframework.checker.nullness.qual.Nullable; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** Utilities for interacting with portability {@link Environment environments}. */ |
| @SuppressWarnings({ |
| "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402) |
| }) |
| public class Environments { |
| private static final Logger LOG = LoggerFactory.getLogger(Environments.class); |
| |
| private static final ObjectMapper MAPPER = |
| new ObjectMapper() |
| .registerModules(ObjectMapper.findModules(ReflectHelpers.findClassLoader())); |
| public static final String ENVIRONMENT_DOCKER = "DOCKER"; |
| public static final String ENVIRONMENT_PROCESS = "PROCESS"; |
| public static final String ENVIRONMENT_EXTERNAL = "EXTERNAL"; |
| public static final String ENVIRONMENT_EMBEDDED = "EMBEDDED"; // Non Public urn for testing |
| public static final String ENVIRONMENT_LOOPBACK = "LOOPBACK"; // Non Public urn for testing |
| |
| /* For development, use the container build by the current user to ensure that the SDK harness and |
| * the SDK agree on how they should interact. This should be changed to a version-specific |
| * container during a release. |
| * |
| * See https://beam.apache.org/contribute/docker-images/ for more information on how to build a |
| * container. |
| */ |
| private static final String JAVA_SDK_HARNESS_CONTAINER_URL = |
| ReleaseInfo.getReleaseInfo().getDefaultDockerRepoRoot() |
| + "/" |
| + ReleaseInfo.getReleaseInfo().getDefaultDockerRepoPrefix() |
| + "java_sdk:" |
| + ReleaseInfo.getReleaseInfo().getSdkVersion(); |
| public static final Environment JAVA_SDK_HARNESS_ENVIRONMENT = |
| createDockerEnvironment(JAVA_SDK_HARNESS_CONTAINER_URL); |
| |
| private Environments() {} |
| |
| public static Environment createOrGetDefaultEnvironment(PortablePipelineOptions options) { |
| String type = options.getDefaultEnvironmentType(); |
| String config = options.getDefaultEnvironmentConfig(); |
| |
| Environment defaultEnvironment; |
| if (Strings.isNullOrEmpty(type)) { |
| defaultEnvironment = JAVA_SDK_HARNESS_ENVIRONMENT; |
| } else { |
| switch (type) { |
| case ENVIRONMENT_EMBEDDED: |
| defaultEnvironment = createEmbeddedEnvironment(config); |
| break; |
| case ENVIRONMENT_EXTERNAL: |
| case ENVIRONMENT_LOOPBACK: |
| defaultEnvironment = createExternalEnvironment(config); |
| break; |
| case ENVIRONMENT_PROCESS: |
| defaultEnvironment = createProcessEnvironment(config); |
| break; |
| case ENVIRONMENT_DOCKER: |
| default: |
| defaultEnvironment = createDockerEnvironment(config); |
| } |
| } |
| return defaultEnvironment |
| .toBuilder() |
| .addAllDependencies(getDeferredArtifacts(options)) |
| .addAllCapabilities(getJavaCapabilities()) |
| .build(); |
| } |
| |
| public static Environment createDockerEnvironment(String dockerImageUrl) { |
| if (Strings.isNullOrEmpty(dockerImageUrl)) { |
| return JAVA_SDK_HARNESS_ENVIRONMENT; |
| } |
| return Environment.newBuilder() |
| .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.DOCKER)) |
| .setPayload( |
| DockerPayload.newBuilder().setContainerImage(dockerImageUrl).build().toByteString()) |
| .build(); |
| } |
| |
| private static Environment createExternalEnvironment(String config) { |
| return Environment.newBuilder() |
| .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.EXTERNAL)) |
| .setPayload( |
| ExternalPayload.newBuilder() |
| .setEndpoint(ApiServiceDescriptor.newBuilder().setUrl(config).build()) |
| .build() |
| .toByteString()) |
| .build(); |
| } |
| |
| private static Environment createProcessEnvironment(String config) { |
| try { |
| ProcessPayloadReferenceJSON payloadReferenceJSON = |
| MAPPER.readValue(config, ProcessPayloadReferenceJSON.class); |
| return createProcessEnvironment( |
| payloadReferenceJSON.getOs(), |
| payloadReferenceJSON.getArch(), |
| payloadReferenceJSON.getCommand(), |
| payloadReferenceJSON.getEnv()); |
| } catch (IOException e) { |
| throw new RuntimeException( |
| String.format("Unable to parse process environment config: %s", config), e); |
| } |
| } |
| |
| private static Environment createEmbeddedEnvironment(String config) { |
| return Environment.newBuilder() |
| .setUrn(ENVIRONMENT_EMBEDDED) |
| .setPayload(ByteString.copyFromUtf8(MoreObjects.firstNonNull(config, ""))) |
| .build(); |
| } |
| |
| public static Environment createProcessEnvironment( |
| String os, String arch, String command, Map<String, String> env) { |
| ProcessPayload.Builder builder = ProcessPayload.newBuilder(); |
| if (!Strings.isNullOrEmpty(os)) { |
| builder.setOs(os); |
| } |
| if (!Strings.isNullOrEmpty(arch)) { |
| builder.setArch(arch); |
| } |
| if (!Strings.isNullOrEmpty(command)) { |
| builder.setCommand(command); |
| } |
| if (env != null) { |
| builder.putAllEnv(env); |
| } |
| return Environment.newBuilder() |
| .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.PROCESS)) |
| .setPayload(builder.build().toByteString()) |
| .build(); |
| } |
| |
| public static Optional<Environment> getEnvironment(String ptransformId, Components components) { |
| PTransform ptransform = components.getTransformsOrThrow(ptransformId); |
| String envId = ptransform.getEnvironmentId(); |
| if (Strings.isNullOrEmpty(envId)) { |
| // Some PTransform payloads may have an unspecified (empty) Environment ID, for example a |
| // WindowIntoPayload with a known WindowFn. Others will never have an Environment ID, such |
| // as a GroupByKeyPayload, and we return null in this case. |
| return Optional.empty(); |
| } else { |
| return Optional.of(components.getEnvironmentsOrThrow(envId)); |
| } |
| } |
| |
| public static Optional<Environment> getEnvironment( |
| PTransform ptransform, RehydratedComponents components) { |
| String envId = ptransform.getEnvironmentId(); |
| if (Strings.isNullOrEmpty(envId)) { |
| return Optional.empty(); |
| } else { |
| // Some PTransform payloads may have an empty (default) Environment ID, for example a |
| // WindowIntoPayload with a known WindowFn. Others will never have an Environment ID, such |
| // as a GroupByKeyPayload, and we return null in this case. |
| return Optional.of(components.getEnvironment(envId)); |
| } |
| } |
| |
| public static List<ArtifactInformation> getArtifacts(List<String> stagingFiles) { |
| ImmutableList.Builder<ArtifactInformation> artifactsBuilder = ImmutableList.builder(); |
| Set<String> deduplicatedStagingFiles = new LinkedHashSet<>(stagingFiles); |
| for (String path : deduplicatedStagingFiles) { |
| File file; |
| String stagedName = null; |
| if (path.contains("=")) { |
| String[] components = path.split("=", 2); |
| file = new File(components[1]); |
| stagedName = components[0]; |
| } else { |
| file = new File(path); |
| } |
| // Spurious items get added to the classpath. Filter by just those that exist. |
| if (file.exists()) { |
| ArtifactInformation.Builder artifactBuilder = ArtifactInformation.newBuilder(); |
| artifactBuilder.setTypeUrn(BeamUrns.getUrn(StandardArtifacts.Types.FILE)); |
| artifactBuilder.setRoleUrn(BeamUrns.getUrn(StandardArtifacts.Roles.STAGING_TO)); |
| HashCode hashCode; |
| if (file.isDirectory()) { |
| File zippedFile; |
| try { |
| zippedFile = zipDirectory(file); |
| hashCode = Files.asByteSource(zippedFile).hash(Hashing.sha256()); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| |
| artifactBuilder.setTypePayload( |
| RunnerApi.ArtifactFilePayload.newBuilder() |
| .setPath(zippedFile.getPath()) |
| .setSha256(hashCode.toString()) |
| .build() |
| .toByteString()); |
| |
| } else { |
| try { |
| hashCode = Files.asByteSource(file).hash(Hashing.sha256()); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| artifactBuilder.setTypePayload( |
| RunnerApi.ArtifactFilePayload.newBuilder() |
| .setPath(file.getPath()) |
| .setSha256(hashCode.toString()) |
| .build() |
| .toByteString()); |
| } |
| if (stagedName == null) { |
| stagedName = createStagingFileName(file, hashCode); |
| } |
| artifactBuilder.setRolePayload( |
| RunnerApi.ArtifactStagingToRolePayload.newBuilder() |
| .setStagedName(stagedName) |
| .build() |
| .toByteString()); |
| artifactsBuilder.add(artifactBuilder.build()); |
| } |
| } |
| return artifactsBuilder.build(); |
| } |
| |
| public static List<ArtifactInformation> getDeferredArtifacts(PipelineOptions options) { |
| List<String> stagingFiles = options.as(PortablePipelineOptions.class).getFilesToStage(); |
| if (stagingFiles == null || stagingFiles.isEmpty()) { |
| return ImmutableList.of(); |
| } |
| |
| String key = UUID.randomUUID().toString(); |
| DefaultArtifactResolver.INSTANCE.register( |
| (info) -> { |
| if (BeamUrns.getUrn(StandardArtifacts.Types.DEFERRED).equals(info.getTypeUrn())) { |
| RunnerApi.DeferredArtifactPayload deferredArtifactPayload; |
| try { |
| deferredArtifactPayload = |
| RunnerApi.DeferredArtifactPayload.parseFrom(info.getTypePayload()); |
| } catch (InvalidProtocolBufferException e) { |
| throw new RuntimeException("Error parsing deferred artifact payload.", e); |
| } |
| if (key.equals(deferredArtifactPayload.getKey())) { |
| return Optional.of(getArtifacts(stagingFiles)); |
| } else { |
| return Optional.empty(); |
| } |
| } else { |
| return Optional.empty(); |
| } |
| }); |
| |
| return ImmutableList.of( |
| ArtifactInformation.newBuilder() |
| .setTypeUrn(BeamUrns.getUrn(StandardArtifacts.Types.DEFERRED)) |
| .setTypePayload( |
| RunnerApi.DeferredArtifactPayload.newBuilder().setKey(key).build().toByteString()) |
| .build()); |
| } |
| |
| public static Set<String> getJavaCapabilities() { |
| ImmutableSet.Builder<String> capabilities = ImmutableSet.builder(); |
| capabilities.addAll(ModelCoders.urns()); |
| capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.MULTI_CORE_BUNDLE_PROCESSING)); |
| capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.PROGRESS_REPORTING)); |
| capabilities.add("beam:version:sdk_base:" + JAVA_SDK_HARNESS_CONTAINER_URL); |
| capabilities.add(BeamUrns.getUrn(SplittableParDoComponents.TRUNCATE_SIZED_RESTRICTION)); |
| return capabilities.build(); |
| } |
| |
| public static String createStagingFileName(File path, HashCode hash) { |
| String encodedHash = Base64Variants.MODIFIED_FOR_URL.encode(hash.asBytes()); |
| String fileName = Files.getNameWithoutExtension(path.getAbsolutePath()); |
| String ext = path.isDirectory() ? "jar" : Files.getFileExtension(path.getAbsolutePath()); |
| String suffix = Strings.isNullOrEmpty(ext) ? "" : "." + ext; |
| return String.format("%s-%s%s", fileName, encodedHash, suffix); |
| } |
| |
| private static File zipDirectory(File directory) throws IOException { |
| File zipFile = File.createTempFile(directory.getName(), ".zip"); |
| try (FileOutputStream fos = new FileOutputStream(zipFile)) { |
| ZipFiles.zipDirectory(directory, fos); |
| } |
| return zipFile; |
| } |
| |
| private static class ProcessPayloadReferenceJSON { |
| private @Nullable String os; |
| private @Nullable String arch; |
| private @Nullable String command; |
| private @Nullable Map<String, String> env; |
| |
| public @Nullable String getOs() { |
| return os; |
| } |
| |
| public @Nullable String getArch() { |
| return arch; |
| } |
| |
| public @Nullable String getCommand() { |
| return command; |
| } |
| |
| public @Nullable Map<String, String> getEnv() { |
| return env; |
| } |
| } |
| } |
