Skip to content

Commit e2a05a4

Browse files
janpfeiferjanpfeifer
authored
Merge pull request #5 from gomlx/platforms
Generic CPU installation code
2 parents 07ddf43 + 34fdd15 commit e2a05a4

27 files changed

Lines changed: 684 additions & 190 deletions

‎.github/workflows/darwin_tests.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ jobs:
3333
shell: bash
3434
run: |
3535
(cd ./cmd/pjrt_installer && go install .)
36-
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./pkg/pjrt/cpuversion.go)"
37-
sudo env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin darwin -version "${CPU_VERSION}" -path /usr/local/lib/go-xla
36+
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./internal/utils/utils.go)"
37+
sudo env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin cpu -version "${CPU_VERSION}" -path /usr/local/lib/go-xla
3838
3939
- name: Test
4040
env:

.github/workflows/linux_tests.yaml renamed to .github/workflows/linux_amd64_tests.yaml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: "Tests(linux)"
1+
name: "Tests(linux/amd64, ubuntu latest)"
22

33
permissions:
44
# read|write|none
@@ -48,8 +48,8 @@ jobs:
4848
shell: bash
4949
run: |
5050
(cd ./cmd/pjrt_installer && go install .)
51-
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./pkg/pjrt/cpuversion.go)"
52-
sudo env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin linux -version "${CPU_VERSION}" -path /usr/local/lib/go-xla
51+
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./internal/utils/utils.go)"
52+
sudo env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin cpu -version "${CPU_VERSION}" -path /usr/local/lib/go-xla
5353
# When adding static linking or dynamic linking at load time:
5454
# sudo ln -sf /usr/local/lib/go-xla/libpjrt* /usr/lib/x86_64-linux-gnu/
5555

‎.github/workflows/linux_arm64_tests.yaml‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
name: "Tests(linux/arm64)"
2+
3+
permissions:
4+
# read|write|none
5+
actions: read
6+
checks: read
7+
contents: write
8+
deployments: read
9+
# id-token: read --> doesn't work
10+
issues: read
11+
discussions: read
12+
packages: read
13+
pages: read
14+
pull-requests: write
15+
repository-projects: read
16+
security-events: read
17+
statuses: read
18+
19+
on:
20+
push:
21+
branches: [main]
22+
pull_request:
23+
branches: [main]
24+
# Allows you to run this workflow manually from the Actions tab
25+
workflow_dispatch:
26+
27+
jobs:
28+
test:
29+
runs-on: ubuntu-24.04-arm
30+
steps:
31+
- uses: actions/checkout@v4
32+
33+
- name: Install Go
34+
uses: actions/setup-go@v5
35+
with:
36+
go-version: "1.25.x"
37+
38+
- name: Test with AutoInstall CPU PJRT plugin
39+
env:
40+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
41+
run: |
42+
go test ./internal/tests/ -test.run=TestEndToEnd/fixed -test.v -test.count=1
43+
44+
# See file pkg/pjrt/cpuversion.go for the default version of the CPU PJRT plugin to use.
45+
- name: Install PJRT plugin
46+
env:
47+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
48+
shell: bash
49+
run: |
50+
(cd ./cmd/pjrt_installer && go install .)
51+
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./internal/utils/utils.go)"
52+
sudo env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin cpu -version "${CPU_VERSION}" -path /usr/local/lib/go-xla
53+
# When adding static linking or dynamic linking at load time:
54+
# sudo ln -sf /usr/local/lib/go-xla/libpjrt* /usr/lib/x86_64-linux-gnu/
55+
56+
- name: Test
57+
env:
58+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
59+
run: |
60+
go test ./pkg/...
61+
go test ./internal/tests/...
62+
63+
# - name: Go Coverage Badge
64+
# if: github.ref == 'refs/heads/main'
65+
# uses: tj-actions/coverage-badge-go@v2
66+
# with:
67+
# green: 80
68+
# filename: docs/coverage.out
69+
70+
# - name: Commit README.md changes
71+
# if: github.ref == 'refs/heads/main'
72+
# run: |
73+
# if git diff --quiet -- 'README.md' ; then
74+
# echo "README.md not modified."
75+
# else
76+
# git config --local user.email "action@github.com"
77+
# git config --local user.name "GitHub Action"
78+
# git add README.md
79+
# git commit -m "chore: Updated coverage badge."
80+
# fi
81+
82+
# - name: Push Changes
83+
# if: github.ref == 'refs/heads/main'
84+
# uses: ad-m/github-push-action@master
85+
# with:
86+
# github_token: ${{ github.token }}
87+
# branch: ${{ github.ref }}

‎.github/workflows/linux_old_amd64_tests.yaml‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
name: "Tests(linux/amd64, ubuntu-22.04)"
2+
3+
permissions:
4+
# read|write|none
5+
actions: read
6+
checks: read
7+
contents: write
8+
deployments: read
9+
# id-token: read --> doesn't work
10+
issues: read
11+
discussions: read
12+
packages: read
13+
pages: read
14+
pull-requests: write
15+
repository-projects: read
16+
security-events: read
17+
statuses: read
18+
19+
on:
20+
push:
21+
branches: [main]
22+
pull_request:
23+
branches: [main]
24+
# Allows you to run this workflow manually from the Actions tab
25+
workflow_dispatch:
26+
27+
jobs:
28+
test:
29+
runs-on: ubuntu-22.04
30+
steps:
31+
- uses: actions/checkout@v4
32+
33+
- name: Install Go
34+
uses: actions/setup-go@v5
35+
with:
36+
go-version: "1.25.x"
37+
38+
- name: Test with AutoInstall CPU PJRT plugin
39+
env:
40+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
41+
run: |
42+
go test ./internal/tests/ -test.run=TestEndToEnd/fixed -test.v -test.count=1
43+
44+
# See file pkg/pjrt/cpuversion.go for the default version of the CPU PJRT plugin to use.
45+
- name: Install PJRT plugin
46+
env:
47+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
48+
shell: bash
49+
run: |
50+
(cd ./cmd/pjrt_installer && go install .)
51+
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./internal/utils/utils.go)"
52+
sudo env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin cpu -version "${CPU_VERSION}" -path /usr/local/lib/go-xla
53+
# When adding static linking or dynamic linking at load time:
54+
# sudo ln -sf /usr/local/lib/go-xla/libpjrt* /usr/lib/x86_64-linux-gnu/
55+
56+
- name: Test
57+
env:
58+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
59+
run: |
60+
go test ./pkg/...
61+
go test ./internal/tests/...
62+
63+
# - name: Go Coverage Badge
64+
# if: github.ref == 'refs/heads/main'
65+
# uses: tj-actions/coverage-badge-go@v2
66+
# with:
67+
# green: 80
68+
# filename: docs/coverage.out
69+
70+
# - name: Commit README.md changes
71+
# if: github.ref == 'refs/heads/main'
72+
# run: |
73+
# if git diff --quiet -- 'README.md' ; then
74+
# echo "README.md not modified."
75+
# else
76+
# git config --local user.email "action@github.com"
77+
# git config --local user.name "GitHub Action"
78+
# git add README.md
79+
# git commit -m "chore: Updated coverage badge."
80+
# fi
81+
82+
# - name: Push Changes
83+
# if: github.ref == 'refs/heads/main'
84+
# uses: ad-m/github-push-action@master
85+
# with:
86+
# github_token: ${{ github.token }}
87+
# branch: ${{ github.ref }}

‎.github/workflows/windows_amd64_tests.yaml‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
name: "Tests(windows/amd64)"
2+
3+
permissions:
4+
# read|write|none
5+
actions: read
6+
checks: read
7+
contents: write
8+
deployments: read
9+
# id-token: read --> doesn't work
10+
issues: read
11+
discussions: read
12+
packages: read
13+
pages: read
14+
pull-requests: write
15+
repository-projects: read
16+
security-events: read
17+
statuses: read
18+
19+
on:
20+
push:
21+
branches: [main]
22+
pull_request:
23+
branches: [main]
24+
# Allows you to run this workflow manually from the Actions tab
25+
workflow_dispatch:
26+
27+
jobs:
28+
test:
29+
runs-on: windows-latest
30+
env:
31+
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
32+
defaults:
33+
run:
34+
shell: bash
35+
36+
steps:
37+
- uses: actions/checkout@v4
38+
39+
- name: Install Go
40+
uses: actions/setup-go@v5
41+
with:
42+
go-version: "1.25.x"
43+
44+
- name: Test with AutoInstall CPU PJRT plugin
45+
run: |
46+
go test ./internal/tests/ -test.run=TestEndToEnd/fixed -test.v -test.count=1
47+
48+
- name: Install PJRT plugin
49+
run: |
50+
(cd ./cmd/pjrt_installer && go work init && go work use . '..\..' && go install .)
51+
export CPU_VERSION="$(sed -n 's/.*var DefaultCPUVersion = "\(.*\)".*/\1/p' ./internal/utils/utils.go)"
52+
env GH_TOKEN=$GH_TOKEN $(go env GOPATH)/bin/pjrt_installer -plugin cpu -version "${CPU_VERSION}" -path '~\AppData\Local\go-xla'
53+
echo "Contents for go-xla:"
54+
ls -la "${HOME}/AppData/Local/go-xla"
55+
echo "-----------"
56+
echo
57+
58+
- name: Test
59+
run: |
60+
go test ./pkg/...
61+
go test ./internal/tests/...

‎cmd/pjrt_installer/cpu.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"path/filepath"
6+
"runtime"
7+
"slices"
8+
9+
"github.com/gomlx/go-xla/pkg/installer"
10+
"k8s.io/klog/v2"
11+
)
12+
13+
func init() {
14+
platform := fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH)
15+
if !slices.Contains(installer.CPUSupportedPlatforms, platform) {
16+
return
17+
}
18+
pluginName := "cpu"
19+
pluginInstallers[pluginName] = func(plugin, version, installPath string) error {
20+
// Platform set to current architecture.
21+
return installer.CPUInstall(platform, version, installPath, *flagCache, installer.VerbosityLevel(*flagVerbosity))
22+
}
23+
pluginValidators[pluginName] = func(plugin, version string) error {
24+
return installer.CPUValidateVersion(platform, version)
25+
}
26+
pluginValues = append(pluginValues, pluginName)
27+
pluginDescriptions = append(pluginDescriptions, fmt.Sprintf("CPU PJRT (%s/%s)", runtime.GOOS, runtime.GOARCH))
28+
pluginPriorities = append(pluginPriorities, 0)
29+
30+
// Install path suggestions.
31+
localLibPath, err := installer.DefaultHomeLibPath()
32+
if err != nil {
33+
klog.Fatalf("Failed to find home directory: %+v", localLibPath)
34+
}
35+
installPathSuggestions = append(installPathSuggestions, filepath.Join(localLibPath, "go-xla"))
36+
if slices.Contains([]string{"darwin", "linux"}, runtime.GOOS) {
37+
installPathSuggestions = append(installPathSuggestions, "/usr/local/lib/go-xla")
38+
}
39+
}

‎cmd/pjrt_installer/cuda.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ func init() {
1515
}
1616
pluginValues = append(pluginValues, "cuda13", "cuda12")
1717
pluginDescriptions = append(pluginDescriptions,
18-
"CUDA PJRT for Linux/amd64, using CUDA 13",
19-
"CUDA PJRT for Linux/amd64, using CUDA 12, deprecated")
18+
"CUDA PJRT (linux/amd64), using CUDA 13",
19+
"CUDA PJRT (linux/amd64), using CUDA 12, deprecated")
2020
pluginPriorities = append(pluginPriorities, 10, 11)
2121
}

‎cmd/pjrt_installer/darwin.go‎

Lines changed: 0 additions & 20 deletions
This file was deleted.

‎cmd/pjrt_installer/go.mod‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,12 @@ require (
99
github.com/charmbracelet/huh/spinner v0.0.0-20251124111010-6575a6e28cb3
1010
github.com/charmbracelet/lipgloss v1.1.0
1111
github.com/charmbracelet/x/term v0.2.2
12-
github.com/gomlx/go-xla v0.1.2
12+
github.com/gomlx/go-xla v0.1.4-0.20251217134619-744d4d4dee0c
1313
github.com/pkg/errors v0.9.1
1414
k8s.io/klog/v2 v2.130.1
1515
)
1616

1717
require (
18-
github.com/AlexsanderHamir/GenPool v1.6.6 // indirect
1918
github.com/atotto/clipboard v0.1.4 // indirect
2019
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
2120
github.com/catppuccin/go v0.3.0 // indirect

‎cmd/pjrt_installer/go.sum‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
github.com/AlexsanderHamir/GenPool v1.6.6 h1:JokjMQ7at0ibhUQ2/SMkFWbfRrclshT5qkWOoKitrz8=
2-
github.com/AlexsanderHamir/GenPool v1.6.6/go.mod h1:xQYqsDsmkUwA92icUHjx3yJEgiDmZztrSMF0MUtFG/c=
31
github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ=
42
github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE=
53
github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
@@ -48,12 +46,10 @@ github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6
4846
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
4947
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
5048
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
51-
github.com/gomlx/go-xla v0.1.2 h1:+OQSvPo5VpGdXcyz0bZj5i/sW/k/N1Q9ME60rHcSM7M=
52-
github.com/gomlx/go-xla v0.1.2/go.mod h1:odQvA8u6ngj7b2rbJjfT/NIN2uQt0/TQvfFpnH25LAs=
49+
github.com/gomlx/go-xla v0.1.4-0.20251217134619-744d4d4dee0c h1:w6cqMvffJWOhxskdpIYWTUXKMtJC61AT35qtBpzVT6o=
50+
github.com/gomlx/go-xla v0.1.4-0.20251217134619-744d4d4dee0c/go.mod h1:6OfAcL+h+es7KHz6q0Ql9Vzi5MTys99MJ7aR3Nhko/o=
5351
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
5452
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
55-
github.com/janpfeifer/go-benchmarks v0.1.1 h1:gLLy07/JrOKSnMWeUxSnjTdhkglgmrNR2IBDnR4kRqw=
56-
github.com/janpfeifer/go-benchmarks v0.1.1/go.mod h1:5AagXCOUzevvmYFQalcgoa4oWPyH1IkZNckolGWfiSM=
5753
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
5854
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
5955
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -75,8 +71,6 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
7571
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
7672
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
7773
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
78-
github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d h1:X4+kt6zM/OVO6gbJdAfJR60MGPsqCzbtXNnjoGqdfAs=
79-
github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU=
8074
github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
8175
github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
8276
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=

0 commit comments

Comments
 (0)