Skip to content

async_app_manager

AsyncSparkAppManager

Bases: LoggingMixin

Manage Spark apps on Kubernetes asynchronously.

Parameters:

Name Type Description Default
k8s_client_manager KubernetesClientManager

Kubernetes client manager. Defaults to None.

 None
logger_name str

logger name. Defaults to "SparkAppManager".

 None
Source code in spark_on_k8s/utils/async_app_manager.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
class AsyncSparkAppManager(LoggingMixin):
    """Manage Spark apps on Kubernetes asynchronously.

    Args:
        k8s_client_manager (KubernetesClientManager, optional): Kubernetes client manager. Defaults to None.
        logger_name (str, optional): logger name. Defaults to "SparkAppManager".
    """

    def __init__(
        self,
        *,
        k8s_client_manager: KubernetesAsyncClientManager | None = None,
        logger_name: str | None = None,
    ):
        super().__init__(logger_name=logger_name or "SparkAppManager")
        self.k8s_client_manager = k8s_client_manager or KubernetesAsyncClientManager()

    async def app_status(
        self,
        *,
        namespace: str,
        pod_name: str | None = None,
        app_id: str | None = None,
        client: k8s_async.CoreV1Api | None = None,
    ) -> SparkAppStatus:
        """Get app status asynchronously.

        Args:
            namespace (str): Namespace.
            pod_name (str): Pod name. Defaults to None.
            app_id (str): App ID. Defaults to None.
            client (k8s.CoreV1Api, optional): Kubernetes client. Defaults to None.

        Returns:
            SparkAppStatus: App status.
        """

        async def _app_status(_client: k8s_async.CoreV1Api) -> SparkAppStatus:
            if pod_name is None and app_id is None:
                raise ValueError("Either pod_name or app_id must be specified")
            if pod_name is not None:
                _pod = await _client.read_namespaced_pod(
                    namespace=namespace,
                    name=pod_name,
                )
            else:
                _pod = (
                    await _client.list_namespaced_pod(
                        namespace=namespace,
                        label_selector=f"spark-app-id={app_id}",
                    )
                ).items[0]
            return get_app_status(_pod)

        if client is None:
            async with self.k8s_client_manager.client() as client:
                api = k8s_async.CoreV1Api(client)
                return await _app_status(api)
        return await _app_status(client)

    async def wait_for_app(
        self,
        *,
        namespace: str,
        pod_name: str | None = None,
        app_id: str | None = None,
        poll_interval: float = 10,
        should_print: bool = False,
    ):
        """Wait for a Spark app to finish asynchronously.

        Args:
            namespace (str): Namespace.
            pod_name (str): Pod name.
            app_id (str): App ID.
            poll_interval (float, optional): Poll interval in seconds. Defaults to 10.
            should_print (bool, optional): Whether to print logs instead of logging them.
        """
        termination_statuses = {SparkAppStatus.Succeeded, SparkAppStatus.Failed, SparkAppStatus.Unknown}
        async with self.k8s_client_manager.client() as client:
            api = k8s_async.CoreV1Api(client)
            while True:
                try:
                    status = await self.app_status(
                        namespace=namespace, pod_name=pod_name, app_id=app_id, client=api
                    )
                    if status in termination_statuses:
                        break
                except ApiException as e:
                    if e.status == 404:
                        self.log(
                            msg=f"Pod {pod_name} was deleted", level=logging.INFO, should_print=should_print
                        )
                        return
                self.log(
                    msg=f"Pod {pod_name} status is {status}, sleep {poll_interval}s",
                    level=logging.INFO,
                    should_print=should_print,
                )
                await asyncio.sleep(poll_interval)
            self.log(
                msg=f"Pod {pod_name} finished with status {status.value}",
                level=logging.INFO,
                should_print=should_print,
            )

    async def logs_streamer(
        self,
        *,
        namespace: str,
        pod_name: str | None = None,
        app_id: str | None = None,
        tail_lines: int = -1,
    ):
        """Stream logs from a Spark app asynchronously.

        Args:
            namespace (str): Namespace.
            pod_name (str): Pod name.
            app_id (str): App ID.
            tail_lines (int, optional): Number of lines to tail. Defaults to -1.
        """
        if pod_name is None and app_id is None:
            raise ValueError("Either pod_name or app_id must be specified")
        if pod_name is None:
            async with self.k8s_client_manager.client() as client:
                api = k8s_async.CoreV1Api(client)
                pods = (
                    await api.list_namespaced_pod(
                        namespace=namespace,
                        label_selector=f"spark-app-id={app_id}",
                    )
                ).items
                if len(pods) == 0:
                    raise ValueError(f"No pods found for app {app_id}")
                pod_name = pods[0].metadata.name

        async with self.k8s_client_manager.client() as client:
            api = k8s_async.CoreV1Api(client)
            while True:
                pod = await api.read_namespaced_pod(
                    namespace=namespace,
                    name=pod_name,
                )
                if pod.status.phase != "Pending":
                    break

            watcher = watch.Watch()
            log_streamer = watcher.stream(
                api.read_namespaced_pod_log,
                namespace=namespace,
                name=pod_name,
                tail_lines=tail_lines if tail_lines > 0 else None,
                follow=True,
            )
            async for line in log_streamer:
                yield line
            watcher.stop()

    async def kill_app(
        self,
        namespace: str,
        pod_name: str | None = None,
        app_id: str | None = None,
    ):
        """Kill an app asynchronously.

        Args:
            namespace (str): Namespace.
            pod_name (str): Pod name.
            app_id (str): App ID.
        """
        if pod_name is None and app_id is None:
            raise ValueError("Either pod_name or app_id must be specified")
        async with self.k8s_client_manager.client() as client:
            api = k8s_async.CoreV1Api(client)
            if pod_name is None:
                pods = (
                    await api.list_namespaced_pod(
                        namespace=namespace,
                        label_selector=f"spark-app-id={app_id}",
                    )
                ).items
                if len(pods) == 0:
                    raise ValueError(f"No pods found for app {app_id}")
                pod = pods[0]
            else:
                pod = await api.read_namespaced_pod(
                    namespace=namespace,
                    name=pod_name,
                )
            container_name = pod.spec.containers[0].name
            if pod.status.phase != "Running":
                raise ValueError(f"Pod {pod.metadata.name} is not running")
            v1_ws = k8s_async.CoreV1Api(api_client=WsApiClient())
            await stream(
                v1_ws.connect_get_namespaced_pod_exec,
                pod.metadata.name,
                namespace,
                command=["/bin/sh", "-c", "kill 1"],
                container=container_name,
                stderr=True,
                stdin=False,
                stdout=True,
                tty=False,
                _preload_content=True,
            )
            await v1_ws.api_client.close()

    async def delete_app(
        self, namespace: str, pod_name: str | None = None, app_id: str | None = None, force: bool = False
    ):
        """Delete an app asynchronously.

        Args:
            namespace (str): Namespace.
            pod_name (str): Pod name.
            app_id (str): App ID.
            force (bool, optional): Whether to force delete the app. Defaults to False.
        """
        if pod_name is None and app_id is None:
            raise ValueError("Either pod_name or app_id must be specified")
        async with self.k8s_client_manager.client() as client:
            api = k8s_async.CoreV1Api(client)
            if app_id:
                # we don't use `delete_collection_namespaced_pod` to know if the app exists or not
                pods = (
                    await api.list_namespaced_pod(
                        namespace=namespace,
                        label_selector=f"spark-app-id={app_id}",
                    )
                ).items
                if len(pods) == 0:
                    raise ValueError(f"No pods found for app {app_id}")
                pod_name = pods[0].metadata.name
            await api.delete_namespaced_pod(
                name=pod_name,
                namespace=namespace,
                body=k8s_async.V1DeleteOptions(
                    grace_period_seconds=0 if force else None,
                    propagation_policy="Foreground",
                ),
            )

app_status(*, namespace, pod_name=None, app_id=None, client=None) async

Get app status asynchronously.

Parameters:

Name Type Description Default
namespace str

Namespace.

 required
pod_name str

Pod name. Defaults to None.

 None
app_id str

App ID. Defaults to None.

 None
client CoreV1Api

Kubernetes client. Defaults to None.

 None

Returns:

Name Type Description
SparkAppStatus SparkAppStatus

App status.
Source code in spark_on_k8s/utils/async_app_manager.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
async def app_status(
    self,
    *,
    namespace: str,
    pod_name: str | None = None,
    app_id: str | None = None,
    client: k8s_async.CoreV1Api | None = None,
) -> SparkAppStatus:
    """Get app status asynchronously.

    Args:
        namespace (str): Namespace.
        pod_name (str): Pod name. Defaults to None.
        app_id (str): App ID. Defaults to None.
        client (k8s.CoreV1Api, optional): Kubernetes client. Defaults to None.

    Returns:
        SparkAppStatus: App status.
    """

    async def _app_status(_client: k8s_async.CoreV1Api) -> SparkAppStatus:
        if pod_name is None and app_id is None:
            raise ValueError("Either pod_name or app_id must be specified")
        if pod_name is not None:
            _pod = await _client.read_namespaced_pod(
                namespace=namespace,
                name=pod_name,
            )
        else:
            _pod = (
                await _client.list_namespaced_pod(
                    namespace=namespace,
                    label_selector=f"spark-app-id={app_id}",
                )
            ).items[0]
        return get_app_status(_pod)

    if client is None:
        async with self.k8s_client_manager.client() as client:
            api = k8s_async.CoreV1Api(client)
            return await _app_status(api)
    return await _app_status(client)

delete_app(namespace, pod_name=None, app_id=None, force=False) async

Delete an app asynchronously.

Parameters:

Name Type Description Default
namespace str

Namespace.

 required
pod_name str

Pod name.

 None
app_id str

App ID.

 None
force bool

Whether to force delete the app. Defaults to False.

 False
Source code in spark_on_k8s/utils/async_app_manager.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
async def delete_app(
    self, namespace: str, pod_name: str | None = None, app_id: str | None = None, force: bool = False
):
    """Delete an app asynchronously.

    Args:
        namespace (str): Namespace.
        pod_name (str): Pod name.
        app_id (str): App ID.
        force (bool, optional): Whether to force delete the app. Defaults to False.
    """
    if pod_name is None and app_id is None:
        raise ValueError("Either pod_name or app_id must be specified")
    async with self.k8s_client_manager.client() as client:
        api = k8s_async.CoreV1Api(client)
        if app_id:
            # we don't use `delete_collection_namespaced_pod` to know if the app exists or not
            pods = (
                await api.list_namespaced_pod(
                    namespace=namespace,
                    label_selector=f"spark-app-id={app_id}",
                )
            ).items
            if len(pods) == 0:
                raise ValueError(f"No pods found for app {app_id}")
            pod_name = pods[0].metadata.name
        await api.delete_namespaced_pod(
            name=pod_name,
            namespace=namespace,
            body=k8s_async.V1DeleteOptions(
                grace_period_seconds=0 if force else None,
                propagation_policy="Foreground",
            ),
        )

kill_app(namespace, pod_name=None, app_id=None) async

Kill an app asynchronously.

Parameters:

Name Type Description Default
namespace str

Namespace.

 required
pod_name str

Pod name.

 None
app_id str

App ID.

 None
Source code in spark_on_k8s/utils/async_app_manager.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
async def kill_app(
    self,
    namespace: str,
    pod_name: str | None = None,
    app_id: str | None = None,
):
    """Kill an app asynchronously.

    Args:
        namespace (str): Namespace.
        pod_name (str): Pod name.
        app_id (str): App ID.
    """
    if pod_name is None and app_id is None:
        raise ValueError("Either pod_name or app_id must be specified")
    async with self.k8s_client_manager.client() as client:
        api = k8s_async.CoreV1Api(client)
        if pod_name is None:
            pods = (
                await api.list_namespaced_pod(
                    namespace=namespace,
                    label_selector=f"spark-app-id={app_id}",
                )
            ).items
            if len(pods) == 0:
                raise ValueError(f"No pods found for app {app_id}")
            pod = pods[0]
        else:
            pod = await api.read_namespaced_pod(
                namespace=namespace,
                name=pod_name,
            )
        container_name = pod.spec.containers[0].name
        if pod.status.phase != "Running":
            raise ValueError(f"Pod {pod.metadata.name} is not running")
        v1_ws = k8s_async.CoreV1Api(api_client=WsApiClient())
        await stream(
            v1_ws.connect_get_namespaced_pod_exec,
            pod.metadata.name,
            namespace,
            command=["/bin/sh", "-c", "kill 1"],
            container=container_name,
            stderr=True,
            stdin=False,
            stdout=True,
            tty=False,
            _preload_content=True,
        )
        await v1_ws.api_client.close()

logs_streamer(*, namespace, pod_name=None, app_id=None, tail_lines=-1) async

Stream logs from a Spark app asynchronously.

Parameters:

Name Type Description Default
namespace str

Namespace.

 required
pod_name str

Pod name.

 None
app_id str

App ID.

 None
tail_lines int

Number of lines to tail. Defaults to -1.

 -1
Source code in spark_on_k8s/utils/async_app_manager.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
async def logs_streamer(
    self,
    *,
    namespace: str,
    pod_name: str | None = None,
    app_id: str | None = None,
    tail_lines: int = -1,
):
    """Stream logs from a Spark app asynchronously.

    Args:
        namespace (str): Namespace.
        pod_name (str): Pod name.
        app_id (str): App ID.
        tail_lines (int, optional): Number of lines to tail. Defaults to -1.
    """
    if pod_name is None and app_id is None:
        raise ValueError("Either pod_name or app_id must be specified")
    if pod_name is None:
        async with self.k8s_client_manager.client() as client:
            api = k8s_async.CoreV1Api(client)
            pods = (
                await api.list_namespaced_pod(
                    namespace=namespace,
                    label_selector=f"spark-app-id={app_id}",
                )
            ).items
            if len(pods) == 0:
                raise ValueError(f"No pods found for app {app_id}")
            pod_name = pods[0].metadata.name

    async with self.k8s_client_manager.client() as client:
        api = k8s_async.CoreV1Api(client)
        while True:
            pod = await api.read_namespaced_pod(
                namespace=namespace,
                name=pod_name,
            )
            if pod.status.phase != "Pending":
                break

        watcher = watch.Watch()
        log_streamer = watcher.stream(
            api.read_namespaced_pod_log,
            namespace=namespace,
            name=pod_name,
            tail_lines=tail_lines if tail_lines > 0 else None,
            follow=True,
        )
        async for line in log_streamer:
            yield line
        watcher.stop()

wait_for_app(*, namespace, pod_name=None, app_id=None, poll_interval=10, should_print=False) async

Wait for a Spark app to finish asynchronously.

Parameters:

Name Type Description Default
namespace str

Namespace.

 required
pod_name str

Pod name.

 None
app_id str

App ID.

 None
poll_interval float

Poll interval in seconds. Defaults to 10.

 10
should_print bool

Whether to print logs instead of logging them.

 False
Source code in spark_on_k8s/utils/async_app_manager.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
async def wait_for_app(
    self,
    *,
    namespace: str,
    pod_name: str | None = None,
    app_id: str | None = None,
    poll_interval: float = 10,
    should_print: bool = False,
):
    """Wait for a Spark app to finish asynchronously.

    Args:
        namespace (str): Namespace.
        pod_name (str): Pod name.
        app_id (str): App ID.
        poll_interval (float, optional): Poll interval in seconds. Defaults to 10.
        should_print (bool, optional): Whether to print logs instead of logging them.
    """
    termination_statuses = {SparkAppStatus.Succeeded, SparkAppStatus.Failed, SparkAppStatus.Unknown}
    async with self.k8s_client_manager.client() as client:
        api = k8s_async.CoreV1Api(client)
        while True:
            try:
                status = await self.app_status(
                    namespace=namespace, pod_name=pod_name, app_id=app_id, client=api
                )
                if status in termination_statuses:
                    break
            except ApiException as e:
                if e.status == 404:
                    self.log(
                        msg=f"Pod {pod_name} was deleted", level=logging.INFO, should_print=should_print
                    )
                    return
            self.log(
                msg=f"Pod {pod_name} status is {status}, sleep {poll_interval}s",
                level=logging.INFO,
                should_print=should_print,
            )
            await asyncio.sleep(poll_interval)
        self.log(
            msg=f"Pod {pod_name} finished with status {status.value}",
            level=logging.INFO,
            should_print=should_print,
        )