flytekitplugins-awsemrserverless 1.16.26__tar.gz
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.
- flytekitplugins_awsemrserverless-1.16.26/PKG-INFO +154 -0
- flytekitplugins_awsemrserverless-1.16.26/README.md +120 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins/awsemrserverless/__init__.py +29 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins/awsemrserverless/_entrypoint.py +203 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins/awsemrserverless/boto_handler.py +389 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins/awsemrserverless/connector.py +825 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins/awsemrserverless/task.py +402 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/PKG-INFO +154 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/SOURCES.txt +18 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/dependency_links.txt +1 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/entry_points.txt +2 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/namespace_packages.txt +1 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/requires.txt +3 -0
- flytekitplugins_awsemrserverless-1.16.26/flytekitplugins_awsemrserverless.egg-info/top_level.txt +1 -0
- flytekitplugins_awsemrserverless-1.16.26/setup.cfg +4 -0
- flytekitplugins_awsemrserverless-1.16.26/setup.py +49 -0
- flytekitplugins_awsemrserverless-1.16.26/tests/test_boto_handler.py +336 -0
- flytekitplugins_awsemrserverless-1.16.26/tests/test_connector.py +1266 -0
- flytekitplugins_awsemrserverless-1.16.26/tests/test_entrypoint.py +338 -0
- flytekitplugins_awsemrserverless-1.16.26/tests/test_task.py +346 -0
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flytekitplugins-awsemrserverless
|
|
3
|
+
Version: 1.16.26
|
|
4
|
+
Summary: Flytekit AWS EMR Serverless Plugin: run Spark and Hive jobs from Flyte tasks
|
|
5
|
+
Author: flyteorg
|
|
6
|
+
Author-email: admin@flyte.org
|
|
7
|
+
License: apache2
|
|
8
|
+
Classifier: Intended Audience :: Science/Research
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Topic :: Scientific/Engineering
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
|
+
Classifier: Topic :: Software Development
|
|
18
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
19
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
20
|
+
Requires-Python: >=3.9
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
Requires-Dist: flytekit>=1.14.6
|
|
23
|
+
Requires-Dist: aioboto3>=12.3.0
|
|
24
|
+
Requires-Dist: boto3>=1.28.0
|
|
25
|
+
Dynamic: author
|
|
26
|
+
Dynamic: author-email
|
|
27
|
+
Dynamic: classifier
|
|
28
|
+
Dynamic: description
|
|
29
|
+
Dynamic: description-content-type
|
|
30
|
+
Dynamic: license
|
|
31
|
+
Dynamic: requires-dist
|
|
32
|
+
Dynamic: requires-python
|
|
33
|
+
Dynamic: summary
|
|
34
|
+
|
|
35
|
+
# Flytekit AWS EMR Serverless Plugin
|
|
36
|
+
|
|
37
|
+
A Flyte connector for [AWS EMR Serverless](https://docs.aws.amazon.com/emr/latest/EMR-Serverless-UserGuide/emr-serverless.html) that submits Spark and Hive jobs to an EMR Serverless application and tracks them through to completion.
|
|
38
|
+
|
|
39
|
+
## Features
|
|
40
|
+
|
|
41
|
+
- **Pythonic Spark mode**: write a Flyte `@task` whose body is regular PySpark; the plugin packages the user code, uploads it to S3, and runs it on EMR Serverless. No long-lived cluster to manage.
|
|
42
|
+
- **Script Spark mode**: point at an existing `main.py` (or JAR) already in S3 and submit it directly.
|
|
43
|
+
- **Hive mode**: submit a Hive query (inline or from S3) against an EMR Serverless application configured for Hive.
|
|
44
|
+
- Async connector lifecycle (`create` / `get` / `delete`) so the connector pod stays light and many jobs can be tracked concurrently.
|
|
45
|
+
- Honours Flyte task retries, timeouts, and cancellation, and surfaces EMR Serverless logs through the Flyte UI when log URIs are available.
|
|
46
|
+
|
|
47
|
+
## Installation
|
|
48
|
+
|
|
49
|
+
```bash
|
|
50
|
+
pip install flytekitplugins-awsemrserverless
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
The connector is registered automatically with `flytekit` via the plugin entry point. Deploy it on a [`flyteconnector`](https://github.com/flyteorg/flyte/tree/master/charts/flyteconnector) pod that has this package installed and an IAM identity allowed to call EMR Serverless `StartJobRun` / `GetJobRun` / `CancelJobRun` and to read/write the script-staging S3 prefix.
|
|
54
|
+
|
|
55
|
+
## Usage
|
|
56
|
+
|
|
57
|
+
### Pythonic Spark task
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
from flytekit import task, workflow
|
|
61
|
+
from flytekitplugins.awsemrserverless import EMRServerless, EMRServerlessSparkJobDriver
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
@task(
|
|
65
|
+
task_config=EMRServerless(
|
|
66
|
+
application_id="00fhabc12345",
|
|
67
|
+
execution_role_arn="arn:aws:iam::123456789012:role/EMRServerlessRole",
|
|
68
|
+
region="us-east-1",
|
|
69
|
+
job_driver=EMRServerlessSparkJobDriver(
|
|
70
|
+
spark_submit_parameters="--conf spark.executor.cores=2 --conf spark.executor.memory=4g",
|
|
71
|
+
),
|
|
72
|
+
),
|
|
73
|
+
)
|
|
74
|
+
def spark_count() -> int:
|
|
75
|
+
from pyspark.sql import SparkSession
|
|
76
|
+
|
|
77
|
+
spark = SparkSession.builder.getOrCreate()
|
|
78
|
+
return spark.range(1_000_000).count()
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
@workflow
|
|
82
|
+
def wf() -> int:
|
|
83
|
+
return spark_count()
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
The plugin serializes the task body, uploads it to S3, and EMR Serverless runs it inside the worker image you have associated with the application.
|
|
87
|
+
|
|
88
|
+
### Script Spark task
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
@task(
|
|
92
|
+
task_config=EMRServerless(
|
|
93
|
+
application_id="00fhabc12345",
|
|
94
|
+
execution_role_arn="arn:aws:iam::123456789012:role/EMRServerlessRole",
|
|
95
|
+
region="us-east-1",
|
|
96
|
+
job_driver=EMRServerlessSparkJobDriver(
|
|
97
|
+
entry_point="s3://my-bucket/scripts/main.py",
|
|
98
|
+
entry_point_arguments=["--date", "2025-01-01"],
|
|
99
|
+
spark_submit_parameters="--conf spark.executor.memory=4g",
|
|
100
|
+
),
|
|
101
|
+
),
|
|
102
|
+
)
|
|
103
|
+
def submit_script():
|
|
104
|
+
...
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
### Hive task
|
|
108
|
+
|
|
109
|
+
```python
|
|
110
|
+
from flytekitplugins.awsemrserverless import EMRServerless, EMRServerlessHiveJobDriver
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@task(
|
|
114
|
+
task_config=EMRServerless(
|
|
115
|
+
application_id="00fhabc12345",
|
|
116
|
+
execution_role_arn="arn:aws:iam::123456789012:role/EMRServerlessRole",
|
|
117
|
+
region="us-east-1",
|
|
118
|
+
job_driver=EMRServerlessHiveJobDriver(
|
|
119
|
+
query="SELECT COUNT(*) FROM my_table",
|
|
120
|
+
),
|
|
121
|
+
),
|
|
122
|
+
)
|
|
123
|
+
def hive_query():
|
|
124
|
+
...
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
## Worker image
|
|
128
|
+
|
|
129
|
+
For Pythonic Spark tasks the worker image must contain `flytekit` and this plugin so the executor can rehydrate the task object on the EMR Serverless side. Script Spark and Hive jobs do not require flytekit on the worker.
|
|
130
|
+
|
|
131
|
+
A reference `Dockerfile` is shipped alongside this plugin; it builds on the public EMR Serverless Spark base image and installs the matching `flytekit` and `flytekitplugins-awsemrserverless` versions:
|
|
132
|
+
|
|
133
|
+
```bash
|
|
134
|
+
docker build \
|
|
135
|
+
--build-arg VERSION=<flytekit-version> \
|
|
136
|
+
-t <registry>/emr-serverless-flytekit:<tag> \
|
|
137
|
+
plugins/flytekit-aws-emr-serverless
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
Override the base with `--build-arg EMR_BASE_IMAGE=...` to track a different EMR release.
|
|
141
|
+
|
|
142
|
+
## IAM
|
|
143
|
+
|
|
144
|
+
The connector pod's IAM principal needs:
|
|
145
|
+
|
|
146
|
+
- `emr-serverless:StartJobRun`, `emr-serverless:GetJobRun`, `emr-serverless:CancelJobRun` on the target application
|
|
147
|
+
- `s3:GetObject` / `s3:PutObject` on the script-staging prefix (Pythonic mode)
|
|
148
|
+
- `iam:PassRole` for the EMR Serverless execution role
|
|
149
|
+
|
|
150
|
+
The execution role attached to the EMR Serverless application is the role the workers run as and needs whatever data-access permissions your jobs require.
|
|
151
|
+
|
|
152
|
+
## Discussion
|
|
153
|
+
|
|
154
|
+
Tracking issue: [flyteorg/flyte#7286](https://github.com/flyteorg/flyte/issues/7286).
|
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
# Flytekit AWS EMR Serverless Plugin
|
|
2
|
+
|
|
3
|
+
A Flyte connector for [AWS EMR Serverless](https://docs.aws.amazon.com/emr/latest/EMR-Serverless-UserGuide/emr-serverless.html) that submits Spark and Hive jobs to an EMR Serverless application and tracks them through to completion.
|
|
4
|
+
|
|
5
|
+
## Features
|
|
6
|
+
|
|
7
|
+
- **Pythonic Spark mode**: write a Flyte `@task` whose body is regular PySpark; the plugin packages the user code, uploads it to S3, and runs it on EMR Serverless. No long-lived cluster to manage.
|
|
8
|
+
- **Script Spark mode**: point at an existing `main.py` (or JAR) already in S3 and submit it directly.
|
|
9
|
+
- **Hive mode**: submit a Hive query (inline or from S3) against an EMR Serverless application configured for Hive.
|
|
10
|
+
- Async connector lifecycle (`create` / `get` / `delete`) so the connector pod stays light and many jobs can be tracked concurrently.
|
|
11
|
+
- Honours Flyte task retries, timeouts, and cancellation, and surfaces EMR Serverless logs through the Flyte UI when log URIs are available.
|
|
12
|
+
|
|
13
|
+
## Installation
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install flytekitplugins-awsemrserverless
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
The connector is registered automatically with `flytekit` via the plugin entry point. Deploy it on a [`flyteconnector`](https://github.com/flyteorg/flyte/tree/master/charts/flyteconnector) pod that has this package installed and an IAM identity allowed to call EMR Serverless `StartJobRun` / `GetJobRun` / `CancelJobRun` and to read/write the script-staging S3 prefix.
|
|
20
|
+
|
|
21
|
+
## Usage
|
|
22
|
+
|
|
23
|
+
### Pythonic Spark task
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
from flytekit import task, workflow
|
|
27
|
+
from flytekitplugins.awsemrserverless import EMRServerless, EMRServerlessSparkJobDriver
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@task(
|
|
31
|
+
task_config=EMRServerless(
|
|
32
|
+
application_id="00fhabc12345",
|
|
33
|
+
execution_role_arn="arn:aws:iam::123456789012:role/EMRServerlessRole",
|
|
34
|
+
region="us-east-1",
|
|
35
|
+
job_driver=EMRServerlessSparkJobDriver(
|
|
36
|
+
spark_submit_parameters="--conf spark.executor.cores=2 --conf spark.executor.memory=4g",
|
|
37
|
+
),
|
|
38
|
+
),
|
|
39
|
+
)
|
|
40
|
+
def spark_count() -> int:
|
|
41
|
+
from pyspark.sql import SparkSession
|
|
42
|
+
|
|
43
|
+
spark = SparkSession.builder.getOrCreate()
|
|
44
|
+
return spark.range(1_000_000).count()
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
@workflow
|
|
48
|
+
def wf() -> int:
|
|
49
|
+
return spark_count()
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
The plugin serializes the task body, uploads it to S3, and EMR Serverless runs it inside the worker image you have associated with the application.
|
|
53
|
+
|
|
54
|
+
### Script Spark task
|
|
55
|
+
|
|
56
|
+
```python
|
|
57
|
+
@task(
|
|
58
|
+
task_config=EMRServerless(
|
|
59
|
+
application_id="00fhabc12345",
|
|
60
|
+
execution_role_arn="arn:aws:iam::123456789012:role/EMRServerlessRole",
|
|
61
|
+
region="us-east-1",
|
|
62
|
+
job_driver=EMRServerlessSparkJobDriver(
|
|
63
|
+
entry_point="s3://my-bucket/scripts/main.py",
|
|
64
|
+
entry_point_arguments=["--date", "2025-01-01"],
|
|
65
|
+
spark_submit_parameters="--conf spark.executor.memory=4g",
|
|
66
|
+
),
|
|
67
|
+
),
|
|
68
|
+
)
|
|
69
|
+
def submit_script():
|
|
70
|
+
...
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
### Hive task
|
|
74
|
+
|
|
75
|
+
```python
|
|
76
|
+
from flytekitplugins.awsemrserverless import EMRServerless, EMRServerlessHiveJobDriver
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
@task(
|
|
80
|
+
task_config=EMRServerless(
|
|
81
|
+
application_id="00fhabc12345",
|
|
82
|
+
execution_role_arn="arn:aws:iam::123456789012:role/EMRServerlessRole",
|
|
83
|
+
region="us-east-1",
|
|
84
|
+
job_driver=EMRServerlessHiveJobDriver(
|
|
85
|
+
query="SELECT COUNT(*) FROM my_table",
|
|
86
|
+
),
|
|
87
|
+
),
|
|
88
|
+
)
|
|
89
|
+
def hive_query():
|
|
90
|
+
...
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
## Worker image
|
|
94
|
+
|
|
95
|
+
For Pythonic Spark tasks the worker image must contain `flytekit` and this plugin so the executor can rehydrate the task object on the EMR Serverless side. Script Spark and Hive jobs do not require flytekit on the worker.
|
|
96
|
+
|
|
97
|
+
A reference `Dockerfile` is shipped alongside this plugin; it builds on the public EMR Serverless Spark base image and installs the matching `flytekit` and `flytekitplugins-awsemrserverless` versions:
|
|
98
|
+
|
|
99
|
+
```bash
|
|
100
|
+
docker build \
|
|
101
|
+
--build-arg VERSION=<flytekit-version> \
|
|
102
|
+
-t <registry>/emr-serverless-flytekit:<tag> \
|
|
103
|
+
plugins/flytekit-aws-emr-serverless
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
Override the base with `--build-arg EMR_BASE_IMAGE=...` to track a different EMR release.
|
|
107
|
+
|
|
108
|
+
## IAM
|
|
109
|
+
|
|
110
|
+
The connector pod's IAM principal needs:
|
|
111
|
+
|
|
112
|
+
- `emr-serverless:StartJobRun`, `emr-serverless:GetJobRun`, `emr-serverless:CancelJobRun` on the target application
|
|
113
|
+
- `s3:GetObject` / `s3:PutObject` on the script-staging prefix (Pythonic mode)
|
|
114
|
+
- `iam:PassRole` for the EMR Serverless execution role
|
|
115
|
+
|
|
116
|
+
The execution role attached to the EMR Serverless application is the role the workers run as and needs whatever data-access permissions your jobs require.
|
|
117
|
+
|
|
118
|
+
## Discussion
|
|
119
|
+
|
|
120
|
+
Tracking issue: [flyteorg/flyte#7286](https://github.com/flyteorg/flyte/issues/7286).
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
"""
|
|
2
|
+
.. currentmodule:: flytekitplugins.awsemrserverless
|
|
3
|
+
|
|
4
|
+
This plugin enables running Spark and Hive jobs on AWS EMR Serverless from
|
|
5
|
+
Flyte workflows. It exposes an async connector that handles the EMR
|
|
6
|
+
Serverless job lifecycle (submit, poll, cancel) and a task config type.
|
|
7
|
+
|
|
8
|
+
.. autosummary::
|
|
9
|
+
:template: custom.rst
|
|
10
|
+
:toctree: generated/
|
|
11
|
+
|
|
12
|
+
EMRServerless
|
|
13
|
+
EMRServerlessSparkJobDriver
|
|
14
|
+
EMRServerlessHiveJobDriver
|
|
15
|
+
EMRServerlessTask
|
|
16
|
+
EMRServerlessConnector
|
|
17
|
+
EMRServerlessJobMetadata
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
from flytekitplugins.awsemrserverless.connector import (
|
|
21
|
+
EMRServerlessConnector,
|
|
22
|
+
EMRServerlessJobMetadata,
|
|
23
|
+
)
|
|
24
|
+
from flytekitplugins.awsemrserverless.task import (
|
|
25
|
+
EMRServerless,
|
|
26
|
+
EMRServerlessHiveJobDriver,
|
|
27
|
+
EMRServerlessSparkJobDriver,
|
|
28
|
+
EMRServerlessTask,
|
|
29
|
+
)
|
|
@@ -0,0 +1,203 @@
|
|
|
1
|
+
"""
|
|
2
|
+
EMR Serverless Pythonic-mode entrypoint script.
|
|
3
|
+
|
|
4
|
+
This file is the canonical source of the bootstrap script that EMR
|
|
5
|
+
Serverless workers execute as ``sparkSubmit.entryPoint`` for Pythonic
|
|
6
|
+
tasks (i.e. tasks that do not provide an explicit ``spark_job_driver``).
|
|
7
|
+
|
|
8
|
+
How it is delivered to EMR
|
|
9
|
+
--------------------------
|
|
10
|
+
|
|
11
|
+
The connector pod:
|
|
12
|
+
|
|
13
|
+
1. reads this file from its own ``site-packages`` install;
|
|
14
|
+
2. computes ``hashlib.sha256(content)[:12]`` and uploads (idempotently)
|
|
15
|
+
to ``s3://<bucket>/flyte/emr-serverless/entrypoint-<hash>.py``;
|
|
16
|
+
3. passes that S3 URI as ``StartJobRun.jobDriver.sparkSubmit.entryPoint``.
|
|
17
|
+
|
|
18
|
+
EMR Serverless then downloads this script onto the Spark driver and
|
|
19
|
+
runs it with ``spark-submit``. ``sys.argv[1:]`` carries the actual
|
|
20
|
+
``pyflyte-fast-execute`` (or ``pyflyte-execute``) invocation that
|
|
21
|
+
should run inside the worker, plus the fast-registration distribution
|
|
22
|
+
arguments.
|
|
23
|
+
|
|
24
|
+
Why a custom entrypoint at all
|
|
25
|
+
------------------------------
|
|
26
|
+
|
|
27
|
+
EMR Serverless's API requires ``sparkSubmit.entryPoint`` to be a
|
|
28
|
+
single Python file URI -- there is no "container as entrypoint"
|
|
29
|
+
escape hatch (cf. SageMaker / Batch / ECS, which run the container
|
|
30
|
+
itself). We need a thin shim that:
|
|
31
|
+
|
|
32
|
+
* downloads the fast-registration tarball from Flyte's blob store,
|
|
33
|
+
* invokes ``pyflyte-fast-execute`` with the right resolver arguments,
|
|
34
|
+
* converts Flytekit's "exit-0-on-user-error" semantics into a non-zero
|
|
35
|
+
exit so EMR reports ``FAILED`` instead of ``SUCCESS``.
|
|
36
|
+
|
|
37
|
+
This file deliberately has only ``flytekit`` as a runtime dependency
|
|
38
|
+
(specifically ``flytekit.tools.fast_registration.download_distribution``)
|
|
39
|
+
because it runs *inside the EMR worker*, not the connector pod.
|
|
40
|
+
|
|
41
|
+
Editing this file
|
|
42
|
+
-----------------
|
|
43
|
+
|
|
44
|
+
Treat this file as part of the connector's *runtime contract* with
|
|
45
|
+
EMR workers, not as plugin internals:
|
|
46
|
+
|
|
47
|
+
* changes here propagate to every Pythonic-mode job on the next
|
|
48
|
+
connector deploy via the content hash in the S3 key;
|
|
49
|
+
* the corresponding unit tests live in ``tests/test_entrypoint.py``
|
|
50
|
+
and exercise this module both as imported Python and as the
|
|
51
|
+
spawned subprocess EMR sees;
|
|
52
|
+
* upstream alignment: this is the EMR analogue of
|
|
53
|
+
``flytetools/flytekitplugins/databricks/entrypoint.py`` --
|
|
54
|
+
same shape, different transport (S3 instead of GitHub).
|
|
55
|
+
"""
|
|
56
|
+
|
|
57
|
+
import os
|
|
58
|
+
import signal
|
|
59
|
+
import subprocess
|
|
60
|
+
import sys
|
|
61
|
+
|
|
62
|
+
from flytekit.tools.fast_registration import download_distribution
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _run_subprocess(cmd, env=None):
|
|
66
|
+
"""Run ``cmd`` and forward SIGTERM, returning ``(returncode, stderr_text)``.
|
|
67
|
+
|
|
68
|
+
stdout streams through to the parent (Spark driver stdout); stderr is
|
|
69
|
+
captured so the caller can inspect it for Flytekit's user-error banner.
|
|
70
|
+
"""
|
|
71
|
+
p = subprocess.Popen(cmd, env=env, stderr=subprocess.PIPE, stdout=None)
|
|
72
|
+
signal.signal(signal.SIGTERM, lambda s, f: p.send_signal(s))
|
|
73
|
+
_, stderr_bytes = p.communicate()
|
|
74
|
+
stderr_text = stderr_bytes.decode("utf-8", errors="replace") if stderr_bytes else ""
|
|
75
|
+
if stderr_text:
|
|
76
|
+
sys.stderr.write(stderr_text)
|
|
77
|
+
return p.returncode, stderr_text
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def _exit_with_code(rc, stderr_text=""):
|
|
81
|
+
"""Translate Flytekit subprocess exit semantics into EMR-correct exits.
|
|
82
|
+
|
|
83
|
+
Flytekit's ``pyflyte-execute`` catches user exceptions, writes the
|
|
84
|
+
error to the Flyte output blob, and exits ``0`` -- by design for
|
|
85
|
+
K8s-based agents where FlytePropeller reads the output.
|
|
86
|
+
|
|
87
|
+
In EMR Serverless the connector only polls EMR job state
|
|
88
|
+
(``SUCCESS`` / ``FAILED``) and cannot read the output blobs. If
|
|
89
|
+
``pyflyte-execute`` exits ``0`` but the user function failed, EMR
|
|
90
|
+
reports ``SUCCESS`` and the connector wrongly reports ``SUCCEEDED``.
|
|
91
|
+
|
|
92
|
+
Detect this by scanning stderr for Flytekit's error banner. When
|
|
93
|
+
found, force a non-zero exit so Spark fails the driver and EMR
|
|
94
|
+
reports ``FAILED``.
|
|
95
|
+
"""
|
|
96
|
+
if rc == 0 and "User Error Captured by Flyte" in stderr_text:
|
|
97
|
+
print(
|
|
98
|
+
"[flyte-entrypoint] pyflyte-execute exited 0 but stderr contains "
|
|
99
|
+
"a user error -- forcing non-zero exit so EMR reports FAILED",
|
|
100
|
+
file=sys.stderr,
|
|
101
|
+
)
|
|
102
|
+
sys.exit(1)
|
|
103
|
+
if rc != 0:
|
|
104
|
+
print(f"[flyte-entrypoint] Task process exited with code {rc}", file=sys.stderr)
|
|
105
|
+
sys.exit(rc)
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def _parse_fast_execute_args(args):
|
|
109
|
+
"""Split a ``pyflyte-fast-execute ...`` argv into its three pieces.
|
|
110
|
+
|
|
111
|
+
Returns ``(additional_distribution, dest_dir, task_cmd_start)``
|
|
112
|
+
where ``task_cmd_start`` is the index in ``args`` where the
|
|
113
|
+
underlying ``pyflyte-execute ...`` command begins.
|
|
114
|
+
|
|
115
|
+
Recognises the two-arg flag forms emitted by Flytekit:
|
|
116
|
+
``--additional-distribution <s3://...>``
|
|
117
|
+
``--dest-dir <path>``
|
|
118
|
+
and the optional ``--`` separator before the inner command.
|
|
119
|
+
"""
|
|
120
|
+
additional_distribution = None
|
|
121
|
+
dest_dir = None
|
|
122
|
+
task_cmd_start = 0
|
|
123
|
+
|
|
124
|
+
i = 1
|
|
125
|
+
while i < len(args):
|
|
126
|
+
if args[i] == "--additional-distribution" and i + 1 < len(args):
|
|
127
|
+
additional_distribution = args[i + 1]
|
|
128
|
+
i += 2
|
|
129
|
+
elif args[i] == "--dest-dir" and i + 1 < len(args):
|
|
130
|
+
dest_dir = args[i + 1]
|
|
131
|
+
i += 2
|
|
132
|
+
elif args[i] == "--":
|
|
133
|
+
task_cmd_start = i + 1
|
|
134
|
+
break
|
|
135
|
+
else:
|
|
136
|
+
task_cmd_start = i
|
|
137
|
+
break
|
|
138
|
+
|
|
139
|
+
return additional_distribution, dest_dir, task_cmd_start
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
def _build_resolver_command(task_execute_cmd, additional_distribution, dest_dir):
|
|
143
|
+
"""Inject the fast-registration distribution args before ``--resolver``.
|
|
144
|
+
|
|
145
|
+
``pyflyte-execute`` resolves task callables via a resolver plugin
|
|
146
|
+
(default ``flytekit.core.python_auto_container.default_task_resolver``).
|
|
147
|
+
For fast-registered code, the resolver needs to know where the
|
|
148
|
+
extracted source tree lives, which we inject as
|
|
149
|
+
``--dynamic-addl-distro`` / ``--dynamic-dest-dir`` immediately
|
|
150
|
+
before ``--resolver``.
|
|
151
|
+
"""
|
|
152
|
+
cmd = []
|
|
153
|
+
for arg in task_execute_cmd:
|
|
154
|
+
if arg == "--resolver":
|
|
155
|
+
cmd.extend(
|
|
156
|
+
[
|
|
157
|
+
"--dynamic-addl-distro",
|
|
158
|
+
additional_distribution or "",
|
|
159
|
+
"--dynamic-dest-dir",
|
|
160
|
+
dest_dir or "",
|
|
161
|
+
]
|
|
162
|
+
)
|
|
163
|
+
cmd.append(arg)
|
|
164
|
+
return cmd
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def main():
|
|
168
|
+
args = sys.argv[1:]
|
|
169
|
+
if not args:
|
|
170
|
+
print("Usage: entrypoint.py pyflyte-fast-execute|pyflyte-execute ...", file=sys.stderr)
|
|
171
|
+
sys.exit(1)
|
|
172
|
+
|
|
173
|
+
if args[0] == "pyflyte-fast-execute":
|
|
174
|
+
additional_distribution, dest_dir, task_cmd_start = _parse_fast_execute_args(args)
|
|
175
|
+
task_execute_cmd = list(args[task_cmd_start:])
|
|
176
|
+
|
|
177
|
+
if additional_distribution:
|
|
178
|
+
if not dest_dir:
|
|
179
|
+
dest_dir = os.getcwd()
|
|
180
|
+
download_distribution(additional_distribution, dest_dir)
|
|
181
|
+
|
|
182
|
+
cmd = _build_resolver_command(task_execute_cmd, additional_distribution, dest_dir)
|
|
183
|
+
|
|
184
|
+
env = os.environ.copy()
|
|
185
|
+
if dest_dir:
|
|
186
|
+
resolved = os.path.realpath(os.path.expanduser(dest_dir))
|
|
187
|
+
env["PYTHONPATH"] = resolved + os.pathsep + env.get("PYTHONPATH", "")
|
|
188
|
+
rc, stderr_text = _run_subprocess(cmd, env)
|
|
189
|
+
_exit_with_code(rc, stderr_text)
|
|
190
|
+
|
|
191
|
+
elif args[0] == "pyflyte-execute":
|
|
192
|
+
env = os.environ.copy()
|
|
193
|
+
env.setdefault("PYTHONPATH", os.getcwd())
|
|
194
|
+
rc, stderr_text = _run_subprocess(args, env)
|
|
195
|
+
_exit_with_code(rc, stderr_text)
|
|
196
|
+
|
|
197
|
+
else:
|
|
198
|
+
print(f"Unrecognized command: {args}", file=sys.stderr)
|
|
199
|
+
sys.exit(1)
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
if __name__ == "__main__":
|
|
203
|
+
main()
|