installer

package
v0.1.4 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Dec 17, 2025 License: Apache-2.0 Imports: 27 Imported by: 0

Documentation

Overview

Package installer provides functionality to install PJRT plugins.

The API exposes several functions to install the different PJRT plugins individually, and it is used the command-line program github.com/gomlx/go-xla/cmd/pjrt_installer.

External users may be interested in the using the AutoInstall function to automatically install the PRJT plugin for the current platform.

By default, the functions are only available to the corresponding platforms (currently only Linux/amd64 and Darwin/arm64). If you use the tag `pjrt_all`, all functions will be available. The AutoInstall function is an exception, it is platform-specific and the tag `pjrt_all` has no effect on it.

Index

Constants

View Source 
const AmazonLinux = "amazonlinux"
View Source 
const BinaryCPUReleasesRepo = "gomlx/pjrt-cpu-binaries"
View Source 
const DeleteToEndOfLine = "\x1b[J"
View Source 
const TPUPJRTPluginName = "pjrt_c_api_tpu_plugin.so"

Variables

View Source 
var CPUSupportedPlatforms = []string{
	"linux_amd64", "linux_arm64",
	"darwin_arm64",
	"windows_amd64",
}
View Source 
var HasNvidiaGPU = sync.OnceValue[bool](func() bool {
	matches, err := filepath.Glob("/dev/nvidia*")
	if err != nil {
		klog.Errorf("Failed to figure out if there is an Nvidia GPU installed while searching for files matching \"/dev/nvidia*\": %v", err)
	} else if len(matches) > 0 {
		return true
	}

	_, lookErr := exec.LookPath("nvidia-smi")
	if lookErr != nil {
		return false
	}
	cmd := exec.Command("nvidia-smi")
	output, cmdErr := cmd.CombinedOutput()
	if cmdErr != nil {
		return false
	}
	return strings.Contains(string(output), "NVIDIA-SMI")
})

HasNvidiaGPU tries to guess if there is an actual Nvidia GPU installed (as opposed to only the drivers/PJRT file installed, but no actual hardware). It does that by checking for the presence of the device files in /dev/nvidia*.

View Source 
var HasTPU = sync.OnceValue(func() bool {

	if _, err := os.Stat("/lib/libtpu.so"); err == nil {
		return true
	}
	if _, err := os.Stat("/usr/lib/libtpu.so"); err == nil {
		return true
	}
	if _, err := os.Stat("/usr/local/lib/libtpu.so"); err == nil {
		return true
	}

	return false
})

HasTPU checks if a TPU is available on the system. This uses sync.OnceValue for efficient repeated calls.

Functions

func AutoInstall 

func AutoInstall(installPath string, useCache bool, verbosity VerbosityLevel) error

AutoInstall automatically installs the PJRT plugin for the current platform, if the corresponding hardware is present.

It returns the first error encountered, or nil if no error occurred.

  • installPath should point to a "lib" directory, under which an "go-xla" subdirectory is created, and the PJRT plugins are installed (e.g.: "~/.local/lib", "/usr/local/lib/"). If it is set to "", it is set to the default user-local library directory ( in Linux is $HOME/.local/lib/, and on MacOS is $HOME/Library/Application Support/).

- useCache: if true, it will use the cache to store downloaded files. Recommended to keep it true.

- verbosity: the verbosity level to use. 0=quiet, 1=normal (1 log line per plugin installed), 2=verbose.

func CPUAutoInstall 

func CPUAutoInstall(goxlaInstallPath string, useCache bool, verbosity VerbosityLevel) error

CPUAutoInstall installs the latest version of the CPU PJRT if not yet installed. goxlaInstallPath is expected to be a "lib/go-xla" directory, under which the PJRT plugin is installed.

func CPUGetDownloadURL 

func CPUGetDownloadURL(platform, version string) (url string, err error)

CPUGetDownloadURL returns the download URL for the given version and plugin.

func CPUInstall 

func CPUInstall(platform, version, installPath string, useCache bool, verbosity VerbosityLevel) error

CPUInstall the assets on the target directory.

func CPUValidateVersion 

func CPUValidateVersion(platform, version string) error

CPUValidateVersion checks whether the linux version selected by "-version" exists.

func CudaAutoInstall 

func CudaAutoInstall(goxlaInstallPath string, useCache bool, verbosity VerbosityLevel) error

CudaAutoInstall installs the latest version of the CUDA PJRT and Nvidia libraries if not yet installed, and there is an actual Nvidia GPU installed.

It uses HasNvidiaGPU to see if there is an actual Nvidia GPU installed.

goxlaInstallPath is expected to be a "lib/go-xla" directory, under which the nvidia/ subdirectory will be (is already) created, and the CUDA PJRT plugin is installed.

func CudaInstall 

func CudaInstall(plugin, version, installPath string, useCache bool, verbosity VerbosityLevel) error

CudaInstall installs the cuda PJRT from the Jax PIP packages, using pypi.org distributed files.

Checks performed: - Version exists - Author email is from the Jax team - Downloaded files sha256 match the ones on pypi.org

If useCache is true, it will save the file in a cache directory and try to reuse it if already downloaded

The installPath parameter should be to the .../lib/go-xla directory. The pjrt plugin itself will be installed under the .../lib/go-xla/nvidia directory -- it needs to be there due to path resolution issues with the plugin/Nvidia libraries.

func CudaInstallNvidiaLibraries 

func CudaInstallNvidiaLibraries(plugin, version, nvidiaSubdir string, useCache bool, verbosity VerbosityLevel) error

CudaInstallNvidiaLibraries installs the required NVIDIA libraries for CUDA.

func CudaInstallPJRT 

func CudaInstallPJRT(plugin, version, installPath string, useCache bool, verbosity VerbosityLevel) (string, error)

CudaInstallPJRT installs the cuda PJRT from the Jax PIP packages, using pypi.org distributed files.

Checks performed: - Version exists - Author email is from the Jax team - Downloaded files sha256 match the ones on pypi.org

If useCache is true, it will save the file in a cache directory and try to reuse it if already downloaded.

Returns the version that was installed -- it can be different if the requested version was "latest", in which case it is translated to the actual version.

func CudaValidateVersion 

func CudaValidateVersion(plugin, version string) error

CudaValidateVersion checks whether the cuda version selected by "-version" exists.

func DefaultHomeLibPath added in v0.1.4 

func DefaultHomeLibPath() (string, error)

DefaultHomeLibPath returns the default user-local library directory ("~/.local/lib" in Linux)

This is the directory used by AutoInstall if the installPath is empty.

func DownloadURLToTemp 

func DownloadURLToTemp(url, fileName, wantSHA256 string, useCache bool, verbosity VerbosityLevel) (
	filePath string, cached bool, err error)

DownloadURLToTemp downloads a file from a given URL to a temporary file.

It displays a spinner while downloading and outputs some information about the download.

If useCache is true, it will save the file in a cache directory and try to reuse it if already downloaded.

If wantSHA256 is not empty, it will verify the hash of the downloaded file.

It returns the path where the file was downloaded, and if the downloaded file is in a cache (so it shouldn't be removed after use).

func ExtractDirFromZip 

func ExtractDirFromZip(zipFilePath, dirInZipFile, outputPath string) error

ExtractDirFromZip extracts from zipFilePath all files and directories under dirInZipFile and saves them with the same directory structure under outputPath.

Notice dirInZipFile is not repeated in outputPath.

func ExtractFileFromZip 

func ExtractFileFromZip(zipFilePath, fileName, outputPath string) error

ExtractFileFromZip searches for a file named fileName within the zipFilePath and extracts the first one found to the outputPath.

The fileName is matched to the full path of the files in the zip archive as well as to the base name: so a base name can be given, and it will find it anywhere in the zip archive.

func GetCachePath 

func GetCachePath(fileName string) (filePath string, cached bool, err error)

GetCachePath finds and prepares the cache directory for gopjrt.

It uses os.UserCacheDir() for portability:

- Linux: $XDG_CACHE_HOME or $HOME/.cache - Darwin: $HOME/Library/Caches - Windows: %LocalAppData% (e.g., C:\Users\user\AppData\Local)

func GitHubDownloadReleaseAssets 

func GitHubDownloadReleaseAssets(repo string, version string) ([]string, error)

GitHubDownloadReleaseAssets downloads the list of assets available for the given repository/release version. E.g.: repo = "gomlx/pjrt-cpu-binaries", version = "v0.83.1"

func GitHubGetLatestVersion 

func GitHubGetLatestVersion() (string, error)

GitHubGetLatestVersion returns the latest version tag from the gomlx/pjrt-cpu-binaries repository.

func GitHubGetVersions 

func GitHubGetVersions(repo string) ([]string, error)

func PipCompareVersion 

func PipCompareVersion(v1, v2 string) int

PipCompareVersion compares two PIP version strings (they include only numbers separated by dots).

func PipPackageLinuxAMD64 

func PipPackageLinuxAMD64() *regexp.Regexp

PipPackageLinuxAMD64 returns the regexp for matching Linux AMD64 packages.

func PipPackageLinuxAMD64Glibc231 

func PipPackageLinuxAMD64Glibc231() *regexp.Regexp

PipPackageLinuxAMD64Glibc231 returns the regexp for matching Linux AMD64 packages with glibc 2.31.

func ReplaceTildeInDir 

func ReplaceTildeInDir(dir string) (string, error)

ReplaceTildeInDir replaces "~" in a directory path with the user's home directory. Returns dir if it doesn't start with "~". It may panic with an error if `dir` has an unknown user (e.g: `~unknown/...`)

func ReportError 

func ReportError(err error)

ReportError prints an error if it is not nil, but otherwise does nothing.

func TPUAutoInstall 

func TPUAutoInstall(goxlaInstallPath string, useCache bool, verbosity VerbosityLevel) error

TPUAutoInstall installs the TPU PJRT if it is available on the system. goxlaInstallPath is expected to be a "lib/go-xla" directory, under which the PJRT plugin is installed.

func TPUInstall 

func TPUInstall(plugin, version, installPath string, useCache bool, verbosity VerbosityLevel) error

TPUInstall installs the TPU PJRT from the "libtpu" PIP packages, using pypi.org distributed files.

Checks performed: - Version exists - Downloaded files sha256 match the ones on pypi.org

func TPUValidateVersion 

func TPUValidateVersion(plugin, version string) error

TPUValidateVersion checks whether the TPU version selected by "-version" exists.

func Untar 

func Untar(tarballPath, outputDirPath string) ([]string, error)

Untar takes a path to a tar/gzip file and an output directory. It returns a list of extracted files and any error encountered.

func Unzip added in v0.1.4 

func Unzip(zipPath, outputDirPath string) ([]string, error)

Unzip takes a path to a zip file and an output directory. It returns a list of extracted files and any error encountered.

Types

type AutoInstaller 

type AutoInstaller func(installPath string, useCache bool, verbosity VerbosityLevel) error

AutoInstaller is a function that automatically installs the PJRT plugin for the current platform. It should be safe to call multiple times, and be a no-op if the plugin is already installed.

It is expected to only be called concurrently.

type PipDependency 

type PipDependency struct {
	Package                                         string
	Minimum, LowerBound, Maximum, UpperBound, Exact string
	Condition                                       string
}

PipDependency represents a parsed dependency from a PIP package.

func (*PipDependency) IsValid 

func (dep *PipDependency) IsValid(version string) bool

IsValid checks if the given version satisfies the dependency constraints.

type PipPackageInfo 

type PipPackageInfo struct {
	// Top-level object returned by the API
	Info struct {
		Name           string `json:"name"`
		Version        string `json:"version"` // This is the LATEST version
		Summary        string `json:"summary"`
		HomePage       string `json:"home_page"`
		Author         string `json:"author"`
		AuthorEmail    string `json:"author_email"`
		License        string `json:"license"`
		ProjectURL     string `json:"project_url"`
		RequiresPython string `json:"requires_python"`

		// This is a list of dependencies/requirements
		RequiresDist []string `json:"requires_dist"`
	} `json:"info"`

	// Releases is a map of all versions, where the key is the version string (e.g., "1.2.3")
	Releases map[string][]PipReleaseInfo `json:"releases"`
}

PipPackageInfo is the JSON response from pypi.org for a given package.

func CudaGetPJRTPipInfo 

func CudaGetPJRTPipInfo(plugin string) (*PipPackageInfo, string, error)

CudaGetPJRTPipInfo returns the JSON info for the PIP package that corresponds to the plugin.

func GetPipInfo 

func GetPipInfo(packageName string) (*PipPackageInfo, error)

GetPipInfo retrieves package information from pypi.org for the given package name.

func TPUGetPJRTPipInfo 

func TPUGetPJRTPipInfo(plugin string) (*PipPackageInfo, string, error)

TPUGetPJRTPipInfo returns the JSON info for the PIP package that corresponds to the plugin.

func (*PipPackageInfo) ParseDependencies 

func (info *PipPackageInfo) ParseDependencies() ([]PipDependency, error)

ParseDependencies parses the RequiresDist field into structured PipDependency objects.

type PipReleaseInfo 

type PipReleaseInfo struct {
	// You can add more fields here if you need file-specific details
	PackageType string            `json:"packagetype"`
	Filename    string            `json:"filename"`
	URL         string            `json:"url"`
	Digests     map[string]string `json:"digests"`
}

PipReleaseInfo is the JSON response from pypi.org for a given package version, for some platform.

func PipSelectRelease 

func PipSelectRelease(releaseInfos []PipReleaseInfo, platform *regexp.Regexp, anyMatchingVersion bool) (*PipReleaseInfo, error)

PipSelectRelease selects the release for the given platform (regexp).

If anyMatchingVersion is true, it will return the first release that matches the platform. Otherwise, it will return an error if there are multiple matches.

type Spinner added in v0.1.2 

type Spinner struct {
	// contains filtered or unexported fields
}

Spinner is a utility to run an action while displaying a rotating text indicator and allowing the action to update the title text.

func NewSpinner added in v0.1.2 

func NewSpinner() *Spinner

NewSpinner creates a new Spinner instance.

func (*Spinner) Action added in v0.1.2 

func (s *Spinner) Action(action func(chan<- string)) *Spinner

Action sets the function to be executed. The function receives a channel it can use to send real-time title updates.

func (*Spinner) Run added in v0.1.2 

func (s *Spinner) Run() error

Run executes the action function and manages the spinner animation and title updates.

func (*Spinner) Title added in v0.1.2 

func (s *Spinner) Title(title string) *Spinner

Title sets the initial title text for the spinner.

type VerbosityLevel 

type VerbosityLevel int
const (
	Quiet VerbosityLevel = iota
	Normal
	Verbose
)

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.