Feature: More robust Python version mismatch handling

Makefile

@@ -41,8 +41,13 @@ proxy:
 	docker push localhost:5001/beta9-proxy:$(tag)
 
 runner:
+	@if [ -z "$(platform)" ]; then \
+		platform_flag=""; \
+	else \
+		platform_flag="--platform=$(platform)"; \
+	fi; \
 	for target in py312 py311 py310 py39 py38; do \
-		docker build . --no-cache --target $$target --platform=linux/amd64 -f ./docker/Dockerfile.runner -t localhost:5001/beta9-runner:$$target-$(runnerTag); \
+		docker build . --no-cache --target $$target $$platform_flag -f ./docker/Dockerfile.runner -t localhost:5001/beta9-runner:$$target-$(runnerTag); \
 		docker push localhost:5001/beta9-runner:$$target-$(runnerTag); \
 	done
 

pkg/abstractions/image/build.go

@@ -264,21 +264,19 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co
 
 	go client.StreamLogs(ctx, containerId, outputChan)
 
-	// Generate the pip install command and prepend it to the commands list
-	if len(opts.PythonPackages) > 0 {
-		pipInstallCmd := b.generatePipInstallCommand(opts.PythonPackages, opts.PythonVersion)
-		opts.Commands = append([]string{pipInstallCmd}, opts.Commands...)
-	}
-
 	log.Printf("container <%v> building with options: %s\n", containerId, opts)
 	startTime := time.Now()
 
-	// Detect if python3.x is installed in the container, if not install it
-	checkPythonVersionCmd := fmt.Sprintf("%s --version", opts.PythonVersion)
-	if resp, err := client.Exec(containerId, checkPythonVersionCmd); err != nil || !resp.Ok {
-		outputChan <- common.OutputMsg{Done: false, Success: false, Msg: fmt.Sprintf("%s not detected, installing it for you...\n", opts.PythonVersion)}
-		installCmd := b.getPythonInstallCommand(opts.PythonVersion)
-		opts.Commands = append([]string{installCmd}, opts.Commands...)
+	// Attempt to install the requested python version and fallback to existing python version on failure
+	err = setupPythonVersion(opts, client, containerId, outputChan)
+	if err != nil {
+		return err
+	}
+
+	// Generate the pip install command and prepend it to the commands list
+	if len(opts.PythonPackages) > 0 {
+		pipInstallCmd := generatePipInstallCommand(opts.PythonPackages, opts.PythonVersion)
+		opts.Commands = append([]string{pipInstallCmd}, opts.Commands...)
 	}
 
 	// Generate the commands to run in the container
@@ -287,7 +285,7 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co
 		case shellCommandType:
 			opts.Commands = append(opts.Commands, step.Command)
 		case pipCommandType:
-			opts.Commands = append(opts.Commands, b.generatePipInstallCommand([]string{step.Command}, opts.PythonVersion))
+			opts.Commands = append(opts.Commands, generatePipInstallCommand([]string{step.Command}, opts.PythonVersion))
 		}
 	}
 
@@ -422,7 +420,7 @@ func (b *Builder) Exists(ctx context.Context, imageId string) bool {
 	return b.registry.Exists(ctx, imageId)
 }
 
-func (b *Builder) getPythonInstallCommand(pythonVersion string) string {
+func getPythonInstallCommand(pythonVersion string) string {
 	baseCmd := "apt-get update -q && apt-get install -q -y software-properties-common gcc curl git"
 	components := []string{
 		"python3-future",
@@ -438,7 +436,7 @@ func (b *Builder) getPythonInstallCommand(pythonVersion string) string {
 	return fmt.Sprintf("%s && add-apt-repository ppa:deadsnakes/ppa && apt-get update && apt-get install -q -y %s && %s", baseCmd, installCmd, postInstallCmd)
 }
 
-func (b *Builder) generatePipInstallCommand(pythonPackages []string, pythonVersion string) string {
+func generatePipInstallCommand(pythonPackages []string, pythonVersion string) string {
 	var flagLines []string
 	var packages []string
 	var flags = []string{"--", "-"}
@@ -465,14 +463,17 @@ func (b *Builder) generatePipInstallCommand(pythonPackages []string, pythonVersi
 var imageNamePattern = regexp.MustCompile(
 	`^` + // Assert position at the start of the string
 		`(?:(?P<Registry>(?:(?:localhost|[\w.-]+(?:\.[\w.-]+)+)(?::\d+)?)|[\w]+:\d+)\/)?` + // Optional registry, which can be localhost, a domain with optional port, or a simple registry with port
-		`(?P<Namespace>(?:[a-z0-9]+(?:(?:[._]|__|[-]*)[a-z0-9]+)*)\/)*` + // Optional namespace, which can contain multiple segments separated by slashes
-		`(?P<Repo>[a-z0-9][-a-z0-9._]+)` + // Required repository name, must start with alphanumeric and can contain alphanumerics, hyphens, dots, and underscores
+		`(?P<Repo>(?:[\w][\w.-]*(?:/[\w][\w.-]*)*))?` + // Full repository path including namespace
 		`(?::(?P<Tag>[\w][\w.-]{0,127}))?` + // Optional tag, which starts with a word character and can contain word characters, dots, and hyphens
 		`(?:@(?P<Digest>[A-Za-z][A-Za-z0-9]*(?:[-_+.][A-Za-z][A-Za-z0-9]*)*:[0-9A-Fa-f]{32,}))?` + // Optional digest, which is a hash algorithm followed by a colon and a hexadecimal hash
 		`$`, // Assert position at the end of the string
 )
 
 func ExtractImageNameAndTag(imageRef string) (BaseImage, error) {
+	if imageRef == "" {
+		return BaseImage{}, errors.New("invalid image URI format")
+	}
+
 	matches := imageNamePattern.FindStringSubmatch(imageRef)
 	if matches == nil {
 		return BaseImage{}, errors.New("invalid image URI format")
@@ -490,7 +491,11 @@ func ExtractImageNameAndTag(imageRef string) (BaseImage, error) {
 		registry = "docker.io"
 	}
 
-	repo := result["Namespace"] + result["Repo"]
+	repo := result["Repo"]
+	if repo == "" {
+		return BaseImage{}, errors.New("invalid image URI format")
+	}
+
 	tag, digest := result["Tag"], result["Digest"]
 	if tag == "" && digest == "" {
 		tag = "latest"
@@ -512,3 +517,40 @@ func hasAnyPrefix(s string, prefixes []string) bool {
 	}
 	return false
 }
+
+func setupPythonVersion(opts *BuildOpts, client *common.RunCClient, containerId string, outputChan chan common.OutputMsg) error {
+	checkPythonVersionCmd := fmt.Sprintf("%s --version", opts.PythonVersion)
+	if resp, err := client.Exec(containerId, checkPythonVersionCmd); err == nil && resp.Ok {
+		return nil
+	}
+
+	outputChan <- common.OutputMsg{Done: false, Success: false, Msg: fmt.Sprintf("request python version (%s) is not detected, attempting to install it for you...\n", opts.PythonVersion)}
+	resp, err := client.Exec(containerId, getPythonInstallCommand(opts.PythonVersion))
+	if err == nil && resp.Ok {
+		return nil
+	}
+
+	// Check if there is a default python version available
+	resp, err = client.Exec(containerId, "python --version")
+	if err != nil || !resp.Ok {
+		outputChan <- common.OutputMsg{Done: true, Success: false, Msg: "Failed to install python and no default python version was detected.\n"}
+		return errors.New("failed to install python and no default python version was detected")
+	}
+
+	outputChan <- common.OutputMsg{Done: false, Success: false, Msg: fmt.Sprintf("Failed to install %s, will continue with image's default version\n", opts.PythonVersion)}
+
+	// Hacky way to avoid needing to update the python version in the stub
+	resp, err = client.Exec(containerId, fmt.Sprintf("ln -s $(which python) /usr/local/bin/%s", opts.PythonVersion))
+	if err != nil || !resp.Ok {
+		outputChan <- common.OutputMsg{Done: true, Success: false, Msg: "Failed to fall back to default python version.\n"}
+		return errors.New("failed to create symlink for python3")
+	}
+
+	// When falling back to the default python version, we might have an old version of pip so we upgrade it
+	resp, err = client.Exec(containerId, fmt.Sprintf("%s -m pip install --upgrade pip", opts.PythonVersion))
+	if err != nil || !resp.Ok {
+		outputChan <- common.OutputMsg{Done: true, Success: false, Msg: "Failed to fall back to default python version.\n"}
+		return errors.New("failed to upgrade pip for default python version")
+	}
+	return nil
+}

pkg/abstractions/image/build_test.go

@@ -80,9 +80,15 @@ func TestExtractImageNameAndTag(t *testing.T) {
 		{
 			ref:          "nvcr.io/nim/meta/llama-3.1-8b-instruct:1.1.0",
 			wantTag:      "1.1.0",
-			wantRepo:     "meta/llama-3.1-8b-instruct",
+			wantRepo:     "nim/meta/llama-3.1-8b-instruct",
 			wantRegistry: "nvcr.io",
 		},
+		{
+			ref:          "ghcr.io/gis-ops/docker-valhalla/valhalla:latest",
+			wantTag:      "latest",
+			wantRepo:     "gis-ops/docker-valhalla/valhalla",
+			wantRegistry: "ghcr.io",
+		},
 	}
 
 	for _, test := range tests {
@@ -145,8 +151,7 @@ func TestGeneratePipInstallCommand(t *testing.T) {
 	}
 
 	for _, tc := range testCases {
-		b := &Builder{}
-		cmd := b.generatePipInstallCommand(tc.opts.PythonPackages, tc.opts.PythonVersion)
+		cmd := generatePipInstallCommand(tc.opts.PythonPackages, tc.opts.PythonVersion)
 		assert.Equal(t, tc.want, cmd)
 	}
 }

pkg/common/config.default.yaml

@@ -53,6 +53,9 @@ imageService:
   registryStore: local
   registryCredentialProvider: docker
   buildContainerPoolSelector: build
+  imageArchitectures:
+    - amd64
+    - arm64
   registries:
     docker:
       username: beamcloud

pkg/types/config.go

@@ -109,6 +109,7 @@ type ImageServiceConfig struct {
 	BuildContainerMemory           int64                 `key:"buildContainerMemory" json:"build_container_memory"`
 	BuildContainerPoolSelector     string                `key:"buildContainerPoolSelector" json:"build_container_pool_selector"`
 	Runner                         RunnerConfig          `key:"runner" json:"runner"`
+	ImageArchitectures             []string              `key:"imageArchitectures" json:"image_architectures"`
 }
 
 type ImageRegistriesConfig struct {

pkg/worker/image.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"slices"
 	"strings"
 	"syscall"
 	"time"
@@ -237,7 +238,7 @@ func (c *ImageClient) InspectAndVerifyImage(ctx context.Context, sourceImage str
 		return err
 	}
 
-	if imageInfo["Architecture"] != "amd64" {
+	if !slices.Contains(c.config.ImageService.ImageArchitectures, imageInfo["Architecture"].(string)) {
 		return &types.ExitCodeError{
 			ExitCode: types.WorkerContainerExitCodeIncorrectImageArch,
 		}