flyte 0.0.1b0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of flyte might be problematic. Click here for more details.
- flyte/__init__.py +62 -0
- flyte/_api_commons.py +3 -0
- flyte/_bin/__init__.py +0 -0
- flyte/_bin/runtime.py +126 -0
- flyte/_build.py +25 -0
- flyte/_cache/__init__.py +12 -0
- flyte/_cache/cache.py +146 -0
- flyte/_cache/defaults.py +9 -0
- flyte/_cache/policy_function_body.py +42 -0
- flyte/_cli/__init__.py +0 -0
- flyte/_cli/_common.py +287 -0
- flyte/_cli/_create.py +42 -0
- flyte/_cli/_delete.py +23 -0
- flyte/_cli/_deploy.py +140 -0
- flyte/_cli/_get.py +235 -0
- flyte/_cli/_run.py +152 -0
- flyte/_cli/main.py +72 -0
- flyte/_code_bundle/__init__.py +8 -0
- flyte/_code_bundle/_ignore.py +113 -0
- flyte/_code_bundle/_packaging.py +187 -0
- flyte/_code_bundle/_utils.py +339 -0
- flyte/_code_bundle/bundle.py +178 -0
- flyte/_context.py +146 -0
- flyte/_datastructures.py +342 -0
- flyte/_deploy.py +202 -0
- flyte/_doc.py +29 -0
- flyte/_docstring.py +32 -0
- flyte/_environment.py +43 -0
- flyte/_group.py +31 -0
- flyte/_hash.py +23 -0
- flyte/_image.py +760 -0
- flyte/_initialize.py +634 -0
- flyte/_interface.py +84 -0
- flyte/_internal/__init__.py +3 -0
- flyte/_internal/controllers/__init__.py +115 -0
- flyte/_internal/controllers/_local_controller.py +118 -0
- flyte/_internal/controllers/_trace.py +40 -0
- flyte/_internal/controllers/pbhash.py +39 -0
- flyte/_internal/controllers/remote/__init__.py +40 -0
- flyte/_internal/controllers/remote/_action.py +141 -0
- flyte/_internal/controllers/remote/_client.py +43 -0
- flyte/_internal/controllers/remote/_controller.py +361 -0
- flyte/_internal/controllers/remote/_core.py +402 -0
- flyte/_internal/controllers/remote/_informer.py +361 -0
- flyte/_internal/controllers/remote/_service_protocol.py +50 -0
- flyte/_internal/imagebuild/__init__.py +11 -0
- flyte/_internal/imagebuild/docker_builder.py +416 -0
- flyte/_internal/imagebuild/image_builder.py +241 -0
- flyte/_internal/imagebuild/remote_builder.py +0 -0
- flyte/_internal/resolvers/__init__.py +0 -0
- flyte/_internal/resolvers/_task_module.py +54 -0
- flyte/_internal/resolvers/common.py +31 -0
- flyte/_internal/resolvers/default.py +28 -0
- flyte/_internal/runtime/__init__.py +0 -0
- flyte/_internal/runtime/convert.py +199 -0
- flyte/_internal/runtime/entrypoints.py +135 -0
- flyte/_internal/runtime/io.py +136 -0
- flyte/_internal/runtime/resources_serde.py +138 -0
- flyte/_internal/runtime/task_serde.py +210 -0
- flyte/_internal/runtime/taskrunner.py +190 -0
- flyte/_internal/runtime/types_serde.py +54 -0
- flyte/_logging.py +124 -0
- flyte/_protos/__init__.py +0 -0
- flyte/_protos/common/authorization_pb2.py +66 -0
- flyte/_protos/common/authorization_pb2.pyi +108 -0
- flyte/_protos/common/authorization_pb2_grpc.py +4 -0
- flyte/_protos/common/identifier_pb2.py +71 -0
- flyte/_protos/common/identifier_pb2.pyi +82 -0
- flyte/_protos/common/identifier_pb2_grpc.py +4 -0
- flyte/_protos/common/identity_pb2.py +48 -0
- flyte/_protos/common/identity_pb2.pyi +72 -0
- flyte/_protos/common/identity_pb2_grpc.py +4 -0
- flyte/_protos/common/list_pb2.py +36 -0
- flyte/_protos/common/list_pb2.pyi +69 -0
- flyte/_protos/common/list_pb2_grpc.py +4 -0
- flyte/_protos/common/policy_pb2.py +37 -0
- flyte/_protos/common/policy_pb2.pyi +27 -0
- flyte/_protos/common/policy_pb2_grpc.py +4 -0
- flyte/_protos/common/role_pb2.py +37 -0
- flyte/_protos/common/role_pb2.pyi +53 -0
- flyte/_protos/common/role_pb2_grpc.py +4 -0
- flyte/_protos/common/runtime_version_pb2.py +28 -0
- flyte/_protos/common/runtime_version_pb2.pyi +24 -0
- flyte/_protos/common/runtime_version_pb2_grpc.py +4 -0
- flyte/_protos/logs/dataplane/payload_pb2.py +96 -0
- flyte/_protos/logs/dataplane/payload_pb2.pyi +168 -0
- flyte/_protos/logs/dataplane/payload_pb2_grpc.py +4 -0
- flyte/_protos/secret/definition_pb2.py +49 -0
- flyte/_protos/secret/definition_pb2.pyi +93 -0
- flyte/_protos/secret/definition_pb2_grpc.py +4 -0
- flyte/_protos/secret/payload_pb2.py +62 -0
- flyte/_protos/secret/payload_pb2.pyi +94 -0
- flyte/_protos/secret/payload_pb2_grpc.py +4 -0
- flyte/_protos/secret/secret_pb2.py +38 -0
- flyte/_protos/secret/secret_pb2.pyi +6 -0
- flyte/_protos/secret/secret_pb2_grpc.py +198 -0
- flyte/_protos/secret/secret_pb2_grpc_grpc.py +198 -0
- flyte/_protos/validate/validate/validate_pb2.py +76 -0
- flyte/_protos/workflow/node_execution_service_pb2.py +26 -0
- flyte/_protos/workflow/node_execution_service_pb2.pyi +4 -0
- flyte/_protos/workflow/node_execution_service_pb2_grpc.py +32 -0
- flyte/_protos/workflow/queue_service_pb2.py +106 -0
- flyte/_protos/workflow/queue_service_pb2.pyi +141 -0
- flyte/_protos/workflow/queue_service_pb2_grpc.py +172 -0
- flyte/_protos/workflow/run_definition_pb2.py +128 -0
- flyte/_protos/workflow/run_definition_pb2.pyi +310 -0
- flyte/_protos/workflow/run_definition_pb2_grpc.py +4 -0
- flyte/_protos/workflow/run_logs_service_pb2.py +41 -0
- flyte/_protos/workflow/run_logs_service_pb2.pyi +28 -0
- flyte/_protos/workflow/run_logs_service_pb2_grpc.py +69 -0
- flyte/_protos/workflow/run_service_pb2.py +133 -0
- flyte/_protos/workflow/run_service_pb2.pyi +175 -0
- flyte/_protos/workflow/run_service_pb2_grpc.py +412 -0
- flyte/_protos/workflow/state_service_pb2.py +58 -0
- flyte/_protos/workflow/state_service_pb2.pyi +71 -0
- flyte/_protos/workflow/state_service_pb2_grpc.py +138 -0
- flyte/_protos/workflow/task_definition_pb2.py +72 -0
- flyte/_protos/workflow/task_definition_pb2.pyi +65 -0
- flyte/_protos/workflow/task_definition_pb2_grpc.py +4 -0
- flyte/_protos/workflow/task_service_pb2.py +44 -0
- flyte/_protos/workflow/task_service_pb2.pyi +31 -0
- flyte/_protos/workflow/task_service_pb2_grpc.py +104 -0
- flyte/_resources.py +226 -0
- flyte/_retry.py +32 -0
- flyte/_reusable_environment.py +25 -0
- flyte/_run.py +411 -0
- flyte/_secret.py +61 -0
- flyte/_task.py +367 -0
- flyte/_task_environment.py +200 -0
- flyte/_timeout.py +47 -0
- flyte/_tools.py +27 -0
- flyte/_trace.py +128 -0
- flyte/_utils/__init__.py +20 -0
- flyte/_utils/asyn.py +119 -0
- flyte/_utils/coro_management.py +25 -0
- flyte/_utils/file_handling.py +72 -0
- flyte/_utils/helpers.py +108 -0
- flyte/_utils/lazy_module.py +54 -0
- flyte/_utils/uv_script_parser.py +49 -0
- flyte/_version.py +21 -0
- flyte/connectors/__init__.py +0 -0
- flyte/errors.py +143 -0
- flyte/extras/__init__.py +5 -0
- flyte/extras/_container.py +273 -0
- flyte/io/__init__.py +11 -0
- flyte/io/_dataframe.py +0 -0
- flyte/io/_dir.py +448 -0
- flyte/io/_file.py +468 -0
- flyte/io/pickle/__init__.py +0 -0
- flyte/io/pickle/transformer.py +117 -0
- flyte/io/structured_dataset/__init__.py +129 -0
- flyte/io/structured_dataset/basic_dfs.py +219 -0
- flyte/io/structured_dataset/structured_dataset.py +1061 -0
- flyte/py.typed +0 -0
- flyte/remote/__init__.py +25 -0
- flyte/remote/_client/__init__.py +0 -0
- flyte/remote/_client/_protocols.py +131 -0
- flyte/remote/_client/auth/__init__.py +12 -0
- flyte/remote/_client/auth/_authenticators/__init__.py +0 -0
- flyte/remote/_client/auth/_authenticators/base.py +397 -0
- flyte/remote/_client/auth/_authenticators/client_credentials.py +73 -0
- flyte/remote/_client/auth/_authenticators/device_code.py +118 -0
- flyte/remote/_client/auth/_authenticators/external_command.py +79 -0
- flyte/remote/_client/auth/_authenticators/factory.py +200 -0
- flyte/remote/_client/auth/_authenticators/pkce.py +516 -0
- flyte/remote/_client/auth/_channel.py +184 -0
- flyte/remote/_client/auth/_client_config.py +83 -0
- flyte/remote/_client/auth/_default_html.py +32 -0
- flyte/remote/_client/auth/_grpc_utils/__init__.py +0 -0
- flyte/remote/_client/auth/_grpc_utils/auth_interceptor.py +288 -0
- flyte/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +151 -0
- flyte/remote/_client/auth/_keyring.py +143 -0
- flyte/remote/_client/auth/_token_client.py +260 -0
- flyte/remote/_client/auth/errors.py +16 -0
- flyte/remote/_client/controlplane.py +95 -0
- flyte/remote/_console.py +18 -0
- flyte/remote/_data.py +155 -0
- flyte/remote/_logs.py +116 -0
- flyte/remote/_project.py +86 -0
- flyte/remote/_run.py +873 -0
- flyte/remote/_secret.py +132 -0
- flyte/remote/_task.py +227 -0
- flyte/report/__init__.py +3 -0
- flyte/report/_report.py +178 -0
- flyte/report/_template.html +124 -0
- flyte/storage/__init__.py +24 -0
- flyte/storage/_remote_fs.py +34 -0
- flyte/storage/_storage.py +251 -0
- flyte/storage/_utils.py +5 -0
- flyte/types/__init__.py +13 -0
- flyte/types/_interface.py +25 -0
- flyte/types/_renderer.py +162 -0
- flyte/types/_string_literals.py +120 -0
- flyte/types/_type_engine.py +2210 -0
- flyte/types/_utils.py +80 -0
- flyte-0.0.1b0.dist-info/METADATA +179 -0
- flyte-0.0.1b0.dist-info/RECORD +390 -0
- flyte-0.0.1b0.dist-info/WHEEL +5 -0
- flyte-0.0.1b0.dist-info/entry_points.txt +3 -0
- flyte-0.0.1b0.dist-info/top_level.txt +1 -0
- union/__init__.py +54 -0
- union/_api_commons.py +3 -0
- union/_bin/__init__.py +0 -0
- union/_bin/runtime.py +113 -0
- union/_build.py +25 -0
- union/_cache/__init__.py +12 -0
- union/_cache/cache.py +141 -0
- union/_cache/defaults.py +9 -0
- union/_cache/policy_function_body.py +42 -0
- union/_cli/__init__.py +0 -0
- union/_cli/_common.py +263 -0
- union/_cli/_create.py +40 -0
- union/_cli/_delete.py +23 -0
- union/_cli/_deploy.py +120 -0
- union/_cli/_get.py +162 -0
- union/_cli/_params.py +579 -0
- union/_cli/_run.py +150 -0
- union/_cli/main.py +72 -0
- union/_code_bundle/__init__.py +8 -0
- union/_code_bundle/_ignore.py +113 -0
- union/_code_bundle/_packaging.py +187 -0
- union/_code_bundle/_utils.py +342 -0
- union/_code_bundle/bundle.py +176 -0
- union/_context.py +146 -0
- union/_datastructures.py +295 -0
- union/_deploy.py +185 -0
- union/_doc.py +29 -0
- union/_docstring.py +26 -0
- union/_environment.py +43 -0
- union/_group.py +31 -0
- union/_hash.py +23 -0
- union/_image.py +760 -0
- union/_initialize.py +585 -0
- union/_interface.py +84 -0
- union/_internal/__init__.py +3 -0
- union/_internal/controllers/__init__.py +77 -0
- union/_internal/controllers/_local_controller.py +77 -0
- union/_internal/controllers/pbhash.py +39 -0
- union/_internal/controllers/remote/__init__.py +40 -0
- union/_internal/controllers/remote/_action.py +131 -0
- union/_internal/controllers/remote/_client.py +43 -0
- union/_internal/controllers/remote/_controller.py +169 -0
- union/_internal/controllers/remote/_core.py +341 -0
- union/_internal/controllers/remote/_informer.py +260 -0
- union/_internal/controllers/remote/_service_protocol.py +44 -0
- union/_internal/imagebuild/__init__.py +11 -0
- union/_internal/imagebuild/docker_builder.py +416 -0
- union/_internal/imagebuild/image_builder.py +243 -0
- union/_internal/imagebuild/remote_builder.py +0 -0
- union/_internal/resolvers/__init__.py +0 -0
- union/_internal/resolvers/_task_module.py +31 -0
- union/_internal/resolvers/common.py +24 -0
- union/_internal/resolvers/default.py +27 -0
- union/_internal/runtime/__init__.py +0 -0
- union/_internal/runtime/convert.py +163 -0
- union/_internal/runtime/entrypoints.py +121 -0
- union/_internal/runtime/io.py +136 -0
- union/_internal/runtime/resources_serde.py +134 -0
- union/_internal/runtime/task_serde.py +202 -0
- union/_internal/runtime/taskrunner.py +179 -0
- union/_internal/runtime/types_serde.py +53 -0
- union/_logging.py +124 -0
- union/_protos/__init__.py +0 -0
- union/_protos/common/authorization_pb2.py +66 -0
- union/_protos/common/authorization_pb2.pyi +106 -0
- union/_protos/common/authorization_pb2_grpc.py +4 -0
- union/_protos/common/identifier_pb2.py +71 -0
- union/_protos/common/identifier_pb2.pyi +82 -0
- union/_protos/common/identifier_pb2_grpc.py +4 -0
- union/_protos/common/identity_pb2.py +48 -0
- union/_protos/common/identity_pb2.pyi +72 -0
- union/_protos/common/identity_pb2_grpc.py +4 -0
- union/_protos/common/list_pb2.py +36 -0
- union/_protos/common/list_pb2.pyi +69 -0
- union/_protos/common/list_pb2_grpc.py +4 -0
- union/_protos/common/policy_pb2.py +37 -0
- union/_protos/common/policy_pb2.pyi +27 -0
- union/_protos/common/policy_pb2_grpc.py +4 -0
- union/_protos/common/role_pb2.py +37 -0
- union/_protos/common/role_pb2.pyi +51 -0
- union/_protos/common/role_pb2_grpc.py +4 -0
- union/_protos/common/runtime_version_pb2.py +28 -0
- union/_protos/common/runtime_version_pb2.pyi +24 -0
- union/_protos/common/runtime_version_pb2_grpc.py +4 -0
- union/_protos/logs/dataplane/payload_pb2.py +96 -0
- union/_protos/logs/dataplane/payload_pb2.pyi +168 -0
- union/_protos/logs/dataplane/payload_pb2_grpc.py +4 -0
- union/_protos/secret/definition_pb2.py +49 -0
- union/_protos/secret/definition_pb2.pyi +93 -0
- union/_protos/secret/definition_pb2_grpc.py +4 -0
- union/_protos/secret/payload_pb2.py +62 -0
- union/_protos/secret/payload_pb2.pyi +94 -0
- union/_protos/secret/payload_pb2_grpc.py +4 -0
- union/_protos/secret/secret_pb2.py +38 -0
- union/_protos/secret/secret_pb2.pyi +6 -0
- union/_protos/secret/secret_pb2_grpc.py +198 -0
- union/_protos/validate/validate/validate_pb2.py +76 -0
- union/_protos/workflow/node_execution_service_pb2.py +26 -0
- union/_protos/workflow/node_execution_service_pb2.pyi +4 -0
- union/_protos/workflow/node_execution_service_pb2_grpc.py +32 -0
- union/_protos/workflow/queue_service_pb2.py +75 -0
- union/_protos/workflow/queue_service_pb2.pyi +103 -0
- union/_protos/workflow/queue_service_pb2_grpc.py +172 -0
- union/_protos/workflow/run_definition_pb2.py +100 -0
- union/_protos/workflow/run_definition_pb2.pyi +256 -0
- union/_protos/workflow/run_definition_pb2_grpc.py +4 -0
- union/_protos/workflow/run_logs_service_pb2.py +41 -0
- union/_protos/workflow/run_logs_service_pb2.pyi +28 -0
- union/_protos/workflow/run_logs_service_pb2_grpc.py +69 -0
- union/_protos/workflow/run_service_pb2.py +133 -0
- union/_protos/workflow/run_service_pb2.pyi +173 -0
- union/_protos/workflow/run_service_pb2_grpc.py +412 -0
- union/_protos/workflow/state_service_pb2.py +58 -0
- union/_protos/workflow/state_service_pb2.pyi +69 -0
- union/_protos/workflow/state_service_pb2_grpc.py +138 -0
- union/_protos/workflow/task_definition_pb2.py +72 -0
- union/_protos/workflow/task_definition_pb2.pyi +65 -0
- union/_protos/workflow/task_definition_pb2_grpc.py +4 -0
- union/_protos/workflow/task_service_pb2.py +44 -0
- union/_protos/workflow/task_service_pb2.pyi +31 -0
- union/_protos/workflow/task_service_pb2_grpc.py +104 -0
- union/_resources.py +226 -0
- union/_retry.py +32 -0
- union/_reusable_environment.py +25 -0
- union/_run.py +374 -0
- union/_secret.py +61 -0
- union/_task.py +354 -0
- union/_task_environment.py +186 -0
- union/_timeout.py +47 -0
- union/_tools.py +27 -0
- union/_utils/__init__.py +11 -0
- union/_utils/asyn.py +119 -0
- union/_utils/file_handling.py +71 -0
- union/_utils/helpers.py +46 -0
- union/_utils/lazy_module.py +54 -0
- union/_utils/uv_script_parser.py +49 -0
- union/_version.py +21 -0
- union/connectors/__init__.py +0 -0
- union/errors.py +128 -0
- union/extras/__init__.py +5 -0
- union/extras/_container.py +263 -0
- union/io/__init__.py +11 -0
- union/io/_dataframe.py +0 -0
- union/io/_dir.py +425 -0
- union/io/_file.py +418 -0
- union/io/pickle/__init__.py +0 -0
- union/io/pickle/transformer.py +117 -0
- union/io/structured_dataset/__init__.py +122 -0
- union/io/structured_dataset/basic_dfs.py +219 -0
- union/io/structured_dataset/structured_dataset.py +1057 -0
- union/py.typed +0 -0
- union/remote/__init__.py +23 -0
- union/remote/_client/__init__.py +0 -0
- union/remote/_client/_protocols.py +129 -0
- union/remote/_client/auth/__init__.py +12 -0
- union/remote/_client/auth/_authenticators/__init__.py +0 -0
- union/remote/_client/auth/_authenticators/base.py +391 -0
- union/remote/_client/auth/_authenticators/client_credentials.py +73 -0
- union/remote/_client/auth/_authenticators/device_code.py +120 -0
- union/remote/_client/auth/_authenticators/external_command.py +77 -0
- union/remote/_client/auth/_authenticators/factory.py +200 -0
- union/remote/_client/auth/_authenticators/pkce.py +515 -0
- union/remote/_client/auth/_channel.py +184 -0
- union/remote/_client/auth/_client_config.py +83 -0
- union/remote/_client/auth/_default_html.py +32 -0
- union/remote/_client/auth/_grpc_utils/__init__.py +0 -0
- union/remote/_client/auth/_grpc_utils/auth_interceptor.py +204 -0
- union/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +144 -0
- union/remote/_client/auth/_keyring.py +154 -0
- union/remote/_client/auth/_token_client.py +258 -0
- union/remote/_client/auth/errors.py +16 -0
- union/remote/_client/controlplane.py +86 -0
- union/remote/_data.py +149 -0
- union/remote/_logs.py +74 -0
- union/remote/_project.py +86 -0
- union/remote/_run.py +820 -0
- union/remote/_secret.py +132 -0
- union/remote/_task.py +193 -0
- union/report/__init__.py +3 -0
- union/report/_report.py +178 -0
- union/report/_template.html +124 -0
- union/storage/__init__.py +24 -0
- union/storage/_remote_fs.py +34 -0
- union/storage/_storage.py +247 -0
- union/storage/_utils.py +5 -0
- union/types/__init__.py +11 -0
- union/types/_renderer.py +162 -0
- union/types/_string_literals.py +120 -0
- union/types/_type_engine.py +2131 -0
- union/types/_utils.py +80 -0
flyte/py.typed
ADDED
|
File without changes
|
flyte/remote/__init__.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Remote Entities that are accessible from the Union Server once deployed or created.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
__all__ = [
|
|
6
|
+
"Action",
|
|
7
|
+
"ActionDetails",
|
|
8
|
+
"ActionInputs",
|
|
9
|
+
"ActionOutputs",
|
|
10
|
+
"Project",
|
|
11
|
+
"Run",
|
|
12
|
+
"RunDetails",
|
|
13
|
+
"Secret",
|
|
14
|
+
"Task",
|
|
15
|
+
"create_channel",
|
|
16
|
+
"upload_dir",
|
|
17
|
+
"upload_file",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
from ._client.auth import create_channel
|
|
21
|
+
from ._data import upload_dir, upload_file
|
|
22
|
+
from ._project import Project
|
|
23
|
+
from ._run import Action, ActionDetails, ActionInputs, ActionOutputs, Run, RunDetails
|
|
24
|
+
from ._secret import Secret
|
|
25
|
+
from ._task import Task
|
|
File without changes
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
from typing import AsyncIterator, Protocol
|
|
2
|
+
|
|
3
|
+
from flyteidl.admin import project_attributes_pb2, project_pb2, version_pb2
|
|
4
|
+
from flyteidl.service import dataproxy_pb2
|
|
5
|
+
from grpc.aio import UnaryStreamCall
|
|
6
|
+
from grpc.aio._typing import RequestType
|
|
7
|
+
|
|
8
|
+
from flyte._protos.secret import payload_pb2
|
|
9
|
+
from flyte._protos.workflow import run_logs_service_pb2, run_service_pb2, task_service_pb2
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class MetadataServiceProtocol(Protocol):
|
|
13
|
+
async def GetVersion(self, request: version_pb2.GetVersionRequest) -> version_pb2.GetVersionResponse: ...
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class ProjectDomainService(Protocol):
|
|
17
|
+
async def RegisterProject(
|
|
18
|
+
self, request: project_pb2.ProjectRegisterRequest
|
|
19
|
+
) -> project_pb2.ProjectRegisterResponse: ...
|
|
20
|
+
|
|
21
|
+
async def UpdateProject(self, request: project_pb2.Project) -> project_pb2.ProjectUpdateResponse: ...
|
|
22
|
+
|
|
23
|
+
async def GetProject(self, request: project_pb2.ProjectGetRequest) -> project_pb2.Project: ...
|
|
24
|
+
|
|
25
|
+
async def ListProjects(self, request: project_pb2.ProjectListRequest) -> project_pb2.Projects: ...
|
|
26
|
+
|
|
27
|
+
async def GetDomains(self, request: project_pb2.GetDomainRequest) -> project_pb2.GetDomainsResponse: ...
|
|
28
|
+
|
|
29
|
+
async def UpdateProjectDomainAttributes(
|
|
30
|
+
self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
|
|
31
|
+
) -> project_pb2.ProjectUpdateResponse: ...
|
|
32
|
+
|
|
33
|
+
async def GetProjectDomainAttributes(
|
|
34
|
+
self, request: project_attributes_pb2.ProjectAttributesGetRequest
|
|
35
|
+
) -> project_attributes_pb2.ProjectAttributes: ...
|
|
36
|
+
|
|
37
|
+
async def DeleteProjectDomainAttributes(
|
|
38
|
+
self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
|
|
39
|
+
) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
|
|
40
|
+
|
|
41
|
+
async def UpdateProjectAttributes(
|
|
42
|
+
self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
|
|
43
|
+
) -> project_attributes_pb2.ProjectAttributesUpdateResponse: ...
|
|
44
|
+
|
|
45
|
+
async def GetProjectAttributes(
|
|
46
|
+
self, request: project_attributes_pb2.ProjectAttributesGetRequest
|
|
47
|
+
) -> project_attributes_pb2.ProjectAttributes: ...
|
|
48
|
+
|
|
49
|
+
async def DeleteProjectAttributes(
|
|
50
|
+
self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
|
|
51
|
+
) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
class TaskService(Protocol):
|
|
55
|
+
async def DeployTask(self, request: task_service_pb2.DeployTaskRequest) -> task_service_pb2.DeployTaskResponse: ...
|
|
56
|
+
|
|
57
|
+
async def GetTaskDetails(
|
|
58
|
+
self, request: task_service_pb2.GetTaskDetailsRequest
|
|
59
|
+
) -> task_service_pb2.GetTaskDetailsResponse: ...
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
class RunService(Protocol):
|
|
63
|
+
async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...
|
|
64
|
+
|
|
65
|
+
async def AbortRun(self, request: run_service_pb2.AbortRunRequest) -> run_service_pb2.AbortRunResponse: ...
|
|
66
|
+
|
|
67
|
+
async def GetRunDetails(
|
|
68
|
+
self, request: run_service_pb2.GetRunDetailsRequest
|
|
69
|
+
) -> run_service_pb2.GetRunDetailsResponse: ...
|
|
70
|
+
|
|
71
|
+
async def WatchRunDetails(
|
|
72
|
+
self, request: run_service_pb2.WatchRunDetailsRequest
|
|
73
|
+
) -> AsyncIterator[run_service_pb2.WatchRunDetailsResponse]: ...
|
|
74
|
+
|
|
75
|
+
async def GetActionDetails(
|
|
76
|
+
self, request: run_service_pb2.GetActionDetailsRequest
|
|
77
|
+
) -> run_service_pb2.GetActionDetailsResponse: ...
|
|
78
|
+
|
|
79
|
+
async def WatchActionDetails(
|
|
80
|
+
self, request: run_service_pb2.WatchActionDetailsRequest
|
|
81
|
+
) -> AsyncIterator[run_service_pb2.WatchActionDetailsResponse]: ...
|
|
82
|
+
|
|
83
|
+
async def GetActionData(
|
|
84
|
+
self, request: run_service_pb2.GetActionDataRequest
|
|
85
|
+
) -> run_service_pb2.GetActionDataResponse: ...
|
|
86
|
+
|
|
87
|
+
async def ListRuns(self, request: run_service_pb2.ListRunsRequest) -> run_service_pb2.ListRunsResponse: ...
|
|
88
|
+
|
|
89
|
+
async def WatchRuns(
|
|
90
|
+
self, request: run_service_pb2.WatchRunsRequest
|
|
91
|
+
) -> AsyncIterator[run_service_pb2.WatchRunsResponse]: ...
|
|
92
|
+
|
|
93
|
+
async def ListActions(self, request: run_service_pb2.ListActionsRequest) -> run_service_pb2.ListActionsResponse: ...
|
|
94
|
+
|
|
95
|
+
async def WatchActions(
|
|
96
|
+
self, request: run_service_pb2.WatchActionsRequest
|
|
97
|
+
) -> AsyncIterator[run_service_pb2.WatchActionsResponse]: ...
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
class DataProxyService(Protocol):
|
|
101
|
+
async def CreateUploadLocation(
|
|
102
|
+
self, request: dataproxy_pb2.CreateUploadLocationRequest
|
|
103
|
+
) -> dataproxy_pb2.CreateUploadLocationResponse: ...
|
|
104
|
+
|
|
105
|
+
async def CreateDownloadLocation(
|
|
106
|
+
self, request: dataproxy_pb2.CreateDownloadLocationRequest
|
|
107
|
+
) -> dataproxy_pb2.CreateDownloadLocationResponse: ...
|
|
108
|
+
|
|
109
|
+
async def CreateDownloadLink(
|
|
110
|
+
self, request: dataproxy_pb2.CreateDownloadLinkRequest
|
|
111
|
+
) -> dataproxy_pb2.CreateDownloadLinkResponse: ...
|
|
112
|
+
|
|
113
|
+
async def GetData(self, request: dataproxy_pb2.GetDataRequest) -> dataproxy_pb2.GetDataResponse: ...
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
class RunLogsService(Protocol):
|
|
117
|
+
def TailLogs(
|
|
118
|
+
self, request: run_logs_service_pb2.TailLogsRequest
|
|
119
|
+
) -> UnaryStreamCall[RequestType, run_logs_service_pb2.TailLogsResponse]: ...
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
class SecretService(Protocol):
|
|
123
|
+
async def CreateSecret(self, request: payload_pb2.CreateSecretRequest) -> payload_pb2.CreateSecretResponse: ...
|
|
124
|
+
|
|
125
|
+
async def UpdateSecret(self, request: payload_pb2.UpdateSecretRequest) -> payload_pb2.UpdateSecretResponse: ...
|
|
126
|
+
|
|
127
|
+
async def GetSecret(self, request: payload_pb2.GetSecretRequest) -> payload_pb2.GetSecretResponse: ...
|
|
128
|
+
|
|
129
|
+
async def ListSecrets(self, request: payload_pb2.ListSecretsRequest) -> payload_pb2.ListSecretsResponse: ...
|
|
130
|
+
|
|
131
|
+
async def DeleteSecret(self, request: payload_pb2.DeleteSecretRequest) -> payload_pb2.DeleteSecretResponse: ...
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
from flyte.remote._client.auth._channel import create_channel
|
|
2
|
+
from flyte.remote._client.auth._client_config import AuthType, ClientConfig
|
|
3
|
+
from flyte.remote._client.auth.errors import AccessTokenNotFoundError, AuthenticationError, AuthenticationPending
|
|
4
|
+
|
|
5
|
+
__all__ = [
|
|
6
|
+
"AccessTokenNotFoundError",
|
|
7
|
+
"AuthType",
|
|
8
|
+
"AuthenticationError",
|
|
9
|
+
"AuthenticationPending",
|
|
10
|
+
"ClientConfig",
|
|
11
|
+
"create_channel",
|
|
12
|
+
]
|
|
File without changes
|
|
@@ -0,0 +1,397 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import dataclasses
|
|
3
|
+
import ssl
|
|
4
|
+
import typing
|
|
5
|
+
from abc import abstractmethod
|
|
6
|
+
from http import HTTPStatus
|
|
7
|
+
|
|
8
|
+
import httpx
|
|
9
|
+
from grpc.aio import Metadata
|
|
10
|
+
|
|
11
|
+
from flyte.remote._client.auth._client_config import ClientConfig, ClientConfigStore
|
|
12
|
+
from flyte.remote._client.auth._keyring import Credentials, KeyringStore
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@dataclasses.dataclass
|
|
16
|
+
class GrpcAuthMetadata:
|
|
17
|
+
creds_id: str
|
|
18
|
+
pairs: Metadata
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class Authenticator(object):
|
|
22
|
+
"""
|
|
23
|
+
Base authenticator for all authentication flows
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
def __init__(
|
|
27
|
+
self,
|
|
28
|
+
endpoint: str,
|
|
29
|
+
*,
|
|
30
|
+
cfg_store: typing.Optional[ClientConfigStore] = None,
|
|
31
|
+
client_config: typing.Optional[ClientConfig] = None,
|
|
32
|
+
credentials: typing.Optional[Credentials] = None,
|
|
33
|
+
http_session: typing.Optional[httpx.AsyncClient] = None,
|
|
34
|
+
http_proxy_url: typing.Optional[str] = None,
|
|
35
|
+
verify: bool = True,
|
|
36
|
+
ca_cert_path: typing.Optional[str] = None,
|
|
37
|
+
**kwargs,
|
|
38
|
+
):
|
|
39
|
+
"""
|
|
40
|
+
Initialize the base authenticator.
|
|
41
|
+
|
|
42
|
+
:param endpoint: The endpoint URL for authentication
|
|
43
|
+
:param cfg_store: Optional client configuration store for retrieving remote configuration
|
|
44
|
+
:param client_config: Optional client configuration containing authentication settings
|
|
45
|
+
:param credentials: Optional credentials to use for authentication
|
|
46
|
+
:param http_session: Optional HTTP session to use for requests
|
|
47
|
+
:param http_proxy_url: Optional HTTP proxy URL
|
|
48
|
+
:param verify: Whether to verify SSL certificates
|
|
49
|
+
:param ca_cert_path: Optional path to CA certificate file
|
|
50
|
+
:param kwargs: Additional keyword arguments passed to get_async_session, which may include:
|
|
51
|
+
- auth: Authentication implementation to use
|
|
52
|
+
- params: Query parameters to include in request URLs
|
|
53
|
+
- headers: HTTP headers to include in requests
|
|
54
|
+
- cookies: Cookies to include in requests
|
|
55
|
+
- cert: SSL client certificate (path or tuple)
|
|
56
|
+
- http1: Whether to enable HTTP/1.1 support
|
|
57
|
+
- http2: Whether to enable HTTP/2 support
|
|
58
|
+
- proxies: Proxy configuration mapping
|
|
59
|
+
- mounts: Mounted transports for specific URL patterns
|
|
60
|
+
- timeout: Request timeout configuration
|
|
61
|
+
- follow_redirects: Whether to follow redirects
|
|
62
|
+
- limits: Connection pool limits
|
|
63
|
+
- max_redirects: Maximum number of redirects to follow
|
|
64
|
+
- event_hooks: Event hooks for request/response lifecycle
|
|
65
|
+
- base_url: Base URL to join with relative URLs
|
|
66
|
+
- transport: Transport implementation to use
|
|
67
|
+
- app: ASGI application to handle requests
|
|
68
|
+
"""
|
|
69
|
+
self._endpoint = endpoint
|
|
70
|
+
self._creds = credentials or KeyringStore.retrieve(endpoint)
|
|
71
|
+
self._http_proxy_url = http_proxy_url
|
|
72
|
+
self._verify = verify
|
|
73
|
+
self._ca_cert_path = ca_cert_path
|
|
74
|
+
self._client_config = client_config
|
|
75
|
+
self._cfg_store = cfg_store
|
|
76
|
+
# Will be populated by _ensure_remote_config
|
|
77
|
+
self._resolved_config: ClientConfig | None = None
|
|
78
|
+
# Lock for coroutine safety
|
|
79
|
+
self._async_lock = asyncio.Lock()
|
|
80
|
+
self._http_session = http_session or get_async_session(**kwargs)
|
|
81
|
+
# Id for tracking credential refresh state
|
|
82
|
+
self._creds_id = self._creds.id if self._creds else None
|
|
83
|
+
|
|
84
|
+
async def _resolve_config(self) -> ClientConfig:
|
|
85
|
+
"""
|
|
86
|
+
Resolves and merges client configuration with remote configuration.
|
|
87
|
+
|
|
88
|
+
This method fetches the remote configuration from the cfg_store and merges it with
|
|
89
|
+
the local client_config, prioritizing local settings over remote ones.
|
|
90
|
+
|
|
91
|
+
This method is thread-safe and coroutine-safe, ensuring the remote config is fetched
|
|
92
|
+
only once regardless of concurrent access from multiple threads or coroutines.
|
|
93
|
+
|
|
94
|
+
:return: A merged ClientConfig object containing resolved configuration settings
|
|
95
|
+
"""
|
|
96
|
+
# First check without locks for performance
|
|
97
|
+
if self._resolved_config is not None:
|
|
98
|
+
return self._resolved_config
|
|
99
|
+
|
|
100
|
+
if self._cfg_store is None:
|
|
101
|
+
raise ValueError("ClientConfigStore is not set. Cannot resolve configuration.")
|
|
102
|
+
|
|
103
|
+
remote_config = await self._cfg_store.get_client_config()
|
|
104
|
+
self._resolved_config = (
|
|
105
|
+
remote_config.with_override(self._client_config) if self._client_config else remote_config
|
|
106
|
+
)
|
|
107
|
+
|
|
108
|
+
return self._resolved_config
|
|
109
|
+
|
|
110
|
+
def get_credentials(self) -> typing.Optional[Credentials]:
|
|
111
|
+
"""
|
|
112
|
+
Get the current credentials.
|
|
113
|
+
|
|
114
|
+
:return: The current credentials or None if not set
|
|
115
|
+
"""
|
|
116
|
+
return self._creds
|
|
117
|
+
|
|
118
|
+
def _set_credentials(self, creds: Credentials):
|
|
119
|
+
"""
|
|
120
|
+
Set the credentials.
|
|
121
|
+
|
|
122
|
+
:param creds: The credentials to set
|
|
123
|
+
"""
|
|
124
|
+
self._creds = creds
|
|
125
|
+
|
|
126
|
+
async def get_grpc_call_auth_metadata(self) -> typing.Optional[GrpcAuthMetadata]:
|
|
127
|
+
"""
|
|
128
|
+
Fetch the authentication metadata for gRPC calls.
|
|
129
|
+
|
|
130
|
+
:return: A tuple of (header_key, header_value) or None if no credentials are available
|
|
131
|
+
"""
|
|
132
|
+
creds = self.get_credentials()
|
|
133
|
+
if creds:
|
|
134
|
+
cfg = await self._resolve_config()
|
|
135
|
+
return GrpcAuthMetadata(
|
|
136
|
+
creds_id=creds.id,
|
|
137
|
+
pairs=Metadata((cfg.header_key, f"Bearer {creds.access_token}")),
|
|
138
|
+
)
|
|
139
|
+
return None
|
|
140
|
+
|
|
141
|
+
async def refresh_credentials(self, creds_id: str | None = None):
|
|
142
|
+
"""
|
|
143
|
+
Refresh the credentials asynchronously with thread and asyncio safety.
|
|
144
|
+
|
|
145
|
+
This method implements a thread-safe and coroutine-safe credential refresh mechanism.
|
|
146
|
+
It uses a timestamp-based approach to prevent redundant credential refreshes when
|
|
147
|
+
multiple threads or coroutines attempt to refresh credentials simultaneously.
|
|
148
|
+
|
|
149
|
+
The caller should capture the current _creds_timestamp before attempting to use credentials.
|
|
150
|
+
If credential usage fails, the caller can pass that timestamp to this method.
|
|
151
|
+
If the timestamp matches the current value, a refresh is needed; otherwise,
|
|
152
|
+
another thread has already refreshed the credentials.
|
|
153
|
+
|
|
154
|
+
:param creds_id: The id of credentials when they were last accessed by the caller.
|
|
155
|
+
If None, force a refresh regardless of id.
|
|
156
|
+
:raises: May raise authentication-related exceptions if the refresh fails
|
|
157
|
+
"""
|
|
158
|
+
# If creds_id is None, force refresh
|
|
159
|
+
# If creds_id matches current value, credentials need refresh
|
|
160
|
+
# If creds_id doesn't match, another thread already refreshed credentials
|
|
161
|
+
if creds_id and creds_id != self._creds_id:
|
|
162
|
+
# Credentials have been refreshed by another thread/coroutine since caller read them
|
|
163
|
+
return
|
|
164
|
+
|
|
165
|
+
# Use the async lock to ensure coroutine safety
|
|
166
|
+
async with self._async_lock:
|
|
167
|
+
# Double-check pattern to avoid unnecessary work
|
|
168
|
+
if creds_id and creds_id != self._creds_id:
|
|
169
|
+
# Another thread/coroutine refreshed credentials while we were waiting for the lock
|
|
170
|
+
return
|
|
171
|
+
|
|
172
|
+
# Perform the actual credential refresh
|
|
173
|
+
try:
|
|
174
|
+
self._creds = await self._do_refresh_credentials()
|
|
175
|
+
KeyringStore.store(self._creds)
|
|
176
|
+
except Exception:
|
|
177
|
+
KeyringStore.delete(self._endpoint)
|
|
178
|
+
raise
|
|
179
|
+
|
|
180
|
+
# Update the timestamp to indicate credentials have been refreshed
|
|
181
|
+
self._creds_id = self._creds.id
|
|
182
|
+
|
|
183
|
+
@abstractmethod
|
|
184
|
+
async def _do_refresh_credentials(self) -> Credentials:
|
|
185
|
+
"""
|
|
186
|
+
Perform the actual credential refresh operation.
|
|
187
|
+
|
|
188
|
+
This method must be implemented by subclasses to handle the specific authentication flow.
|
|
189
|
+
It should update the internal credentials object (_creds) with a new access token.
|
|
190
|
+
|
|
191
|
+
Implementations typically use the resolved configuration from _resolve_config() to
|
|
192
|
+
determine authentication endpoints, scopes, audience, and other parameters needed for
|
|
193
|
+
the specific authentication flow.
|
|
194
|
+
|
|
195
|
+
:raises: May raise authentication-related exceptions if the refresh fails
|
|
196
|
+
"""
|
|
197
|
+
...
|
|
198
|
+
|
|
199
|
+
|
|
200
|
+
class AsyncAuthenticatedClient(httpx.AsyncClient):
|
|
201
|
+
"""
|
|
202
|
+
An httpx.AsyncClient that automatically adds authentication headers to requests.
|
|
203
|
+
This class extends httpx.AsyncClient which is inherently async for network operations.
|
|
204
|
+
"""
|
|
205
|
+
|
|
206
|
+
def __init__(self, authenticator: Authenticator, **kwargs):
|
|
207
|
+
"""
|
|
208
|
+
Initialize the authenticated client.
|
|
209
|
+
|
|
210
|
+
:param authenticator: The authenticator to use for authentication
|
|
211
|
+
:param kwargs: Additional arguments passed to the httpx.AsyncClient constructor
|
|
212
|
+
"""
|
|
213
|
+
super().__init__(**kwargs)
|
|
214
|
+
self.auth_adapter = AsyncAuthenticationHTTPAdapter(authenticator)
|
|
215
|
+
self.authenticator = authenticator
|
|
216
|
+
|
|
217
|
+
async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
|
|
218
|
+
"""
|
|
219
|
+
Sends the request with added authentication headers.
|
|
220
|
+
Must be async because it performs network IO operations and may need to refresh credentials.
|
|
221
|
+
If the response returns a 401 status code, refreshes the credentials and retries the request.
|
|
222
|
+
|
|
223
|
+
:param request: The request object to send.
|
|
224
|
+
:param kwargs: Additional keyword arguments passed to the parent httpx.AsyncClient.send method, which may
|
|
225
|
+
include:
|
|
226
|
+
- auth: Authentication implementation to use for this request
|
|
227
|
+
- follow_redirects: Whether to follow redirects for this request
|
|
228
|
+
- timeout: Request timeout configuration for this request
|
|
229
|
+
:return: The response object.
|
|
230
|
+
"""
|
|
231
|
+
|
|
232
|
+
creds_id = await self.auth_adapter.add_auth_header(request)
|
|
233
|
+
response = await super().send(request, **kwargs)
|
|
234
|
+
|
|
235
|
+
if response.status_code == HTTPStatus.UNAUTHORIZED:
|
|
236
|
+
await self.authenticator.refresh_credentials(creds_id=creds_id)
|
|
237
|
+
await self.auth_adapter.add_auth_header(request)
|
|
238
|
+
response = await super().send(request, **kwargs)
|
|
239
|
+
|
|
240
|
+
return response
|
|
241
|
+
|
|
242
|
+
|
|
243
|
+
class AsyncAuthenticationHTTPAdapter:
|
|
244
|
+
"""
|
|
245
|
+
A custom async HTTP adapter that adds authentication headers to requests of an httpx.AsyncClient.
|
|
246
|
+
This is the async equivalent of AuthenticationHTTPAdapter for requests.
|
|
247
|
+
"""
|
|
248
|
+
|
|
249
|
+
def __init__(self, authenticator: Authenticator):
|
|
250
|
+
"""
|
|
251
|
+
Initialize the authentication HTTP adapter.
|
|
252
|
+
|
|
253
|
+
:param authenticator: The authenticator to use for authentication
|
|
254
|
+
"""
|
|
255
|
+
self.authenticator = authenticator
|
|
256
|
+
|
|
257
|
+
async def add_auth_header(self, request: httpx.Request) -> typing.Optional[str]:
|
|
258
|
+
"""
|
|
259
|
+
Adds authentication headers to the request.
|
|
260
|
+
Must be async because it may call refresh_credentials which performs IO operations.
|
|
261
|
+
|
|
262
|
+
:param request: The request object to add headers to.
|
|
263
|
+
:return: The credentials ID (creds_id) used for tracking credential refresh state
|
|
264
|
+
"""
|
|
265
|
+
if self.authenticator.get_credentials() is None:
|
|
266
|
+
await self.authenticator.refresh_credentials()
|
|
267
|
+
|
|
268
|
+
metadata = await self.authenticator.get_grpc_call_auth_metadata()
|
|
269
|
+
if metadata is None:
|
|
270
|
+
return None
|
|
271
|
+
for key, value in metadata.pairs.keys():
|
|
272
|
+
request.headers[key] = value
|
|
273
|
+
return metadata.creds_id
|
|
274
|
+
|
|
275
|
+
|
|
276
|
+
def upgrade_async_session_to_proxy_authenticated(
|
|
277
|
+
http_session: httpx.AsyncClient, proxy_authenticator: typing.Optional[Authenticator] = None, **kwargs
|
|
278
|
+
) -> httpx.AsyncClient:
|
|
279
|
+
"""
|
|
280
|
+
Given an httpx.AsyncClient, it returns a new session that uses AsyncAuthenticationHTTPAdapter
|
|
281
|
+
to perform authentication with a proxy in front of Flyte
|
|
282
|
+
|
|
283
|
+
:param http_session: httpx.AsyncClient Precreated session
|
|
284
|
+
:param proxy_authenticator: Optional authenticator for proxy authentication
|
|
285
|
+
:param kwargs: Additional arguments passed to AsyncAuthenticatedClient, which may include:
|
|
286
|
+
- auth: Authentication implementation to use
|
|
287
|
+
- params: Query parameters to include in request URLs
|
|
288
|
+
- headers: HTTP headers to include in requests
|
|
289
|
+
- cookies: Cookies to include in requests
|
|
290
|
+
- verify: SSL verification mode (True/False/path to certificate)
|
|
291
|
+
- cert: SSL client certificate (path or tuple)
|
|
292
|
+
- http1: Whether to enable HTTP/1.1 support
|
|
293
|
+
- http2: Whether to enable HTTP/2 support
|
|
294
|
+
- proxies: Proxy configuration mapping
|
|
295
|
+
- mounts: Mounted transports for specific URL patterns
|
|
296
|
+
- timeout: Request timeout configuration
|
|
297
|
+
- follow_redirects: Whether to follow redirects
|
|
298
|
+
- limits: Connection pool limits
|
|
299
|
+
- max_redirects: Maximum number of redirects to follow
|
|
300
|
+
- event_hooks: Event hooks for request/response lifecycle
|
|
301
|
+
- base_url: Base URL to join with relative URLs
|
|
302
|
+
- transport: Transport implementation to use
|
|
303
|
+
- app: ASGI application to handle requests
|
|
304
|
+
:return: httpx.AsyncClient with authentication
|
|
305
|
+
"""
|
|
306
|
+
if proxy_authenticator:
|
|
307
|
+
return AsyncAuthenticatedClient(proxy_authenticator, **kwargs)
|
|
308
|
+
else:
|
|
309
|
+
return http_session
|
|
310
|
+
|
|
311
|
+
|
|
312
|
+
def get_async_session(
|
|
313
|
+
proxy_authenticator: Authenticator | None = None,
|
|
314
|
+
ca_cert_path: str | None = None,
|
|
315
|
+
verify: bool | None = None,
|
|
316
|
+
**kwargs,
|
|
317
|
+
) -> httpx.AsyncClient:
|
|
318
|
+
"""
|
|
319
|
+
Returns a new httpx.AsyncClient with proxy authentication if proxy_authenticator is provided.
|
|
320
|
+
|
|
321
|
+
This function creates a new httpx.AsyncClient and optionally configures it with proxy authentication
|
|
322
|
+
if a proxy authenticator is provided.
|
|
323
|
+
|
|
324
|
+
:param proxy_authenticator: Optional authenticator for proxy authentication
|
|
325
|
+
:param ca_cert_path: Optional path to CA certificate file for SSL verification
|
|
326
|
+
:param verify: Optional SSL verification mode (True/False/path to certificate)
|
|
327
|
+
:param kwargs: Additional keyword arguments passed to httpx.AsyncClient constructor and AsyncAuthenticatedClient,
|
|
328
|
+
which may include:
|
|
329
|
+
- auth: Authentication implementation to use
|
|
330
|
+
- params: Query parameters to include in request URLs
|
|
331
|
+
- headers: HTTP headers to include in requests
|
|
332
|
+
- cookies: Cookies to include in requests
|
|
333
|
+
- cert: SSL client certificate (path or tuple)
|
|
334
|
+
- http1: Whether to enable HTTP/1.1 support
|
|
335
|
+
- http2: Whether to enable HTTP/2 support
|
|
336
|
+
- proxies: Proxy configuration mapping
|
|
337
|
+
- mounts: Mounted transports for specific URL patterns
|
|
338
|
+
- timeout: Request timeout configuration
|
|
339
|
+
- follow_redirects: Whether to follow redirects
|
|
340
|
+
- limits: Connection pool limits
|
|
341
|
+
- max_redirects: Maximum number of redirects to follow
|
|
342
|
+
- event_hooks: Event hooks for request/response lifecycle
|
|
343
|
+
- base_url: Base URL to join with relative URLs
|
|
344
|
+
- transport: Transport implementation to use
|
|
345
|
+
- app: ASGI application to handle requests
|
|
346
|
+
- proxy_env: Environment variables for proxy command
|
|
347
|
+
- proxy_timeout: Timeout for proxy command execution
|
|
348
|
+
- header_key: Header key to use for authentication
|
|
349
|
+
- endpoint: The endpoint URL for authentication
|
|
350
|
+
- client_id: Client ID for authentication
|
|
351
|
+
- client_secret: Client secret for authentication
|
|
352
|
+
- scopes: List of scopes to request during authentication
|
|
353
|
+
- audience: Audience for the token
|
|
354
|
+
- http_proxy_url: HTTP proxy URL
|
|
355
|
+
:return: An httpx.AsyncClient instance, optionally configured with proxy authentication
|
|
356
|
+
"""
|
|
357
|
+
|
|
358
|
+
# Extract known httpx.AsyncClient parameters from kwargs
|
|
359
|
+
client_kwargs = {
|
|
360
|
+
k: v
|
|
361
|
+
for k, v in kwargs.items()
|
|
362
|
+
if k
|
|
363
|
+
in [
|
|
364
|
+
"auth",
|
|
365
|
+
"params",
|
|
366
|
+
"headers",
|
|
367
|
+
"cookies",
|
|
368
|
+
"verify",
|
|
369
|
+
"cert",
|
|
370
|
+
"http1",
|
|
371
|
+
"http2",
|
|
372
|
+
"proxies",
|
|
373
|
+
"mounts",
|
|
374
|
+
"timeout",
|
|
375
|
+
"follow_redirects",
|
|
376
|
+
"limits",
|
|
377
|
+
"max_redirects",
|
|
378
|
+
"event_hooks",
|
|
379
|
+
"base_url",
|
|
380
|
+
"transport",
|
|
381
|
+
"app",
|
|
382
|
+
]
|
|
383
|
+
}
|
|
384
|
+
|
|
385
|
+
if ca_cert_path:
|
|
386
|
+
context = ssl.create_default_context(capath=ca_cert_path)
|
|
387
|
+
verify = True if context is not None else False
|
|
388
|
+
|
|
389
|
+
if verify is not None:
|
|
390
|
+
client_kwargs["verify"] = verify
|
|
391
|
+
|
|
392
|
+
http_session = httpx.AsyncClient(**client_kwargs)
|
|
393
|
+
if proxy_authenticator:
|
|
394
|
+
http_session = upgrade_async_session_to_proxy_authenticated(
|
|
395
|
+
http_session, proxy_authenticator=proxy_authenticator, **kwargs
|
|
396
|
+
)
|
|
397
|
+
return http_session
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
from flyte._logging import logger
|
|
2
|
+
from flyte.remote._client.auth import _token_client as token_client
|
|
3
|
+
from flyte.remote._client.auth._authenticators.base import Authenticator
|
|
4
|
+
from flyte.remote._client.auth._keyring import Credentials
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class ClientCredentialsAuthenticator(Authenticator):
|
|
8
|
+
"""
|
|
9
|
+
This Authenticator uses ClientId and ClientSecret to authenticate
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
def __init__(
|
|
13
|
+
self,
|
|
14
|
+
client_id: str,
|
|
15
|
+
client_credentials_secret: str,
|
|
16
|
+
**kwargs,
|
|
17
|
+
):
|
|
18
|
+
"""
|
|
19
|
+
Initialize the client credentials authenticator.
|
|
20
|
+
|
|
21
|
+
:param client_id: The client ID for authentication
|
|
22
|
+
:param client_credentials_secret: The client secret for authentication
|
|
23
|
+
:param kwargs: Additional keyword arguments passed to the base Authenticator
|
|
24
|
+
|
|
25
|
+
**Keyword Arguments passed to base Authenticator**:
|
|
26
|
+
:param endpoint: The endpoint URL for authentication (required)
|
|
27
|
+
:param cfg_store: Optional client configuration store for retrieving remote configuration
|
|
28
|
+
:param client_config: Optional client configuration containing authentication settings
|
|
29
|
+
:param credentials: Optional credentials to use for authentication
|
|
30
|
+
:param http_session: Optional HTTP session to use for requests
|
|
31
|
+
:param http_proxy_url: Optional HTTP proxy URL
|
|
32
|
+
:param verify: Whether to verify SSL certificates (default: True)
|
|
33
|
+
:param ca_cert_path: Optional path to CA certificate file
|
|
34
|
+
:param scopes: List of scopes to request during authentication
|
|
35
|
+
:param audience: Audience for the token
|
|
36
|
+
"""
|
|
37
|
+
if not client_id or not client_credentials_secret:
|
|
38
|
+
raise ValueError("both client_id and client_credentials_secret are required.")
|
|
39
|
+
self._client_id = client_id
|
|
40
|
+
self._client_credentials_secret = client_credentials_secret
|
|
41
|
+
super().__init__(**kwargs)
|
|
42
|
+
|
|
43
|
+
async def _do_refresh_credentials(self) -> Credentials:
|
|
44
|
+
"""
|
|
45
|
+
Refreshes the authentication credentials using client credentials flow.
|
|
46
|
+
|
|
47
|
+
This function is used by the _handle_rpc_error() decorator, depending on the AUTH_MODE config object.
|
|
48
|
+
This handler is meant for SDK use-cases of auth (like pyflyte, or when users call SDK functions that require
|
|
49
|
+
access to Admin, like when waiting for another workflow to complete from within a task). This function uses
|
|
50
|
+
basic auth, which means the credentials for basic auth must be present from wherever this code is running.
|
|
51
|
+
"""
|
|
52
|
+
cfg = await self._resolve_config()
|
|
53
|
+
|
|
54
|
+
# Note that unlike the Pkce flow, the client ID does not come from Admin.
|
|
55
|
+
logger.debug(f"Basic authorization flow with client id {self._client_id} scope {cfg.scopes}")
|
|
56
|
+
authorization_header = token_client.get_basic_authorization_header(
|
|
57
|
+
self._client_id, self._client_credentials_secret
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
token, refresh_token, expires_in = await token_client.get_token(
|
|
61
|
+
token_endpoint=cfg.token_endpoint,
|
|
62
|
+
authorization_header=authorization_header,
|
|
63
|
+
http_proxy_url=self._http_proxy_url,
|
|
64
|
+
verify=self._verify,
|
|
65
|
+
scopes=cfg.scopes,
|
|
66
|
+
audience=cfg.audience,
|
|
67
|
+
http_session=self._http_session,
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
logger.info("Retrieved new token, expires in {}".format(expires_in))
|
|
71
|
+
return Credentials(
|
|
72
|
+
for_endpoint=self._endpoint, access_token=token, refresh_token=refresh_token, expires_in=expires_in
|
|
73
|
+
)
|