From f37e36a51bb7f0d53dc1f59e9ec77e59eabae8ca Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Sat, 30 Apr 2022 20:34:02 -0700 Subject: [PATCH 1/4] Squashed commit of the following: commit a53457bad1c8bf42bde61e4f43b4cd08e3b3a996 Author: Valay Dave Date: Sat Apr 30 20:33:14 2022 -0700 Fixed bugs in the remote-debuggers support. commit 3589d6ae110bd7b50bb5ea73895dfeeafcafc70f Author: Valay Dave Date: Sat Apr 30 20:18:25 2022 -0700 Added a debug decorator - support for Ngrok based remote debugging. --- metaflow/metaflow_config.py | 3 + metaflow/plugins/__init__.py | 3 + metaflow/plugins/debug_decorator.py | 147 ++++++++++++++++++++++++++++ metaflow/plugins/remote_pdb.py | 137 ++++++++++++++++++++++++++ 4 files changed, 290 insertions(+) create mode 100644 metaflow/plugins/debug_decorator.py create mode 100644 metaflow/plugins/remote_pdb.py diff --git a/metaflow/metaflow_config.py b/metaflow/metaflow_config.py index 9d7601870cb..76d86809935 100644 --- a/metaflow/metaflow_config.py +++ b/metaflow/metaflow_config.py @@ -134,6 +134,9 @@ def from_conf(name, default=None): # Default container registry DEFAULT_CONTAINER_REGISTRY = from_conf("METAFLOW_DEFAULT_CONTAINER_REGISTRY") +# Ngrok Debug Decorator Configuration +NGROK_KEY = from_conf("METAFLOW_NGROK_KEY", None) + ### # AWS Batch configuration ### diff --git a/metaflow/plugins/__init__.py b/metaflow/plugins/__init__.py index a7862fa8aa2..d9333dcb6d8 100644 --- a/metaflow/plugins/__init__.py +++ b/metaflow/plugins/__init__.py @@ -113,6 +113,7 @@ def get_plugin_cli(): from .conda.conda_step_decorator import CondaStepDecorator from .cards.card_decorator import CardDecorator from .frameworks.pytorch import PytorchParallelDecorator +from .debug_decorator import NgrokDebugDecorator, DebugDecorator STEP_DECORATORS = [ @@ -130,6 +131,8 @@ def get_plugin_cli(): PytorchParallelDecorator, InternalTestUnboundedForeachDecorator, ArgoWorkflowsInternalDecorator, + NgrokDebugDecorator, + DebugDecorator, ] _merge_lists(STEP_DECORATORS, _ext_plugins["STEP_DECORATORS"], "name") diff --git a/metaflow/plugins/debug_decorator.py b/metaflow/plugins/debug_decorator.py new file mode 100644 index 00000000000..d5fd4e06de9 --- /dev/null +++ b/metaflow/plugins/debug_decorator.py @@ -0,0 +1,147 @@ +import functools +import time +from metaflow.decorators import StepDecorator +from metaflow.metaflow_config import NGROK_KEY +from .remote_pdb import RemotePdb +from pdb import set_trace +from .remote_pdb import cry as log_message +import os + + +class DebugDecorator(StepDecorator): + name = "debugger" + + defaults = dict(port=8292, host="localhost") + + def __init__(self, attributes=None, statically_defined=False): + super().__init__(attributes, statically_defined) + self._isvalid = True + + def task_pre_step( + self, + step_name, + task_datastore, + metadata, + run_id, + task_id, + flow, + graph, + retry_count, + max_user_code_retries, + ubf_context, + inputs, + ): + from metaflow import current + + breakpoint = BreakPoint( + self.attributes["host"], + self.attributes["port"], + is_remote=False, + is_active=self._isvalid, + ) + current._update_env({"debug": breakpoint}) + + +class BreakPoint(object): + def __init__(self, host, port, is_remote=False, is_active=False): + self._host, self._port, self._is_remote, self._is_active = ( + host, + port, + is_remote, + is_active, + ) + + self._ngrok_tunnel, self._debug_host, self._debug_port = None, None, None + self._activated = False + self._remote_pdb = None + + def _setup_ngrok_tunnel(self): + try: + from pyngrok import ngrok + except: + return None + if not NGROK_KEY: + return None + ngrok.set_auth_token(NGROK_KEY) + debugger_terminal = ngrok.connect(self._port, "tcp") + debug_host, debug_port = debugger_terminal.public_url.split("tcp://")[1].split( + ":" + ) + return debugger_terminal, debug_host, debug_port + + @property + def breakpoint(self): + if self._activated: + return self._remote_pdb.set_trace() + + if not self._is_active: + return None + + if self._is_remote: + if not self._ngrok_tunnel: + ( + self._ngrok_tunnel, + self._debug_host, + self._debug_port, + ) = self._setup_ngrok_tunnel() + log_message( + "Starting a Remote Debug Tunnel Using Ngrok on tcp://%s:%s. " + "Connect to this job from your local machine using : `telnet %s %s`" + % ( + self._debug_host, + self._debug_port, + self._debug_host, + self._debug_port, + ) + ) + else: + log_message( + "Connect to this job's debugger using : `telnet %s %s`" + % (self._host, self._port) + ) + self._activated = True + self._remote_pdb = RemotePdb(self._host, self._port, quiet=True) + return self._remote_pdb.set_trace() + + +class NgrokDebugDecorator(DebugDecorator): + name = "remote_debugger" + + def _validate_ngrok(self): + try: + from pyngrok import ngrok + except: + return False + ngrok_key = os.environ.get("METAFLOW_NGROK_KEY", None) + if not ngrok_key: + return False + return True + + def task_pre_step( + self, + step_name, + task_datastore, + metadata, + run_id, + task_id, + flow, + graph, + retry_count, + max_user_code_retries, + ubf_context, + inputs, + ): + self._isvalid = self._validate_ngrok() + return super().task_pre_step( + step_name, + task_datastore, + metadata, + run_id, + task_id, + flow, + graph, + retry_count, + max_user_code_retries, + ubf_context, + inputs, + ) diff --git a/metaflow/plugins/remote_pdb.py b/metaflow/plugins/remote_pdb.py new file mode 100644 index 00000000000..166cb502efd --- /dev/null +++ b/metaflow/plugins/remote_pdb.py @@ -0,0 +1,137 @@ +# Code Plucked From https://github.com/ionelmc/python-remote-pdb/blob/master/src/remote_pdb.py +# Thank you to the author. +# TODO: Add license later. +from __future__ import print_function + +import errno +import logging +import os +import re +import socket +import sys +from pdb import Pdb + +__version__ = "2.1.0" + +PY3 = sys.version_info[0] == 3 +log = logging.getLogger(__name__) + + +def cry(message, stderr=sys.__stderr__): + log.critical(message) + print(message, file=stderr) + stderr.flush() + + +class LF2CRLF_FileWrapper(object): + def __init__(self, connection): + self.connection = connection + self.stream = fh = connection.makefile("rw") + self.read = fh.read + self.readline = fh.readline + self.readlines = fh.readlines + self.close = fh.close + self.flush = fh.flush + self.fileno = fh.fileno + if hasattr(fh, "encoding"): + self._send = lambda data: connection.sendall(data.encode(fh.encoding)) + else: + self._send = connection.sendall + + @property + def encoding(self): + return self.stream.encoding + + def __iter__(self): + return self.stream.__iter__() + + def write(self, data, nl_rex=re.compile("\r?\n")): + data = nl_rex.sub("\r\n", data) + self._send(data) + + def writelines(self, lines, nl_rex=re.compile("\r?\n")): + for line in lines: + self.write(line, nl_rex) + + +class RemotePdb(Pdb): + """ + This will run pdb as a ephemeral telnet service. Once you connect no one + else can connect. On construction this object will block execution till a + client has connected. + Based on https://github.com/tamentis/rpdb I think ... + To use this:: + RemotePdb(host='0.0.0.0', port=4444).set_trace() + Then run: telnet 127.0.0.1 4444 + """ + + active_instance = None + + def __init__(self, host, port, patch_stdstreams=False, quiet=False): + self._quiet = quiet + listen_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listen_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True) + listen_socket.bind((host, port)) + if not self._quiet: + cry( + "RemotePdb session open at %s:%s, waiting for connection ..." + % listen_socket.getsockname() + ) + listen_socket.listen(1) + connection, address = listen_socket.accept() + if not self._quiet: + cry("RemotePdb accepted connection from %s." % repr(address)) + self.handle = LF2CRLF_FileWrapper(connection) + Pdb.__init__(self, completekey="tab", stdin=self.handle, stdout=self.handle) + self.backup = [] + if patch_stdstreams: + for name in ( + "stderr", + "stdout", + "__stderr__", + "__stdout__", + "stdin", + "__stdin__", + ): + self.backup.append((name, getattr(sys, name))) + setattr(sys, name, self.handle) + RemotePdb.active_instance = self + + def __restore(self): + if self.backup and not self._quiet: + cry("Restoring streams: %s ..." % self.backup) + for name, fh in self.backup: + setattr(sys, name, fh) + self.handle.close() + RemotePdb.active_instance = None + + def do_quit(self, arg): + self.__restore() + return Pdb.do_quit(self, arg) + + do_q = do_exit = do_quit + + def set_trace(self, frame=None): + if frame is None: + frame = sys._getframe().f_back + try: + Pdb.set_trace(self, frame) + except IOError as exc: + if exc.errno != errno.ECONNRESET: + raise + + +def set_trace(host=None, port=None, patch_stdstreams=False, quiet=None): + """ + Opens a remote PDB on first available port. + """ + if host is None: + host = os.environ.get("REMOTE_PDB_HOST", "127.0.0.1") + if port is None: + port = int(os.environ.get("REMOTE_PDB_PORT", "0")) + if quiet is None: + quiet = bool(os.environ.get("REMOTE_PDB_QUIET", "")) + rdb = RemotePdb( + host=host, port=port, patch_stdstreams=patch_stdstreams, quiet=quiet + ) + rdb.set_trace(frame=sys._getframe().f_back) From f566c51ae1c98abbf51319d2da089898726e10dd Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Sat, 30 Apr 2022 21:15:15 -0700 Subject: [PATCH 2/4] bug fixes. --- metaflow/plugins/debug_decorator.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/metaflow/plugins/debug_decorator.py b/metaflow/plugins/debug_decorator.py index d5fd4e06de9..ebada848a32 100644 --- a/metaflow/plugins/debug_decorator.py +++ b/metaflow/plugins/debug_decorator.py @@ -11,7 +11,7 @@ class DebugDecorator(StepDecorator): name = "debugger" - defaults = dict(port=8292, host="localhost") + defaults = dict(port=8292, host="localhost", auth_token=None) def __init__(self, attributes=None, statically_defined=False): super().__init__(attributes, statically_defined) @@ -38,17 +38,19 @@ def task_pre_step( self.attributes["port"], is_remote=False, is_active=self._isvalid, + auth_token=self.attributes["auth_token"], ) current._update_env({"debug": breakpoint}) class BreakPoint(object): - def __init__(self, host, port, is_remote=False, is_active=False): - self._host, self._port, self._is_remote, self._is_active = ( + def __init__(self, host, port, is_remote=False, is_active=False, auth_token=None): + self._host, self._port, self._is_remote, self._is_active, self._auth_token = ( host, port, is_remote, is_active, + auth_token, ) self._ngrok_tunnel, self._debug_host, self._debug_port = None, None, None @@ -58,11 +60,11 @@ def __init__(self, host, port, is_remote=False, is_active=False): def _setup_ngrok_tunnel(self): try: from pyngrok import ngrok - except: + except ImportError: return None - if not NGROK_KEY: + if not self._auth_token: return None - ngrok.set_auth_token(NGROK_KEY) + ngrok.set_auth_token(self._auth_token) debugger_terminal = ngrok.connect(self._port, "tcp") debug_host, debug_port = debugger_terminal.public_url.split("tcp://")[1].split( ":" @@ -112,11 +114,17 @@ def _validate_ngrok(self): from pyngrok import ngrok except: return False - ngrok_key = os.environ.get("METAFLOW_NGROK_KEY", None) - if not ngrok_key: + if not self.attributes["auth_token"]: return False return True + def runtime_task_created( + self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context + ): + if not self.attributes["auth_token"]: + if NGROK_KEY: + self.attributes["auth_token"] = NGROK_KEY + def task_pre_step( self, step_name, From 6f01b7db472d3aa66394843d780d36e82dbbf175 Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Sun, 1 May 2022 04:55:10 +0000 Subject: [PATCH 3/4] Vendoring py-ngrok. - Fixing bugs related to py-ngrok. - Decorator works with remote execution. --- metaflow/_vendor/pyngrok.LICENSE | 21 ++ metaflow/_vendor/pyngrok/__init__.py | 0 metaflow/_vendor/pyngrok/conf.py | 115 ++++++ metaflow/_vendor/pyngrok/exception.py | 91 +++++ metaflow/_vendor/pyngrok/installer.py | 269 ++++++++++++++ metaflow/_vendor/pyngrok/ngrok.py | 510 ++++++++++++++++++++++++++ metaflow/_vendor/pyngrok/process.py | 470 ++++++++++++++++++++++++ metaflow/_vendor/vendor_any.txt | 1 + metaflow/plugins/debug_decorator.py | 41 +-- 9 files changed, 1491 insertions(+), 27 deletions(-) create mode 100644 metaflow/_vendor/pyngrok.LICENSE create mode 100644 metaflow/_vendor/pyngrok/__init__.py create mode 100644 metaflow/_vendor/pyngrok/conf.py create mode 100644 metaflow/_vendor/pyngrok/exception.py create mode 100644 metaflow/_vendor/pyngrok/installer.py create mode 100644 metaflow/_vendor/pyngrok/ngrok.py create mode 100644 metaflow/_vendor/pyngrok/process.py diff --git a/metaflow/_vendor/pyngrok.LICENSE b/metaflow/_vendor/pyngrok.LICENSE new file mode 100644 index 00000000000..ea37a951ba7 --- /dev/null +++ b/metaflow/_vendor/pyngrok.LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Alex Laird + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/metaflow/_vendor/pyngrok/__init__.py b/metaflow/_vendor/pyngrok/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/metaflow/_vendor/pyngrok/conf.py b/metaflow/_vendor/pyngrok/conf.py new file mode 100644 index 00000000000..375c659c76a --- /dev/null +++ b/metaflow/_vendor/pyngrok/conf.py @@ -0,0 +1,115 @@ +import os + +from .installer import get_ngrok_bin + +__author__ = "Alex Laird" +__copyright__ = "Copyright 2021, Alex Laird" +__version__ = "5.1.0" + +BIN_DIR = os.path.normpath(os.path.join(os.path.abspath(os.path.dirname(__file__)), "bin")) +DEFAULT_NGROK_PATH = os.path.join(BIN_DIR, get_ngrok_bin()) +DEFAULT_CONFIG_PATH = None + +DEFAULT_NGROK_CONFIG_PATH = os.path.join(os.path.expanduser("~"), ".ngrok2", "ngrok.yml") + +_default_pyngrok_config = None + + +class PyngrokConfig: + """ + An object containing ``pyngrok``'s configuration for interacting with the ``ngrok`` binary. All values are + optional when it is instantiated, and default values will be used for parameters not passed. + + Use :func:`~pyngrok.conf.get_default` and :func:`~pyngrok.conf.set_default` to interact with the default + ``pyngrok_config``, or pass another instance of this object as the ``pyngrok_config`` keyword arg to most + methods in the :mod:`~pyngrok.ngrok` module to override the default. + + .. code-block:: python + + from pyngrok import conf, ngrok + + # Here we update the entire default config + pyngrok_config = conf.PyngrokConfig(ngrok_path="/usr/local/bin/ngrok") + conf.set_default(pyngrok_config) + + # Here we update just one variable in the default config + conf.get_default().ngrok_path = "/usr/local/bin/ngrok" + + # Here we leave the default config as-is and pass an override + pyngrok_config = conf.PyngrokConfig(ngrok_path="/usr/local/bin/ngrok") + ngrok.connect(pyngrok_config=pyngrok_config) + + :var ngrok_path: The path to the ``ngrok`` binary, defaults to the value in + `conf.DEFAULT_NGROK_PATH `_ + :vartype ngrok_path: str + :var config_path: The path to the ``ngrok`` config, defaults to ``None`` and ``ngrok`` manages it. + :vartype config_path: str + :var auth_token: An authtoken to pass to commands (overrides what is in the config). + :vartype auth_token: str + :var region: The region in which ``ngrok`` should start. + :vartype region: str + :var monitor_thread: Whether ``ngrok`` should continue to be monitored (for logs, etc.) after startup + is complete. + :vartype monitor_thread: bool + :var log_event_callback: A callback that will be invoked each time ``ngrok`` emits a log. ``monitor_thread`` + must be set to ``True`` or the function will stop being called after ``ngrok`` finishes starting. + :vartype log_event_callback: types.FunctionType + :var startup_timeout: The max number of seconds to wait for ``ngrok`` to start before timing out. + :vartype startup_timeout: int + :var max_logs: The max number of logs to store in :class:`~pyngrok.process.NgrokProcess`'s ``logs`` variable. + :vartype max_logs: int + :var request_timeout: The max timeout when making requests to ``ngrok``'s API. + :vartype request_timeout: float + :var start_new_session: Passed to :py:class:`subprocess.Popen` when launching ``ngrok``. (Python 3 and POSIX only) + :vartype start_new_session: bool + """ + + def __init__(self, + ngrok_path=None, + config_path=None, + auth_token=None, + region=None, + monitor_thread=True, + log_event_callback=None, + startup_timeout=15, + max_logs=100, + request_timeout=4, + start_new_session=False): + self.ngrok_path = DEFAULT_NGROK_PATH if ngrok_path is None else ngrok_path + self.config_path = DEFAULT_CONFIG_PATH if config_path is None else config_path + self.auth_token = auth_token + self.region = region + self.monitor_thread = monitor_thread + self.log_event_callback = log_event_callback + self.startup_timeout = startup_timeout + self.max_logs = max_logs + self.request_timeout = request_timeout + self.start_new_session = start_new_session + + +def get_default(): + """ + Get the default config to be used with methods in the :mod:`~pyngrok.ngrok` module. To override the + default individually, the ``pyngrok_config`` keyword arg can also be passed to most of these methods, + or set a new default config with :func:`~pyngrok.conf.set_default`. + + :return: The default ``pyngrok_config``. + :rtype: PyngrokConfig + """ + if _default_pyngrok_config is None: + set_default(PyngrokConfig()) + + return _default_pyngrok_config + + +def set_default(pyngrok_config): + """ + Set a new default config to be used with methods in the :mod:`~pyngrok.ngrok` module. To override the + default individually, the ``pyngrok_config`` keyword arg can also be passed to most of these methods. + + :param pyngrok_config: The new ``pyngrok_config`` to be used by default. + :type pyngrok_config: PyngrokConfig + """ + global _default_pyngrok_config + + _default_pyngrok_config = pyngrok_config diff --git a/metaflow/_vendor/pyngrok/exception.py b/metaflow/_vendor/pyngrok/exception.py new file mode 100644 index 00000000000..b55c844ddb1 --- /dev/null +++ b/metaflow/_vendor/pyngrok/exception.py @@ -0,0 +1,91 @@ +__author__ = "Alex Laird" +__copyright__ = "Copyright 2020, Alex Laird" +__version__ = "4.1.0" + + +class PyngrokError(Exception): + """ + Raised when a general ``pyngrok`` error has occurred. + """ + pass + + +class PyngrokSecurityError(PyngrokError): + """ + Raised when a ``pyngrok`` security error has occurred. + """ + pass + + +class PyngrokNgrokInstallError(PyngrokError): + """ + Raised when an error has occurred while downloading and installing the ``ngrok`` binary. + """ + pass + + +class PyngrokNgrokError(PyngrokError): + """ + Raised when an error occurs interacting directly with the ``ngrok`` binary. + + :var error: A description of the error being thrown. + :vartype error: str + :var ngrok_logs: The ``ngrok`` logs, which may be useful for debugging the error. + :vartype ngrok_logs: list[NgrokLog] + :var ngrok_error: The error that caused the ``ngrok`` process to fail. + :vartype ngrok_error: str + """ + + def __init__(self, error, ngrok_logs=None, ngrok_error=None): + super(PyngrokNgrokError, self).__init__(error) + + if ngrok_logs is None: + ngrok_logs = [] + + self.ngrok_logs = ngrok_logs + self.ngrok_error = ngrok_error + + +class PyngrokNgrokHTTPError(PyngrokNgrokError): + """ + Raised when an error occurs making a request to the ``ngrok`` web interface. The ``body`` + contains the error response received from ``ngrok``. + + :var error: A description of the error being thrown. + :vartype error: str + :var url: The request URL that failed. + :vartype url: str + :var status_code: The response status code from ``ngrok``. + :vartype status_code: int + :var message: The response message from ``ngrok``. + :vartype message: str + :var headers: The request headers sent to ``ngrok``. + :vartype headers: dict[str, str] + :var body: The response body from ``ngrok``. + :vartype body: str + """ + + def __init__(self, error, url, status_code, message, headers, body): + super(PyngrokNgrokHTTPError, self).__init__(error) + + self.url = url + self.status_code = status_code + self.message = message + self.headers = headers + self.body = body + + +class PyngrokNgrokURLError(PyngrokNgrokError): + """ + Raised when an error occurs when trying to initiate an API request. + + :var error: A description of the error being thrown. + :vartype error: str + :var reason: The reason for the URL error. + :vartype reason: str + """ + + def __init__(self, error, reason): + super(PyngrokNgrokURLError, self).__init__(error) + + self.reason = reason diff --git a/metaflow/_vendor/pyngrok/installer.py b/metaflow/_vendor/pyngrok/installer.py new file mode 100644 index 00000000000..77cfbe46047 --- /dev/null +++ b/metaflow/_vendor/pyngrok/installer.py @@ -0,0 +1,269 @@ +import logging +import os +import platform +import socket +import sys +import tempfile +import time +import zipfile +from http import HTTPStatus +from urllib.request import urlopen + +import yaml + +from .exception import PyngrokNgrokInstallError, PyngrokSecurityError, PyngrokError + +__author__ = "Alex Laird" +__copyright__ = "Copyright 2021, Alex Laird" +__version__ = "5.1.0" + +logger = logging.getLogger(__name__) + +CDN_URL_PREFIX = "https://bin.equinox.io/c/4VmDzA7iaHb/" +PLATFORMS = { + "darwin_x86_64": CDN_URL_PREFIX + "ngrok-stable-darwin-amd64.zip", + "darwin_x86_64_arm": CDN_URL_PREFIX + "ngrok-stable-darwin-arm64.zip", + "windows_x86_64": CDN_URL_PREFIX + "ngrok-stable-windows-amd64.zip", + "windows_i386": CDN_URL_PREFIX + "ngrok-stable-windows-386.zip", + "linux_x86_64_arm": CDN_URL_PREFIX + "ngrok-stable-linux-arm64.zip", + "linux_i386_arm": CDN_URL_PREFIX + "ngrok-stable-linux-arm.zip", + "linux_i386": CDN_URL_PREFIX + "ngrok-stable-linux-386.zip", + "linux_x86_64": CDN_URL_PREFIX + "ngrok-stable-linux-amd64.zip", + "freebsd_x86_64": CDN_URL_PREFIX + "ngrok-stable-freebsd-amd64.zip", + "freebsd_i386": CDN_URL_PREFIX + "ngrok-stable-freebsd-386.zip", + "cygwin_x86_64": CDN_URL_PREFIX + "ngrok-stable-windows-amd64.zip", +} +DEFAULT_DOWNLOAD_TIMEOUT = 6 +DEFAULT_RETRY_COUNT = 0 + +_config_cache = None +_print_progress_enabled = True + + +def get_ngrok_bin(): + """ + Get the ``ngrok`` executable for the current system. + + :return: The name of the ``ngrok`` executable. + :rtype: str + """ + system = platform.system().lower() + if system in ["darwin", "linux", "freebsd"]: + return "ngrok" + elif system in ["windows", "cygwin"]: # pragma: no cover + return "ngrok.exe" + else: # pragma: no cover + raise PyngrokNgrokInstallError("\"{}\" is not a supported platform".format(system)) + + +def install_ngrok(ngrok_path, **kwargs): + """ + Download and install the latest ``ngrok`` for the current system, overwriting any existing contents + at the given path. + + :param ngrok_path: The path to where the ``ngrok`` binary will be downloaded. + :type ngrok_path: str + :param kwargs: Remaining ``kwargs`` will be passed to :func:`_download_file`. + :type kwargs: dict, optional + """ + logger.debug( + "Installing ngrok to {}{} ...".format(ngrok_path, ", overwriting" if os.path.exists(ngrok_path) else "")) + + ngrok_dir = os.path.dirname(ngrok_path) + + if not os.path.exists(ngrok_dir): + os.makedirs(ngrok_dir) + + arch = "x86_64" if sys.maxsize > 2 ** 32 else "i386" + if platform.uname()[4].startswith("arm") or \ + platform.uname()[4].startswith("aarch64"): + arch += "_arm" + system = platform.system().lower() + if "cygwin" in system: + system = "cygwin" + + plat = system + "_" + arch + try: + url = PLATFORMS[plat] + + logger.debug("Platform to download: {}".format(plat)) + except KeyError: + raise PyngrokNgrokInstallError("\"{}\" is not a supported platform".format(plat)) + + try: + download_path = _download_file(url, **kwargs) + + _install_ngrok_zip(ngrok_path, download_path) + except Exception as e: + raise PyngrokNgrokInstallError("An error occurred while downloading ngrok from {}: {}".format(url, e)) + + +def _install_ngrok_zip(ngrok_path, zip_path): + """ + Extract the ``ngrok`` zip file to the given path. + + :param ngrok_path: The path where ``ngrok`` will be installed. + :type ngrok_path: str + :param zip_path: The path to the ``ngrok`` zip file to be extracted. + :type zip_path: str + """ + _print_progress("Installing ngrok ... ") + + with zipfile.ZipFile(zip_path, "r") as zip_ref: + logger.debug("Extracting ngrok binary from {} to {} ...".format(zip_path, ngrok_path)) + zip_ref.extractall(os.path.dirname(ngrok_path)) + + os.chmod(ngrok_path, int("777", 8)) + + _clear_progress() + + +def get_ngrok_config(config_path, use_cache=True): + """ + Get the ``ngrok`` config from the given path. + + :param config_path: The ``ngrok`` config path to read. + :type config_path: str + :param use_cache: Use the cached version of the config (if populated). + :type use_cache: bool + :return: The ``ngrok`` config. + :rtype: dict + """ + global _config_cache + + if not _config_cache or not use_cache: + with open(config_path, "r") as config_file: + config = yaml.safe_load(config_file) + if config is None: + config = {} + + _config_cache = config + + return _config_cache + + +def install_default_config(config_path, data=None): + """ + Install the given data to the ``ngrok`` config. If a config is not already present for the given path, create one. + Before saving new data to the default config, validate that they are compatible with ``pyngrok``. + + :param config_path: The path to where the ``ngrok`` config should be installed. + :type config_path: str + :param data: A dictionary of things to add to the default config. + :type data: dict, optional + """ + if data is None: + data = {} + + config_dir = os.path.dirname(config_path) + if not os.path.exists(config_dir): + os.makedirs(config_dir) + if not os.path.exists(config_path): + open(config_path, "w").close() + + config = get_ngrok_config(config_path, use_cache=False) + + config.update(data) + + validate_config(config) + + with open(config_path, "w") as config_file: + logger.debug("Installing default ngrok config to {} ...".format(config_path)) + + yaml.dump(config, config_file) + + +def validate_config(data): + """ + Validate that the given dict of config items are valid for ``ngrok`` and ``pyngrok``. + + :param data: A dictionary of things to be validated as config items. + :type data: dict + """ + if data.get("web_addr", None) is False: + raise PyngrokError("\"web_addr\" cannot be False, as the ngrok API is a dependency for pyngrok") + elif data.get("log_format") == "json": + raise PyngrokError("\"log_format\" must be \"term\" to be compatible with pyngrok") + elif data.get("log_level", "info") not in ["info", "debug"]: + raise PyngrokError("\"log_level\" must be \"info\" to be compatible with pyngrok") + + +def _download_file(url, retries=0, **kwargs): + """ + Download a file to a temporary path and emit a status to stdout (if possible) as the download progresses. + + :param url: The URL to download. + :type url: str + :param retries: The retry attempt index, if download fails. + :type retries: int, optional + :param kwargs: Remaining ``kwargs`` will be passed to :py:func:`urllib.request.urlopen`. + :type kwargs: dict, optional + :return: The path to the downloaded temporary file. + :rtype: str + """ + kwargs["timeout"] = kwargs.get("timeout", DEFAULT_DOWNLOAD_TIMEOUT) + + if not url.lower().startswith("http"): + raise PyngrokSecurityError("URL must start with \"http\": {}".format(url)) + + try: + _print_progress("Downloading ngrok ...") + + logger.debug("Download ngrok from {} ...".format(url)) + + local_filename = url.split("/")[-1] + response = urlopen(url, **kwargs) + + status_code = response.getcode() + + if status_code != HTTPStatus.OK: + logger.debug("Response status code: {}".format(status_code)) + + return None + + length = response.getheader("Content-Length") + if length: + length = int(length) + chunk_size = max(4096, length // 100) + else: + chunk_size = 64 * 1024 + + download_path = os.path.join(tempfile.gettempdir(), local_filename) + with open(download_path, "wb") as f: + size = 0 + while True: + buffer = response.read(chunk_size) + + if not buffer: + break + + f.write(buffer) + size += len(buffer) + + if length: + percent_done = int((float(size) / float(length)) * 100) + _print_progress("Downloading ngrok: {}%".format(percent_done)) + + _clear_progress() + + return download_path + except socket.timeout as e: + if retries < DEFAULT_RETRY_COUNT: + logger.warning("ngrok download failed, retrying in 0.5 seconds ...") + time.sleep(0.5) + + return _download_file(url, retries + 1, **kwargs) + else: + raise e + + +def _print_progress(line): + if _print_progress_enabled: + sys.stdout.write("{}\r".format(line)) + sys.stdout.flush() + + +def _clear_progress(spaces=100): + if _print_progress_enabled: + sys.stdout.write((" " * spaces) + "\r") + sys.stdout.flush() diff --git a/metaflow/_vendor/pyngrok/ngrok.py b/metaflow/_vendor/pyngrok/ngrok.py new file mode 100644 index 00000000000..c3487b71657 --- /dev/null +++ b/metaflow/_vendor/pyngrok/ngrok.py @@ -0,0 +1,510 @@ +import json +import logging +import os +import socket +import sys +import uuid +from http import HTTPStatus +from urllib.error import HTTPError, URLError +from urllib.parse import urlencode +from urllib.request import urlopen, Request + +from . import process, conf, installer +from .exception import PyngrokNgrokHTTPError, PyngrokNgrokURLError, PyngrokSecurityError, PyngrokError + +__author__ = "Alex Laird" +__copyright__ = "Copyright 2021, Alex Laird" +__version__ = "5.1.0" + +logger = logging.getLogger(__name__) + +_current_tunnels = {} + + +class NgrokTunnel: + """ + An object containing information about a ``ngrok`` tunnel. + + :var data: The original tunnel data. + :vartype data: dict + :var name: The name of the tunnel. + :vartype name: str + :var proto: The protocol of the tunnel. + :vartype proto: str + :var uri: The tunnel URI, a relative path that can be used to make requests to the ``ngrok`` web interface. + :vartype uri: str + :var public_url: The public ``ngrok`` URL. + :vartype public_url: str + :var config: The config for the tunnel. + :vartype config: dict + :var metrics: Metrics for `the tunnel `_. + :vartype metrics: dict + :var pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok``. + :vartype pyngrok_config: PyngrokConfig + :var api_url: The API URL for the ``ngrok`` web interface. + :vartype api_url: str + """ + + def __init__(self, data, pyngrok_config, api_url): + self.data = data + + self.name = data.get("name") + self.proto = data.get("proto") + self.uri = data.get("uri") + self.public_url = data.get("public_url") + self.config = data.get("config", {}) + self.metrics = data.get("metrics", {}) + + self.pyngrok_config = pyngrok_config + self.api_url = api_url + + def __repr__(self): + return " \"{}\">".format(self.public_url, self.config["addr"]) if self.config.get( + "addr", None) else "" + + def __str__(self): # pragma: no cover + return "NgrokTunnel: \"{}\" -> \"{}\"".format(self.public_url, self.config["addr"]) if self.config.get( + "addr", None) else "" + + def refresh_metrics(self): + """ + Get the latest metrics for the tunnel and update the ``metrics`` variable. + """ + logger.info("Refreshing metrics for tunnel: {}".format(self.public_url)) + + data = api_request("{}{}".format(self.api_url, self.uri), method="GET", + timeout=self.pyngrok_config.request_timeout) + + if "metrics" not in data: + raise PyngrokError("The ngrok API did not return \"metrics\" in the response") + + self.data["metrics"] = data["metrics"] + self.metrics = self.data["metrics"] + + +def install_ngrok(pyngrok_config=None): + """ + Download, install, and initialize ``ngrok`` for the given config. If ``ngrok`` and its default + config is already installed, calling this method will do nothing. + + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + if not os.path.exists(pyngrok_config.ngrok_path): + installer.install_ngrok(pyngrok_config.ngrok_path) + + # If no config_path is set, ngrok will use its default path + if pyngrok_config.config_path is not None: + config_path = pyngrok_config.config_path + else: + config_path = conf.DEFAULT_NGROK_CONFIG_PATH + + # Install the config to the requested path + if not os.path.exists(config_path): + installer.install_default_config(config_path) + + # Install the default config, even if we don't need it this time, if it doesn't already exist + if conf.DEFAULT_NGROK_CONFIG_PATH != config_path and \ + not os.path.exists(conf.DEFAULT_NGROK_CONFIG_PATH): + installer.install_default_config(conf.DEFAULT_NGROK_CONFIG_PATH) + + +def set_auth_token(token, pyngrok_config=None): + """ + Set the ``ngrok`` auth token in the config file, enabling authenticated features (for instance, + more concurrent tunnels, custom subdomains, etc.). + + If ``ngrok`` is not installed at :class:`~pyngrok.conf.PyngrokConfig`'s ``ngrok_path``, calling this method + will first download and install ``ngrok``. + + :param token: The auth token to set. + :type token: str + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + install_ngrok(pyngrok_config) + + process.set_auth_token(pyngrok_config, token) + + +def get_ngrok_process(pyngrok_config=None): + """ + Get the current ``ngrok`` process for the given config's ``ngrok_path``. + + If ``ngrok`` is not installed at :class:`~pyngrok.conf.PyngrokConfig`'s ``ngrok_path``, calling this method + will first download and install ``ngrok``. + + If ``ngrok`` is not running, calling this method will first start a process with + :class:`~pyngrok.conf.PyngrokConfig`. + + Use :func:`~pyngrok.process.is_process_running` to check if a process is running without also implicitly + installing and starting it. + + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + :return: The ``ngrok`` process. + :rtype: NgrokProcess + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + install_ngrok(pyngrok_config) + + return process.get_process(pyngrok_config) + + +def connect(addr=None, proto=None, name=None, pyngrok_config=None, **options): + """ + Establish a new ``ngrok`` tunnel for the given protocol to the given port, returning an object representing + the connected tunnel. + + If a `tunnel definition in ngrok's config file `_ matches the given + ``name``, it will be loaded and used to start the tunnel. When ``name`` is ``None`` and a "pyngrok-default" tunnel + definition exists in ``ngrok``'s config, it will be loaded and use. Any ``kwargs`` passed as ``options`` will + override properties from the loaded tunnel definition. + + If ``ngrok`` is not installed at :class:`~pyngrok.conf.PyngrokConfig`'s ``ngrok_path``, calling this method + will first download and install ``ngrok``. + + If ``ngrok`` is not running, calling this method will first start a process with + :class:`~pyngrok.conf.PyngrokConfig`. + + .. note:: + + ``ngrok``'s default behavior for ``http`` when no additional properties are passed is to open *two* tunnels, + one ``http`` and one ``https``. This method will return a reference to the ``http`` tunnel in this case. If + only a single tunnel is needed, pass ``bind_tls=True`` and a reference to the ``https`` tunnel will be returned. + + :param addr: The local port to which the tunnel will forward traffic, or a + `local directory or network address `_, defaults to "80". + :type addr: str, optional + :param proto: A valid `tunnel protocol `_, defaults to "http". + :type proto: str, optional + :param name: A friendly name for the tunnel, or the name of a `ngrok tunnel definition `_ + to be used. + :type name: str, optional + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + :param options: Remaining ``kwargs`` are passed as `configuration for the ngrok + tunnel `_. + :type options: dict, optional + :return: The created ``ngrok`` tunnel. + :rtype: NgrokTunnel + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + if pyngrok_config.config_path is not None: + config_path = pyngrok_config.config_path + else: + config_path = conf.DEFAULT_NGROK_CONFIG_PATH + + if os.path.exists(config_path): + config = installer.get_ngrok_config(config_path) + else: + config = {} + + # If a "pyngrok-default" tunnel definition exists in the ngrok config, use that + tunnel_definitions = config.get("tunnels", {}) + if not name and "pyngrok-default" in tunnel_definitions: + name = "pyngrok-default" + + # Use a tunnel definition for the given name, if it exists + if name and name in tunnel_definitions: + tunnel_definition = tunnel_definitions[name] + + addr = tunnel_definition.get("addr") if not addr else addr + proto = tunnel_definition.get("proto") if not proto else proto + # Use the tunnel definition as the base, but override with any passed in options + tunnel_definition.update(options) + options = tunnel_definition + + addr = str(addr) if addr else "80" + if not proto: + proto = "http" + + if not name: + if not addr.startswith("file://"): + name = "{}-{}-{}".format(proto, addr, uuid.uuid4()) + else: + name = "{}-file-{}".format(proto, uuid.uuid4()) + + logger.info("Opening tunnel named: {}".format(name)) + + config = { + "name": name, + "addr": addr, + "proto": proto + } + options.update(config) + + api_url = get_ngrok_process(pyngrok_config).api_url + + logger.debug("Creating tunnel with options: {}".format(options)) + + tunnel = NgrokTunnel(api_request("{}/api/tunnels".format(api_url), method="POST", data=options, + timeout=pyngrok_config.request_timeout), + pyngrok_config, api_url) + + if proto == "http" and options.get("bind_tls", "both") == "both": + tunnel = NgrokTunnel(api_request("{}{}%20%28http%29".format(api_url, tunnel.uri), method="GET", + timeout=pyngrok_config.request_timeout), + pyngrok_config, api_url) + + _current_tunnels[tunnel.public_url] = tunnel + + return tunnel + + +def disconnect(public_url, pyngrok_config=None): + """ + Disconnect the ``ngrok`` tunnel for the given URL, if open. + + :param public_url: The public URL of the tunnel to disconnect. + :type public_url: str + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + # If ngrok is not running, there are no tunnels to disconnect + if not process.is_process_running(pyngrok_config.ngrok_path): + return + + api_url = get_ngrok_process(pyngrok_config).api_url + + if public_url not in _current_tunnels: + get_tunnels(pyngrok_config) + + # One more check, if the given URL is still not in the list of tunnels, it is not active + if public_url not in _current_tunnels: + return + + tunnel = _current_tunnels[public_url] + + logger.info("Disconnecting tunnel: {}".format(tunnel.public_url)) + + api_request("{}{}".format(api_url, tunnel.uri), method="DELETE", + timeout=pyngrok_config.request_timeout) + + _current_tunnels.pop(public_url, None) + + +def get_tunnels(pyngrok_config=None): + """ + Get a list of active ``ngrok`` tunnels for the given config's ``ngrok_path``. + + If ``ngrok`` is not installed at :class:`~pyngrok.conf.PyngrokConfig`'s ``ngrok_path``, calling this method + will first download and install ``ngrok``. + + If ``ngrok`` is not running, calling this method will first start a process with + :class:`~pyngrok.conf.PyngrokConfig`. + + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + :return: The active ``ngrok`` tunnels. + :rtype: list[NgrokTunnel] + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + api_url = get_ngrok_process(pyngrok_config).api_url + + _current_tunnels.clear() + for tunnel in api_request("{}/api/tunnels".format(api_url), method="GET", + timeout=pyngrok_config.request_timeout)["tunnels"]: + ngrok_tunnel = NgrokTunnel(tunnel, pyngrok_config, api_url) + _current_tunnels[ngrok_tunnel.public_url] = ngrok_tunnel + + return list(_current_tunnels.values()) + + +def kill(pyngrok_config=None): + """ + Terminate the ``ngrok`` processes, if running, for the given config's ``ngrok_path``. This method will not + block, it will just issue a kill request. + + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + process.kill_process(pyngrok_config.ngrok_path) + + _current_tunnels.clear() + + +def get_version(pyngrok_config=None): + """ + Get a tuple with the ``ngrok`` and ``pyngrok`` versions. + + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + :return: A tuple of ``(ngrok_version, pyngrok_version)``. + :rtype: tuple + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + ngrok_version = process.capture_run_process(pyngrok_config.ngrok_path, ["--version"]).split("version ")[1] + + return ngrok_version, __version__ + + +def update(pyngrok_config=None): + """ + Update ``ngrok`` for the given config's ``ngrok_path``, if an update is available. + + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + :return: The result from the ``ngrok`` update. + :rtype: str + """ + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + return process.capture_run_process(pyngrok_config.ngrok_path, ["update"]) + + +def api_request(url, method="GET", data=None, params=None, timeout=4): + """ + Invoke an API request to the given URL, returning JSON data from the response. + + One use for this method is making requests to ``ngrok`` tunnels: + + .. code-block:: python + + from pyngrok import ngrok + + public_url = ngrok.connect() + response = ngrok.api_request("{}/some-route".format(public_url), + method="POST", data={"foo": "bar"}) + + Another is making requests to the ``ngrok`` API itself: + + .. code-block:: python + + from pyngrok import ngrok + + api_url = ngrok.get_ngrok_process().api_url + response = ngrok.api_request("{}/api/requests/http".format(api_url), + params={"tunnel_name": "foo"}) + + :param url: The request URL. + :type url: str + :param method: The HTTP method. + :type method: str, optional + :param data: The request body. + :type data: dict, optional + :param params: The URL parameters. + :type params: dict, optional + :param timeout: The request timeout, in seconds. + :type timeout: float, optional + :return: The response from the request. + :rtype: dict + """ + if params is None: + params = [] + + if not url.lower().startswith("http"): + raise PyngrokSecurityError("URL must start with \"http\": {}".format(url)) + + data = json.dumps(data).encode("utf-8") if data else None + + if params: + url += "?{}".format(urlencode([(x, params[x]) for x in params])) + + request = Request(url, method=method.upper()) + request.add_header("Content-Type", "application/json") + + logger.debug("Making {} request to {} with data: {}".format(method, url, data)) + + try: + response = urlopen(request, data, timeout) + response_data = response.read().decode("utf-8") + + status_code = response.getcode() + logger.debug("Response {}: {}".format(status_code, response_data.strip())) + + if str(status_code)[0] != "2": + raise PyngrokNgrokHTTPError("ngrok client API returned {}: {}".format(status_code, response_data), url, + status_code, None, request.headers, response_data) + elif status_code == HTTPStatus.NO_CONTENT: + return None + + return json.loads(response_data) + except socket.timeout: + raise PyngrokNgrokURLError("ngrok client exception, URLError: timed out", "timed out") + except HTTPError as e: + response_data = e.read().decode("utf-8") + + status_code = e.getcode() + logger.debug("Response {}: {}".format(status_code, response_data.strip())) + + raise PyngrokNgrokHTTPError("ngrok client exception, API returned {}: {}".format(status_code, response_data), + e.url, + status_code, e.msg, e.hdrs, response_data) + except URLError as e: + raise PyngrokNgrokURLError("ngrok client exception, URLError: {}".format(e.reason), e.reason) + + +def run(args=None, pyngrok_config=None): + """ + Ensure ``ngrok`` is installed at the default path, then call :func:`~pyngrok.process.run_process`. + + This method is meant for interacting with ``ngrok`` from the command line and is not necessarily + compatible with non-blocking API methods. For that, use :mod:`~pyngrok.ngrok`'s interface methods (like + :func:`~pyngrok.ngrok.connect`), or use :func:`~pyngrok.process.get_process`. + + :param args: Arguments to be passed to the ``ngrok`` process. + :type args: list[str], optional + :param pyngrok_config: A ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary, + overriding :func:`~pyngrok.conf.get_default()`. + :type pyngrok_config: PyngrokConfig, optional + """ + if args is None: + args = [] + if pyngrok_config is None: + pyngrok_config = conf.get_default() + + install_ngrok(pyngrok_config) + + process.run_process(pyngrok_config.ngrok_path, args) + + +def main(): + """ + Entry point for the package's ``console_scripts``. This initializes a call from the command + line and invokes :func:`~pyngrok.ngrok.run`. + + This method is meant for interacting with ``ngrok`` from the command line and is not necessarily + compatible with non-blocking API methods. For that, use :mod:`~pyngrok.ngrok`'s interface methods (like + :func:`~pyngrok.ngrok.connect`), or use :func:`~pyngrok.process.get_process`. + """ + run(sys.argv[1:]) + + if len(sys.argv) == 1 or len(sys.argv) == 2 and sys.argv[1].lstrip("-").lstrip("-") == "help": + print("\nPYNGROK VERSION:\n {}".format(__version__)) + elif len(sys.argv) == 2 and sys.argv[1].lstrip("-").lstrip("-") in ["v", "version"]: + print("pyngrok version {}".format(__version__)) + + +if __name__ == "__main__": + main() diff --git a/metaflow/_vendor/pyngrok/process.py b/metaflow/_vendor/pyngrok/process.py new file mode 100644 index 00000000000..509424933bd --- /dev/null +++ b/metaflow/_vendor/pyngrok/process.py @@ -0,0 +1,470 @@ +import atexit +import logging +import os +import shlex +import subprocess +import threading +import time +from http import HTTPStatus +from urllib.request import Request, urlopen + +import yaml + +from . import conf, installer +from .exception import PyngrokNgrokError, PyngrokSecurityError + +__author__ = "Alex Laird" +__copyright__ = "Copyright 2021, Alex Laird" +__version__ = "5.1.0" + +logger = logging.getLogger(__name__) +ngrok_logger = logging.getLogger("{}.ngrok".format(__name__)) + +_current_processes = {} + + +class NgrokProcess: + """ + An object containing information about the ``ngrok`` process. + + :var proc: The child process that is running ``ngrok``. + :vartype proc: subprocess.Popen + :var pyngrok_config: The ``pyngrok`` configuration to use with ``ngrok``. + :vartype pyngrok_config: PyngrokConfig + :var api_url: The API URL for the ``ngrok`` web interface. + :vartype api_url: str + :var logs: A list of the most recent logs from ``ngrok``, limited in size to ``max_logs``. + :vartype logs: list[NgrokLog] + :var startup_error: If ``ngrok`` startup fails, this will be the log of the failure. + :vartype startup_error: str + """ + + def __init__(self, proc, pyngrok_config): + self.proc = proc + self.pyngrok_config = pyngrok_config + + self.api_url = None + self.logs = [] + self.startup_error = None + + self._tunnel_started = False + self._client_connected = False + self._monitor_thread = None + + def __repr__(self): + return "".format(self.api_url) + + def __str__(self): # pragma: no cover + return "NgrokProcess: \"{}\"".format(self.api_url) + + @staticmethod + def _line_has_error(log): + return log.lvl in ["ERROR", "CRITICAL"] + + def _log_startup_line(self, line): + """ + Parse the given startup log line and use it to manage the startup state + of the ``ngrok`` process. + + :param line: The line to be parsed and logged. + :type line: str + :return: The parsed log. + :rtype: NgrokLog + """ + log = self._log_line(line) + + if log is None: + return + elif self._line_has_error(log): + self.startup_error = log.err + elif log.msg: + # Log ngrok startup states as they come in + if "starting web service" in log.msg and log.addr is not None: + self.api_url = "http://{}".format(log.addr) + elif "tunnel session started" in log.msg: + self._tunnel_started = True + elif "client session established" in log.msg: + self._client_connected = True + + return log + + def _log_line(self, line): + """ + Parse, log, and emit (if ``log_event_callback`` in :class:`~pyngrok.conf.PyngrokConfig` is registered) the + given log line. + + :param line: The line to be processed. + :type line: str + :return: The parsed log. + :rtype: NgrokLog + """ + log = NgrokLog(line) + + if log.line == "": + return None + + ngrok_logger.log(getattr(logging, log.lvl), log.line) + self.logs.append(log) + if len(self.logs) > self.pyngrok_config.max_logs: + self.logs.pop(0) + + if self.pyngrok_config.log_event_callback is not None: + self.pyngrok_config.log_event_callback(log) + + return log + + def healthy(self): + """ + Check whether the ``ngrok`` process has finished starting up and is in a running, healthy state. + + :return: ``True`` if the ``ngrok`` process is started, running, and healthy. + :rtype: bool + """ + if self.api_url is None or \ + not self._tunnel_started or \ + not self._client_connected: + return False + + if not self.api_url.lower().startswith("http"): + raise PyngrokSecurityError("URL must start with \"http\": {}".format(self.api_url)) + + # Ensure the process is available for requests before registering it as healthy + request = Request("{}/api/tunnels".format(self.api_url)) + response = urlopen(request) + if response.getcode() != HTTPStatus.OK: + return False + + return self.proc.poll() is None + + def _monitor_process(self): + thread = threading.current_thread() + + thread.alive = True + while thread.alive and self.proc.poll() is None: + self._log_line(self.proc.stdout.readline()) + + self._monitor_thread = None + + def start_monitor_thread(self): + """ + Start a thread that will monitor the ``ngrok`` process and its logs until it completes. + + If a monitor thread is already running, nothing will be done. + """ + if self._monitor_thread is None: + logger.debug("Monitor thread will be started") + + self._monitor_thread = threading.Thread(target=self._monitor_process) + self._monitor_thread.daemon = True + self._monitor_thread.start() + + def stop_monitor_thread(self): + """ + Set the monitor thread to stop monitoring the ``ngrok`` process after the next log event. This will not + necessarily terminate the thread immediately, as the thread may currently be idle, rather it sets a flag + on the thread telling it to terminate the next time it wakes up. + + This has no impact on the ``ngrok`` process itself, only ``pyngrok``'s monitor of the process and + its logs. + """ + if self._monitor_thread is not None: + logger.debug("Monitor thread will be stopped") + + self._monitor_thread.alive = False + + +class NgrokLog: + """ + An object containing a parsed log from the ``ngrok`` process. + + :var line: The raw, unparsed log line. + :vartype line: str + :var t: The log's ISO 8601 timestamp. + :vartype t: str + :var lvl: The log's level. + :vartype lvl: str + :var msg: The log's message. + :vartype msg: str + :var err: The log's error, if applicable. + :vartype err: str + :var addr: The URL, if ``obj`` is "web". + :vartype addr: str + """ + + def __init__(self, line): + self.line = line.strip() + self.t = None + self.lvl = "NOTSET" + self.msg = None + self.err = None + self.addr = None + + for i in shlex.split(self.line): + if "=" not in i: + continue + + key, value = i.split("=", 1) + + if key == "lvl": + if not value: + value = self.lvl + + value = value.upper() + if value == "CRIT": + value = "CRITICAL" + elif value in ["ERR", "EROR"]: + value = "ERROR" + elif value == "WARN": + value = "WARNING" + + if not hasattr(logging, value): + value = self.lvl + + setattr(self, key, value) + + def __repr__(self): + return "".format(self.t, self.lvl, self.msg) + + def __str__(self): # pragma: no cover + attrs = [attr for attr in dir(self) if not attr.startswith("_") and getattr(self, attr) is not None] + attrs.remove("line") + + return " ".join("{}=\"{}\"".format(attr, getattr(self, attr)) for attr in attrs) + + +def set_auth_token(pyngrok_config, token): + """ + Set the ``ngrok`` auth token in the config file, enabling authenticated features (for instance, + more concurrent tunnels, custom subdomains, etc.). + + :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary. + :type pyngrok_config: PyngrokConfig + :param token: The auth token to set. + :type token: str + """ + start = [pyngrok_config.ngrok_path, "authtoken", token, "--log=stdout"] + if pyngrok_config.config_path: + logger.info("Updating authtoken for \"config_path\": {}".format(pyngrok_config.config_path)) + start.append("--config={}".format(pyngrok_config.config_path)) + else: + logger.info( + "Updating authtoken for default \"config_path\" of \"ngrok_path\": {}".format(pyngrok_config.ngrok_path)) + + result = subprocess.check_output(start) + + if "Authtoken saved" not in str(result): + raise PyngrokNgrokError("An error occurred when saving the auth token: {}".format(result)) + + +def is_process_running(ngrok_path): + """ + Check if the ``ngrok`` process is currently running. + + :param ngrok_path: The path to the ``ngrok`` binary. + :type ngrok_path: str + :return: ``True`` if ``ngrok`` is running from the given path. + """ + if ngrok_path in _current_processes: + # Ensure the process is still running and hasn't been killed externally, otherwise cleanup + if _current_processes[ngrok_path].proc.poll() is None: + return True + else: + logger.debug( + "Removing stale process for \"ngrok_path\" {}".format(ngrok_path)) + + _current_processes.pop(ngrok_path, None) + + return False + + +def get_process(pyngrok_config): + """ + Get the current ``ngrok`` process for the given config's ``ngrok_path``. + + If ``ngrok`` is not running, calling this method will first start a process with + :class:`~pyngrok.conf.PyngrokConfig`. + + :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary. + :type pyngrok_config: PyngrokConfig + :return: The ``ngrok`` process. + :rtype: NgrokProcess + """ + if is_process_running(pyngrok_config.ngrok_path): + return _current_processes[pyngrok_config.ngrok_path] + + return _start_process(pyngrok_config) + + +def kill_process(ngrok_path): + """ + Terminate the ``ngrok`` processes, if running, for the given path. This method will not block, it will just + issue a kill request. + + :param ngrok_path: The path to the ``ngrok`` binary. + :type ngrok_path: str + """ + if is_process_running(ngrok_path): + ngrok_process = _current_processes[ngrok_path] + + logger.info("Killing ngrok process: {}".format(ngrok_process.proc.pid)) + + try: + ngrok_process.proc.kill() + ngrok_process.proc.wait() + except OSError as e: # pragma: no cover + # If the process was already killed, nothing to do but cleanup state + if e.errno != 3: + raise e + + _current_processes.pop(ngrok_path, None) + else: + logger.debug("\"ngrok_path\" {} is not running a process".format(ngrok_path)) + + +def run_process(ngrok_path, args): + """ + Start a blocking ``ngrok`` process with the binary at the given path and the passed args. + + This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not + necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`. + + :param ngrok_path: The path to the ``ngrok`` binary. + :type ngrok_path: str + :param args: The args to pass to ``ngrok``. + :type args: list[str] + """ + _validate_path(ngrok_path) + + start = [ngrok_path] + args + subprocess.call(start) + + +def capture_run_process(ngrok_path, args): + """ + Start a blocking ``ngrok`` process with the binary at the given path and the passed args. When the process + returns, so will this method, and the captured output from the process along with it. + + This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not + necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`. + + :param ngrok_path: The path to the ``ngrok`` binary. + :type ngrok_path: str + :param args: The args to pass to ``ngrok``. + :type args: list[str] + :return: The output from the process. + :rtype: str + """ + _validate_path(ngrok_path) + + start = [ngrok_path] + args + output = subprocess.check_output(start) + + return output.decode("utf-8").strip() + + +def _validate_path(ngrok_path): + """ + Validate the given path exists, is a ``ngrok`` binary, and is ready to be started, otherwise raise a + relevant exception. + + :param ngrok_path: The path to the ``ngrok`` binary. + :type ngrok_path: str + """ + if not os.path.exists(ngrok_path): + raise PyngrokNgrokError( + "ngrok binary was not found. Be sure to call \"ngrok.install_ngrok()\" first for " + "\"ngrok_path\": {}".format(ngrok_path)) + + if ngrok_path in _current_processes: + raise PyngrokNgrokError("ngrok is already running for the \"ngrok_path\": {}".format(ngrok_path)) + + +def _validate_config(config_path): + with open(config_path, "r") as config_file: + config = yaml.safe_load(config_file) + + if config is not None: + installer.validate_config(config) + + +def _terminate_process(process): + if process is None: + return + + try: + process.terminate() + except OSError: # pragma: no cover + logger.debug("ngrok process already terminated: {}".format(process.pid)) + + +def _start_process(pyngrok_config): + """ + Start a ``ngrok`` process with no tunnels. This will start the ``ngrok`` web interface, against + which HTTP requests can be made to create, interact with, and destroy tunnels. + + :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary. + :type pyngrok_config: PyngrokConfig + :return: The ``ngrok`` process. + :rtype: NgrokProcess + """ + if pyngrok_config.config_path is not None: + config_path = pyngrok_config.config_path + else: + config_path = conf.DEFAULT_NGROK_CONFIG_PATH + + _validate_path(pyngrok_config.ngrok_path) + _validate_config(config_path) + + start = [pyngrok_config.ngrok_path, "start", "--none", "--log=stdout"] + if pyngrok_config.config_path: + logger.info("Starting ngrok with config file: {}".format(pyngrok_config.config_path)) + start.append("--config={}".format(pyngrok_config.config_path)) + if pyngrok_config.auth_token: + logger.info("Overriding default auth token") + start.append("--authtoken={}".format(pyngrok_config.auth_token)) + if pyngrok_config.region: + logger.info("Starting ngrok in region: {}".format(pyngrok_config.region)) + start.append("--region={}".format(pyngrok_config.region)) + + popen_kwargs = {"stdout": subprocess.PIPE, "universal_newlines": True} + if os.name == "posix": + popen_kwargs.update(start_new_session=pyngrok_config.start_new_session) + elif pyngrok_config.start_new_session: + logger.warning("Ignoring start_new_session=True, which requires POSIX") + proc = subprocess.Popen(start, **popen_kwargs) + atexit.register(_terminate_process, proc) + + logger.debug("ngrok process starting with PID: {}".format(proc.pid)) + + ngrok_process = NgrokProcess(proc, pyngrok_config) + _current_processes[pyngrok_config.ngrok_path] = ngrok_process + + timeout = time.time() + pyngrok_config.startup_timeout + while time.time() < timeout: + line = proc.stdout.readline() + ngrok_process._log_startup_line(line) + + if ngrok_process.healthy(): + logger.debug("ngrok process has started with API URL: {}".format(ngrok_process.api_url)) + + ngrok_process.startup_error = None + + if pyngrok_config.monitor_thread: + ngrok_process.start_monitor_thread() + + break + elif ngrok_process.proc.poll() is not None: + break + + if not ngrok_process.healthy(): + # If the process did not come up in a healthy state, clean up the state + kill_process(pyngrok_config.ngrok_path) + + if ngrok_process.startup_error is not None: + raise PyngrokNgrokError("The ngrok process errored on start: {}.".format(ngrok_process.startup_error), + ngrok_process.logs, + ngrok_process.startup_error) + else: + raise PyngrokNgrokError("The ngrok process was unable to start.", ngrok_process.logs) + + return ngrok_process diff --git a/metaflow/_vendor/vendor_any.txt b/metaflow/_vendor/vendor_any.txt index f7fc59bf3bb..73721edc659 100644 --- a/metaflow/_vendor/vendor_any.txt +++ b/metaflow/_vendor/vendor_any.txt @@ -1 +1,2 @@ click==7.1.2 +pyngrok=5.1.0 diff --git a/metaflow/plugins/debug_decorator.py b/metaflow/plugins/debug_decorator.py index ebada848a32..6101068c894 100644 --- a/metaflow/plugins/debug_decorator.py +++ b/metaflow/plugins/debug_decorator.py @@ -58,10 +58,8 @@ def __init__(self, host, port, is_remote=False, is_active=False, auth_token=None self._remote_pdb = None def _setup_ngrok_tunnel(self): - try: - from pyngrok import ngrok - except ImportError: - return None + from metaflow._vendor.pyngrok import ngrok + if not self._auth_token: return None ngrok.set_auth_token(self._auth_token) @@ -110,21 +108,10 @@ class NgrokDebugDecorator(DebugDecorator): name = "remote_debugger" def _validate_ngrok(self): - try: - from pyngrok import ngrok - except: - return False if not self.attributes["auth_token"]: return False return True - def runtime_task_created( - self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context - ): - if not self.attributes["auth_token"]: - if NGROK_KEY: - self.attributes["auth_token"] = NGROK_KEY - def task_pre_step( self, step_name, @@ -139,17 +126,17 @@ def task_pre_step( ubf_context, inputs, ): + if not self.attributes["auth_token"]: + if NGROK_KEY: + self.attributes["auth_token"] = NGROK_KEY self._isvalid = self._validate_ngrok() - return super().task_pre_step( - step_name, - task_datastore, - metadata, - run_id, - task_id, - flow, - graph, - retry_count, - max_user_code_retries, - ubf_context, - inputs, + from metaflow import current + + breakpoint = BreakPoint( + self.attributes["host"], + self.attributes["port"], + is_remote=True, + is_active=self._isvalid, + auth_token=self.attributes["auth_token"], ) + current._update_env({"debug": breakpoint}) From ab6c18253ad873e3dd7afc2c4db165137f59f5e7 Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Tue, 3 May 2022 19:01:06 +0000 Subject: [PATCH 4/4] Made one decorator for debugging. - removed remote_debugger decorator --- metaflow/plugins/__init__.py | 3 +- metaflow/plugins/debug_decorator.py | 78 ++++++++++++++--------------- 2 files changed, 38 insertions(+), 43 deletions(-) diff --git a/metaflow/plugins/__init__.py b/metaflow/plugins/__init__.py index d9333dcb6d8..94166947135 100644 --- a/metaflow/plugins/__init__.py +++ b/metaflow/plugins/__init__.py @@ -113,7 +113,7 @@ def get_plugin_cli(): from .conda.conda_step_decorator import CondaStepDecorator from .cards.card_decorator import CardDecorator from .frameworks.pytorch import PytorchParallelDecorator -from .debug_decorator import NgrokDebugDecorator, DebugDecorator +from .debug_decorator import DebugDecorator STEP_DECORATORS = [ @@ -131,7 +131,6 @@ def get_plugin_cli(): PytorchParallelDecorator, InternalTestUnboundedForeachDecorator, ArgoWorkflowsInternalDecorator, - NgrokDebugDecorator, DebugDecorator, ] _merge_lists(STEP_DECORATORS, _ext_plugins["STEP_DECORATORS"], "name") diff --git a/metaflow/plugins/debug_decorator.py b/metaflow/plugins/debug_decorator.py index 6101068c894..045afdd6626 100644 --- a/metaflow/plugins/debug_decorator.py +++ b/metaflow/plugins/debug_decorator.py @@ -1,6 +1,8 @@ import functools +import random import time from metaflow.decorators import StepDecorator +from metaflow.exception import MetaflowException from metaflow.metaflow_config import NGROK_KEY from .remote_pdb import RemotePdb from pdb import set_trace @@ -11,11 +13,35 @@ class DebugDecorator(StepDecorator): name = "debugger" - defaults = dict(port=8292, host="localhost", auth_token=None) + defaults = dict( + host="localhost", + port=9983, + auth_token=None, + ) + + def _validate_ngrok(self): + if not self.attributes["auth_token"]: + raise MetaflowException( + "Ngrok `auth_token` required when calling @debugger with @batch/@kubernetes. " + "Set the token via `auth_token` in @debugger or set the `METAFLOW_NGROK_KEY` environment variable." + ) + return True def __init__(self, attributes=None, statically_defined=False): super().__init__(attributes, statically_defined) self._isvalid = True + self.backend = None + self._is_remote = False + + def step_init( + self, flow, graph, step_name, decorators, environment, flow_datastore, logger + ): + remote_decos = [ + deco for deco in decorators if deco.name in ["kubernetes", "batch"] + ] + if len(remote_decos) > 0: + self.backend = "ngrok" + self._is_remote = True def task_pre_step( self, @@ -31,12 +57,20 @@ def task_pre_step( ubf_context, inputs, ): + self.port = self.attributes["port"] + # print("Using Port ",self.port) + if self.backend == "ngrok": + if not self.attributes["auth_token"]: + if NGROK_KEY: + self.attributes["auth_token"] = NGROK_KEY + self._isvalid = self._validate_ngrok() + from metaflow import current breakpoint = BreakPoint( self.attributes["host"], - self.attributes["port"], - is_remote=False, + self.port, + is_remote=self._is_remote, is_active=self._isvalid, auth_token=self.attributes["auth_token"], ) @@ -102,41 +136,3 @@ def breakpoint(self): self._activated = True self._remote_pdb = RemotePdb(self._host, self._port, quiet=True) return self._remote_pdb.set_trace() - - -class NgrokDebugDecorator(DebugDecorator): - name = "remote_debugger" - - def _validate_ngrok(self): - if not self.attributes["auth_token"]: - return False - return True - - def task_pre_step( - self, - step_name, - task_datastore, - metadata, - run_id, - task_id, - flow, - graph, - retry_count, - max_user_code_retries, - ubf_context, - inputs, - ): - if not self.attributes["auth_token"]: - if NGROK_KEY: - self.attributes["auth_token"] = NGROK_KEY - self._isvalid = self._validate_ngrok() - from metaflow import current - - breakpoint = BreakPoint( - self.attributes["host"], - self.attributes["port"], - is_remote=True, - is_active=self._isvalid, - auth_token=self.attributes["auth_token"], - ) - current._update_env({"debug": breakpoint})