[WIP] feat: Implement controller support for polling and tracking training job progression

.flake8

@@ -1,3 +1,3 @@
 [flake8]
 max-line-length = 100
-extend-ignore = W503
+extend-ignore = W503, E203

.gitattributes

@@ -0,0 +1,3 @@
+**/zz_generated.*.go linguist-generated=true
+api/** linguist-generated=true
+pkg/client/** linguist-generated=true

.github/ISSUE_TEMPLATE/bug_report.yaml

@@ -33,7 +33,7 @@ body:
         ```
         Kubeflow Trainer version:
         ```bash
-        $ kubectl get pods -n kubeflow -l app.kubernetes.io/name=trainer -o jsonpath="{.items[*].spec.containers[*].image}"
+        $ kubectl get pods -n kubeflow-system -l app.kubernetes.io/name=kubeflow-trainer -o jsonpath="{.items[*].spec.containers[*].image}"
 
         ```
         Kubeflow Python SDK version:

.github/workflows/build-and-push-images.yaml

@@ -4,9 +4,9 @@ on:
   push:
     branches:
       - master
-      - 'release-*'
+      - "release-*"
     tags:
-      - 'v*'
+      - "v*"
   pull_request:
 
 jobs:
@@ -34,12 +34,16 @@ jobs:
           - component-name: deepspeed-runtime
             dockerfile: cmd/runtimes/deepspeed/Dockerfile
             platforms: linux/amd64,linux/arm64
+          # TODO (andreyvelich): mlx[cuda] doesn't support arm at the moment: https://github.com/ml-explore/mlx/issues/2469
           - component-name: mlx-runtime
             dockerfile: cmd/runtimes/mlx/Dockerfile
-            platforms: linux/arm64
+            platforms: linux/amd64
           - component-name: torchtune-trainer
             dockerfile: cmd/trainers/torchtune/Dockerfile
             platforms: linux/amd64,linux/arm64
+          - component-name: data-cache
+            dockerfile: cmd/data_cache/Dockerfile
+            platforms: linux/amd64,linux/arm64
     steps:
       - name: Checkout
         uses: actions/checkout@v4

.github/workflows/check-pr-title.yaml

@@ -0,0 +1,44 @@
+name: Check PR Title
+
+on:
+  pull_request_target:
+    branches:
+      - master
+    types:
+      - edited
+      - opened
+      - reopened
+      - synchronize
+
+jobs:
+  check:
+    name: Check PR Title
+    runs-on: ubuntu-latest
+    steps:
+      - uses: amannn/action-semantic-pull-request@v5.5.3
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          types: |
+            chore
+            fix
+            feat
+            revert
+
+          scopes: |
+            api
+            ci
+            docs
+            examples
+            exporter
+            initializer
+            manifests
+            operator
+            runtimes
+            scripts
+            test
+            cache
+          requireScope: false
+          ignoreLabels: |
+            do-not-merge/work-in-progress
+            dependencies

.github/workflows/gh-workflow-approve.yaml

@@ -0,0 +1,72 @@
+name: Approve Workflow Runs
+
+permissions:
+  actions: write
+  contents: read
+
+on:
+  pull_request_target:
+    types:
+      - labeled
+      - synchronize
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event.number }}
+  cancel-in-progress: true
+
+jobs:
+  ok-to-test:
+    if: contains(github.event.pull_request.labels.*.name, 'ok-to-test') || github.event_name == 'pull_request_target'
+    runs-on: ubuntu-latest
+    continue-on-error: true
+
+    steps:
+      - name: Check if author is a Kubeflow GitHub member
+        id: membership-check
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const username = context.payload.pull_request.user.login;
+            const org = context.repo.owner;
+            try {
+              const res = await github.rest.orgs.checkMembershipForUser({
+                org,
+                username
+              });
+              core.setOutput("is_member", true);
+            } catch (error) {
+              if (error.status === 404) {
+                // User is not a member
+                core.setOutput("is_member", false);
+              } else {
+                throw error;
+              }
+            }
+
+      - name: Approve Pending Workflow Runs
+        if: steps.membership-check.outputs.is_member == 'true' || contains(github.event.pull_request.labels.*.name, 'ok-to-test')
+        uses: actions/github-script@v7
+        with:
+          retries: 3
+          script: |
+            const request = {
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              event: "pull_request",
+              status: "action_required",
+              head_sha: context.payload.pull_request.head.sha,
+            }
+
+            core.info(`Getting workflow runs that need approval for commit ${request.head_sha}`)
+            const runs = await github.paginate(github.rest.actions.listWorkflowRunsForRepo, request)
+
+            core.info(`Found ${runs.length} workflow runs that need approval`)
+            for (const run of runs) {
+              core.info(`Approving workflow run ${run.id}`)
+              const request = {
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                run_id: run.id,
+              }
+              await github.rest.actions.approveWorkflowRun(request)
+            }

.github/workflows/publish-helm-charts.yaml

@@ -0,0 +1,53 @@
+name: Publish Helm Charts
+
+on:
+  push:
+    branches:
+      - master
+    tags:
+      - "v*"
+
+env:
+  CHART_PATH: charts/kubeflow-trainer
+
+jobs:
+  publish:
+    permissions:
+      contents: read
+      packages: write
+
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Update appVersion in Chart.yaml
+        run: |
+          if [[ "${{ github.ref }}" == refs/tags/* ]]; then
+            APP_VERSION="${GITHUB_REF#refs/tags/}"
+          else
+            APP_VERSION="sha-$(git rev-parse --short=7 HEAD)"
+          fi
+
+          echo "appVersion: ${APP_VERSION}" >> ${{ env.CHART_PATH }}/Chart.yaml
+          echo "Updated Chart.yaml with appVersion: ${APP_VERSION}"
+
+      - name: Build, package, and push Helm chart
+        run: |
+          if [[ "${{ github.ref }}" == refs/tags/* ]]; then
+            VERSION=$(grep '^version:' ${{ env.CHART_PATH }}/Chart.yaml | awk '{print $2}')
+          else
+            VERSION="0.0.0-sha-$(git rev-parse --short=7 HEAD)"
+          fi
+          echo "Chart version set to: ${VERSION}"
+
+          helm dependency build ${{ env.CHART_PATH }}
+          helm package ${{ env.CHART_PATH }} --version ${VERSION}
+          helm push kubeflow-trainer-${VERSION}.tgz oci://ghcr.io/kubeflow/charts

.github/workflows/template-publish-image/action.yaml

@@ -82,7 +82,7 @@ runs:
         images: ${{ inputs.image }}
         tags: |
           type=ref,event=tag
-          type=raw,latest
+          type=raw,value=latest,enable={{is_default_branch}}
           type=sha
 
     - name: Build and Push

.github/workflows/test-e2e-gpu.yaml

@@ -0,0 +1,93 @@
+name: GPU E2E Test
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize, labeled]
+
+permissions:
+  contents: read
+  pull-requests: read
+
+jobs:
+  gpu-e2e-test:
+    name: GPU E2E Test
+    runs-on: oracle-vm-16cpu-a10gpu-240gb
+
+    env:
+      GOPATH: ${{ github.workspace }}/go
+
+    defaults:
+      run:
+        working-directory: ${{ env.GOPATH }}/src/github.com/kubeflow/trainer
+
+    strategy:
+      fail-fast: false
+      matrix:
+        kubernetes-version: ["1.33.1"]
+
+    steps:
+      - name: Check GPU label
+        id: check-label
+        run: |
+          if [[ "${{ join(github.event.pull_request.labels.*.name, ',') }}" != *"ok-to-test-gpu-runner"* ]]; then
+            echo "✅ Skipping GPU E2E tests (label not present)."
+            echo "skip=true" >> $GITHUB_OUTPUT
+            exit 0
+          else
+            echo "Label found. Requesting environment approval to run GPU tests."
+            echo "skip=false" >> $GITHUB_OUTPUT
+          fi
+
+      - name: Check out code
+        if: steps.check-label.outputs.skip == 'false'
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+          path: ${{ env.GOPATH }}/src/github.com/kubeflow/trainer
+
+      - name: Setup Go
+        if: steps.check-label.outputs.skip == 'false'
+        uses: actions/setup-go@v5
+        with:
+          go-version-file: ${{ env.GOPATH }}/src/github.com/kubeflow/trainer/go.mod
+
+      - name: Setup Python
+        if: steps.check-label.outputs.skip == 'false'
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.11
+
+      - name: Install dependencies
+        if: steps.check-label.outputs.skip == 'false'
+        run: |
+          pip install papermill==2.6.0 jupyter==1.1.1 ipykernel==6.29.5
+          pip install git+https://github.com/kubeflow/sdk.git@main
+
+      - name: Setup cluster with GPU support using nvidia/kind
+        if: steps.check-label.outputs.skip == 'false'
+        run: |
+          make test-e2e-setup-gpu-cluster K8S_VERSION=${{ matrix.kubernetes-version }}
+
+      - name: Run e2e test on GPU cluster
+        if: steps.check-label.outputs.skip == 'false'
+        run: |
+          mkdir -p artifacts/notebooks
+          make test-e2e-notebook NOTEBOOK_INPUT=./examples/torchtune/qwen2_5/qwen2.5-1.5B-with-alpaca.ipynb NOTEBOOK_OUTPUT=./artifacts/notebooks/${{ matrix.kubernetes-version }}_qwen2_5_with_alpaca-trainjob-yaml.ipynb TIMEOUT=900
+
+      - name: Upload Artifacts to GitHub
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.kubernetes-version }}
+          path: ${{ env.GOPATH }}/src/github.com/kubeflow/trainer/artifacts/*
+          retention-days: 1
+
+  delete-kind-cluster:
+    name: Delete kind Cluster
+    runs-on: oracle-vm-16cpu-a10gpu-240gb
+    needs: [gpu-e2e-test]
+    if: always()
+    steps:
+      - name: Delete any existing kind cluster
+        run: |
+          sudo kind delete cluster --name kind-gpu && echo "kind cluster has been deleted" || echo "kind cluster doesn't exist"

.github/workflows/test-e2e.yaml

@@ -17,7 +17,7 @@ jobs:
       fail-fast: false
       matrix:
         # Kubernetes versions for e2e tests on Kind cluster.
-        kubernetes-version: ["1.29.14", "1.30.0", "1.31.0", "1.32.3"]
+        kubernetes-version: ["1.31.0", "1.32.3", "1.33.1", "1.34.0"]
 
     steps:
       - name: Check out code
@@ -40,8 +40,8 @@ jobs:
           echo "Install Papermill"
           pip install papermill==2.6.0 jupyter==1.1.1 ipykernel==6.29.5
 
-          echo "Install Kubeflow SDK"
-          pip install git+https://github.com/kubeflow/sdk.git@main#subdirectory=python
+          echo "Install Kubeflow SDK with Docker support for local Notebooks"
+          pip install "kubeflow[docker] @ git+https://github.com/kubeflow/sdk.git@main"
 
       - name: Setup cluster
         run: |
@@ -56,6 +56,8 @@ jobs:
           mkdir -p artifacts/notebooks
           make test-e2e-notebook NOTEBOOK_INPUT=./examples/pytorch/image-classification/mnist.ipynb NOTEBOOK_OUTPUT=./artifacts/notebooks/${{ matrix.kubernetes-version }}_mnist.ipynb TIMEOUT=900
           make test-e2e-notebook NOTEBOOK_INPUT=./examples/pytorch/question-answering/fine-tune-distilbert.ipynb NOTEBOOK_OUTPUT=./artifacts/notebooks/${{ matrix.kubernetes-version }}_fine-tune-distilbert.ipynb TIMEOUT=900
+          make test-e2e-notebook NOTEBOOK_INPUT=./examples/local/local-training-mnist.ipynb NOTEBOOK_OUTPUT=./artifacts/notebooks/${{ matrix.kubernetes-version }}_local-training-mnist.ipynb TIMEOUT=900
+          make test-e2e-notebook NOTEBOOK_INPUT=./examples/local/local-container-mnist.ipynb NOTEBOOK_OUTPUT=./artifacts/notebooks/${{ matrix.kubernetes-version }}_local-container-mnist.ipynb TIMEOUT=900
 
       # TODO (andreyvelich): Discuss how we can upload artifacts for multiple Notebooks.
       - name: Upload Artifacts to GitHub

.github/workflows/test-go.yaml

@@ -8,6 +8,7 @@ jobs:
   generate:
     name: Generate
     runs-on: ubuntu-latest
+    if: ${{ github.repository == 'kubeflow/trainer' }}
     env:
       GOPATH: ${{ github.workspace }}/go
     defaults:
@@ -47,6 +48,7 @@ jobs:
   test:
     name: Test
     runs-on: ubuntu-latest
+    if: ${{ github.repository == 'kubeflow/trainer' }}
     env:
       GOPATH: ${{ github.workspace }}/go
     defaults:
@@ -77,4 +79,3 @@ jobs:
         with:
           path-to-profile: cover.out
           working-directory: ${{ env.GOPATH }}/src/github.com/kubeflow/trainer
-          parallel: true