<commit_before># Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

import os
import json

from tornado import gen, web
from tornado.escape import json_encode, json_decode, url_escape
from tornado.httpclient import HTTPClient, AsyncHTTPClient, HTTPError

from notebook.services.kernels.kernelmanager import MappingKernelManager
from notebook.services.sessions.sessionmanager import (
    SessionManager as BaseSessionManager
)
from jupyter_client.kernelspec import KernelSpecManager
from notebook.utils import url_path_join

from traitlets import Instance, Unicode, default

# TODO: Find a better way to specify global configuration options 
# for a server extension.
KG_URL = os.getenv('KG_URL', 'http://127.0.0.1:8888/')
KG_HEADERS = json.loads(os.getenv('KG_HEADERS', '{}'))
KG_HEADERS.update({
    'Authorization': 'token {}'.format(os.getenv('KG_AUTH_TOKEN', ''))
})
VALIDATE_KG_CERT = os.getenv('VALIDATE_KG_CERT') not in ['no', 'false']

@gen.coroutine
def fetch_kg(endpoint, **kwargs):
    """Make an async request to kernel gateway endpoint."""
    client = AsyncHTTPClient()
    url = url_path_join(KG_URL, endpoint)
    response = yield client.fetch(url, headers=KG_HEADERS, validate_cert=VALIDATE_KG_CERT, **kwargs)
    raise gen.Return(response)


class RemoteKernelManager(MappingKernelManager):
    """Kernel manager that supports remote kernels hosted by Jupyter 
    kernel gateway."""

    kernels_endpoint_env = 'KG_KERNELS_ENDPOINT'
    kernels_endpoint = Unicode(config=True,
        help="""The kernel gateway API endpoint for accessing kernel resources 
        (KG_KERNELS_ENDPOINT env var)""")
    @default('kernels_endpoint')
    def kernels_endpoint_default(self):
        return os.getenv(self.kernels_endpoint_env, '/api/kernels')

    # TODO: The notebook code base assumes a sync operation to determine if
    # kernel manager has a kernel_id (existing kernel manager stores kernels
    # in dictionary).  Keeping such a dictionary in sync with remote KG is
    # NOT something we want to do, is it?
    #
    # options:
    #  - update internal dictionary on every /api/kernels request
    #  - replace `__contains__` with more formal async get_kernel() API
    #    (requires notebook code base changes)
    _kernels = {}
    
    def __contains__(self, kernel_id):
        self.log.debug('RemoteKernelManager.__contains__ {}'.format(kernel_id))
        return kernel_id in self._kernels

    def _remove_kernel(self, kernel_id):
        """Remove a kernel from our mapping, mainly so that a dead kernel can be 
        removed without having to call shutdown_kernel.

        The kernel object is returned.

        Parameters
        ----------
        kernel_id: kernel UUID
        """
        try:
            return self._kernels.pop(kernel_id)
        except KeyError:
            pass

    def _kernel_id_to_url(self, kernel_id):
        """Builds a url for the given kernel UUID.

        Parameters
        ----------
        kernel_id: kernel UUID
        """
        return url_path_join(self.kernels_endpoint, url_escape(str(kernel_id)))

    @gen.coroutine
    def start_kernel(self, kernel_id=None, path=None, **kwargs):
        """Start a kernel for a session and return its kernel_id.

        Parameters
        ----------
        kernel_id : uuid
            The uuid to associate the new kernel with. If this
            is not None, this kernel will be persistent whenever it is
            requested.
        path : API path
            The API path (unicode, '/' delimited) for the cwd.
            Will be transformed to an OS path relative to root_dir.
        """
        self.log.info(
            'Request start kernel: kernel_id=%s, path="%s"',
            kernel_id, path
        )

        if kernel_id is None:
            kernel_name = kwargs.get('kernel_name', 'python3')
            self.log.info("Request new kernel at: %s" % self.kernels_endpoint)
            response = yield fetch_kg(
                self.kernels_endpoint,
                method='POST',
                body=json_encode({
                    'name': kernel_name,
                    'env': {k:v for (k,v) in dict(os.environ).items() if k.startswith('KERNEL_') or k in os.environ['KG_WHITELIST'].split(":")}
                })
            )
            kernel = json_decode(response.body)
            kernel_id = kernel['id']
            self.log.info("Kernel started: %s" % kernel_id)
        else:
            kernel = yield self.get_kernel(kernel_id)
            kernel_id = kernel['id']
            self.log.info("Using existing kernel: %s" % kernel_id)
        self._kernels[kernel_id] = kernel<commit_msg>Refactoring KG_WHITELIST to KG_ENV_WHITELIST<commit_after># Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.

import os
import json

from tornado import gen, web
from tornado.escape import json_encode, json_decode, url_escape
from tornado.httpclient import HTTPClient, AsyncHTTPClient, HTTPError

from notebook.services.kernels.kernelmanager import MappingKernelManager
from notebook.services.sessions.sessionmanager import (
    SessionManager as BaseSessionManager
)
from jupyter_client.kernelspec import KernelSpecManager
from notebook.utils import url_path_join

from traitlets import Instance, Unicode, default

# TODO: Find a better way to specify global configuration options 
# for a server extension.
KG_URL = os.getenv('KG_URL', 'http://127.0.0.1:8888/')
KG_HEADERS = json.loads(os.getenv('KG_HEADERS', '{}'))
KG_HEADERS.update({
    'Authorization': 'token {}'.format(os.getenv('KG_AUTH_TOKEN', ''))
})
VALIDATE_KG_CERT = os.getenv('VALIDATE_KG_CERT') not in ['no', 'false']

@gen.coroutine
def fetch_kg(endpoint, **kwargs):
    """Make an async request to kernel gateway endpoint."""
    client = AsyncHTTPClient()
    url = url_path_join(KG_URL, endpoint)
    response = yield client.fetch(url, headers=KG_HEADERS, validate_cert=VALIDATE_KG_CERT, **kwargs)
    raise gen.Return(response)


class RemoteKernelManager(MappingKernelManager):
    """Kernel manager that supports remote kernels hosted by Jupyter 
    kernel gateway."""

    kernels_endpoint_env = 'KG_KERNELS_ENDPOINT'
    kernels_endpoint = Unicode(config=True,
        help="""The kernel gateway API endpoint for accessing kernel resources 
        (KG_KERNELS_ENDPOINT env var)""")
    @default('kernels_endpoint')
    def kernels_endpoint_default(self):
        return os.getenv(self.kernels_endpoint_env, '/api/kernels')

    # TODO: The notebook code base assumes a sync operation to determine if
    # kernel manager has a kernel_id (existing kernel manager stores kernels
    # in dictionary).  Keeping such a dictionary in sync with remote KG is
    # NOT something we want to do, is it?
    #
    # options:
    #  - update internal dictionary on every /api/kernels request
    #  - replace `__contains__` with more formal async get_kernel() API
    #    (requires notebook code base changes)
    _kernels = {}
    
    def __contains__(self, kernel_id):
        self.log.debug('RemoteKernelManager.__contains__ {}'.format(kernel_id))
        return kernel_id in self._kernels

    def _remove_kernel(self, kernel_id):
        """Remove a kernel from our mapping, mainly so that a dead kernel can be 
        removed without having to call shutdown_kernel.

        The kernel object is returned.

        Parameters
        ----------
        kernel_id: kernel UUID
        """
        try:
            return self._kernels.pop(kernel_id)
        except KeyError:
            pass

    def _kernel_id_to_url(self, kernel_id):
        """Builds a url for the given kernel UUID.

        Parameters
        ----------
        kernel_id: kernel UUID
        """
        return url_path_join(self.kernels_endpoint, url_escape(str(kernel_id)))

    @gen.coroutine
    def start_kernel(self, kernel_id=None, path=None, **kwargs):
        """Start a kernel for a session and return its kernel_id.

        Parameters
        ----------
        kernel_id : uuid
            The uuid to associate the new kernel with. If this
            is not None, this kernel will be persistent whenever it is
            requested.
        path : API path
            The API path (unicode, '/' delimited) for the cwd.
            Will be transformed to an OS path relative to root_dir.
        """
        self.log.info(
            'Request start kernel: kernel_id=%s, path="%s"',
            kernel_id, path
        )

        if kernel_id is None:
            kernel_name = kwargs.get('kernel_name', 'python3')
            self.log.info("Request new kernel at: %s" % self.kernels_endpoint)
            response = yield fetch_kg(
                self.kernels_endpoint,
                method='POST',
                body=json_encode({
                    'name': kernel_name,
                    'env': {k:v for (k,v) in dict(os.environ).items() if k.startswith('KERNEL_') or k in os.environ['KG_ENV_WHITELIST'].split(",")}
                })
            )
            kernel = json_decode(response.body)
            kernel_id = kernel['id']
            self.log.info("Kernel started: %s" % kernel_id)
        else:
            kernel = yield self.get_kernel(kernel_id)
            kernel_id = kernel['id']
            self.log.info("Using existing kernel: %s" % kernel_id)
        self._kernels[kernel_id] = kernel
{'fraction_non_alphanumeric': 0.09604579667126047, 'fraction_numerical': 0.002332237888264603, 'mean_word_length': 2.6679626749611196, 'pattern_counts': {'":': 1, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1037763', 'n_tokens_mistral': 2733, 'n_tokens_neox': 2640, 'n_words': 806}