merged master

Signed-off-by: Kevin Su <[email protected]>

.github/workflows/build_image.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11"]
     steps:
       - uses: actions/checkout@v4
         with:

.github/workflows/pythonbuild.yml

@@ -15,6 +15,41 @@ concurrency:
 
 jobs:
   build:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+        python-version: ["3.8", "3.11", "3.12"]
+    steps:
+      - uses: insightsengineering/disk-space-reclaimer@v1
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Cache pip
+        uses: actions/cache@v3
+        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 uninstall -y pandas
+          pip freeze
+      - name: Test with coverage
+        run: |
+          make unit_test_codecov
+      - name: Codecov
+        uses: codecov/[email protected]
+        with:
+          fail_ci_if_error: false
+          files: coverage.xml
+
+  build-with-extras:
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -46,9 +81,6 @@ jobs:
         if: ${{ matrix.python-version != '3.12' }}
         run: |
           make unit_test_extras_codecov
-      - name: Test with coverage
-        run: |
-          make unit_test_codecov
       - name: Codecov
         uses: codecov/[email protected]
         with:
@@ -207,6 +239,7 @@ jobs:
           # onnx-tensorflow needs a version of tensorflow that does not work with protobuf>4.
           # The issue is being tracked on the tensorflow side in https://github.com/tensorflow/tensorflow/issues/53234#issuecomment-1330111693
           # flytekit-onnx-tensorflow
+          - flytekit-openai
           - flytekit-pandera
           - flytekit-papermill
           - flytekit-polars
@@ -250,6 +283,7 @@ jobs:
             plugin-names: "flytekit-whylogs"
     steps:
       - uses: insightsengineering/disk-space-reclaimer@v1
+        if: ${{ matrix.plugin-names == 'flytekit-envd' }}
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
@@ -264,11 +298,12 @@ jobs:
           key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('plugins/{0}/requirements.txt', matrix.plugin-names ))) }}
       - name: Install dependencies
         run: |
+          export SETUPTOOLS_SCM_PRETEND_VERSION="2.0.0"
           make setup
           cd plugins/${{ matrix.plugin-names }}
-          pip install --pre .
-          if [ -f dev-requirements.txt ]; then pip install -r dev-requirements.txt; fi
-          pip install --pre -U $GITHUB_WORKSPACE
+          pip install .
+          if [ -f dev-requirements.in ]; then pip install -r dev-requirements.in; fi
+          pip install -U $GITHUB_WORKSPACE
           pip freeze
       - name: Test with coverage
         run: |

.github/workflows/pythonpublish.yml

@@ -24,7 +24,7 @@ jobs:
         run: |
           # from refs/tags/v1.2.3 get 1.2.3
           VERSION=$(echo $GITHUB_REF | sed 's#.*/v##')
-          echo "::set-output name=version::$VERSION"
+          echo "version=$VERSION" >> $GITHUB_OUTPUT
         shell: bash
       - name: Build and publish
         env:
@@ -167,6 +167,16 @@ jobs:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"
           password: "${{ secrets.FLYTE_BOT_PAT }}"
+      - name: Prepare Flyte Agent Slim Image Names
+        id: flyteagent-slim-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/flyteagent-slim
+          tags: |
+            latest
+            ${{ github.sha }}
+            ${{ needs.deploy.outputs.version }}
       - name: Prepare Flyte Agent Image Names
         id: flyteagent-names
         uses: docker/metadata-action@v3
@@ -177,11 +187,25 @@ jobs:
             latest
             ${{ github.sha }}
             ${{ needs.deploy.outputs.version }}
-      - name: Push External Plugin Service Image to GitHub Registry
+      - name: Push flyteagent-slim Image to GitHub Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: "."
+          platforms: linux/arm64, linux/amd64
+          target: agent-slim
+          push: ${{ github.event_name == 'release' }}
+          tags: ${{ steps.flyteagent-slim-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ needs.deploy.outputs.version }}
+          file: ./Dockerfile.agent
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+      - name: Push flyteagent-all Image to GitHub Registry
         uses: docker/build-push-action@v2
         with:
           context: "."
           platforms: linux/arm64, linux/amd64
+          target: agent-all
           push: ${{ github.event_name == 'release' }}
           tags: ${{ steps.flyteagent-names.outputs.tags }}
           build-args: |

Dockerfile.agent

@@ -1,4 +1,4 @@
-FROM python:3.9-slim-bookworm
+FROM python:3.9-slim-bookworm as agent-slim
 
 MAINTAINER Flyte Team <[email protected]>
 LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytekit
@@ -7,16 +7,22 @@ ARG VERSION
 
 RUN apt-get update && apt-get install build-essential -y
 
-RUN pip install prometheus-client
+RUN pip install prometheus-client grpcio-health-checking
 RUN pip install --no-cache-dir -U flytekit==$VERSION \
-  flytekitplugins-bigquery==$VERSION \
   flytekitplugins-airflow==$VERSION \
-  flytekitplugins-mmcloud==$VERSION \
-  flytekitplugins-spark==$VERSION \
+  flytekitplugins-bigquery==$VERSION \
+  flytekitplugins-chatgpt==$VERSION \
   flytekitplugins-snowflake==$VERSION \
   && apt-get clean autoclean \
   && apt-get autoremove --yes \
   && rm -rf /var/lib/{apt,dpkg,cache,log}/ \
   && :
 
 CMD pyflyte serve agent --port 8000
+
+FROM agent-slim as agent-all
+ARG VERSION
+
+RUN pip install --no-cache-dir -U \
+  flytekitplugins-mmcloud==$VERSION \
+  flytekitplugins-spark==$VERSION

Dockerfile.dev

@@ -28,8 +28,8 @@ COPY . /flytekit
 # 3. Clean up the apt cache to reduce image size. Reference: https://gist.github.com/marvell/7c812736565928e602c4
 # 4. Create a non-root user 'flytekit' and set appropriate permissions for directories.
 RUN apt-get update && apt-get install build-essential vim libmagic1 git -y \
+    && pip install "git+https://github.com/flyteorg/flyte.git@master#subdirectory=flyteidl" \
     && pip install --no-cache-dir -U --pre \
-        flyteidl \
         -e /flytekit \
         -e /flytekit/plugins/flytekit-k8s-pod \
         -e /flytekit/plugins/flytekit-deck-standard \
@@ -43,6 +43,7 @@ RUN apt-get update && apt-get install build-essential vim libmagic1 git -y \
     && chown flytekit: /home \
     && :
 
+
 ENV PYTHONPATH "/flytekit:/flytekit/plugins/flytekit-k8s-pod:/flytekit/plugins/flytekit-deck-standard:"
 
 # Switch to the 'flytekit' user for better security.

Makefile

@@ -25,8 +25,7 @@ update_boilerplate:
 
 .PHONY: setup
 setup: install-piptools ## Install requirements
-	pip install --pre -r dev-requirements.in
-
+	pip install -r dev-requirements.in
 
 .PHONY: fmt
 fmt:
@@ -63,14 +62,14 @@ unit_test_extras_codecov:
 unit_test:
 	# Skip all extra tests and run them with the necessary env var set so that a working (albeit slower)
 	# library is used to serialize/deserialize protobufs is used.
-	$(PYTEST_AND_OPTS) -m "not (serial or sandbox_test)" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models ${CODECOV_OPTS}
+	$(PYTEST_AND_OPTS) -m "not (serial or sandbox_test)" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
 	# Run serial tests without any parallelism
-	$(PYTEST) -m "serial" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models ${CODECOV_OPTS}
+	$(PYTEST) -m "serial" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
 
 
 .PHONY: unit_test_extras
 unit_test_extras:
-	PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python $(PYTEST_AND_OPTS) tests/flytekit/unit/extras ${CODECOV_OPTS}
+	PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python $(PYTEST_AND_OPTS) tests/flytekit/unit/extras tests/flytekit/unit/extend ${CODECOV_OPTS}
 
 .PHONY: test_serialization_codecov
 test_serialization_codecov:

dev-requirements.in

@@ -1,4 +1,5 @@
 -e file:.#egg=flytekit
+git+https://github.com/flyteorg/flyte.git@master#subdirectory=flyteidl
 
 coverage[toml]
 hypothesis

flytekit/clis/sdk_in_container/run.py

@@ -595,7 +595,7 @@ def _get_params(
         defaults: typing.Optional[typing.Dict[str, Parameter]] = None,
     ) -> typing.List["click.Parameter"]:
         params = []
-        flyte_ctx = context_manager.FlyteContextManager.current_context()
+        flyte_ctx = ctx.obj.remote_instance().context
         for name, var in inputs.items():
             if fixed and name in fixed:
                 continue
@@ -609,6 +609,7 @@ def _get_params(
         return params
 
     def get_params(self, ctx: click.Context) -> typing.List["click.Parameter"]:
+        ctx.obj.remote = True
         if not self.params:
             self.params = []
             entity = self._fetch_entity(ctx)

flytekit/clis/sdk_in_container/serve.py

@@ -1,11 +1,13 @@
 from concurrent import futures
 
+import grpc
 import rich_click as click
+from flyteidl.service import agent_pb2
 from flyteidl.service.agent_pb2_grpc import (
     add_AgentMetadataServiceServicer_to_server,
     add_AsyncAgentServiceServicer_to_server,
+    add_SyncAgentServiceServicer_to_server,
 )
-from grpc import aio
 
 
 @click.group("serve")
@@ -51,21 +53,46 @@ def agent(_: click.Context, port, worker, timeout):
 
 
 async def _start_grpc_server(port: int, worker: int, timeout: int):
-    click.secho("Starting up the server to expose the prometheus metrics...", fg="blue")
-    from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService
+    from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService, SyncAgentService
 
-    try:
-        from prometheus_client import start_http_server
-
-        start_http_server(9090)
-    except ImportError as e:
-        click.secho(f"Failed to start the prometheus server with error {e}", fg="red")
+    _start_http_server()
     click.secho("Starting the agent service...", fg="blue")
-    server = aio.server(futures.ThreadPoolExecutor(max_workers=worker))
+    server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=worker))
 
     add_AsyncAgentServiceServicer_to_server(AsyncAgentService(), server)
+    add_SyncAgentServiceServicer_to_server(SyncAgentService(), server)
     add_AgentMetadataServiceServicer_to_server(AgentMetadataService(), server)
+    _start_health_check_server(server, worker)
 
     server.add_insecure_port(f"[::]:{port}")
     await server.start()
     await server.wait_for_termination(timeout)
+
+
+def _start_http_server():
+    try:
+        from prometheus_client import start_http_server
+
+        click.secho("Starting up the server to expose the prometheus metrics...", fg="blue")
+        start_http_server(9090)
+    except ImportError as e:
+        click.secho(f"Failed to start the prometheus server with error {e}", fg="red")
+
+
+def _start_health_check_server(server: grpc.Server, worker: int):
+    try:
+        from grpc_health.v1 import health, health_pb2, health_pb2_grpc
+
+        health_servicer = health.HealthServicer(
+            experimental_non_blocking=True,
+            experimental_thread_pool=futures.ThreadPoolExecutor(max_workers=worker),
+        )
+
+        for service in agent_pb2.DESCRIPTOR.services_by_name.values():
+            health_servicer.set(service.full_name, health_pb2.HealthCheckResponse.SERVING)
+        health_servicer.set(health.SERVICE_NAME, health_pb2.HealthCheckResponse.SERVING)
+
+        health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
+    except ImportError as e:
+        click.secho(f"Failed to start the health check servicer with error {e}", fg="red")

flytekit/core/artifact.py

@@ -147,7 +147,13 @@ def __init__(
         self.time_partition = time_partition
         self.partitions = partitions
         self.tag = tag
-        self.bindings = bindings
+        if len(bindings) > 0:
+            b = set(bindings)
+            if len(b) > 1:
+                raise ValueError(f"Multiple bindings found in query {self}")
+            self.binding: Optional[Artifact] = bindings[0]
+        else:
+            self.binding = None
 
     def to_flyte_idl(
         self,
@@ -391,23 +397,19 @@ def concrete_artifact_id(self) -> art_id.ArtifactID:
 
     def embed_as_query(
         self,
-        bindings: typing.List[Artifact],
         partition: Optional[str] = None,
         bind_to_time_partition: Optional[bool] = None,
         expr: Optional[str] = None,
     ) -> art_id.ArtifactQuery:
         """
         This should only be called in the context of a Trigger
-        :param bindings: The list of artifacts in trigger_on
         :param partition: Can embed a time partition
         :param bind_to_time_partition: Set to true if you want to bind to a time partition
         :param expr: Only valid if there's a time partition.
         """
         # Find self in the list, raises ValueError if not there.
-        idx = bindings.index(self)
         aq = art_id.ArtifactQuery(
             binding=art_id.ArtifactBindingData(
-                index=idx,
                 partition_key=partition,
                 bind_to_time_partition=bind_to_time_partition,
                 transform=str(expr) if expr and (partition or bind_to_time_partition) else None,