PR into #488

.github/workflows/lint.yml

@@ -27,12 +27,51 @@ jobs:
         run: |
           make fmt
           git diff --exit-code
+  check-flyteidl2-versions:
+    name: check flyteidl2 versions
+    runs-on: ubuntu-latest
+    steps:
+      - name: Fetch the code
+        uses: actions/checkout@v4
+      - name: Check flyteidl2 version consistency
+        run: |
+          # Extract flyteidl2 version from root pyproject.toml
+          ROOT_VER=$(grep 'flyteidl2==' pyproject.toml | head -1 | sed 's/.*flyteidl2==\([^"]*\).*/\1/')
+          echo "Root pyproject.toml: flyteidl2==$ROOT_VER"
+
+          # Extract flyteidl2 version from rs_controller/Cargo.toml
+          CARGO_VER=$(grep 'flyteidl2' rs_controller/Cargo.toml | grep -v '^#' | sed 's/.*"=\(.*\)".*/\1/')
+          echo "rs_controller/Cargo.toml: flyteidl2=$CARGO_VER"
+
+          # Extract flyteidl2 version from rs_controller/pyproject.toml
+          RS_VER=$(grep 'flyteidl2==' rs_controller/pyproject.toml | head -1 | sed 's/.*flyteidl2==\([^"]*\).*/\1/')
+          echo "rs_controller/pyproject.toml: flyteidl2==$RS_VER"
+
+          # Compare all three
+          if [ "$ROOT_VER" != "$CARGO_VER" ] || [ "$ROOT_VER" != "$RS_VER" ]; then
+            echo "ERROR: flyteidl2 versions do not match!"
+            echo "  pyproject.toml:                $ROOT_VER"
+            echo "  rs_controller/Cargo.toml:      $CARGO_VER"
+            echo "  rs_controller/pyproject.toml:  $RS_VER"
+            exit 1
+          fi
+          echo "All flyteidl2 versions match: $ROOT_VER"
   rs-fmt:
     name: rust fmt
     runs-on: ubuntu-latest
     steps:
       - name: Fetch the code
         uses: actions/checkout@v4
+      - name: Cache Cargo registry and build
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            rs_controller/target
+          key: ${{ runner.os }}-cargo-fmt-${{ hashFiles('rs_controller/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-fmt-
       - name: Install nightly toolchain
         run: |
           rustup toolchain install nightly
@@ -46,6 +85,16 @@ jobs:
     steps:
       - name: Fetch the code
         uses: actions/checkout@v4
+      - name: Cache Cargo registry and build
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            rs_controller/target
+          key: ${{ runner.os }}-cargo-lint-${{ hashFiles('rs_controller/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-lint-
       - name: Install toolchain
         run: |
           rustup toolchain install

.github/workflows/publish.yml

@@ -3,14 +3,95 @@ name: Publish
 on:
   release:
     types: [published]
+  push:
+    branches: [revisit-protos-ci]
   pull_request:
     paths:
       - ".github/workflows/publish.yml"
       - "maint_tools/build_default_image.py"
 
 jobs:
+  rs-controller-wheels:
+    name: Build RS controller wheel (${{ matrix.target }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - target: x86_64
+            os: ubuntu-latest
+          - target: aarch64
+            os: ubuntu-24.04-arm
+          - target: aarch64-apple-darwin
+            os: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: "0"
+      - name: Set version from tag
+        run: |
+          if [[ "$GITHUB_REF" == refs/tags/v* ]]; then
+            VERSION=$(echo "$GITHUB_REF" | sed 's|refs/tags/v||')
+          else
+            VERSION="0.0.0.dev0"
+          fi
+          echo "VERSION=$VERSION" >> $GITHUB_ENV
+          sed "s/^version = .*/version = \"$VERSION\"/" rs_controller/pyproject.toml > tmp && mv tmp rs_controller/pyproject.toml
+          echo "Set version to $VERSION"
+          cat rs_controller/pyproject.toml
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: ${{ matrix.target }}
+          manylinux: auto
+          args: --release --out dist -m rs_controller/Cargo.toml
+          sccache: true
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: rs-controller-wheel-${{ matrix.target }}
+          path: dist/*.whl
+
+  rs-controller-publish:
+    name: Publish RS controller to PyPI
+    needs: rs-controller-wheels
+    runs-on: ubuntu-latest
+    if: github.event_name == 'release'
+    steps:
+      - name: Download all wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          pattern: rs-controller-wheel-*
+          merge-multiple: true
+          path: dist/
+      - name: Install twine
+        run: pip install twine
+      - name: Publish to PyPI
+        env:
+          TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+          TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+        run: |
+          twine upload --verbose dist/*
+      - name: Wait for PyPI availability
+        run: |
+          VERSION=$(echo "$GITHUB_REF" | sed 's|refs/tags/v||')
+          LINK="https://pypi.org/project/flyte_controller_base/${VERSION}/"
+          echo "Waiting for $LINK"
+          for i in $(seq 1 60); do
+            if curl -L -I -s -f "$LINK"; then
+              echo "Found on PyPI: $LINK"
+              exit 0
+            else
+              echo "Attempt $i: not yet available, retrying in 10s..."
+              sleep 10
+            fi
+          done
+          echo "ERROR: timed out waiting for PyPI"
+          exit 1
+
   flyte-pypi:
     name: PyPI package
+    needs: rs-controller-publish
+    if: always() && (needs.rs-controller-publish.result == 'success' || needs.rs-controller-publish.result == 'skipped')
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -27,18 +108,25 @@ jobs:
         run: |
           uv venv
           uv pip install build twine setuptools wheel
+      - name: Pin flyte_controller_base version (release only)
+        if: github.event_name == 'release'
+        run: |
+          VERSION=$(echo "$GITHUB_REF" | sed 's|refs/tags/v||')
+          sed -i "s/flyte_controller_base>=2.0.0b0/flyte_controller_base==${VERSION}/" pyproject.toml
+          echo "Pinned flyte_controller_base==${VERSION}"
+          grep flyte_controller_base pyproject.toml
       - name: Build and publish
         run: |
           uv run python -m build --wheel --installer uv
       - name: Publish
         if: ${{ github.event_name == 'release' }}
-        working-directory: ${{ matrix.workdir }}
         env:
           TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
           TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
         run: |
           uv run python -m twine upload --verbose dist/*
       - name: Sleep until pypi is available
+        if: ${{ github.event_name == 'release' }}
         id: pypiwait
         run: |
           # from 'refs/tags/v1.2.3 get 1.2.3' and make sure it's not an empty string

examples/basics/large_slow_run.py

@@ -1,8 +1,18 @@
 import asyncio
+from pathlib import Path
 
 import flyte
+from flyte._image import PythonWheels
 
-env = flyte.TaskEnvironment("large-slow-run")
+controller_dist_folder = Path("/Users/ytong/go/src/github.com/flyteorg/sdk-rust/rs_controller/dist")
+wheel_layer = PythonWheels(wheel_dir=controller_dist_folder, package_name="flyte_controller_base")
+base = flyte.Image.from_debian_base()
+rs_controller_image = base.clone(addl_layer=wheel_layer)
+
+env = flyte.TaskEnvironment(
+    "large-slow-run",
+    image=rs_controller_image,
+)
 
 
 @env.task

examples/stress/duplicate_action_id.py

@@ -14,20 +14,29 @@
 """
 
 import asyncio
+from pathlib import Path
 
 import flyte
+from flyte._image import PythonWheels
+
+controller_dist_folder = Path("/Users/ytong/go/src/github.com/flyteorg/sdk-rust/rs_controller/dist")
+wheel_layer = PythonWheels(wheel_dir=controller_dist_folder, package_name="flyte_controller_base")
+base = flyte.Image.from_debian_base()
+rs_controller_image = base.clone(addl_layer=wheel_layer)
 
 env_worker = flyte.TaskEnvironment(
     name="worker",
     resources=flyte.Resources(cpu=1, memory="250Mi"),
     cache="disable",
+    image=rs_controller_image,
 )
 
 env_main = flyte.TaskEnvironment(
     name="main",
     resources=flyte.Resources(cpu=1, memory="250Mi"),
     depends_on=[env_worker],
     cache="disable",
+    image=rs_controller_image,
 )
 
 

pyproject.toml

@@ -24,7 +24,8 @@ dependencies = [
     "async-lru>=2.0.5",
     "mashumaro",
     "aiolimiter>=1.2.1",
-    "flyteidl2==2.0.2",
+    "flyteidl2==2.0.6",
+    "flyte_controller_base>=2.0.0b0",
     "packaging",
 ]
 
@@ -166,6 +167,9 @@ ignore = ["PGH003", "PLC0415", "ASYNC210", "ASYNC240", "FURB122"]
 disable_error_code = ["import-untyped"]
 ignore_missing_imports = true
 
+[tool.uv.sources]
+flyte_controller_base = { path = "rs_controller" }
+
 [tool.ty.rules]
 unused-ignore-comment = "warn"
 redundant-cast = "ignore"