JAX inference offloading bridge

.github/container/git-clone.sh

@@ -10,6 +10,8 @@ Usage: $0 [OPTION]... GIT_URL#REF PATH
   -h, --help             Print usage.
   -m, --manifest FILE    The manifest yaml file to which the downloaded library is documented.
                          Default is /opt/manifest.d/git-clone.yaml
+  --sparse-path PATH     If specified, use sparse checkout to only fetch the given path.
+                         PATH is relative to the repo root.
 
 Example:
   # clone JAX's main branch at /opt/jax and update manifest file at the default location
@@ -20,14 +22,15 @@ EOF
     exit $1
 }
 
-args=$(getopt -o hm: --long help,manifest: -- "$@")
+args=$(getopt -o hm: --long help,manifest:,sparse-path: -- "$@")
 if [[ $? -ne 0 ]]; then
     exit 1
 fi
 
 ## Set default arguments
 
 MANIFEST="/opt/manifest.d/git-clone.yaml"
+SPARSE_PATH=""
 
 eval set -- "$args"
 while [ : ]; do
@@ -39,6 +42,10 @@ while [ : ]; do
         MANIFEST="$2"
         shift 2
         ;;
+    --sparse-path)
+        SPARSE_PATH="$2"
+        shift 2
+        ;;
     --)
         shift;
         break
@@ -72,7 +79,15 @@ echo "Fetching $GIT_REPO#$GIT_REF to $DESTINATION"
 
 set -ex -o pipefail
 
+if [[ -n "${SPARSE_PATH}" ]]; then
+    git clone --no-checkout ${GIT_REPO} ${DESTINATION}
+    pushd ${DESTINATION}
+    git sparse-checkout init --cone
+    git sparse-checkout set ${SPARSE_PATH}
+    popd
+else
 git clone ${GIT_REPO} ${DESTINATION}
+fi
 pushd ${DESTINATION}
 if [[ -n "${GIT_REF}" ]]; then
     git checkout ${GIT_REF}

.github/workflows/jio.yaml

@@ -0,0 +1,91 @@
+name: JAX Inference Offloading
+
+on:
+  schedule:
+    - cron: '30 9 * * *'  # Pacific Time 01:30 AM in UTC
+  pull_request:
+    types:
+      - opened
+      - reopened
+      - ready_for_review
+      - synchronize
+    paths:
+      - '.github/workflows/jio.yaml'
+      - 'jax-inference-offloading/**'
+  workflow_dispatch:
+    inputs:
+      PUBLISH:
+        type: boolean
+        description: Publish dated images and update the 'latest' tag?
+        default: false
+        required: false      
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read        # to fetch code
+  actions:  write       # to cancel previous workflows
+  packages: write       # to upload containers
+
+jobs:
+  metadata:
+    runs-on: ubuntu-22.04
+    outputs:
+      BUILD_DATE: ${{ steps.date.outputs.BUILD_DATE }}
+      PUBLISH: ${{ steps.if-publish.outputs.PUBLISH }}
+    steps:
+      - name: Cancel workflow run if the trigger is a draft PR
+        id: cancel-if-draft
+        if: github.event_name == 'pull_request' && github.event.pull_request.draft == true
+        run: |
+          echo "Cancelling workflow for draft PR"
+          curl -X POST -H "Authorization: token ${{ github.token }}" \
+                -H "Accept: application/vnd.github.v3+json" \
+                "https://api.github.com/repos/${{ github.repository }}/actions/runs/${{ github.run_id }}/cancel"
+          while true; do sleep 1; done  # blocks execution in case workflow cancellation takes time
+
+      - name: Set build date
+        id: date
+        shell: bash -x -e {0}
+        run: |
+          BUILD_DATE=$(TZ='US/Los_Angeles' date '+%Y-%m-%d')
+          echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT
+
+      - name: Determine whether results will be 'published'
+        id: if-publish
+        shell: bash -x -e {0}
+        run: |
+          echo "PUBLISH=${{ github.event_name == 'schedule' || inputs.PUBLISH }}" >> $GITHUB_OUTPUT
+
+  build:
+    needs: metadata
+    strategy:
+      fail-fast: true
+      matrix:
+        ARCHITECTURE: [amd64, arm64]
+    runs-on: [self-hosted, "${{ matrix.ARCHITECTURE }}", "small"]
+    steps:
+        - name: Checkout repository
+          uses: actions/checkout@v4
+        - name: Build container
+          id: build-container
+          uses: ./.github/actions/build-container
+          with:
+            ARCHITECTURE: ${{ matrix.ARCHITECTURE }}
+            ARTIFACT_NAME: artifact-jio-build
+            BADGE_FILENAME: badge-jio-build
+            BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+            CONTAINER_NAME: jio
+            DOCKERFILE: jax-inference-offloading/dockerfile/oss.dockerfile
+            RUNNER_SIZE: small
+            ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
+            ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
+            github-token: ${{ secrets.GITHUB_TOKEN }}
+            # EXTRA_BUILD_ARGS: |
+            #   URLREF_TUNIX=${{ fromJson(inputs.SOURCE_URLREFS).TUNIX }}
+
+    outputs:
+      DOCKER_TAG_MEALKIT: ${{ steps.build-container.outputs.DOCKER_TAG_MEALKIT }}
+      DOCKER_TAG_FINAL:   ${{ steps.build-container.outputs.DOCKER_TAG_FINAL }}

jax-inference-offloading/.flake8

@@ -0,0 +1,26 @@
+[flake8]
+indent-size = 2
+max-line-length = 120
+extend-ignore = E203,W503  # match Black / NumPy convention
+
+# -------- Plugins -------------
+# Activate popular extra rule-sets; install alongside flake8:
+#   pip install flake8-bugbear flake8-docstrings flake8-import-order flake8-builtins
+select =  C,E,F,W,B,B950,D,I
+# C (mccabe), E/W (pycodestyle), F (pyflakes), B (bugbear), B950 (long line hint),
+# D (pydocstyle), I (import order)
+
+# -------- Per‑file tweaks ------
+per-file-ignores =
+  __init__.py:F401       # allow re‑exported symbols
+  tests/*:D              # relax docstring rules in test helpers
+
+# -------- Exclusions -----------
+exclude =
+  .git,
+  __pycache__,
+  build,
+  dist,
+  .venv,
+  notebooks,
+

jax-inference-offloading/.gitignore

@@ -0,0 +1,8 @@
+*.swp
+*.swo
+*.swn
+__pycache__
+.nfs*
+*.egg-info
+build
+.env

jax-inference-offloading/CONTRIBUTING.md

@@ -0,0 +1,16 @@
+# Signing Your Work
+
+* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
+
+  * Any contribution which contains commits that are not Signed-Off will not be accepted.
+
+* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes:
+  ```bash
+  $ git commit -s -m "Add cool feature."
+  ```
+  This will append the following to your commit message:
+  ```
+  Signed-off-by: Your Name <[email protected]>
+  ```
+
+* Full text of the DCO can be found [here](DCO)

jax-inference-offloading/DCO

@@ -0,0 +1,34 @@
+Developer Certificate of Origin
+Version 1.1
+
+Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
+
+Everyone is permitted to copy and distribute verbatim copies of this
+license document, but changing it is not allowed.
+
+
+Developer's Certificate of Origin 1.1
+
+By making a contribution to this project, I certify that:
+
+(a) The contribution was created in whole or in part by me and I
+    have the right to submit it under the open source license
+    indicated in the file; or
+
+(b) The contribution is based upon previous work that, to the best
+    of my knowledge, is covered under an appropriate open source
+    license and I have the right under that license to submit that
+    work with modifications, whether created in whole or in part
+    by me, under the same open source license (unless I am
+    permitted to submit under a different license), as indicated
+    in the file; or
+
+(c) The contribution was provided directly to me by some other
+    person who certified (a), (b) or (c) and I have not modified
+    it.
+
+(d) I understand and agree that this project and the contribution
+    are public and that a record of the contribution (including all
+    personal information I submit with it, including my sign-off) is
+    maintained indefinitely and may be redistributed consistent with
+    this project or the open source license(s) involved.