[5/6 Deps Update] Update all workflows and actions to use UV

.github/actions/setup-python-tools/action.yml

@@ -37,28 +37,16 @@ runs:
     using: "composite"
     # needs: comment-link-to-workflow  # Ensure that a comment is posted with workflow id
     steps:
-      # Step 1: Set up Python environment (Python 3.9.13).
-      - name: Set up Python
-        uses: actions/setup-python@v4
+      # Step 1: Set up Python environment
+      - name: "Set up Python"
+        uses: actions/setup-python@v6
         with:
-          # Available versions: https://raw.githubusercontent.com/actions/python-versions/main/versions-manifest.json
-          # Ensure to use a version that has support for arm64-darwin so we can build for Apple Silicon (macOS 14).
-          python-version: '3.9.13'
+          python-version-file: ".python-version"
 
-      # Step 2: Install pip-tools, which is used to generate hashed requirements.
-      # Note_1: pip 25.1 has a bug that causes pip-tools to fail with the following error:
-      # File ".../python3.9/site-packages/piptools/repositories/pypi.py", line 452, in allow_all_wheels
-      #   self.finder.find_all_candidates.cache_clear()
-      #   AttributeError: 'function' object has no attribute 'cache_clear'
-      # Note_2: Even though some wheels are guarded behind conditionals i.e. only use this if platform = linux;
-      # pip-tools 7.5.0 fails with the following error:
-      # pip._internal.exceptions.UnsupportedWheel: pyg_lib-0.4....linux_x86_64.whl is not a supported wheel on this platform.
-      # Thus, we fix the pip version to 25.0.1 and pip-tools version to 7.4.1.
-      - name: Install pip-tools
-        shell: bash
-        run: |
-          python -m pip install "pip==25.0.1"
-          python -m pip install "pip-tools==7.4.1"
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.9.5" # Matches the version in install_py_deps.sh
 
       # Step 3: Set up Gcloud AUTH using Workload Identity Federation
       # See following for context: https://cloud.google.com/blog/products/identity-security/enabling-keyless-authentication-from-github-actions

.github/cloud_builder/run_command_on_active_checkout.yaml

@@ -3,7 +3,7 @@ substitutions:
 options:
   logging: CLOUD_LOGGING_ONLY
 steps:
-  - name: us-central1-docker.pkg.dev/external-snap-ci-github-gigl/gigl-base-images/gigl-builder:6a94ae7cad3ec0c633246b0c9340a5095527deb9.63.2
+  - name: us-central1-docker.pkg.dev/external-snap-ci-github-gigl/gigl-base-images/gigl-builder:51af343c1c298ab465a96ecffd4e50ea6dffacb7.88.1
     entrypoint: /bin/bash
     args:
       - -c
@@ -18,15 +18,18 @@ steps:
 
         echo "Setting up environment..."
         # gcloud runner will run as a non-root user, but all paths/profiles, etc are set up for root
+        mkdir -p /builder/home/.local/bin
+        cp -r /root/.local/bin/ /builder/home/.local/
         echo "source /root/.bashrc" >> ~/.bashrc
         echo "source /root/.profile" >> ~/.profile
 
         source ~/.profile
+        docker version
         docker buildx create --driver=docker-container --use
         docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
         gcloud auth configure-docker us-central1-docker.pkg.dev
         # Install GiGL
-        pip install -e ./python/
+        uv pip install -e .
         # The builder operates in its own user dir, usually /workspace,
         # so we need to copy the gigl tools dir to the current cloud_builder's user dir.
         # See: containers/Dockerfile.builder.

.github/scripts/update_docker_image_refs.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Script to update dep_vars.env and cloud builder config with new Docker image references
+
+set -e
+
+echo "Writing new image names to dep_vars.env:"
+echo "  DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CUDA_IMAGE}"
+echo "  DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CPU_IMAGE}"
+echo "  DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=${GIGL_BASE_DATAFLOW_IMAGE}"
+echo "  DOCKER_LATEST_BUILDER_IMAGE_NAME_WITH_TAG=${GIGL_BUILDER_IMAGE}"
+
+sed -i "s|^DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=.*|DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CUDA_IMAGE}|" dep_vars.env
+sed -i "s|^DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=.*|DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CPU_IMAGE}|" dep_vars.env
+sed -i "s|^DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=.*|DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=${GIGL_BASE_DATAFLOW_IMAGE}|" dep_vars.env
+sed -i "s|name: us-central1-docker\.pkg\.dev.*|name: ${GIGL_BUILDER_IMAGE}|" .github/cloud_builder/run_command_on_active_checkout.yaml

.github/workflows/build-base-docker-images.yml

@@ -6,6 +6,7 @@ on:
       pr_number:
         description: 'PR to run the workflow on'
         required: true
+
 env:
   DOCKER_BUILDKIT: 1
   GIGL_BASE_CUDA_IMAGE: us-central1-docker.pkg.dev/${{ vars.GCP_PROJECT_ID }}/public-gigl/gigl-cuda-base:${{ github.sha }}.${{ github.run_number }}.${{ github.run_attempt }}
@@ -16,6 +17,7 @@ env:
 
 jobs:
   comment-workflow-started:
+
     runs-on: ubuntu-latest
     steps:
     - name: Comment on PR
@@ -29,7 +31,7 @@ jobs:
           Once done, the workflow will update the `dep_vars.env` file with the new image names.
 
   build-cuda-base-image:
-    runs-on: gigl-large-instances # x64 Ubuntu:latest w/ 4 cores, 16GB RAM, 150 GB SSD
+    runs-on: gigl-large-instances # x64 Ubuntu:latest w/ 8-cores, 32GB RAM, 300 GB SSD
     permissions:
       # Needed for gcloud auth: https://github.com/google-github-actions/auth
       contents: 'read'
@@ -41,7 +43,7 @@ jobs:
         github-token: ${{ secrets.GITHUB_TOKEN }}
         pr_number: ${{ inputs.pr_number }}
     - name: Setup Machine for building Docker images
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
         setup_gcloud: "true"
         try_cleaning_disk_space: "true"
@@ -56,8 +58,8 @@ jobs:
         docker push ${GIGL_BASE_CUDA_IMAGE}
         echo "Pushed CUDA base image to ${GIGL_BASE_CUDA_IMAGE}"
 
-  build-cpu-base-images:
-    runs-on: gigl-large-instances # x64 Ubuntu:latest w/ 4 cores, 16GB RAM, 150 GB SSD
+  build-cpu-base-image:
+    runs-on: ubuntu-latest
     permissions:
       # Needed for gcloud auth: https://github.com/google-github-actions/auth
       contents: 'read'
@@ -69,14 +71,13 @@ jobs:
         github-token: ${{ secrets.GITHUB_TOKEN }}
         pr_number: ${{ inputs.pr_number }}
     - name: Setup Machine for building Docker images
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
         setup_gcloud: "true"
         try_cleaning_disk_space: "true"
         gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
         workload_identity_provider: ${{ secrets.WORKLOAD_IDENTITY_PROVIDER }}
         gcp_service_account_email: ${{ secrets.GCP_SERVICE_ACCOUNT_EMAIL }}
-
     - name: Build and Push CPU Base Image and Docker CPU Image
       run: |
         gcloud auth configure-docker us-central1-docker.pkg.dev
@@ -85,8 +86,30 @@ jobs:
         docker push ${GIGL_BASE_CPU_IMAGE}
         echo "Pushed CPU base image to ${GIGL_BASE_CPU_IMAGE}"
 
-        echo "Will use CPU image ${GIGL_BASE_CPU_IMAGE} as base image for Dataflow image."
-        docker build -f ./containers/Dockerfile.dataflow.base --build-arg BASE_IMAGE=${GIGL_BASE_CPU_IMAGE} -t ${GIGL_BASE_DATAFLOW_IMAGE} .
+  build-dataflow-base-image:
+    runs-on: ubuntu-latest
+    permissions:
+      # Needed for gcloud auth: https://github.com/google-github-actions/auth
+      contents: 'read'
+      id-token: 'write'
+    steps:
+    - name: Checkout PR Branch
+      uses: snapchat/gigl/.github/actions/checkout-pr-branch@main
+      with:
+        github-token: ${{ secrets.GITHUB_TOKEN }}
+        pr_number: ${{ inputs.pr_number }}
+    - name: Setup Machine for building Docker images
+      uses: ./.github/actions/setup-python-tools
+      with:
+        setup_gcloud: "true"
+        try_cleaning_disk_space: "true"
+        gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
+        workload_identity_provider: ${{ secrets.WORKLOAD_IDENTITY_PROVIDER }}
+        gcp_service_account_email: ${{ secrets.GCP_SERVICE_ACCOUNT_EMAIL }}
+    - name: Build and Push Dataflow Base Image
+      run: |
+        gcloud auth configure-docker us-central1-docker.pkg.dev
+        docker build -f ./containers/Dockerfile.dataflow.base -t ${GIGL_BASE_DATAFLOW_IMAGE} .
         docker push ${GIGL_BASE_DATAFLOW_IMAGE}
         echo "Pushed Dataflow base image to ${GIGL_BASE_DATAFLOW_IMAGE}"
 
@@ -103,7 +126,7 @@ jobs:
         github-token: ${{ secrets.GITHUB_TOKEN }}
         pr_number: ${{ inputs.pr_number }}
     - name: Setup Machine for building Docker images
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
         setup_gcloud: "true"
         try_cleaning_disk_space: "true"
@@ -124,7 +147,8 @@ jobs:
   build-and-commit-base-images:
     needs:
       - build-cuda-base-image
-      - build-cpu-base-images
+      - build-cpu-base-image
+      - build-dataflow-base-image
       - build-builder-image
     runs-on: ubuntu-latest
     steps:
@@ -134,23 +158,12 @@ jobs:
         github-token: ${{ secrets.GITHUB_TOKEN }}
         pr_number: ${{ inputs.pr_number }}
         should_leave_progress_comments: "false"
-        command: |
-          echo "Writing new image names to dep_vars.env:"
-          echo "  DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CUDA_IMAGE}"
-          echo "  DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CPU_IMAGE}"
-          echo "  DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=${GIGL_BASE_DATAFLOW_IMAGE}"
-          echo "  DOCKER_LATEST_BUILDER_IMAGE_NAME_WITH_TAG=${GIGL_BUILDER_IMAGE}"
-          sed -i "s|^DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=.*|DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CUDA_IMAGE}|" dep_vars.env
-          sed -i "s|^DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=.*|DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=${GIGL_BASE_CPU_IMAGE}|" dep_vars.env
-          sed -i "s|^DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=.*|DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=${GIGL_BASE_DATAFLOW_IMAGE}|" dep_vars.env
-          sed -i "s|name: us-central1-docker\.pkg\.dev.*|name: ${GIGL_BUILDER_IMAGE}|" .github/cloud_builder/run_command_on_active_checkout.yaml
-
+        command: bash .github/scripts/update_docker_image_refs.sh
     - name: Commit and Push Dep Vars
       uses: snapchat/gigl/.github/actions/commit-and-push@main
       with:
         commit_message: "[AUTOMATED] Update dep.vars, and other relevant files with new image names"
         github_token: ${{ secrets.GITHUB_TOKEN }}
-
     - uses: snapchat/gigl/.github/actions/comment-on-pr@main
       with:
         pr_number: ${{ inputs.pr_number }}

.github/workflows/on-pr-comment.yml

@@ -24,7 +24,7 @@ jobs:
     - name: Setup Python
       uses: actions/setup-python@v4
       with:
-        python-version: '3.x'
+        python-version-file: ".python-version"
 
     - name: Install PyYAML
       run: pip install PyYAML
@@ -164,6 +164,7 @@ jobs:
         workload_identity_provider: ${{ secrets.WORKLOAD_IDENTITY_PROVIDER }}
         gcp_service_account_email: ${{ secrets.GCP_SERVICE_ACCOUNT_EMAIL }}
         command: |
+          # sourcing .profile is important to resolve paths for java, sbt, et al.
+          # It is setup in the setup-python-tools action.
           source ~/.profile
-          make check_format
-          make assert_yaml_configs_parse
+          make lint_test

.github/workflows/on-pr-merge.yml

@@ -26,7 +26,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Setup development environment
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
           setup_gcloud: "true"
           gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
@@ -37,7 +37,7 @@ jobs:
       # using GFile library (a.k.a anything that does IO w/ Tensorflow). GFile does not understand
       # how to leverage Workload Identity Federation to read assets from GCS, et al. See:
       # https://github.com/tensorflow/tensorflow/issues/57104
-      uses: snapchat/gigl/.github/actions/run-cloud-run-command-on-active-checkout@main
+      uses: ./.github/actions/run-cloud-run-command-on-active-checkout
       with:
         cmd: "make unit_test_py"
         service_account:  ${{ secrets.gcp_service_account_email }}
@@ -53,14 +53,18 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Setup development environment
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
           setup_gcloud: "true"
           gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
           workload_identity_provider: ${{ secrets.workload_identity_provider }}
           gcp_service_account_email: ${{ secrets.gcp_service_account_email }}
     - name: Run Scala Unit Tests
-      uses: snapchat/gigl/.github/actions/run-cloud-run-command-on-active-checkout@main
+      # We use cloud run here instead of using github hosted runners because of limitation of tests
+      # using GFile library (a.k.a anything that does IO w/ Tensorflow). GFile does not understand
+      # how to leverage Workload Identity Federation to read assets from GCS, et al. See:
+      # https://github.com/tensorflow/tensorflow/issues/57104
+      uses: ./.github/actions/run-cloud-run-command-on-active-checkout
       with:
         cmd: "make unit_test_scala"
         service_account:  ${{ secrets.gcp_service_account_email }}
@@ -72,14 +76,14 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Setup development environment
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
           setup_gcloud: "true"
           gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
           workload_identity_provider: ${{ secrets.workload_identity_provider }}
           gcp_service_account_email: ${{ secrets.gcp_service_account_email }}
     - name: Run Integration Tests
-      uses: snapchat/gigl/.github/actions/run-cloud-run-command-on-active-checkout@main
+      uses: ./.github/actions/run-cloud-run-command-on-active-checkout
       with:
         cmd: "make integration_test"
         service_account:  ${{ secrets.gcp_service_account_email }}
@@ -91,14 +95,14 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Setup development environment
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
           setup_gcloud: "true"
           gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
           workload_identity_provider: ${{ secrets.workload_identity_provider }}
           gcp_service_account_email: ${{ secrets.gcp_service_account_email }}
     - name: Run E2E Tests
-      uses: snapchat/gigl/.github/actions/run-cloud-run-command-on-active-checkout@main
+      uses: ./.github/actions/run-cloud-run-command-on-active-checkout
       with:
         cmd: "make run_all_e2e_tests"
         service_account:  ${{ secrets.gcp_service_account_email }}
@@ -111,14 +115,14 @@ jobs:
   #   steps:
   #   - uses: actions/checkout@v4
   #   - name: Setup development environment
-  #     uses: snapchat/gigl/.github/actions/setup-python-tools@main
+  #     uses: ./.github/actions/setup-python-tools
   #     with:
   #         setup_gcloud: "true"
   #         gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
   #         workload_identity_provider: ${{ secrets.workload_identity_provider }}
   #         gcp_service_account_email: ${{ secrets.gcp_service_account_email }}
   #   - name: Run Example Notebook E2E Tests
-  #     uses: snapchat/gigl/.github/actions/run-cloud-run-command-on-active-checkout@main
+  #     uses: ./.github/actions/run-cloud-run-command-on-active-checkout
   #     with:
   #       cmd: "make notebooks_test"
   #       service_account:  ${{ secrets.gcp_service_account_email }}
@@ -130,15 +134,17 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Setup development environment
-      uses: snapchat/gigl/.github/actions/setup-python-tools@main
+      uses: ./.github/actions/setup-python-tools
       with:
           install_dev_deps: "true"
           setup_gcloud: "true"
           gcp_project_id: ${{ vars.GCP_PROJECT_ID }}
           workload_identity_provider: ${{ secrets.workload_identity_provider }}
           gcp_service_account_email: ${{ secrets.gcp_service_account_email }}
     - name: Run Lint Tests
+      shell: bash
       run: |
+        # sourcing .profile is important to resolve paths for java, sbt, et al.
+        # It is setup in the setup-python-tools action.
         source ~/.profile
-        make check_format
-        make assert_yaml_configs_parse
+        make lint_test

.github/workflows/release-documentation.yml

@@ -44,7 +44,7 @@ jobs:
       # We also make gigl available w/ editable install `-e` so that autodoc can find it.
       - name: Install necessary doc dependencies
         run: |
-          pip install -e "./python[docs]"
+          uv sync --extra docs
       - name: Sphinx build
         run: |
           make build_docs

Makefile

@@ -151,7 +151,7 @@ format: format_py format_scala format_md
 type_check:
 	uv run mypy ${PYTHON_DIRS} --check-untyped-defs
 
-lint_test: check_format assert_yaml_config_parse
+lint_test: check_format assert_yaml_configs_parse
 	@echo "Lint checks pass!"
 
 # compiles current working state of scala projects to local jars

dep_vars.env

@@ -1,7 +1,7 @@
 # Note this file only supports static key value pairs so it can be loaded by make, bash, python, and sbt without any additional parsing.
-DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cuda-base:6a94ae7cad3ec0c633246b0c9340a5095527deb9.63.2
-DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cpu-base:6a94ae7cad3ec0c633246b0c9340a5095527deb9.63.2
-DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-dataflow-base:6a94ae7cad3ec0c633246b0c9340a5095527deb9.63.2
+DOCKER_LATEST_BASE_CUDA_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cuda-base:51af343c1c298ab465a96ecffd4e50ea6dffacb7.88.1
+DOCKER_LATEST_BASE_CPU_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-cpu-base:51af343c1c298ab465a96ecffd4e50ea6dffacb7.88.1
+DOCKER_LATEST_BASE_DATAFLOW_IMAGE_NAME_WITH_TAG=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/gigl-dataflow-base:51af343c1c298ab465a96ecffd4e50ea6dffacb7.88.1
 
 DEFAULT_GIGL_RELEASE_SRC_IMAGE_CUDA=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/src-cuda:0.0.11
 DEFAULT_GIGL_RELEASE_SRC_IMAGE_CPU=us-central1-docker.pkg.dev/external-snap-ci-github-gigl/public-gigl/src-cpu:0.0.11

deployment/configs/e2e_cicd_resource_config.yaml

@@ -44,7 +44,7 @@ split_generator_config:
 trainer_config:
   vertex_ai_trainer_config:
     machine_type: "n1-highmem-8" # set to `ACCELERATOR_TYPE_UNSPECIFIED` for cpu training
-    gpu_type: nvidia-tesla-p100
+    gpu_type: NVIDIA_TESLA_T4
     gpu_limit: 1  # set to 0 for cpu training
     num_replicas: 2
 inferencer_config:

deployment/configs/unittest_resource_config.yaml

@@ -46,7 +46,7 @@ split_generator_config:
 trainer_config:
   vertex_ai_trainer_config:
     machine_type: "n1-highmem-8"
-    gpu_type: nvidia-tesla-p100  # set to `ACCELERATOR_TYPE_UNSPECIFIED` for cpu training
+    gpu_type: NVIDIA_TESLA_T4  # set to `ACCELERATOR_TYPE_UNSPECIFIED` for cpu training
     gpu_limit: 1  # set to 0 for cpu training
     num_replicas: 2
 inferencer_config: