TPC: Pull gpu installer from in universe repo
BUG=b/322205785
TEST=Built cos-extensions, copied binary to VM, ran the binary with
different values of /etc/cloud-api-domains set.
Change-Id: I7be5add99bbc91fcc62224cd7aa5c619d8d598a9
Reviewed-on: https://cos-review.googlesource.com/c/cos/cos-extensions/+/84824
Reviewed-by: Vaibhav Rustagi <vaibhavrustagi@google.com>
Reviewed-by: Arnav Kansal <rnv@google.com>
Tested-by: Dexter Rivera <riverade@google.com>
diff --git a/extensions/gpu/gpu.go b/extensions/gpu/gpu.go
index aa74b81..0be4e37 100644
--- a/extensions/gpu/gpu.go
+++ b/extensions/gpu/gpu.go
@@ -2,6 +2,7 @@
package gpu
import (
+ "errors"
"fmt"
"io"
"net/http"
@@ -15,18 +16,24 @@
)
const (
- osReleasePath = "/etc/os-release"
- versionID = "VERSION_ID"
- buildID = "BUILD_ID"
- instanceURL = "http://metadata.google.internal/computeMetadata/v1/instance/zone?alt=text"
- extensionCache = "/var/lib/cos-extensions"
+ osReleasePath = "/etc/os-release"
+ versionID = "VERSION_ID"
+ buildID = "BUILD_ID"
+ apiDomain = "API_DOMAIN"
+ artifactRegistryDomain = "ARTIFACT_REGISTRY_DOMAIN"
+ projectPrefix = "PROJECT_PREFIX"
+ instanceURL = "http://metadata.google.internal/computeMetadata/v1/instance/zone?alt=text"
+ extensionCache = "/var/lib/cos-extensions"
+ gcrPrefix = "gcr.io/cos-cloud"
)
var (
- run = osutils.Run
- parse = osutils.Parse
- artifactRegions = []string{"us", "eu", "asia"}
- DefaultInstaller = "gcr.io/cos-cloud/cos-gpu-installer:v2.3.2"
+ run = osutils.Run
+ parse = osutils.Parse
+ artifactRegions = []string{"us", "eu", "asia"}
+ apiDomainsPath = "/etc/cloud-api-domains"
+ tpcGPUInstallerRepo = "docker.%s/%s/cos-cloud/cos-gpu-installer"
+ DefaultInstaller = "gcr.io/cos-cloud/cos-gpu-installer:v2.3.2"
)
// Installers represent the names of the possible installers to be used during the program execution.
@@ -149,15 +156,38 @@
// GetCosInstaller retrieves the installers to be used in execution of the program's commands.
func GetCosInstaller(c Client) (Installers, error) {
var installers Installers
+ defaultInstaller := DefaultInstaller
+
+ // Modify DefaultInstaller based on APIDomainsPath. If
+ // APIDomainsPath does not exist, just use DefaultInstaller. This
+ // adjustment is required when running in a TPC environment,
+ // where the installer needs to be fetched from an
+ // internal repository.
+ if _, err := os.Stat(apiDomainsPath); err == nil {
+ domains, err := parse(apiDomainsPath)
+ if err != nil {
+ return installers, fmt.Errorf("failed to parse %q, err: %v", apiDomainsPath, err)
+ }
+ if domains[apiDomain] != "googleapis.com" {
+ // We are running in TPC, set default installer to in-universe
+ // repo.
+ universeRepo := fmt.Sprintf(tpcGPUInstallerRepo, domains[artifactRegistryDomain], strings.TrimSuffix(domains[projectPrefix], ":"))
+ defaultInstaller = strings.Replace(defaultInstaller, gcrPrefix, universeRepo, -1)
+ glog.Infof("Using TPC Default installer: %q", defaultInstaller)
+ }
+ } else if !errors.Is(err, os.ErrNotExist) {
+ return installers, fmt.Errorf("failed to stat %q, err: %v", apiDomainsPath, err)
+ }
// Checks if it is in the environmental variables and returns the value.
if installer := os.Getenv("COS_GPU_INSTALLER"); installer != "" {
installers = Installers{
- DefaultInstaller: DefaultInstaller,
+ DefaultInstaller: defaultInstaller,
SelectedInstaller: installer,
}
return installers, nil
}
+
// Change gpu installer container image name if location of the vm is available
// and it's from cos-cloud project.
zone, err := retrieveInstanceZone(c, instanceURL)
@@ -168,13 +198,14 @@
}
if zone != "" && slices.Contains(artifactRegions, zone) {
installers = Installers{
- DefaultInstaller: fmt.Sprintf("%s.%s", zone, DefaultInstaller),
+ DefaultInstaller: fmt.Sprintf("%s.%s", zone, defaultInstaller),
}
return installers, nil
}
installers = Installers{
- DefaultInstaller: DefaultInstaller,
+ DefaultInstaller: defaultInstaller,
}
+ glog.Infof("Using installers: %v", installers)
return installers, nil
}
diff --git a/extensions/gpu/gpu_test.go b/extensions/gpu/gpu_test.go
index 85e677c..09e7b12 100644
--- a/extensions/gpu/gpu_test.go
+++ b/extensions/gpu/gpu_test.go
@@ -107,6 +107,7 @@
wantInstaller Installers
setEnv bool
setInstaller string
+ apiDomains string
}{
{
desc: "Vm instance zone in artifact regions",
@@ -128,10 +129,38 @@
httpRes: []byte("us"),
wantInstaller: Installers{DefaultInstaller: "gcr.io/cos-cloud/cos-gpu-installer:v2.3.2", SelectedInstaller: "cloud-installer-101"},
},
+ {
+ desc: "Running in TPC",
+ wantErr: false,
+ httpRes: []byte("u"),
+ wantInstaller: Installers{DefaultInstaller: "docker.fake-ar-domain.goog/fake-project-prefix/cos-cloud/cos-gpu-installer/cos-gpu-installer:v2.3.2", SelectedInstaller: ""},
+ apiDomains: "API_DOMAIN=fake-api-domain.goog\nARTIFACT_REGISTRY_DOMAIN=fake-ar-domain.goog\nPROJECT_PREFIX=fake-project-prefix",
+ },
+ {
+ desc: "Running in GDU with domain file",
+ wantErr: false,
+ httpRes: []byte("us"),
+ wantInstaller: Installers{DefaultInstaller: "us.gcr.io/cos-cloud/cos-gpu-installer:v2.3.2", SelectedInstaller: ""},
+ apiDomains: "API_DOMAIN=googleapis.com\nARTIFACT_REGISTRY_DOMAIN=pkg.dev\nPROJECT_PREFIX=",
+ },
}
for _, test := range tests {
t.Run(test.desc, func(t *testing.T) {
c := &FakeHttpClient{Body: test.httpRes}
+ if test.apiDomains != "" {
+ tmpFile, err := os.CreateTemp("", "")
+ if err != nil {
+ t.Fatalf("failed to create tmpFile: %v", err)
+ }
+ defer os.Remove(tmpFile.Name())
+ _, err = tmpFile.WriteString(test.apiDomains)
+ if err != nil {
+ t.Fatalf("failed to write to tmpFile: %v", err)
+ }
+ realAPIDomainsPath := apiDomainsPath
+ apiDomainsPath = tmpFile.Name()
+ defer func() { apiDomainsPath = realAPIDomainsPath }()
+ }
if test.setEnv == true {
t.Setenv("COS_GPU_INSTALLER", test.setInstaller)
}
diff --git a/go.mod b/go.mod
index 9c3441f..0226d69 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,6 @@
github.com/fsouza/fake-gcs-server v1.49.2
github.com/golang/glog v1.2.1
github.com/spf13/cobra v1.8.1
- github.com/spf13/pflag v1.0.5
google.golang.org/api v0.183.0
)
@@ -32,6 +31,7 @@
github.com/gorilla/mux v1.8.1 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/pkg/xattr v0.4.9 // indirect
+ github.com/spf13/pflag v1.0.5 // indirect
go.opencensus.io v0.24.0 // indirect
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
