akshaychitneni commented on code in PR #34840:
URL: https://github.com/apache/airflow/pull/34840#discussion_r1385350302
##########
airflow/providers/papermill/operators/papermill.py:
##########
@@ -17,17 +17,102 @@
# under the License.
from __future__ import annotations
+from functools import cached_property
from typing import TYPE_CHECKING, ClassVar, Collection, Sequence
import attr
+from papermill.utils import remove_args, merge_kwargs
+from pydantic import typing
+
import papermill as pm
+from papermill.engines import NBClientEngine
+from papermill.clientwrap import PapermillNotebookClient
+from jupyter_client.manager import AsyncKernelManager
+from jupyter_client.client import KernelClient
+from traitlets import Unicode
from airflow.lineage.entities import File
from airflow.models import BaseOperator
+from airflow.providers.papermill.hooks.kernel import KernelHook
if TYPE_CHECKING:
from airflow.utils.context import Context
+REMOTE_KERNEL_ENGINE = "remote_kernel_engine"
+
+
+class RemoteKernelManager(AsyncKernelManager):
+ """
+ Jupyter kernel manager that connects to a remote kernel.
+ """
+ session_key = Unicode('', config=True, help="Session key to connect to
remote kernel")
+
+ @property
+ def has_kernel(self) -> bool:
+ return True
+
+ async def _async_is_alive(self) -> bool:
+ return True
+
+ def shutdown_kernel(self, now: bool = False, restart: bool = False) ->
None:
+ pass
+
+ def client(self, **kwargs: typing.Any) -> KernelClient:
+ """Create a client configured to connect to our kernel"""
+ kernel_client = super().client(**kwargs)
+ # load connection info to set session_key
+ config = dict(
+ ip=self.ip,
+ shell_port=self.shell_port,
+ iopub_port=self.iopub_port,
+ stdin_port=self.stdin_port,
+ control_port=self.control_port,
+ hb_port=self.hb_port,
+ key=self.session_key,
+ transport="tcp",
+ signature_scheme="hmac-sha256",
+ )
+ kernel_client.load_connection_info(config)
+ return kernel_client
+
+
+class RemoteKernelEngine(NBClientEngine):
Review Comment:
Updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
