ctrs-rs new

.github/workflows/lint.yml

@@ -27,6 +27,31 @@ jobs:
         run: |
           make fmt
           git diff --exit-code
+  rs-fmt:
+    name: rust fmt
+    runs-on: ubuntu-latest
+    steps:
+      - name: Fetch the code
+        uses: actions/checkout@v4
+      - name: Install nightly toolchain
+        run: |
+          rustup toolchain install nightly
+          rustup component add --toolchain nightly-x86_64-unknown-linux-gnu rustfmt
+      - name: fmt
+        run: |
+          make -C rs_controller check-fmt
+  rs-lint:
+    name: rust lint
+    runs-on: ubuntu-latest
+    steps:
+      - name: Fetch the code
+        uses: actions/checkout@v4
+      - name: Install toolchain
+        run: |
+          rustup toolchain install
+      - name: lint
+        run: |
+          make -C rs_controller lint
   mypy:
     name: make mypy
     runs-on: ubuntu-latest

.gitignore

@@ -10,9 +10,6 @@ __pycache__/
 # C extensions
 *.so
 
-# Temporary
-rs_controller/
-
 # Distribution / packaging
 .Python
 build/
@@ -194,5 +191,9 @@ examples/local_debug.ipynb
 examples/remote_example.py
 config.yaml
 .claude/
+tests/flyte/internal/bin/flyte-inputs.json
 tests/flyte/internal/bin/flyte-parameters.json
+.cargo/
+.run/
 .flyteignore
+

Makefile

@@ -1,3 +1,8 @@
+# Default registry for image builds
+REGISTRY ?= ghcr.io/flyteorg
+# Default name for connector image
+CONNECTOR_IMAGE_NAME ?= flyte-connector
+
 # Default target: show all available targets
 .PHONY: help
 help:
@@ -92,6 +97,12 @@ unit_test_plugins:
 		fi \
 	done
 
+.PHONY: dev-rs-dist
+dev-rs-dist:
+	cd rs_controller && $(MAKE) build-wheels
+	$(MAKE) dist
+	uv run python maint_tools/build_default_image.py --registry $(REGISTRY) --name $(CONNECTOR_IMAGE_NAME)
+	uv pip install --find-links ./rs_controller/dist --no-index --force-reinstall --no-deps flyte_controller_base
 
 .PHONY: cli-docs-gen
 cli-docs-gen: ## Generate CLI documentation

README.md

@@ -142,3 +142,74 @@ pip install flyte[tui]
 ## License
 
 Apache 2.0 — see [LICENSE](LICENSE).
+
+## Developing the Rust Core Controller
+
+The Rust core controller (`flyte_controller_base`) ships as a separate PyPI wheel from the main
+SDK. Keeping it separate avoids forcing the whole SDK build toolchain to become maturin-based;
+we may revisit that decision once the Rust controller is the default execution path.
+
+When iterating, keep `flyteidl2` pinned to the same version on both the Python and Rust sides
+(see `pyproject.toml` and `rs_controller/Cargo.toml`).
+
+### One-time builder setup
+
+`cd` into `rs_controller` and run:
+
+```bash
+make build-builders
+```
+
+This builds the manylinux builder images once. They can be reused as the Rust source changes.
+
+### Iteration cycle
+
+Build the multi-arch wheels (linux/amd64 and linux/arm64):
+
+```bash
+cd rs_controller
+make build-wheels
+```
+
+`make build-wheel-local` builds a macOS wheel for local Rust development.
+
+`cd` back up to the repo root, then:
+
+```bash
+make dist
+python maint_tools/build_default_image.py
+```
+
+Install the freshly-built wheel into your venv:
+
+```bash
+uv pip install --find-links ./rs_controller/dist --no-index --force-reinstall --no-deps flyte_controller_base
+```
+
+Repeat this loop — build new wheels, force-reinstall — to iterate.
+
+### Build configuration
+
+To support both Rust crate use and Python wheel distribution, the crate toggles the
+`pyo3/extension-module` feature via Cargo features:
+
+```toml
+# rs_controller/Cargo.toml
+[features]
+default = ["pyo3/auto-initialize"]              # For local cargo run / cargo test (links libpython)
+extension-module = ["pyo3/extension-module"]    # For Python wheels (no libpython linking)
+
+[lib]
+crate-type = ["rlib", "cdylib"]                 # Both Rust and Python consumers
+```
+
+`auto-initialize` embeds Python in Rust binaries — convenient on macOS where libpython is
+linkable, but not available in manylinux. `extension-module` extends Python from Rust and
+must *not* link libpython for wheel portability. The maturin pyproject opts into the
+extension-module feature explicitly:
+
+```toml
+# rs_controller/pyproject.toml
+[tool.maturin]
+features = ["extension-module"]
+```

examples/advanced/hybrid_mode.py

@@ -0,0 +1,83 @@
+import asyncio
+import os
+from pathlib import Path
+from typing import List
+
+import flyte
+import flyte.storage
+from flyte.storage import S3
+
+env = flyte.TaskEnvironment(name="hello_world", cache="disable")
+
+
+@env.task
+async def say_hello_hybrid(data: str, lt: List[int]) -> str:
+    print(f"Hello, world! - {flyte.ctx().action}")
+    return f"Hello {data} {lt}"
+
+
+@env.task
+async def squared(i: int = 3) -> int:
+    print(flyte.ctx().action)
+    return i * i
+
+
+@env.task
+async def squared_2(i: int = 3) -> int:
+    print(flyte.ctx().action)
+    return i * i
+
+
+@env.task
+async def say_hello_hybrid_nested(data: str = "default string") -> str:
+    print(f"Hello, nested! - {flyte.ctx().action}")
+    coros = []
+    for i in range(3):
+        coros.append(squared(i=i))
+
+    vals = await asyncio.gather(*coros)
+    return await say_hello_hybrid(data=data, lt=vals)
+
+
+@env.task
+async def hybrid_parent_placeholder():
+    import sys
+    import time
+
+    print(f"Hello, hybrid parent placeholder - Task Context: {flyte.ctx()}")
+    print(f"Run command: {sys.argv}")
+    print("Environment Variables:")
+    for k, value in sorted(os.environ.items()):
+        if k.startswith("FLYTE_") or k.startswith("_U"):  # noqa: PIE810
+            print(f"{k}: {value}")
+
+    print("Sleeping for 24 hours to simulate a long-running task...", flush=True)
+    time.sleep(86400)  # noqa: ASYNC251
+
+
+if __name__ == "__main__":
+    # Configurable via env vars so the example doesn't bake in machine-local paths.
+    #   FLYTE_CONFIG       -> path to flyte config yaml (default: ~/.flyte/config.yaml)
+    #   FLYTE_HYBRID_RUN_NAME -> the long-running parent run id to attach hybrid sub-actions to
+    #   FLYTE_HYBRID_RUN_BASE -> s3 base dir for the hybrid run
+    current_directory = Path(os.getcwd())
+    repo_root = current_directory.parent.parent
+    s3_sandbox = S3.for_sandbox()
+
+    config_path = os.getenv("FLYTE_CONFIG") or str(Path.home() / ".flyte" / "config.yaml")
+    flyte.init_from_config(config_path, root_dir=repo_root, storage=s3_sandbox)
+
+    run_name = os.getenv("FLYTE_HYBRID_RUN_NAME")
+    run_base_dir = os.getenv("FLYTE_HYBRID_RUN_BASE")
+    if not run_name or not run_base_dir:
+        raise SystemExit(
+            "Set FLYTE_HYBRID_RUN_NAME and FLYTE_HYBRID_RUN_BASE to the parent run id "
+            "and s3 base dir of an already-launched hybrid run before invoking this example."
+        )
+
+    outputs = flyte.with_runcontext(
+        mode="hybrid",
+        name=run_name,
+        run_base_dir=run_base_dir,
+    ).run(say_hello_hybrid_nested, data="hello world")
+    print("Output:", outputs)

examples/advanced/rust_controller.py

@@ -0,0 +1,139 @@
+"""
+Run a small fanout against the Rust-backed RemoteController.
+
+The Rust controller (`flyte_controller_base`) is opt-in via the
+`_F_USE_RUST_CONTROLLER=1` environment variable. With the variable unset,
+the SDK uses the pure-Python `RemoteController` — this example will still
+work, just on the Python path.
+
+Selection rules:
+
+    _F_USE_RUST_CONTROLLER  _U_USE_ACTIONS   wheel installed   selected
+    ----------------------  --------------   ---------------   -----------------
+    unset / 0               any              any               Python controller
+    1                       1                any               Python controller (Actions service is Python-only)
+    1                       unset / 0        yes               Rust controller
+    1                       unset / 0        no                Python controller (warning logged)
+
+Quickstart (remote cluster)::
+
+    # 1. Build the wheel for the cluster's architecture (multi-arch by default):
+    cd rs_controller && make build-wheels && cd ..
+
+    # 2. Install the wheel locally so the parent process can import it:
+    uv pip install --find-links ./rs_controller/dist --no-index \\
+        --force-reinstall --no-deps flyte_controller_base
+
+    # 3. Opt in and run. The example automatically discovers
+    #    ./rs_controller/dist and bakes the wheel into the task image
+    #    via a PythonWheels layer, so child task containers also get
+    #    the Rust controller wheel.
+    _F_USE_RUST_CONTROLLER=1 python examples/advanced/rust_controller.py
+
+To override the wheel location (e.g. shared CI artifacts dir)::
+
+    FLYTE_RS_CONTROLLER_WHEELS=/path/to/wheels \\
+        _F_USE_RUST_CONTROLLER=1 python examples/advanced/rust_controller.py
+
+If no wheel directory is found the example falls back to the default
+debian-base image — the Python controller is used end-to-end and the
+opt-in becomes a no-op.
+
+The example does a 50-way fanout of a trivial task to give the Rust
+informer/submit loop something to chew on. Bump `FAN_OUT` in your
+environment to push it harder.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+from pathlib import Path
+
+import flyte
+from flyte._image import PythonWheels
+
+
+def _maybe_rust_controller_image() -> flyte.Image | None:
+    """Build a debian-base image with the locally-built flyte_controller_base
+    wheel installed, so task containers can pick up the Rust controller path.
+
+    Resolution order for the wheel directory:
+
+    1. ``FLYTE_RS_CONTROLLER_WHEELS`` env var — explicit override.
+    2. ``rs_controller/dist`` relative to this file (works for in-tree
+       development with ``make build-wheels`` or ``make build-wheel-local``).
+
+    Returns ``None`` if no wheel directory is found, in which case we fall
+    back to the default image. The example will still run on the Python
+    controller path; the Rust opt-in will simply log a warning and use
+    Python in the parent task.
+    """
+    override = os.getenv("FLYTE_RS_CONTROLLER_WHEELS")
+    if override:
+        wheel_dir = Path(override).expanduser().resolve()
+    else:
+        # examples/advanced/rust_controller.py -> repo_root/rs_controller/dist
+        repo_root = Path(__file__).resolve().parent.parent.parent
+        wheel_dir = repo_root / "rs_controller" / "dist"
+
+    if not wheel_dir.is_dir() or not list(wheel_dir.glob("*.whl")):
+        return None
+
+    base = flyte.Image.from_debian_base()
+    return base.clone(
+        addl_layer=PythonWheels(wheel_dir=wheel_dir, package_name="flyte_controller_base")
+    )
+
+
+_image = _maybe_rust_controller_image()
+
+env = flyte.TaskEnvironment(
+    name="rust_controller_demo",
+    resources=flyte.Resources(cpu=1, memory="500Mi"),
+    # Only attach the custom image when we found a wheel; otherwise let
+    # TaskEnvironment fall back to its default.
+    **({"image": _image} if _image is not None else {}),
+)
+
+
+@env.task
+async def square(i: int) -> int:
+    """Trivial leaf task — purpose is to be cheap so the controller's
+    submit/watch loop is the bottleneck, not the work."""
+    return i * i
+
+
+@env.task
+async def fanout(n: int) -> int:
+    """Submit `n` parallel sub-actions and return the sum of squares."""
+    print(f"[fanout] submitting {n} sub-actions")
+    t0 = time.monotonic()
+    results = await asyncio.gather(*[square(i=i) for i in range(n)])
+    elapsed = time.monotonic() - t0
+    total = sum(results)
+    print(f"[fanout] {n} sub-actions completed in {elapsed:.2f}s, sum={total}")
+    return total
+
+
+if __name__ == "__main__":
+    fan_out = int(os.getenv("FAN_OUT", "50"))
+    using_rust = os.getenv("_F_USE_RUST_CONTROLLER") == "1"
+    wheel_attached = _image is not None
+
+    print(f"Controller path:   {'Rust' if using_rust else 'Python'}")
+    print(f"Wheel layer:       {'attached' if wheel_attached else 'not found (using default image)'}")
+    print(f"FAN_OUT:           {fan_out}")
+    if using_rust and not wheel_attached:
+        print(
+            "WARNING: _F_USE_RUST_CONTROLLER=1 was set but no wheel directory was "
+            "found. Build wheels with `cd rs_controller && make build-wheels` or set "
+            "FLYTE_RS_CONTROLLER_WHEELS to the directory containing flyte_controller_base*.whl."
+        )
+
+    flyte.init_from_config()
+    run = flyte.run(fanout, n=fan_out)
+    print(f"run name: {run.name}")
+    print(f"run url:  {run.url}")
+    run.wait()

rs_controller/.gitignore

@@ -0,0 +1 @@
+docker_cargo_cache/