Backport 1.7 to 1.2 (#1715)

* FLYTECTL_CONFIG env var higher precedence, config flag respected in pyflyte package (#1662) Signed-off-by: Yee Hing Tong <[email protected]> * Change flytekit Pytorch, TFJob and MPI plugins to use new kubeflow config (#1627) * upgrade tensorflow plugin to v1 Signed-off-by: Yubo Wang <[email protected]> * minor fix Signed-off-by: Yubo Wang <[email protected]> * fix tests and lints Signed-off-by: Yubo Wang <[email protected]> * move models file into task make backward compatible Signed-off-by: Yubo Wang <[email protected]> Signed-off-by: Yubo Wang <[email protected]> * add code example in README Signed-off-by: Yubo Wang <[email protected]> * bump flyteidl Signed-off-by: Yubo Wang <[email protected]> * add pytorch Signed-off-by: Yubo Wang <[email protected]> * add mpi and fix requirements.txt Signed-off-by: Yubo Wang <[email protected]> * lint and fmt Signed-off-by: Yubo Wang <[email protected]> * Regenerate requirements files using python 3.8 Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: Yubo Wang <[email protected]> Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: Yubo Wang <[email protected]> Co-authored-by: eduardo apolinario <[email protected]> * Root cert should be byte string when loading from caCertFilePath (#1669) Signed-off-by: Yee Hing Tong <[email protected]> * Explicitly set the content type for flyte deck (#1658) * Set content type for flyte deck Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]> * Use protos of new kubeflow.pytorch plugin instead of legacy pytorch plugin (#1678) Signed-off-by: Fabio Grätz <[email protected]> Co-authored-by: Fabio Grätz <[email protected]> * More time info for time line deck (#1680) * more visualization Signed-off-by: Yicheng-Lu-llll <[email protected]> * more visualization Signed-off-by: Yicheng-Lu-llll <[email protected]> --------- Signed-off-by: Yicheng-Lu-llll <[email protected]> * Add http_proxy to client & Fix deviceflow (#1611) * Add http_proxy to client & Fix deviceflow RB=3890720 Signed-off-by: byhsu <[email protected]> * nit Signed-off-by: byhsu <[email protected]> * lint! Signed-off-by: byhsu <[email protected]> --------- Signed-off-by: byhsu <[email protected]> Co-authored-by: byhsu <[email protected]> * Pass verify flag to all authenticators (#1641) Signed-off-by: byhsu <[email protected]> * feat: Add Auth0/audience support for ClientCredentials flow (#1639) * feat: Add Auth0/audience support for ClientCredentials flow Signed-off-by: tnam <[email protected]> * refactor: Remove unneeded variables & condense code Signed-off-by: tnam <[email protected]> * refactor: Reduce verbosity of code Signed-off-by: tnam <[email protected]> * refactor(chore): Remove unused commented code Signed-off-by: tnam <[email protected]> * fix: Missing comma in input args - authenticator.py 213 Signed-off-by: tnam <[email protected]> * style: Run pre-commit against all files Signed-off-by: tnam <[email protected]> --------- Signed-off-by: tnam <[email protected]> Co-authored-by: tnam <[email protected]> * pyflyte run remote file (#1670) Signed-off-by: ChungYujoyce <[email protected]> Signed-off-by: Kevin Su <[email protected]> Co-authored-by: Kevin Su <[email protected]> * upload deck.html only with deck enable (#1693) Signed-off-by: Yicheng-Lu-llll <[email protected]> * Add dask plugin #patch (#1366) * Add dummy task type to test backend plugin Signed-off-by: Bernhard Stadlbauer <[email protected]> * Add docs page Signed-off-by: Bernhard Stadlbauer <[email protected]> * Add dask models Signed-off-by: Bernhard Stadlbauer <[email protected]> * Add function to convert resources Signed-off-by: Bernhard Stadlbauer <[email protected]> * Add tests to `dask` task Signed-off-by: Bernhard Stadlbauer <[email protected]> * Remove namespace Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update setup.py Signed-off-by: Bernhard Stadlbauer <[email protected]> * Add dask to `plugin/README.md` Signed-off-by: Bernhard Stadlbauer <[email protected]> * Add README.md for `dask` Signed-off-by: Bernhard Stadlbauer <[email protected]> * Top level export of `JopPodSpec` and `DaskCluster` Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update docs for images Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update README.md Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update models after `flyteidl` change Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update task after `flyteidl` change Signed-off-by: Bernhard Stadlbauer <[email protected]> * Raise error when less than 1 worker Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update flyteidl to >= 1.3.2 Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update doc requirements Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update doc-requirements.txt Signed-off-by: Bernhard Stadlbauer <[email protected]> * Re-lock dependencies on linux Signed-off-by: Bernhard Stadlbauer <[email protected]> * Update dask API docs Signed-off-by: Bernhard Stadlbauer <[email protected]> * Fix documentation links Signed-off-by: Bernhard Stadlbauer <[email protected]> * Default optional model constructor arguments to `None` Signed-off-by: Bernhard Stadlbauer <[email protected]> * Refactor `convert_resources_to_resource_model` to `core.resources` Signed-off-by: Bernhard Stadlbauer <[email protected]> * Use `convert_resources_to_resource_model` in `core.node` Signed-off-by: Bernhard Stadlbauer <[email protected]> * Incorporate review feedback Signed-off-by: Eduardo Apolinario <[email protected]> * Lint Signed-off-by: Eduardo Apolinario <[email protected]> Signed-off-by: Bernhard Stadlbauer <[email protected]> Signed-off-by: Bernhard Stadlbauer <[email protected]> Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]> * Do not bring dask changes Signed-off-by: eduardo apolinario <[email protected]> * Remove readthedocs Signed-off-by: eduardo apolinario <[email protected]> * Linting Signed-off-by: eduardo apolinario <[email protected]> * Force scipy<1.11.0 is on windows Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: Yee Hing Tong <[email protected]> Signed-off-by: Yubo Wang <[email protected]> Signed-off-by: eduardo apolinario <[email protected]> Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Fabio Grätz <[email protected]> Signed-off-by: Yicheng-Lu-llll <[email protected]> Signed-off-by: byhsu <[email protected]> Signed-off-by: tnam <[email protected]> Signed-off-by: ChungYujoyce <[email protected]> Signed-off-by: Bernhard Stadlbauer <[email protected]> Signed-off-by: Bernhard Stadlbauer <[email protected]> Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Yee Hing Tong <[email protected]> Co-authored-by: Yubo Wang <[email protected]> Co-authored-by: Yubo Wang <[email protected]> Co-authored-by: eduardo apolinario <[email protected]> Co-authored-by: Kevin Su <[email protected]> Co-authored-by: Fabio M. Graetz, Ph.D <[email protected]> Co-authored-by: Fabio Grätz <[email protected]> Co-authored-by: Yicheng-Lu-llll <[email protected]> Co-authored-by: ByronHsu <[email protected]> Co-authored-by: byhsu <[email protected]> Co-authored-by: TomNam <[email protected]> Co-authored-by: tnam <[email protected]> Co-authored-by: ChungYujoyce <[email protected]> Co-authored-by: bstadlbauer <[email protected]>
flyteorg · Jun 30, 2023 · 7c1a354 · 7c1a354
1 parent 0b4cc5e
commit 7c1a354
Show file tree

Hide file tree

Showing 47 changed files with 2,025 additions and 488 deletions.
diff --git a/.readthedocs.yml b/.readthedocs.yml
diff --git a/dev-requirements.in b/dev-requirements.in
@@ -22,6 +22,7 @@ tensorflow==2.8.1; platform_machine!='arm64' or platform_system!='Darwin'
 # we put this constraint while we do not have per-environment requirements files
 torch<=1.12.1
 scikit-learn
+scipy<1.11.0; platform_system=='Windows'
 types-protobuf<4
 types-croniter
 types-mock

diff --git a/docs/source/plugins/dask.rst b/docs/source/plugins/dask.rst
@@ -0,0 +1,12 @@
+.. _dask:
+
+###################################################
+Dask API reference
+###################################################
+
+.. tags:: Integration, DistributedComputing, KubernetesOperator
+
+.. automodule:: flytekitplugins.dask
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:
diff --git a/docs/source/plugins/index.rst b/docs/source/plugins/index.rst
@@ -9,6 +9,7 @@ Plugin API reference
 * :ref:`AWS Sagemaker <awssagemaker>` - AWS Sagemaker plugin reference
 * :ref:`Google Bigquery <bigquery>` - Google Bigquery plugin reference
 * :ref:`FS Spec <fsspec>` - FS Spec API reference
+* :ref:`Dask <dask>` - Dask standard API reference
 * :ref:`Deck standard <deck>` - Deck standard API reference
 * :ref:`Dolt standard <dolt>` - Dolt standard API reference
 * :ref:`Great expectations <greatexpectations>` - Great expectations API reference
@@ -41,6 +42,7 @@ Plugin API reference
    AWS Sagemaker <awssagemaker>
    Google Bigquery <bigquery>
    FS Spec <fsspec>
+   Dask <dask>
    Deck standard <deck>
    Dolt standard <dolt>
    Great expectations <greatexpectations>

diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py
@@ -24,6 +24,7 @@
 from flytekit.core.data_persistence import FileAccessProvider
 from flytekit.core.map_task import MapTaskResolver
 from flytekit.core.promise import VoidPromise
+from flytekit.deck.deck import _output_deck
 from flytekit.exceptions import scopes as _scoped_exceptions
 from flytekit.exceptions import scopes as _scopes
 from flytekit.interfaces.stats.taggable import get_stats as _get_stats
@@ -158,6 +159,10 @@ def _dispatch_execute(
 
     ctx.file_access.put_data(ctx.execution_state.engine_dir, output_prefix, is_multipart=True)
     logger.info(f"Engine folder written successfully to the output prefix {output_prefix}")
+
+    if not task_def.disable_deck:
+        _output_deck(task_def.name.split(".")[-1], ctx.user_space_params)
+
     logger.debug("Finished _dispatch_execute")
 
     if os.environ.get("FLYTE_FAIL_ON_ERROR", "").lower() == "true" and _constants.ERROR_FILE_NAME in output_file_dict:
@@ -183,14 +188,16 @@ def get_one_of(*args) -> str:
 @contextlib.contextmanager
 def setup_execution(
     raw_output_data_prefix: str,
+    output_metadata_prefix: Optional[str] = None,
     checkpoint_path: Optional[str] = None,
     prev_checkpoint: Optional[str] = None,
     dynamic_addl_distro: Optional[str] = None,
     dynamic_dest_dir: Optional[str] = None,
 ):
     """
 
-    :param raw_output_data_prefix:
+    :param raw_output_data_prefix: Where to write offloaded data (files, directories, dataframes).
+    :param output_metadata_prefix: Where to write primitive outputs.
     :param checkpoint_path:
     :param prev_checkpoint:
     :param dynamic_addl_distro: Works in concert with the other dynamic arg. If present, indicates that if a dynamic
@@ -247,6 +254,7 @@ def setup_execution(
         logging=user_space_logger,
         tmp_dir=user_workspace_dir,
         raw_output_prefix=raw_output_data_prefix,
+        output_metadata_prefix=output_metadata_prefix,
         checkpoint=checkpointer,
         task_id=_identifier.Identifier(_identifier.ResourceType.TASK, tk_project, tk_domain, tk_name, tk_version),
     )
@@ -337,6 +345,7 @@ def _execute_task(
 
     with setup_execution(
         raw_output_data_prefix,
+        output_prefix,
         checkpoint_path,
         prev_checkpoint,
         dynamic_addl_distro,

diff --git a/flytekit/clients/auth/authenticator.py b/flytekit/clients/auth/authenticator.py
@@ -23,6 +23,7 @@ class ClientConfig:
     device_authorization_endpoint: typing.Optional[str] = None
     scopes: typing.List[str] = None
     header_key: str = "authorization"
+    audience: typing.Optional[str] = None
 
 
 class ClientConfigStore(object):
@@ -48,10 +49,19 @@ 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,
+        verify: typing.Optional[typing.Union[bool, 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
+        self._verify = verify
 
     def get_credentials(self) -> Credentials:
         return self._creds
@@ -87,10 +97,9 @@ def __init__(
         """
         Initialize with default creds from KeyStore using the endpoint name
         """
-        super().__init__(endpoint, header_key, KeyringStore.retrieve(endpoint))
+        super().__init__(endpoint, header_key, KeyringStore.retrieve(endpoint), verify=verify)
         self._cfg_store = cfg_store
         self._auth_client = None
-        self._verify = verify
 
     def _initialize_auth_client(self):
         if not self._auth_client:
@@ -162,6 +171,9 @@ 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,
+        verify: typing.Optional[typing.Union[bool, str]] = None,
+        audience: typing.Optional[str] = None,
     ):
         if not client_id or not client_secret:
             raise ValueError("Client ID and Client SECRET both are required.")
@@ -171,7 +183,8 @@ 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)
+        self._audience = audience or cfg.audience
+        super().__init__(endpoint, cfg.header_key or header_key, http_proxy_url=http_proxy_url, verify=verify)
 
     def refresh_credentials(self):
         """
@@ -183,11 +196,21 @@ def refresh_credentials(self):
         """
         token_endpoint = self._token_endpoint
         scopes = self._scopes
+        audience = self._audience
 
         # 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=token_endpoint,
+            authorization_header=authorization_header,
+            http_proxy_url=self._http_proxy_url,
+            verify=self._verify,
+            scopes=scopes,
+            audience=audience,
+        )
+
         logging.info("Retrieved new token, expires in {}".format(expires_in))
         self._creds = Credentials(token)
 
@@ -207,6 +230,8 @@ def __init__(
         cfg_store: ClientConfigStore,
         header_key: typing.Optional[str] = None,
         audience: typing.Optional[str] = None,
+        http_proxy_url: typing.Optional[str] = None,
+        verify: typing.Optional[typing.Union[bool, str]] = None,
     ):
         self._audience = audience
         cfg = cfg_store.get_client_config()
@@ -219,21 +244,32 @@ 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,
+            verify=verify,
         )
 
     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, self._verify
+        )
         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,
+                verify=self._verify,
+            )
             self._creds = Credentials(access_token=token, expires_in=expires_in, for_endpoint=self._endpoint)
             KeyringStore.store(self._creds)
         except Exception:

diff --git a/flytekit/clients/auth/token_client.py b/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"],
         )
@@ -76,7 +74,10 @@ def get_token(
     authorization_header: typing.Optional[str] = None,
     client_id: typing.Optional[str] = None,
     device_code: typing.Optional[str] = None,
+    audience: typing.Optional[str] = None,
     grant_type: GrantType = GrantType.CLIENT_CREDS,
+    http_proxy_url: typing.Optional[str] = None,
+    verify: typing.Optional[typing.Union[bool, 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
@@ -98,8 +99,12 @@ def get_token(
         body["device_code"] = device_code
     if scopes is not None:
         body["scope"] = ",".join(scopes)
+    if audience:
+        body["audience"] = audience
+
+    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, verify=verify)
 
-    response = requests.post(token_endpoint, data=body, headers=headers)
     if not response.ok:
         j = response.json()
         if "error" in j:
@@ -118,19 +123,29 @@ 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,
+    verify: typing.Optional[typing.Union[bool, 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, verify=verify)
     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,
+    verify: typing.Optional[typing.Union[bool, str]] = None,
+) -> typing.Tuple[str, int]:
     tick = datetime.now()
     interval = timedelta(seconds=resp.interval)
     end_time = tick + timedelta(seconds=resp.expires_in)
@@ -141,13 +156,16 @@ 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,
+                verify=verify,
             )
             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