merge master

Signed-off-by: Yee Hing Tong <[email protected]>

.github/workflows/docs_build.yml

@@ -23,4 +23,4 @@ jobs:
         run: |
           sudo apt-get install python3-sphinx
           pip install -r doc-requirements.txt
-          SPHINXOPTS="-W" cd docs && make html
+          cd docs && SPHINXOPTS="-W" make html

.github/workflows/pythonpublish.yml

@@ -141,3 +141,87 @@ jobs:
           file: ./plugins/flytekit-sqlalchemy/Dockerfile
           cache-from: type=gha
           cache-to: type=gha,mode=max
+
+  build-and-push-external-plugin-service-images:
+    runs-on: ubuntu-latest
+    needs: deploy
+    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
+        if: ${{ github.event_name == 'release' }}
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: "${{ secrets.FLYTE_BOT_USERNAME }}"
+          password: "${{ secrets.FLYTE_BOT_PAT }}"
+      - name: Prepare External Plugin Service Image Names
+        id: external-plugin-service-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/external-plugin-service
+          tags: |
+            latest
+            ${{ github.sha }}
+            ${{ needs.deploy.outputs.version }}
+      - name: Push External Plugin Service Image to GitHub Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: "."
+          platforms: linux/arm64, linux/amd64
+          push: ${{ github.event_name == 'release' }}
+          tags: ${{ steps.external-plugin-service-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ needs.deploy.outputs.version }}
+          file: ./Dockerfile.external-plugin-service
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+  build-and-push-spark-images:
+    runs-on: ubuntu-latest
+    needs: deploy
+    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
+        if: ${{ github.event_name == 'release' }}
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: "${{ secrets.FLYTE_BOT_USERNAME }}"
+          password: "${{ secrets.FLYTE_BOT_PAT }}"
+      - name: Prepare Spark Image Names
+        id: spark-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/flytekit
+          tags: |
+            spark-latest
+            spark-${{ github.sha }}
+            spark-${{ needs.deploy.outputs.version }}
+      - name: Push Spark Image to GitHub Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: "./plugins/flytekit-spark/"
+          platforms: linux/arm64, linux/amd64
+          push: ${{ github.event_name == 'release' }}
+          tags: ${{ steps.spark-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ needs.deploy.outputs.version }}
+          file: ./plugins/flytekit-spark/Dockerfile
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

Dockerfile.external-plugin-service

@@ -0,0 +1,9 @@
+FROM python:3.9-slim-buster
+
+MAINTAINER Flyte Team <[email protected]>
+LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytekit
+
+ARG VERSION
+RUN pip install -U flytekit==$VERSION flytekitplugins-bigquery==$VERSION
+
+CMD pyflyte serve --port 8000

doc-requirements.in

@@ -1,6 +1,7 @@
 .
 -e file:.#egg=flytekit
 
+grpcio<=1.49.1
 git+https://github.com/flyteorg/furo@main
 sphinx
 sphinx-gallery
@@ -11,7 +12,7 @@ sphinx-autoapi
 sphinx-copybutton
 sphinx_fontawesome
 sphinx-panels
-sphinxcontrib-yt
+sphinxcontrib-youtube
 cryptography
 google-api-core[grpc]
 scikit-learn

doc-requirements.txt

@@ -244,7 +244,7 @@ flask==2.2.3
     # via mlflow
 flatbuffers==23.1.21
     # via tensorflow
-flyteidl==1.5.0
+flyteidl==1.5.4
     # via flytekit
 fonttools==4.38.0
     # via matplotlib
@@ -1101,7 +1101,7 @@ sphinxcontrib-qthelp==1.0.3
     # via sphinx
 sphinxcontrib-serializinghtml==1.1.5
     # via sphinx
-sphinxcontrib-yt==0.2.2
+sphinxcontrib-youtube==1.2.0
     # via -r doc-requirements.in
 sqlalchemy==1.4.46
     # via

docs/source/conf.py

@@ -57,7 +57,7 @@
     "sphinx-prompt",
     "sphinx_copybutton",
     "sphinx_panels",
-    "sphinxcontrib.yt",
+    "sphinxcontrib.youtube",
     "sphinx_tags",
     "sphinx_click",
 ]

flytekit/__init__.py

@@ -197,6 +197,8 @@
 
 from rich import traceback
 
+from flytekit.lazy_import.lazy_module import lazy_module
+
 if sys.version_info < (3, 10):
     from importlib_metadata import entry_points
 else:
@@ -224,15 +226,15 @@
 from flytekit.core.workflow import ImperativeWorkflow as Workflow
 from flytekit.core.workflow import WorkflowFailurePolicy, reference_workflow, workflow
 from flytekit.deck import Deck
-from flytekit.extras import pytorch, sklearn, tensorflow
+from flytekit.image_spec import ImageSpec
 from flytekit.loggers import logger
 from flytekit.models.common import Annotations, AuthRole, Labels
 from flytekit.models.core.execution import WorkflowExecutionPhase
 from flytekit.models.core.types import BlobType
 from flytekit.models.documentation import Description, Documentation, SourceCode
 from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar
 from flytekit.models.types import LiteralType
-from flytekit.types import directory, file, numpy, schema
+from flytekit.types import directory, file
 from flytekit.types.structured.structured_dataset import (
     StructuredDataset,
     StructuredDatasetFormat,

flytekit/clients/auth/authenticator.py

@@ -48,10 +48,17 @@ class Authenticator(object):
     Base authenticator for all authentication flows
     """
 
-    def __init__(self, endpoint: str, header_key: str, credentials: Credentials = None):
+    def __init__(
+        self,
+        endpoint: str,
+        header_key: str,
+        credentials: Credentials = None,
+        http_proxy_url: typing.Optional[str] = None,
+    ):
         self._endpoint = endpoint
         self._creds = credentials
         self._header_key = header_key if header_key else "authorization"
+        self._http_proxy_url = http_proxy_url
 
     def get_credentials(self) -> Credentials:
         return self._creds
@@ -162,6 +169,7 @@ def __init__(
         cfg_store: ClientConfigStore,
         header_key: typing.Optional[str] = None,
         scopes: typing.Optional[typing.List[str]] = None,
+        http_proxy_url: typing.Optional[str] = None,
     ):
         if not client_id or not client_secret:
             raise ValueError("Client ID and Client SECRET both are required.")
@@ -171,7 +179,7 @@ def __init__(
         self._scopes = scopes or cfg.scopes
         self._client_id = client_id
         self._client_secret = client_secret
-        super().__init__(endpoint, cfg.header_key or header_key)
+        super().__init__(endpoint, cfg.header_key or header_key, http_proxy_url=http_proxy_url)
 
     def refresh_credentials(self):
         """
@@ -187,7 +195,9 @@ def refresh_credentials(self):
         # Note that unlike the Pkce flow, the client ID does not come from Admin.
         logging.debug(f"Basic authorization flow with client id {self._client_id} scope {scopes}")
         authorization_header = token_client.get_basic_authorization_header(self._client_id, self._client_secret)
-        token, expires_in = token_client.get_token(token_endpoint, scopes, authorization_header)
+        token, expires_in = token_client.get_token(
+            token_endpoint, scopes, authorization_header, http_proxy_url=self._http_proxy_url
+        )
         logging.info("Retrieved new token, expires in {}".format(expires_in))
         self._creds = Credentials(token)
 
@@ -207,6 +217,7 @@ def __init__(
         cfg_store: ClientConfigStore,
         header_key: typing.Optional[str] = None,
         audience: typing.Optional[str] = None,
+        http_proxy_url: typing.Optional[str] = None,
     ):
         self._audience = audience
         cfg = cfg_store.get_client_config()
@@ -219,21 +230,27 @@ def __init__(
                 "Device Authentication is not available on the Flyte backend / authentication server"
             )
         super().__init__(
-            endpoint=endpoint, header_key=header_key or cfg.header_key, credentials=KeyringStore.retrieve(endpoint)
+            endpoint=endpoint,
+            header_key=header_key or cfg.header_key,
+            credentials=KeyringStore.retrieve(endpoint),
+            http_proxy_url=http_proxy_url,
         )
 
     def refresh_credentials(self):
-        resp = token_client.get_device_code(self._device_auth_endpoint, self._client_id, self._audience, self._scope)
+        resp = token_client.get_device_code(
+            self._device_auth_endpoint, self._client_id, self._audience, self._scope, self._http_proxy_url
+        )
         print(
             f"""
 To Authenticate navigate in a browser to the following URL: {resp.verification_uri} and enter code: {resp.user_code}
-OR copy paste the following URL: {resp.verification_uri_complete}
         """
         )
         try:
             # Currently the refresh token is not retreived. We may want to add support for refreshTokens so that
             # access tokens can be refreshed for once authenticated machines
-            token, expires_in = token_client.poll_token_endpoint(resp, self._token_endpoint, client_id=self._client_id)
+            token, expires_in = token_client.poll_token_endpoint(
+                resp, self._token_endpoint, client_id=self._client_id, http_proxy_url=self._http_proxy_url
+            )
             self._creds = Credentials(access_token=token, expires_in=expires_in, for_endpoint=self._endpoint)
             KeyringStore.store(self._creds)
         except Exception:

flytekit/clients/auth/token_client.py

@@ -9,6 +9,7 @@
 
 import requests
 
+from flytekit import logger
 from flytekit.clients.auth.exceptions import AuthenticationError, AuthenticationPending
 
 utf_8 = "utf-8"
@@ -31,15 +32,13 @@ class DeviceCodeResponse:
     {'device_code': 'code',
          'user_code': 'BNDJJFXL',
          'verification_uri': 'url',
-         'verification_uri_complete': 'url',
          'expires_in': 600,
          'interval': 5}
     """
 
     device_code: str
     user_code: str
     verification_uri: str
-    verification_uri_complete: str
     expires_in: int
     interval: int
 
@@ -49,7 +48,6 @@ def from_json_response(cls, j: typing.Dict) -> "DeviceCodeResponse":
             device_code=j["device_code"],
             user_code=j["user_code"],
             verification_uri=j["verification_uri"],
-            verification_uri_complete=j["verification_uri_complete"],
             expires_in=j["expires_in"],
             interval=j["interval"],
         )
@@ -77,6 +75,7 @@ def get_token(
     client_id: typing.Optional[str] = None,
     device_code: typing.Optional[str] = None,
     grant_type: GrantType = GrantType.CLIENT_CREDS,
+    http_proxy_url: typing.Optional[str] = None,
 ) -> typing.Tuple[str, int]:
     """
     :rtype: (Text,Int) The first element is the access token retrieved from the IDP, the second is the expiration
@@ -99,7 +98,8 @@ def get_token(
     if scopes is not None:
         body["scope"] = ",".join(scopes)
 
-    response = requests.post(token_endpoint, data=body, headers=headers)
+    proxies = {"https": http_proxy_url, "http": http_proxy_url} if http_proxy_url else None
+    response = requests.post(token_endpoint, data=body, headers=headers, proxies=proxies)
     if not response.ok:
         j = response.json()
         if "error" in j:
@@ -118,19 +118,24 @@ def get_device_code(
     client_id: str,
     audience: typing.Optional[str] = None,
     scope: typing.Optional[typing.List[str]] = None,
+    http_proxy_url: typing.Optional[str] = None,
 ) -> DeviceCodeResponse:
     """
     Retrieves the device Authentication code that can be done to authenticate the request using a browser on a
     separate device
     """
-    payload = {"client_id": client_id, "scope": scope, "audience": audience}
-    resp = requests.post(device_auth_endpoint, payload)
+    _scope = " ".join(scope) if scope is not None else ""
+    payload = {"client_id": client_id, "scope": _scope, "audience": audience}
+    proxies = {"https": http_proxy_url, "http": http_proxy_url} if http_proxy_url else None
+    resp = requests.post(device_auth_endpoint, payload, proxies=proxies)
     if not resp.ok:
         raise AuthenticationError(f"Unable to retrieve Device Authentication Code for {payload}, Reason {resp.reason}")
     return DeviceCodeResponse.from_json_response(resp.json())
 
 
-def poll_token_endpoint(resp: DeviceCodeResponse, token_endpoint: str, client_id: str) -> typing.Tuple[str, int]:
+def poll_token_endpoint(
+    resp: DeviceCodeResponse, token_endpoint: str, client_id: str, http_proxy_url: typing.Optional[str] = None
+) -> typing.Tuple[str, int]:
     tick = datetime.now()
     interval = timedelta(seconds=resp.interval)
     end_time = tick + timedelta(seconds=resp.expires_in)
@@ -141,13 +146,15 @@ def poll_token_endpoint(resp: DeviceCodeResponse, token_endpoint: str, client_id
                 grant_type=GrantType.DEVICE_CODE,
                 client_id=client_id,
                 device_code=resp.device_code,
+                http_proxy_url=http_proxy_url,
             )
             print("Authentication successful!")
             return access_token, expires_in
         except AuthenticationPending:
             ...
-        except Exception:
-            raise
+        except Exception as e:
+            logger.error("Authentication attempt failed: ", e)
+            raise e
         print("Authentication Pending...")
         time.sleep(interval.total_seconds())
         tick = tick + interval

flytekit/clients/auth_helper.py

@@ -72,6 +72,7 @@ def get_authenticator(cfg: PlatformConfig, cfg_store: ClientConfigStore) -> Auth
             client_secret=cfg.client_credentials_secret,
             cfg_store=cfg_store,
             scopes=cfg.scopes,
+            http_proxy_url=cfg.http_proxy_url,
         )
     elif cfg_auth == AuthType.EXTERNAL_PROCESS or cfg_auth == AuthType.EXTERNALCOMMAND:
         client_cfg = None
@@ -82,7 +83,9 @@ def get_authenticator(cfg: PlatformConfig, cfg_store: ClientConfigStore) -> Auth
             header_key=client_cfg.header_key if client_cfg else None,
         )
     elif cfg_auth == AuthType.DEVICEFLOW:
-        return DeviceCodeAuthenticator(endpoint=cfg.endpoint, cfg_store=cfg_store, audience=cfg.audience)
+        return DeviceCodeAuthenticator(
+            endpoint=cfg.endpoint, cfg_store=cfg_store, audience=cfg.audience, http_proxy_url=cfg.http_proxy_url
+        )
     else:
         raise ValueError(
             f"Invalid auth mode [{cfg_auth}] specified." f"Please update the creds config to use a valid value"

flytekit/clients/friendly.py

@@ -1007,7 +1007,7 @@ def get_upload_signed_url(
 
     def get_download_signed_url(
         self, native_url: str, expires_in: datetime.timedelta = None
-    ) -> _data_proxy_pb2.CreateUploadLocationResponse:
+    ) -> _data_proxy_pb2.CreateDownloadLocationRequest:
         expires_in_pb = None
         if expires_in:
             expires_in_pb = Duration()