From 5342a960af2d8b16241c7f7e882c79c760c29a35 Mon Sep 17 00:00:00 2001 From: tvalentyn Date: Thu, 7 Mar 2024 15:08:05 -0800 Subject: [PATCH] Attempt to install staged Beam SDK if one is available, before checking if Beam is installed. (#30565) --- sdks/python/container/boot.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/sdks/python/container/boot.go b/sdks/python/container/boot.go index 5c84ad6bc7e4..3da31ad04b23 100644 --- a/sdks/python/container/boot.go +++ b/sdks/python/container/boot.go @@ -385,17 +385,17 @@ func installSetupPackages(ctx context.Context, logger *tools.Logger, files []str bufLogger.Printf(ctx, "Failed to setup acceptable wheel specs, leave it as empty: %v", err) } + // Install the Dataflow Python SDK if one was staged. In released + // container images, SDK is already installed, but can be overriden + // using the --sdk_location pipeline option. + if err := installSdk(ctx, logger, files, workDir, sdkSrcFile, acceptableWhlSpecs, false); err != nil { + return fmt.Errorf("failed to install SDK: %v", err) + } pkgName := "apache-beam" isSdkInstalled := isPackageInstalled(pkgName) if !isSdkInstalled { return fmt.Errorf("Apache Beam is not installed in the runtime environment. If you use a custom container image, you must install apache-beam package in the custom image using same version of Beam as in the pipeline submission environment. For more information, see: the https://beam.apache.org/documentation/runtime/environments/") } - // Install the Dataflow Python SDK and worker packages. - // We install the extra requirements in case of using the beam sdk. These are ignored by pip - // if the user is using an SDK that does not provide these. - if err := installSdk(ctx, logger, files, workDir, sdkSrcFile, acceptableWhlSpecs, false); err != nil { - return fmt.Errorf("failed to install SDK: %v", err) - } // The staged files will not disappear due to restarts because workDir is a // folder that is mapped to the host (and therefore survives restarts). for _, f := range requirementsFiles {