Attempt to install staged Beam SDK if one is available, before checki…

…ng if Beam is installed. (#30565)

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 {