Flytekit add async agent test, fix delete task agent service and fix python3.8 asyncio to_thread issue

.github/workflows/build_image.yml

@@ -0,0 +1,86 @@
+name: Publish Python Package
+
+on:
+  workflow_dispatch:
+
+jobs:
+  build-and-push-docker-images-manual:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.9", "3.10", "3.11"]
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: "0"
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v1
+      - name: Set up Docker Buildx
+        id: buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: "${{ secrets.FLYTE_BOT_USERNAME }}"
+          password: "${{ secrets.FLYTE_BOT_PAT }}"
+      - name: Prepare Flytekit Image Names
+        id: flytekit-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/flytekit
+          tags: |
+            py${{ matrix.python-version }}-${{ github.sha }}
+      - name: Build & Push Flytekit Python${{ matrix.python-version }} Docker Image to Github Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          platforms: linux/arm64, linux/amd64
+          push: true
+          tags: ${{ steps.flytekit-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ github.sha }}
+            DOCKER_IMAGE=ghcr.io/${{ github.repository_owner }}/flytekit:py${{ matrix.python-version }}-${{ github.sha }}
+            PYTHON_VERSION=${{ matrix.python-version }}
+          file: Dockerfile
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+  build-and-push-flyteagent-images-manual:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: "0"
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v1
+      - name: Set up Docker Buildx
+        id: buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: "${{ secrets.FLYTE_BOT_USERNAME }}"
+          password: "${{ secrets.FLYTE_BOT_PAT }}"
+      - name: Prepare Flyte Agent Image Names
+        id: flyteagent-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/flyteagent
+          tags: |
+            ${{ github.sha }}
+      - name: Push External Plugin Service Image to GitHub Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: "."
+          platforms: linux/arm64, linux/amd64
+          push: true
+          tags: ${{ steps.flyteagent-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ github.sha }}
+          file: ./Dockerfile.agent
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

.github/workflows/pythonbuild.yml

@@ -16,13 +16,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
-        exclude:
-          # Ignore this test because we failed to install docker-py
-          # docker-py will install pywin32==227, whereas pywin only added support for python 3.10 in version 301.
-          # For more detail, see https://github.com/flyteorg/flytekit/pull/856#issuecomment-1067152855
-          - python-version: 3.10
-            os: windows-latest
+        python-version: ["3.8", "3.11"]
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
@@ -47,13 +41,56 @@ jobs:
         with:
           fail_ci_if_error: false
 
+  build-integration:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        # python 3.11 has intermittent issues with the docker build + push step
+        # https://github.com/flyteorg/flytekit/actions/runs/5800978835/job/15724237979?pr=1579
+        python-version: ["3.8", "3.11"]
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Cache pip
+        uses: actions/cache@v2
+        with:
+          # This path is specific to Ubuntu
+          path: ~/.cache/pip
+          # Look to see if there is a cache hit for the corresponding requirements files
+          key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.in', 'requirements.in')) }}
+      - name: Install dependencies
+        run: make setup && pip freeze
+      - name: Install FlyteCTL
+        uses: unionai-oss/flytectl-setup-action@master
+      - name: Setup Flyte Sandbox
+        run: |
+          flytectl demo start
+          flytectl config init
+      - name: Build and push to local registry
+        run: |
+          docker build . -f Dockerfile.dev -t localhost:30000/flytekit:dev --build-arg PYTHON_VERSION=${{ matrix.python-version }}
+      - name: Integration Test with coverage
+        env:
+          FLYTEKIT_IMAGE: localhost:30000/flytekit:dev
+          FLYTEKIT_CI: 1
+        run: make integration_test_codecov
+      - name: Codecov
+        uses: codecov/[email protected]
+        with:
+          fail_ci_if_error: false
+
   build-plugins:
     needs: build
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.11"]
         plugin-names:
           # Please maintain an alphabetical order in the following list
           - flytekit-aws-athena
@@ -104,12 +141,6 @@ jobs:
             plugin-names: "flytekit-greatexpectations"
           # onnxruntime does not support python 3.10 yet
           # https://github.com/microsoft/onnxruntime/issues/9782
-          - python-version: 3.10
-            plugin-names: "flytekit-onnx-pytorch"
-          - python-version: 3.10
-            plugin-names: "flytekit-onnx-scikitlearn"
-          - python-version: 3.10
-            plugin-names: "flytekit-onnx-tensorflow"
           - python-version: 3.11
             plugin-names: "flytekit-onnx-pytorch"
           - python-version: 3.11
@@ -158,10 +189,14 @@ jobs:
           # onnx plugins does not support protobuf>4 yet (in fact it is tensorflow that
           # does not support that yet). More details in https://github.com/onnx/onnx/issues/4239.
           if [[ ${{ matrix.plugin-names }} == *"onnx"* || ${{ matrix.plugin-names }} == "flytekit-whylogs" || ${{ matrix.plugin-names }} == "flytekit-mlflow" ]]; then
-             PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python coverage run -m pytest tests
+             PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python coverage run -m pytest tests --cov=./ --cov-report=xml --cov-append
           else
-             coverage run -m pytest tests
+             coverage run -m pytest tests --cov=./ --cov-report=xml --cov-append
           fi
+      - name: Codecov
+        uses: codecov/[email protected]
+        with:
+          fail_ci_if_error: false
   lint:
     runs-on: ubuntu-latest
     steps:

.pre-commit-config.yaml

@@ -1,28 +1,28 @@
 repos:
-- repo: https://github.com/PyCQA/flake8
-  rev:  4.0.1
-  hooks:
-    - id: flake8
-- repo: https://github.com/psf/black
-  rev:  22.3.0
-  hooks:
-    - id: black
-- repo: https://github.com/PyCQA/isort
-  rev:  5.12.0
-  hooks:
-    - id: isort
-      args: ["--profile", "black"]
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.2.0
-  hooks:
-    - id: check-yaml
-    - id: end-of-file-fixer
-    - id: trailing-whitespace
-- repo: https://github.com/shellcheck-py/shellcheck-py
-  rev: v0.8.0.4
-  hooks:
-  - id: shellcheck
-- repo: https://github.com/conorfalvey/check_pdb_hook
-  rev: 0.0.9
-  hooks:
-    - id: check_pdb_hook
+  - repo: https://github.com/PyCQA/flake8
+    rev: 4.0.1
+    hooks:
+      - id: flake8
+  - repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+      - id: black
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.2.0
+    hooks:
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/shellcheck-py/shellcheck-py
+    rev: v0.8.0.4
+    hooks:
+      - id: shellcheck
+  - repo: https://github.com/conorfalvey/check_pdb_hook
+    rev: 0.0.9
+    hooks:
+      - id: check_pdb_hook

.readthedocs.yml

@@ -20,4 +20,3 @@ sphinx:
 python:
   install:
     - requirements: doc-requirements.txt
-    - requirements: docs/requirements.txt

Dockerfile

@@ -16,7 +16,6 @@ RUN apt-get update && apt-get install build-essential -y
 RUN pip install -U flytekit==$VERSION \
 	flytekitplugins-pod==$VERSION \
 	flytekitplugins-deck-standard==$VERSION \
-	flytekitplugins-envd==$VERSION \
 	scikit-learn
 
 RUN useradd -u 1000 flytekit

Dockerfile.agent

@@ -4,6 +4,7 @@ MAINTAINER Flyte Team <[email protected]>
 LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytekit
 
 ARG VERSION
+RUN pip install prometheus-client
 RUN pip install -U flytekit==$VERSION flytekitplugins-bigquery==$VERSION
 
 CMD pyflyte serve --port 8000

Makefile

@@ -1,6 +1,6 @@
 export REPOSITORY=flytekit
 
-PIP_COMPILE = pip-compile --upgrade --verbose
+PIP_COMPILE = pip-compile --upgrade --verbose --resolver=backtracking
 MOCK_FLYTE_REPO=tests/flytekit/integration/remote/mock_flyte_repo/workflows
 
 .SILENT: help
@@ -49,8 +49,6 @@ test: lint unit_test
 
 .PHONY: unit_test_codecov
 unit_test_codecov:
-	# Ensure coverage file
-	rm coverage.xml || true
 	$(MAKE) CODECOV_OPTS="--cov=./ --cov-report=xml --cov-append" unit_test
 
 .PHONY: unit_test
@@ -60,9 +58,17 @@ unit_test:
 	pytest -m "not sandbox_test" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} && \
 		PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS}
 
+.PHONY: integration_test_codecov
+integration_test_codecov:
+	$(MAKE) CODECOV_OPTS="--cov=./ --cov-report=xml --cov-append" integration_test
+
+.PHONY: integration_test
+integration_test:
+	pytest tests/flytekit/integration/experimental ${CODECOV_OPTS}
+
 doc-requirements.txt: export CUSTOM_COMPILE_COMMAND := make doc-requirements.txt
 doc-requirements.txt: doc-requirements.in install-piptools
-	docker run --platform linux/amd64  --rm -it --volume .:/root python:3.9-slim-buster sh -c "cd /root && apt-get update && apt-get install git -y && pip install pip-tools && pip-compile --upgrade --verbose doc-requirements.in"
+	$(PIP_COMPILE) $<
 
 ${MOCK_FLYTE_REPO}/requirements.txt: export CUSTOM_COMPILE_COMMAND := make ${MOCK_FLYTE_REPO}/requirements.txt
 ${MOCK_FLYTE_REPO}/requirements.txt: ${MOCK_FLYTE_REPO}/requirements.in install-piptools

README.md

@@ -16,9 +16,10 @@
 [![PyPI version fury.io](https://badge.fury.io/py/flytekit.svg)](https://pypi.python.org/pypi/flytekit/)
 [![PyPI download day](https://img.shields.io/pypi/dd/flytekit.svg)](https://pypi.python.org/pypi/flytekit/)
 [![PyPI download month](https://img.shields.io/pypi/dm/flytekit.svg)](https://pypi.python.org/pypi/flytekit/)
+[![PyPI total download](https://static.pepy.tech/badge/flytekit)](https://static.pepy.tech/badge/flytekit)
 [![PyPI format](https://img.shields.io/pypi/format/flytekit.svg)](https://pypi.python.org/pypi/flytekit/)
 [![PyPI implementation](https://img.shields.io/pypi/implementation/flytekit.svg)](https://pypi.python.org/pypi/flytekit/)
-![Codecov](https://img.shields.io/codecov/c/github/flyteorg/flytekit?style=plastic)
+[![Codecov](https://img.shields.io/codecov/c/github/flyteorg/flytekit?style=plastic)](https://app.codecov.io/gh/flyteorg/flytekit)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/flytekit.svg)](https://pypi.python.org/pypi/flytekit/)
 [![Docs](https://readthedocs.org/projects/flytekit/badge/?version=latest&style=plastic)](https://flytekit.rtfd.io)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

codecov.yml

@@ -1,7 +1,9 @@
-ignore_paths:
+ignore:
   - "flytekit/bin"
+  - "flytekit/clis/**/*"
   - "test_*.py"
-  - "flytekit/__init__.py"
-  - "flytekit/extend/__init__.py"
-  - "flytekit/testing/__init__.py"
-  - "tests/*"
+  - "tests/**/*"
+  - "setup.py"
+  - "plugins/tests/**/*"
+  - "plugins/setup.py"
+  - "plugins/**/setup.py"

dev-requirements.in

@@ -5,6 +5,7 @@ hypothesis
 joblib
 mock
 pytest
+pytest-asyncio
 pytest-cov
 mypy
 pre-commit
@@ -17,7 +18,7 @@ keyrings.alt
 # Only install tensorflow if not running on an arm Mac.
 tensorflow==2.8.1; python_version<'3.11' and (platform_machine!='arm64' or platform_system!='Darwin')
 # Tensorflow release candidate supports python 3.11
-tensorflow==2.12.0rc1; python_version>='3.11' and (platform_machine!='arm64' or platform_system!='Darwin')
+tensorflow==2.13.0; python_version>='3.11' and (platform_machine!='arm64' or platform_system!='Darwin')
 
 # Newer versions of torch bring in nvidia dependencies that are not present in windows, so
 # we put this constraint while we do not have per-environment requirements files
@@ -30,3 +31,4 @@ types-protobuf
 types-croniter
 types-mock
 autoflake
+prometheus-client

doc-requirements.in

@@ -11,12 +11,14 @@ sphinx-autoapi
 sphinx-copybutton
 sphinx_fontawesome
 sphinx-panels
-sphinxcontrib-youtube
+sphinxcontrib-youtube==1.2.0
 cryptography
 google-api-core[grpc]
 scikit-learn
 sphinx-tags
 sphinx-click
+retry
+mashumaro
 
 # Packages for Plugin docs
 # Package name              Plugin needing it
@@ -29,23 +31,25 @@ plotly                      # deck
 pandas_profiling            # deck
 dolt_integrations           # dolt
 great-expectations          # greatexpectations
+datasets                    # huggingface
 kubernetes                  # k8s-pod
 modin                       # modin
 pandera                     # pandera
 papermill                   # papermill
 jupyter                     # papermill
+polars                      # polars
 pyspark                     # spark
 sqlalchemy                  # sqlalchemy
 torch                       # pytorch
-# TODO: Remove after buf migration is done and packages updated
-# skl2onnx                    # onnxscikitlearn
-# tf2onnx                     # onnxtensorflow
+skl2onnx                    # onnxscikitlearn
+tf2onnx                     # onnxtensorflow
 tensorflow                  # onnxtensorflow
-whylogs                     # whylogs
+whylogs==1.3.3              # whylogs
 whylabs-client              # whylogs
-ray                         # ray
+ray==2.6.3                  # ray
 scikit-learn                # scikit-learn
 dask[distributed]           # dask
 vaex                        # vaex
-mlflow                      # mlflow
+mlflow==2.7.0               # mlflow
 duckdb                      # duckdb
+snowflake-connector-python  # snowflake