PR into #488 (#675)

* Bump flyteidl2 version to latest and get rid of the manually copied
code
* Fix the pem issue where we had to manually download the cert from
amazon before creating the channel (with tonic connections you have to
call `with_native_roots()` to pick up root certs).
* Fix maturin setting so it'll build.
* Pull in some changes that have been made to the main controller.
(#607, #521, #621)
* Add in some CI as a separate PR -
#676. But we need to play with
this experience a bit more first. Need to go through the whole
experience and see how things feel.

* More changes described in
#675 (comment)

---------

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Signed-off-by: machichima <nary12321@gmail.com>
Co-authored-by: machichima <nary12321@gmail.com>

Author: wild-endeavor

.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,6 +108,13 @@ 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

README.md

@@ -129,68 +129,108 @@ pip install flyte[tui]
 
 Flyte 2 is licensed under the [Apache 2.0 License](LICENSE).
 
-## Developing the Core Controller
+## Rust Controller (experimental)
 
-Create a separate virtual environment for the Rust contoller inside the rs_controller folder. The reason for this is
-because the rust controller should be a separate pypi package. The reason it should be a separate pypi package is that
-including it into the main SDK as a core component means the entire build toolchain for the SDK will need to become
-rust/maturin based. We should probably move to this model in the future though.
+The Rust controller is an alternative implementation of the remote controller written in Rust and exposed
+to Python via maturin / pyo3. Distributed as a separate `flyte_controller_base` wheel so the main SDK does
+not need to switch its build toolchain to rust/maturin. Keep important dependencies (notably `flyteidl2`)
+in lockstep between `pyproject.toml`, `rs_controller/pyproject.toml`, and `rs_controller/Cargo.toml`.
 
-Keep important dependencies the same though, namely flyteidl2.
+### Running with the Rust controller
 
-The following instructions are for helping to build the default multi-arch image. Each architecture needs a different wheel. Each wheel needs to be built by a different docker image.
+The Rust controller is gated behind an env var. Set it to `1` (also accepts `true` / `yes`):
 
-### Setup Builders
-`cd` into `rs_controller` and run `make build-builders`. This will build the builder images once, so you can keep using them as the rust code changes.
+```bash
+_F_USE_RUST_CONTROLLER=1 python examples/basics/hello_v2.py
+```
+
+The driver propagates this env var to all sub-task pods, so both the driver and child actions use the
+Rust controller for that run.
+
+> **v1 limitations.** The Rust controller currently supports only the legacy
+> QueueService + StateService path. Do **not** combine `_F_USE_RUST_CONTROLLER=1` with
+> `_U_USE_ACTIONS=1` until ActionsService support lands. Other gaps tracked as follow-ups:
+> abort RPC on cancel, trace-action enqueue, `Code.ABORTED` fast-fail, tunable retries / QPS,
+> graceful `stop()`. See PR #675.
+
+> Dev iteration requires the local image builder. The `flyte_controller_base` wheel is not
+> on PyPI until release, and the remote image builder installs all wheels in a layer at once,
+> so it cannot resolve `flyte_controller_base` from a sibling layer. Use the local image
+> builder while developing the Rust controller:
+>
+> ```yaml
+> # .flyte/config.yaml
+> image:
+>   builder: local
+> ```
+
+### Developing the Rust controller
+
+#### One-time setup
 
-### Iteration Cycle
-Run `make build-wheels` to actually build the multi-arch wheels. This command should probably be updated to build all three,
-currently it only builds for linux/amd64 and linux/arm64... the `make build-wheel-local` command builds a macosx wheel,
-unclear what the difference is between that and the arm64 one, and unclear if both are present, which one pip chooses.
+Build the manylinux builder images. They are cached, so you only need to rebuild them when the
+build tooling itself changes:
 
-`cd` back up to the root folder of this project and proceed with
 ```bash
-make dist
-python maint_tools/build_default_image.py
+cd rs_controller
+make build-builders
+cd ..
 ```
 
-To install the wheel locally for testing, use the following command with your venv active.
+#### Iteration loop
+
+After every Rust change, run the all-in-one dev target from the repo root:
+
 ```bash
-uv pip install --find-links ./rs_controller/dist --no-index --force-reinstall --no-deps flyte_controller_base
+REGISTRY=<your-registry> make dev-rs-dist
 ```
-Repeat this process to iterate - build new wheels, force reinstall the controller package.
 
-### Build Configuration Summary
+`dev-rs-dist` does four things:
 
-In order to support both Rust crate publication and Python wheel distribution, we have
-to sometimes use and sometimes not use the 'pyo3/extension-module' feature. To do this, this
-project's Cargo.toml itself can toggle this on and off.
+1. `cd rs_controller && make build-wheels` — build manylinux x86_64 + aarch64 wheels (use
+   `make build-wheel-local` if you only need a macOS wheel for the driver).
+2. `make dist` — build the main `flyte` SDK wheel.
+3. `uv run python maint_tools/build_default_image.py --registry $(REGISTRY)` — build the default
+   image with both wheels baked in and push it to your registry.
+4. `uv pip install --find-links ./rs_controller/dist --no-index --force-reinstall --no-deps flyte_controller_base` —
+   refresh the wheel in your local venv so the driver picks up the new build.
 
-  [features]
-  default = ["pyo3/auto-initialize"]     # For Rust crate users (links to libpython)
-  extension-module = ["pyo3/extension-module"]  # For Python wheels (no libpython linking)
+After this, any `flyte.TaskEnvironment` that does not pass an explicit `image=` will resolve to the default
+debian image and automatically have the Rust wheel layered in. If you do pass an explicit `image=`, the
+auto-bake is skipped; in that case, chain `.with_local_rs_controller()` onto the image to bake the Rust wheel
+manually.
 
-The cargo file contains
+If you only changed Python (not Rust), you can skip the wheel rebuild and just run `make dist` plus
+the rebuild image step. The Rust wheel is reused.
 
-  # Cargo.toml
-  [lib]
-  crate-type = ["rlib", "cdylib"]  # Support both Rust and Python usage
+### Build configuration summary
 
-When using 'default', 'auto-initialize' is turned on, which requires linking to libpython, which exists on local Mac so
-this works nicely. It is not available in manylinux however, so trying to build with this feature in a manylinux docker
-image will fail. But that's okay, because the purpose of the manylinux container is to build wheels,
-and for wheels, we need the 'extension-module' feature, which disables linking to libpython.
+The Rust crate ships with two cargo features so the same project can produce a Rust rlib and a
+Python extension wheel:
 
-The key insight: auto-initialize is for embedding Python in Rust (needs libpython), while
-extension-module is for extending Python with Rust (must NOT link libpython for portability).
+```toml
+[features]
+default = ["pyo3/auto-initialize"]            # Rust crate users; links libpython
+extension-module = ["pyo3/extension-module"]  # Python wheels; no libpython linking
+
+[lib]
+crate-type = ["rlib", "cdylib"]               # Both Rust and Python usage
+```
 
-This setup makes it possible to build wheels and also run Rust binaries with `cargo run --bin`. 
+- `pyo3/auto-initialize` embeds Python into Rust (works locally on macOS, fails inside the manylinux
+  builder because libpython is unavailable there).
+- `pyo3/extension-module` extends Python with Rust (must not link libpython for portable wheels).
+
+So local `cargo run --bin <name>` uses `default` features, and the manylinux builder explicitly
+disables defaults and turns on `extension-module`:
+
+```toml
+# rs_controller/pyproject.toml
+[tool.maturin]
+no-default-features = true
+features = ["extension-module"]
+```
 
-(not sure if this is needed)
-  # pyproject.toml
-  [tool.maturin]
-  features = ["extension-module"]  # Tells maturin to use extension-module feature
-  
 ## Learn More
 
 - **[Live Demo](https://flyte2intro.apps.demo.hosted.unionai.cloud/)** — Try Flyte 2 in your browser

examples/advanced/cancel_tasks.py

@@ -1,17 +1,8 @@
 import asyncio
-from pathlib import Path
 
-import flyte
 import flyte.errors
-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 = flyte.TaskEnvironment("cancel", image=rs_controller_image)
+env = flyte.TaskEnvironment("cancel")
 
 
 @env.task

examples/basics/devbox_one.py

@@ -1,20 +1,12 @@
 import asyncio
 import logging
-from pathlib import Path
 from typing import List
 
 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 = flyte.TaskEnvironment(
     name="hello_world",
     resources=flyte.Resources(cpu=1, memory="1Gi"),
-    image=rs_controller_image,
 )
 
 

examples/basics/hello.py

@@ -1,18 +1,9 @@
-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)
 
 # TaskEnvironments provide a simple way of grouping configuration used by tasks (more later).
 env = flyte.TaskEnvironment(
     name="hello_world",
     resources=flyte.Resources(memory="250Mi"),
-    image=rs_controller_image,
 )