diff --git a/dev-values.yaml b/dev-values.yaml index 6ff2c58..4769943 100644 --- a/dev-values.yaml +++ b/dev-values.yaml @@ -91,22 +91,21 @@ jupyterhub: config: base_url: "https://minikube.lsst.codes" - options_form: - images_url: "http://cachemachine.cachemachine.svc.cluster.local/cachemachine/jupyter/available" - images: [] - sizes: - - name: Tiny - cpu: .5 - ram: 1536M - - name: Small - cpu: 1 - ram: 3072M - - name: Medium - cpu: 2 - ram: 6144M - - name: Large - cpu: 4 - ram: 12288M + images_url: "http://cachemachine.cachemachine.svc.cluster.local/cachemachine/jupyter/available" + pinned_images: [] + sizes: + - name: Tiny + cpu: 0.5 + ram: 1536M + - name: Small + cpu: 1 + ram: 3072M + - name: Medium + cpu: 2 + ram: 6144M + - name: Large + cpu: 4 + ram: 12288M user_resources: - apiVersion: v1 kind: Namespace @@ -127,7 +126,7 @@ config: SODA_ROUTE: /api/image/soda WORKFLOW_ROUTE: /wf NO_SUDO: 'TRUE' - AUTO_REPO_URLS: https://github.com/lsst-sqre/notebook-demo + AUTO_REPO_URLS: "https://github.com/lsst-sqre/notebook-demo" EXTERNAL_GROUPS: "{{ external_groups }}" EXTERNAL_USER: "{{ user }}" EXTERNAL_UID: "{{ uid }}" diff --git a/src/nublado2/auth.py b/src/nublado2/auth.py index f9fce64..ddc7a9c 100644 --- a/src/nublado2/auth.py +++ b/src/nublado2/auth.py @@ -9,8 +9,8 @@ from jupyterhub.utils import url_path_join from tornado import web +from nublado2.http import get_session from nublado2.nublado_config import NubladoConfig -from nublado2.options import session if TYPE_CHECKING: from typing import Any, Dict, List, Optional, Tuple, Type, Union @@ -151,10 +151,11 @@ async def _build_auth_info(headers: HTTPHeaders) -> Dict[str, Any]: raise web.HTTPError(401, "No request token") # Retrieve the token metadata. - base_url = NubladoConfig().get().get("base_url") + base_url = NubladoConfig().base_url if not base_url: raise web.HTTPError(500, "base_url not set in configuration") api_url = url_path_join(base_url, "/auth/analyze") + session = await get_session() resp = await session.post(api_url, data={"token": token}) if resp.status != 200: raise web.HTTPError(500, "Cannot reach token analysis API") diff --git a/src/nublado2/hooks.py b/src/nublado2/hooks.py index e570bad..fa77dc4 100644 --- a/src/nublado2/hooks.py +++ b/src/nublado2/hooks.py @@ -20,8 +20,8 @@ async def pre_spawn(self, spawner: Spawner) -> None: ) spawner.image = options.image_info.reference - spawner.mem_limit = options.ram - spawner.cpu_limit = options.cpu + spawner.mem_limit = options.size.ram + spawner.cpu_limit = options.size.cpu auth_state = await spawner.user.get_auth_state() diff --git a/src/nublado2/http.py b/src/nublado2/http.py new file mode 100644 index 0000000..fd6bdf6 --- /dev/null +++ b/src/nublado2/http.py @@ -0,0 +1,17 @@ +from aiohttp import ClientSession + +_session = None + + +async def get_session() -> ClientSession: + """This is the way to retrieve a ClientSession to make HTTP requests. + + ClientSession needs to be created inside an async function, so by + calling this, you ensure it exists, or create it if it doesn't. + + Since there are some connection pools, we don't want to be creating + these all the time. Better to just reuse one.""" + global _session + if not _session: + _session = ClientSession() + return _session diff --git a/src/nublado2/hub_config.py b/src/nublado2/hub_config.py index 4a4708d..8e38f66 100644 --- a/src/nublado2/hub_config.py +++ b/src/nublado2/hub_config.py @@ -16,7 +16,7 @@ def configure(self, c: JupyterHub) -> None: self.log.info("Configuring JupyterHub Nublado2 style") self.log.debug(f"JupyterHub configuration starting as: {c}") - nc = NubladoConfig().get() + nc = NubladoConfig() self.log.debug(f"Nublado Config is:\n{nc}") c.JupyterHub.hub_connect_url = self._get_hub_connect_url() diff --git a/src/nublado2/imageinfo.py b/src/nublado2/imageinfo.py index 9a7234c..cbb622f 100644 --- a/src/nublado2/imageinfo.py +++ b/src/nublado2/imageinfo.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Dict FIELD_DELIMITER = "|" @@ -31,14 +31,18 @@ class ImageInfo: Example: "sha256:419c4b7e14603711b25fa9e0569460a753c4b2449fe275bb5f89743b01794a30" # noqa: E501 """ - packed_string: str = field(init=False, default="") - """packed_string is the form in which the image info is packed into the - JupyterHub options form and in which is is returned as the form - selection. It is specification, display_name, and digest - concatenated with the pipe character. + @property + def packed_string(self) -> str: + """packed_string is the form in which the image info is packed into the + JupyterHub options form and in which is is returned as the form + selection. It is specification, display_name, and digest + concatenated with the pipe character. - Example: "registry.hub.docker.com/lsstsqre/sciplat-lab:w_2021_13|Weekly 13|sha256:419c4b7e14603711b25fa9e0569460a753c4b2449fe275bb5f89743b01794a30" # noqa: E501 - """ + Example: "registry.hub.docker.com/lsstsqre/sciplat-lab:w_2021_13|Weekly 13|sha256:419c4b7e14603711b25fa9e0569460a753c4b2449fe275bb5f89743b01794a30" # noqa: E501 + """ + return FIELD_DELIMITER.join( + [self.reference, self.display_name, self.digest] + ) @classmethod def from_cachemachine_entry(cls, entry: CachemachineEntry) -> "ImageInfo": @@ -67,11 +71,3 @@ def from_packed_string(cls, packed_string: str) -> "ImageInfo": return cls( reference=fields[0], display_name=fields[1], digest=fields[2] ) - - def __post_init__(self) -> None: - object.__setattr__(self, "packed_string", self._pack()) - - def _pack(self) -> str: - return FIELD_DELIMITER.join( - [self.reference, self.display_name, self.digest] - ) diff --git a/src/nublado2/labsize.py b/src/nublado2/labsize.py new file mode 100644 index 0000000..ac96e30 --- /dev/null +++ b/src/nublado2/labsize.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class LabSize: + """The cpu and ram settings for a lab container.""" + + cpu: float + """Number of virtual CPUs to allocate for this lab. + + This can be a partial number, such as 2.5 or .5 vCPUs.""" + + name: str + """The name referring to this pairing of cpu and ram.""" + + ram: str + """Amount of memory to allocate for this lab. + + This is a string with special characters for units, such as + 2048M, or 2G.""" diff --git a/src/nublado2/nublado_config.py b/src/nublado2/nublado_config.py index 7c8f576..2e6a84a 100644 --- a/src/nublado2/nublado_config.py +++ b/src/nublado2/nublado_config.py @@ -2,26 +2,64 @@ __all__ = ["NubladoConfig"] -from typing import Any, Dict, Tuple +from typing import Any, Dict, List, Tuple from ruamel import yaml from ruamel.yaml import RoundTripLoader -from traitlets.config import LoggingConfigurable +from nublado2.imageinfo import ImageInfo +from nublado2.labsize import LabSize -class NubladoConfig(LoggingConfigurable): - def get(self) -> Dict[str, Any]: + +class NubladoConfig: + def __init__(self) -> None: + """Load the nublado_config.yaml file from disk. + + This file normally comes from mounting a configmap with the + nublado_config.yaml mounted into the hub container.""" with open("/etc/jupyterhub/nublado_config.yaml") as f: - nc = yaml.load(f.read(), Loader=RoundTripLoader) + self._config = yaml.load(f.read(), Loader=RoundTripLoader) + + self._sizes = { + s.name: s + for s in [ + LabSize(float(s["cpu"]), s["name"], s["ram"]) + for s in self._config["options_form"]["sizes"] + ] + } + + @property + def base_url(self) -> str: + """Base URL for the environment, like https://data.lsst.cloud""" + return self._config["base_url"] + + @property + def images_url(self) -> str: + """URL to fetch list of images to show in options form. + + Generally, this is a link to the cachemachine service.""" + return self._config["images_url"] - self.log.debug(f"Loaded Nublado Config:\n{nc}") - return nc + @property + def pinned_images(self) -> List[ImageInfo]: + """List of images to keep pinned in the options form.""" + return [ + ImageInfo.from_cachemachine_entry(i) + for i in self._config["pinned_images"] + ] - def lookup_size(self, name: str) -> Tuple[float, str]: - sizes = self.get()["options_form"]["sizes"] + @property + def signing_key(self) -> str: + """Retrieve the gafaelfawr signing key to mint tokens.""" + with open("/etc/keys/signing_key.pem", "r") as f: + return f.read() - for s in sizes: - if s["name"] == name: - return (float(s["cpu"]), s["ram"]) + @property + def sizes(self) -> Dict[str, LabSize]: + """Retrieve a copy of the sizes a lab can spawn as.""" + return dict(self._sizes) - raise ValueError(f"Size {name} not found") + @property + def user_resources(self) -> Tuple[Any, ...]: + """Retrieve a copy of the lab resources templates.""" + return tuple(self._config.get("user_resources", [])) diff --git a/src/nublado2/options.py b/src/nublado2/options.py index d55d18d..02f8500 100644 --- a/src/nublado2/options.py +++ b/src/nublado2/options.py @@ -1,10 +1,10 @@ -from typing import Dict, List, Optional +from typing import List, Optional, Tuple -from aiohttp import ClientSession from jinja2 import Template from jupyterhub.spawner import Spawner from traitlets.config import LoggingConfigurable +from nublado2.http import get_session from nublado2.imageinfo import ImageInfo from nublado2.nublado_config import NubladoConfig @@ -73,47 +73,42 @@ """ ) -# Don't have this be a member of NubladoOptions, we should -# share this connection pool. Also the LoggingConfigurable -# will try to pickle it to json, and it can't pickle a session. -session = ClientSession() - class NubladoOptions(LoggingConfigurable): async def show_options_form(self, spawner: Spawner) -> str: - options_config = NubladoConfig().get()["options_form"] - sizes = options_config["sizes"] - - images_url = options_config.get("images_url") + nc = NubladoConfig() - cachemachine_response = await self._get_images_from_url(images_url) + (cached_images, all_images) = await self._get_images_from_url( + nc.images_url + ) + cached_images.extend(nc.pinned_images) - all_imageinfos = [ - ImageInfo.from_cachemachine_entry(img) - for img in cachemachine_response["all"] - ] - # Start with the cachemachine response, then extend it with - # contents of options_config - cached_images = cachemachine_response["images"] - cached_images.extend(options_config["images"]) - cached_imageinfos = [ - ImageInfo.from_cachemachine_entry(img) for img in cached_images - ] return options_template.render( dropdown_sentinel=DROPDOWN_SENTINEL_VALUE, - cached_images=cached_imageinfos, - all_images=all_imageinfos, - sizes=sizes, + cached_images=cached_images, + all_images=all_images, + sizes=nc.sizes.values(), ) async def _get_images_from_url( self, url: Optional[str] - ) -> Dict[str, List[Dict[str, str]]]: + ) -> Tuple[List[ImageInfo], List[ImageInfo]]: if not url: - return {"all": [], "images": []} + return ([], []) + session = await get_session() r = await session.get(url) if r.status != 200: raise Exception(f"Error {r.status} from {url}") - return await r.json() + body = await r.json() + + cached_images = [ + ImageInfo.from_cachemachine_entry(img) for img in body["images"] + ] + + all_images = [ + ImageInfo.from_cachemachine_entry(img) for img in body["all"] + ] + + return (cached_images, all_images) diff --git a/src/nublado2/resourcemgr.py b/src/nublado2/resourcemgr.py index 929b1eb..82f5c39 100644 --- a/src/nublado2/resourcemgr.py +++ b/src/nublado2/resourcemgr.py @@ -43,7 +43,7 @@ async def create_user_resources( auth_state = await spawner.user.get_auth_state() self.log.debug(f"Auth state={auth_state}") - nc = NubladoConfig().get() + nc = NubladoConfig() groups = auth_state["groups"] # Build a comma separated list of group:gid @@ -62,15 +62,13 @@ async def create_user_resources( "token": auth_state["token"], "groups": groups, "external_groups": external_groups, - "base_url": nc.get("base_url"), + "base_url": nc.base_url, "dask_yaml": await self._build_dask_template(spawner), - "auto_repo_urls": nc.get("auto_repo_urls"), "options": options, } self.log.debug(f"Template values={template_values}") - resources = nc.get("user_resources", []) - for r in resources: + for r in nc.user_resources: t_yaml = yaml.dump(r, Dumper=RoundTripDumper) self.log.debug(f"Resource template:\n{t_yaml}") t = Template(t_yaml) @@ -86,9 +84,8 @@ async def create_user_resources( async def _request_homedir_provisioning(self, spawner: Spawner) -> None: """Submit a request for provisioning via Moneypenny.""" - nc = NubladoConfig().get() + nc = NubladoConfig() hc = self.http_client - base_url = nc.get("base_url") uname = spawner.user.name auth_state = await spawner.user.get_auth_state() dossier = { @@ -97,13 +94,13 @@ async def _request_homedir_provisioning(self, spawner: Spawner) -> None: "groups": auth_state["groups"], } token = await self._mint_admin_token() - endpt = f"{base_url}/moneypenny/commission" + endpt = f"{nc.base_url}/moneypenny/commission" auth = {"Authorization": f"Bearer {token}"} self.log.debug(f"Posting dossier {dossier} to {endpt}") resp = await hc.post(endpt, json=dossier, headers=auth) self.log.debug(f"POST got {resp.status}") resp.raise_for_status() - route = f"{base_url}/moneypenny/{uname}" + route = f"{nc.base_url}/moneypenny/{uname}" count = 0 async def _check_moneypenny_completion() -> bool: @@ -132,22 +129,18 @@ async def _mint_admin_token(self) -> str: """Create a token with exec:admin scope, signed as if Gafaelfawr had created it, in order to submit orders to Moneypenny. """ - nc = NubladoConfig().get() + nc = NubladoConfig() template_file = os.path.join( os.path.dirname(__file__), "static/moneypenny-jwt-template.json" ) - base_url = nc.get("base_url") - signing_key_path = nc.get("signing_key_path") - assert isinstance(signing_key_path, str) - with open(signing_key_path, "r") as f: - signing_key = f.read() - current_time = int( - datetime.datetime.now(tz=datetime.timezone.utc).timestamp() - ) + current_time = int( + datetime.datetime.now(tz=datetime.timezone.utc).timestamp() + ) with open(template_file, "r") as f: token_template = Template(f.read()) + token_data = { - "environment_url": base_url, + "environment_url": nc.base_url, "username": "moneypenny", "uidnumber": 1007, "issue_time": current_time, @@ -157,7 +150,7 @@ async def _mint_admin_token(self) -> str: token_dict = json.loads(rendered_token) token = jwt.encode( token_dict, - key=signing_key, + key=nc.signing_key, headers={"kid": "reissuer"}, algorithm="RS256", ) diff --git a/src/nublado2/selectedoptions.py b/src/nublado2/selectedoptions.py index 4f020a7..d245503 100644 --- a/src/nublado2/selectedoptions.py +++ b/src/nublado2/selectedoptions.py @@ -1,6 +1,7 @@ from typing import Any, Dict from nublado2.imageinfo import ImageInfo +from nublado2.labsize import LabSize from nublado2.nublado_config import NubladoConfig from nublado2.options import DROPDOWN_SENTINEL_VALUE @@ -27,7 +28,7 @@ def __init__(self, options: Dict[str, Any]) -> None: self._image_info = ImageInfo.from_packed_string(image_list) nc = NubladoConfig() - (self._cpu, self._ram) = nc.lookup_size(size_name) + self._size = nc.sizes[size_name] self._debug = "TRUE" if "enable_debug" in options else "" self._clear_dotlocal = "TRUE" if "clear_dotlocal" in options else "" @@ -54,14 +55,5 @@ def image_info(self) -> ImageInfo: return self._image_info @property - def cpu(self) -> float: - """Number of vCPUs for the lab pod. Comes from the size.""" - return self._cpu - - @property - def ram(self) -> str: - """Amount of RAM for the lab pod. - - This is in kubernetes format, like 2g or 2048M, and comes - from the size.""" - return self._ram + def size(self) -> LabSize: + return self._size diff --git a/tests/auth_test.py b/tests/auth_test.py index f574411..d875d19 100644 --- a/tests/auth_test.py +++ b/tests/auth_test.py @@ -16,10 +16,18 @@ from tornado.httputil import HTTPHeaders from nublado2.auth import GafaelfawrAuthenticator, GafaelfawrLoginHandler -from nublado2.nublado_config import NubladoConfig if TYPE_CHECKING: - from typing import Any, Callable, Dict + from typing import Any, AsyncGenerator, Callable, Dict + + +@pytest.fixture(autouse=True) +async def use_config_mock() -> AsyncGenerator: + """Use a mock NubladoConfig object.""" + with patch("nublado2.auth.NubladoConfig") as mock: + mock.return_value = MagicMock() + mock.return_value.base_url = "https://data.example.com/" + yield def test_authenticator() -> None: @@ -48,115 +56,103 @@ def handler(url: str, **kwargs: Any) -> CallbackResult: @pytest.mark.asyncio async def test_login_handler() -> None: - # Using patch.object as a decorator doesn't work with Python 3.7. - with patch.object(NubladoConfig, "get") as mock_config_get: - mock_config_get.return_value = {} - headers = HTTPHeaders({"X-Auth-Request-Token": "some-token"}) + # No headers. + with aioresponses() as m: + with pytest.raises(web.HTTPError): + await GafaelfawrLoginHandler._build_auth_info(HTTPHeaders()) + + headers = HTTPHeaders({"X-Auth-Request-Token": "some-token"}) + + # Invalid token. + with aioresponses() as m: + handler = build_handler({"uid": "foo"}, valid=False) + m.post("https://data.example.com/auth/analyze", callback=handler) with pytest.raises(web.HTTPError): await GafaelfawrLoginHandler._build_auth_info(headers) - mock_config_get.return_value = { - "base_url": "https://data.example.com/" - } + # Bad API status. + with aioresponses() as m: + m.post("https://data.example.com/auth/analyze", payload={}, status=500) + with pytest.raises(web.HTTPError): + await GafaelfawrLoginHandler._build_auth_info(headers) - headers = HTTPHeaders() + # Invalid response. + with aioresponses() as m: + m.post("https://data.example.com/auth/analyze", payload={}) with pytest.raises(web.HTTPError): await GafaelfawrLoginHandler._build_auth_info(headers) - headers.add("X-Auth-Request-Token", "some-token") - - # Invalid token. - with aioresponses() as m: - handler = build_handler({"uid": "foo"}, valid=False) - m.post("https://data.example.com/auth/analyze", callback=handler) - with pytest.raises(web.HTTPError): - await GafaelfawrLoginHandler._build_auth_info(headers) - - # Bad API status. - with aioresponses() as m: - m.post( - "https://data.example.com/auth/analyze", payload={}, status=500 - ) - with pytest.raises(web.HTTPError): - await GafaelfawrLoginHandler._build_auth_info(headers) - - # Invalid response. - with aioresponses() as m: - m.post("https://data.example.com/auth/analyze", payload={}) - with pytest.raises(web.HTTPError): - await GafaelfawrLoginHandler._build_auth_info(headers) - - # Test minimum data. - with aioresponses() as m: - handler = build_handler({"uid": "foo"}) - m.post("https://data.example.com/auth/analyze", callback=handler) - assert await GafaelfawrLoginHandler._build_auth_info(headers) == { - "name": "foo", - "auth_state": { - "uid": None, - "token": "some-token", - "groups": [], - }, - } + # Test minimum data. + with aioresponses() as m: + handler = build_handler({"uid": "foo"}) + m.post("https://data.example.com/auth/analyze", callback=handler) + assert await GafaelfawrLoginHandler._build_auth_info(headers) == { + "name": "foo", + "auth_state": { + "uid": None, + "token": "some-token", + "groups": [], + }, + } - # Test full data. - with aioresponses() as m: - handler = build_handler( - { - "uid": "bar", - "uidNumber": "4510", - "isMemberOf": [ - {"name": "group-one", "id": 1726}, - {"name": "group-two", "id": "1618"}, - {"name": "another", "id": 6789, "foo": "bar"}, - ], - } - ) - m.post("https://data.example.com/auth/analyze", callback=handler) - assert await GafaelfawrLoginHandler._build_auth_info(headers) == { - "name": "bar", - "auth_state": { - "uid": 4510, - "token": "some-token", - "groups": [ - {"name": "group-one", "id": 1726}, - {"name": "group-two", "id": 1618}, - {"name": "another", "id": 6789}, - ], - }, + # Test full data. + with aioresponses() as m: + handler = build_handler( + { + "uid": "bar", + "uidNumber": "4510", + "isMemberOf": [ + {"name": "group-one", "id": 1726}, + {"name": "group-two", "id": "1618"}, + {"name": "another", "id": 6789, "foo": "bar"}, + ], } + ) + m.post("https://data.example.com/auth/analyze", callback=handler) + assert await GafaelfawrLoginHandler._build_auth_info(headers) == { + "name": "bar", + "auth_state": { + "uid": 4510, + "token": "some-token", + "groups": [ + {"name": "group-one", "id": 1726}, + {"name": "group-two", "id": 1618}, + {"name": "another", "id": 6789}, + ], + }, + } - # Check invalid format of isMemberOf. - with aioresponses() as m: - handler = build_handler( - {"uid": "bar", "isMemberOf": [{"name": "foo", "id": ["foo"]}]} - ) - m.post("https://data.example.com/auth/analyze", callback=handler) - with pytest.raises(web.HTTPError): - await GafaelfawrLoginHandler._build_auth_info(headers) - - # Test groups without GIDs. - with aioresponses() as m: - handler = build_handler( - { - "uid": "bar", - "uidNumber": "4510", - "isMemberOf": [ - {"name": "group-one", "id": 1726}, - {"name": "group-two"}, - {"name": "another", "id": 6789}, - ], - } - ) - m.post("https://data.example.com/auth/analyze", callback=handler) - assert await GafaelfawrLoginHandler._build_auth_info(headers) == { - "name": "bar", - "auth_state": { - "uid": 4510, - "token": "some-token", - "groups": [ - {"name": "group-one", "id": 1726}, - {"name": "another", "id": 6789}, - ], - }, + # Check invalid format of isMemberOf. + with aioresponses() as m: + handler = build_handler( + {"uid": "bar", "isMemberOf": [{"name": "foo", "id": ["foo"]}]} + ) + m.post("https://data.example.com/auth/analyze", callback=handler) + with pytest.raises(web.HTTPError): + await GafaelfawrLoginHandler._build_auth_info(headers) + + # Test groups without GIDs. + with aioresponses() as m: + handler = build_handler( + { + "uid": "bar", + "uidNumber": "4510", + "isMemberOf": [ + {"name": "group-one", "id": 1726}, + {"name": "group-two"}, + {"name": "another", "id": 6789}, + ], } + ) + m.post("https://data.example.com/auth/analyze", callback=handler) + assert await GafaelfawrLoginHandler._build_auth_info(headers) == { + "name": "bar", + "auth_state": { + "uid": 4510, + "token": "some-token", + "groups": [ + {"name": "group-one", "id": 1726}, + {"name": "another", "id": 6789}, + ], + }, + }