[Internal] Rework runner/.justfile as a module (#4139)

Author: un-def

.justfile

@@ -11,12 +11,15 @@
 # * website/.justfile – Building and previewing the React landing page
 # * .tox.justfile – Running Python tests via tox
 
-# Run tests via tox
-mod tox '.tox.justfile'
+set minimum-version := '1.55.0'
 
-set allow-duplicate-recipes
+[doc("Building and uploading dstack runner and shim")]
+mod runner "runner/.justfile"
+
+[doc("Running Python tests via tox")]
+mod tox ".tox.justfile"
 
-import "runner/.justfile"
+set allow-duplicate-recipes
 
 import "frontend/.justfile"
 

runner/.justfile

@@ -9,7 +9,7 @@
 #
 # Build Process:
 # - Runner and shim are always built for linux (GOOS=linux is the only supported OS)
-# - The target architecture is configurable via DSTACK_SHIM_BUILD_ARCH (or `just --set arch ...`)
+# - The target architecture is configurable via DSTACK_SHIM_BUILD_ARCH (or `just build --arch ...`)
 # - CGO is enabled only for native builds (Linux host with a matching architecture);
 #   otherwise it is disabled and DCGM support is dropped
 #
@@ -24,114 +24,115 @@
 #   * See README.md for instructions on running dstack server with uploaded binaries
 #   * Upload is required for testing with standard backends (including SSH fleets)
 
-default:
-    @just --list
-
 # Version of the runner and shim to upload
-export version := env("DSTACK_SHIM_UPLOAD_VERSION", "0.0.0")
+version := env("DSTACK_SHIM_UPLOAD_VERSION", "0.0.0")
 
 # S3 bucket to upload binaries to
-export s3_bucket := env("DSTACK_SHIM_UPLOAD_S3_BUCKET", "dstack-runner-downloads-stgn")
+s3_bucket := env("DSTACK_SHIM_UPLOAD_S3_BUCKET", "dstack-runner-downloads-stgn")
 
 # Target architecture for runner and shim (GOOS is always linux)
-export arch := env("DSTACK_SHIM_BUILD_ARCH", "amd64")
-
-# Download URLs
-export runner_download_url := "s3://" + s3_bucket + "/" + version + "/binaries/dstack-runner-linux-" + arch
-export shim_download_url := "s3://" + s3_bucket + "/" + version + "/binaries/dstack-shim-linux-" + arch
+arch := env("DSTACK_SHIM_BUILD_ARCH", "amd64")
 
 # Go toolchain image for running tests in a container (keep in sync with go.mod)
-export go_version := env("DSTACK_GO_VERSION", "1.25")
+go_version := env("DSTACK_GO_VERSION", "1.25")
+
+[doc("Build both runner and shim")]
+[arg("arch", long)]
+build arch=arch: (build-runner-binary arch) (build-shim-binary arch)
+    @echo "Build complete! linux/{{arch}} binaries are in their respective cmd directories."
+
+[doc("Clean build artifacts")]
+clean:
+    rm -f ./cmd/runner/runner
+    rm -f ./cmd/shim/shim
+    @echo "Build artifacts cleaned!"
+
+[doc("Run tests for runner and shim (native; requires a Linux host)")]
+test:
+    go test -v ./...
+
+# Examples:
+#   just test-in-container  # short suite, all packages
+#   just test-in-container -run TestPullImage ./internal/shim/
+[doc("Run tests for runner and shim in a Linux container (use on macOS/Windows, where native builds are not available)")]
+test-in-container *args="-short ./...":
+    docker run --rm -t \
+        -v .:/src -w /src \
+        -v dstack-go-mod:/go/pkg/mod \
+        -v dstack-go-build:/root/.cache/go-build \
+        -v /var/run/docker.sock:/var/run/docker.sock \
+        golang:{{go_version}} \
+        go test -race {{args}}
+
+[doc("Upload both runner and shim to S3")]
+[arg("arch", long)]
+upload arch=arch: (upload-runner-binary arch) (upload-shim-binary arch)
 
-# Build runner
 [private]
-build-runner-binary:
-    #!/usr/bin/env bash
-    set -e
-    echo "Building runner for linux/$arch"
-    cd {{source_directory()}}/cmd/runner && CGO_ENABLED=0 GOOS=linux GOARCH=$arch go build -ldflags "-X 'main.Version=$version' -extldflags '-static'"
-    echo "Runner build complete!"
+[doc("Build runner")]
+[arg("arch", long)]
+[working-directory: "./cmd/runner"]
+build-runner-binary arch=arch:
+    @echo "Building runner for linux/{{arch}}"
+    CGO_ENABLED=0 GOOS=linux GOARCH={{arch}} go build -ldflags "-X 'main.Version={{version}}' -extldflags '-static'"
+    @echo "Runner build (version: {{version}}) complete!"
 
-# Build shim
 [private]
-build-shim-binary:
+[doc("Build shim")]
+[arg("arch", long)]
+[working-directory: "./cmd/shim"]
+build-shim-binary arch=arch:
     #!/usr/bin/env bash
     set -e
-    cd {{source_directory()}}/cmd/shim
-    echo "Building shim for linux/$arch"
+    echo "Building shim for linux/{{arch}}"
     host_arch=$(uname -m)
     case "$host_arch" in
         x86_64) host_arch=amd64 ;;
         aarch64 | arm64) host_arch=arm64 ;;
     esac
-    if [ "$(uname -s)" = "Linux" ] && [ "$host_arch" = "$arch" ]; then
-        CGO_ENABLED=1 GOOS=linux GOARCH=$arch go build -ldflags "-X 'main.Version=$version'"
+    if [ "$(uname -s)" = "Linux" ] && [ "$host_arch" = "{{arch}}" ]; then
+        CGO_ENABLED=1 GOOS=linux GOARCH={{arch}} go build -ldflags "-X 'main.Version={{version}}'"
     else
-        echo "WARNING: Cross-compiling to linux/$arch, disabling CGO (DCGM unavailable)"
-        CGO_ENABLED=0 GOOS=linux GOARCH=$arch go build -ldflags "-X 'main.Version=$version' -extldflags '-static'"
+        echo "WARNING: Cross-compiling to linux/{{arch}}, disabling CGO (DCGM unavailable)"
+        CGO_ENABLED=0 GOOS=linux GOARCH={{arch}} go build -ldflags "-X 'main.Version={{version}}' -extldflags '-static'"
     fi
-    echo "Shim build (version: $version) complete!"
-
-# Build both runner and shim
-build-runner: build-runner-binary build-shim-binary
-    echo "Build complete! linux/$arch binaries are in their respective cmd directories."
+    echo "Shim build (version: {{version}}) complete!"
 
-# Clean build artifacts
-clean-runner:
-    rm -f {{source_directory()}}/cmd/runner/runner
-    rm -f {{source_directory()}}/cmd/shim/shim
-    echo "Build artifacts cleaned!"
-
-# Run tests for runner and shim (native; requires a Linux host)
-test-runner:
-    cd {{source_directory()}} && go test -v ./...
-
-# Run tests for runner and shim in a Linux container (use on macOS/Windows, where native builds are not available)
-# Examples:
-#   just test-runner-in-container                              # short suite, all packages
-#   just test-runner-in-container -run TestPullImage ./internal/shim/
-test-runner-in-container *args="-short ./...":
-    docker run --rm -t \
-        -v {{source_directory()}}:/src -w /src \
-        -v dstack-go-mod:/go/pkg/mod \
-        -v dstack-go-build:/root/.cache/go-build \
-        -v /var/run/docker.sock:/var/run/docker.sock \
-        golang:{{go_version}} \
-        go test -race {{args}}
-
-# Validate shim is built for the configured linux architecture
 [private]
-validate-shim-binary:
+[doc("Validate shim is built for the configured linux architecture")]
+[arg("arch", long)]
+validate-shim-binary arch=arch:
     #!/usr/bin/env bash
     set -e
-    case "$arch" in
+    case "{{arch}}" in
         amd64) expected="x86-64" ;;
         arm64) expected="ARM aarch64" ;;
-        *) echo "Error: Unsupported arch '$arch'"; exit 1 ;;
+        *) echo "Error: Unsupported arch '{{arch}}'"; exit 1 ;;
     esac
-    if ! file {{source_directory()}}/cmd/shim/shim | grep -q "ELF 64-bit LSB executable, $expected"; then
-        echo "Error: Shim must be built for linux/$arch for upload"
+    if [[ ! -f ./cmd/shim/shim ]]; then
+        echo "Error: Shim binary not found"
+        exit 1
+    fi
+    if ! file ./cmd/shim/shim | grep -q "ELF 64-bit LSB executable, $expected"; then
+        echo "Error: Shim must be built for linux/{{arch}} for upload"
         exit 1
     fi
 
-# Upload both runner and shim to S3
-upload-runner: upload-runner-binary upload-shim-binary
+[private]
+[doc("Upload runner to S3")]
+[arg("arch", long)]
+upload-runner-binary arch=arch: (build-runner-binary arch)
+    aws s3 cp ./cmd/runner/runner s3://{{s3_bucket}}/{{version}}/binaries/dstack-runner-linux-{{arch}} --acl public-read
+    @echo "Uploaded runner to S3"
 
-# Upload runner to S3
 [private]
-upload-runner-binary:
-    #!/usr/bin/env bash
-    set -e
-    just build-runner-binary
-    aws s3 cp {{source_directory()}}/cmd/runner/runner "{{runner_download_url}}" --acl public-read
-    echo "Uploaded runner to S3"
+[doc("Upload shim to S3")]
+[arg("arch", long)]
+upload-shim-binary arch=arch: (build-shim-binary arch) (validate-shim-binary arch)
+    aws s3 cp ./cmd/shim/shim s3://{{s3_bucket}}/{{version}}/binaries/dstack-shim-linux-{{arch}} --acl public-read
+    @echo "Uploaded shim to S3"
 
-# Upload shim to S3
+[default]
 [private]
-upload-shim-binary:
-    #!/usr/bin/env bash
-    set -e
-    just build-shim-binary
-    just validate-shim-binary
-    aws s3 cp {{source_directory()}}/cmd/shim/shim "{{shim_download_url}}" --acl public-read
-    echo "Uploaded shim to S3"
+default:
+    @just --list --unsorted

runner/README.md

@@ -9,7 +9,7 @@ For overview of `dstack-shim` and `dstack-runner`, see [/contributing/RUNNER-AND
 Run shim and runner tests on any OS inside a Docker container:
 
 ```shell
-just test-runner-in-container
+just test-in-container
 ```
 
 ## Running locally (standalone)
@@ -55,27 +55,33 @@ You can test the built shim and runner with `dstack` using standard backends (in
 > [!NOTE]
 > To run with standard backends, both the runner and shim must be built for linux.
 
-Build the runner and shim and upload them to S3 using `just` (see [`justfile`](justfile)).
+Build the runner and shim and upload them to S3 using `just` (see [`.justfile`](.justfile)).
 
 > [!IMPORTANT]
-> Before running any `just` commands that upload to S3, you must set the following environment variables:
+> Before running any `just` commands that upload to S3, configure the upload via environment variables:
 >
 > ```shell
 > export DSTACK_SHIM_UPLOAD_VERSION="your-version"
 > export DSTACK_SHIM_UPLOAD_S3_BUCKET="your-bucket"
+> export DSTACK_SHIM_BUILD_ARCH="arm64"  # Defaults to amd64 if not set
 > ```
 >
-> These variables are required and must be set before running any upload commands.
+> `DSTACK_SHIM_UPLOAD_VERSION` and `DSTACK_SHIM_UPLOAD_S3_BUCKET` are required and must be set before
+> running any upload commands. `DSTACK_SHIM_BUILD_ARCH` is optional.
+>
+> Set the target architecture via `DSTACK_SHIM_BUILD_ARCH`, not via `just upload --arch ...` — the
+> download URLs below are derived from the environment variable, so `--arch` would upload to one
+> architecture while the URLs point at another.
 
 ```shell
-just upload-runner
+just upload
 ```
 
 To use the built shim and runner with the `dstack` server, pass the URLs via `DSTACK_SHIM_DOWNLOAD_URL` and `DSTACK_RUNNER_DOWNLOAD_URL`:
 
 ```shell
-export DSTACK_SHIM_DOWNLOAD_URL="https://${DSTACK_SHIM_UPLOAD_S3_BUCKET}.s3.amazonaws.com/${DSTACK_SHIM_UPLOAD_VERSION}/binaries/dstack-shim-linux-amd64"
-export DSTACK_RUNNER_DOWNLOAD_URL="https://${DSTACK_SHIM_UPLOAD_S3_BUCKET}.s3.amazonaws.com/${DSTACK_SHIM_UPLOAD_VERSION}/binaries/dstack-runner-linux-amd64"
+export DSTACK_SHIM_DOWNLOAD_URL="https://${DSTACK_SHIM_UPLOAD_S3_BUCKET}.s3.amazonaws.com/${DSTACK_SHIM_UPLOAD_VERSION}/binaries/dstack-shim-linux-${DSTACK_SHIM_BUILD_ARCH:-amd64}"
+export DSTACK_RUNNER_DOWNLOAD_URL="https://${DSTACK_SHIM_UPLOAD_S3_BUCKET}.s3.amazonaws.com/${DSTACK_SHIM_UPLOAD_VERSION}/binaries/dstack-runner-linux-${DSTACK_SHIM_BUILD_ARCH:-amd64}"
 
 dstack server --log-level=debug
 ```