From ddf160ec29c98181b75e994ea2c0e0648693284a Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 24 Jul 2026 14:03:23 -0700 Subject: [PATCH 1/3] feat(computing-unit): out-of-pod dataset mount infrastructure Perform the FUSE mount for dataset repositories outside the (unprivileged) computing-unit pod. A per-node privileged `texera-mounter` DaemonSet runs GeeseFS on the pod's behalf and the read-only mount is exposed back into the pod via mount propagation, scoped to that computing unit. A JWT-authenticated S3 proxy in file-service fronts the LakeFS S3 gateway: it verifies the pod's JWT, checks the user's read access, and re-signs to LakeFS with credentials held only server-side, so no global credential ever enters the pod. - `texera-mounter` DaemonSet: mounter.py (+ tests), dockerfile, helm daemonset/rbac/values. - Unprivileged CU pod + mount-propagation wiring (KubernetesClient) and mounter config/env. - File-service JWT S3 proxy (S3ProxyServlet). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01H27sHgS29u2JcLYPPR6Hdx --- bin/dockerfiles/mounter.dockerfile | 38 ++ .../base/mounter/mounter-daemonset.yaml | 75 +++ .../templates/base/mounter/mounter-rbac.yaml | 48 ++ ...low-computing-unit-manager-deployment.yaml | 2 + bin/k8s/values.yaml | 13 + bin/mounter/mounter.py | 442 ++++++++++++++++++ bin/mounter/tests/conftest.py | 111 +++++ bin/mounter/tests/test_mounter.py | 397 ++++++++++++++++ .../config/src/main/resources/kubernetes.conf | 13 + .../common/config/EnvironmentalVariable.scala | 10 + .../common/config/KubernetesConfig.scala | 5 + .../service/util/KubernetesClient.scala | 67 ++- .../apache/texera/service/FileService.scala | 7 + .../texera/service/util/S3ProxyServlet.scala | 260 +++++++++++ .../service/util/S3ProxyServletSpec.scala | 86 ++++ 15 files changed, 1562 insertions(+), 12 deletions(-) create mode 100644 bin/dockerfiles/mounter.dockerfile create mode 100644 bin/k8s/templates/base/mounter/mounter-daemonset.yaml create mode 100644 bin/k8s/templates/base/mounter/mounter-rbac.yaml create mode 100644 bin/mounter/mounter.py create mode 100644 bin/mounter/tests/conftest.py create mode 100644 bin/mounter/tests/test_mounter.py create mode 100644 file-service/src/main/scala/org/apache/texera/service/util/S3ProxyServlet.scala create mode 100644 file-service/src/test/scala/org/apache/texera/service/util/S3ProxyServletSpec.scala diff --git a/bin/dockerfiles/mounter.dockerfile b/bin/dockerfiles/mounter.dockerfile new file mode 100644 index 00000000000..d385e2aa57a --- /dev/null +++ b/bin/dockerfiles/mounter.dockerfile @@ -0,0 +1,38 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# The texera-mounter runs as a per-node privileged DaemonSet. It performs the GeeseFS +# FUSE mount that computing-unit pods used to do themselves, so that CU pods (which run +# untrusted user code) can be unprivileged. It needs python3 (the mounter), fuse3 + geesefs +# (to mount), and util-linux (umount) — all part of a minimal Debian base. +FROM debian:bookworm-slim + +ARG GEESEFS_VERSION=v0.43.8 +RUN apt-get update \ + && apt-get install -y --no-install-recommends python3 fuse3 mount ca-certificates curl \ + && curl -fsSL -o /usr/local/bin/geesefs \ + "https://github.com/yandex-cloud/geesefs/releases/download/${GEESEFS_VERSION}/geesefs-linux-$(dpkg --print-architecture)" \ + && chmod 755 /usr/local/bin/geesefs \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* \ + # allow FUSE mounts to be accessible by other users (the unprivileged CU pod's UID) + && echo "user_allow_other" >> /etc/fuse.conf + +COPY bin/mounter/mounter.py /opt/mounter/mounter.py + +EXPOSE 8100 +ENTRYPOINT ["python3", "/opt/mounter/mounter.py"] diff --git a/bin/k8s/templates/base/mounter/mounter-daemonset.yaml b/bin/k8s/templates/base/mounter/mounter-daemonset.yaml new file mode 100644 index 00000000000..e846cc282d1 --- /dev/null +++ b/bin/k8s/templates/base/mounter/mounter-daemonset.yaml @@ -0,0 +1,75 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Per-node privileged mounter that performs the GeeseFS FUSE mount on behalf of +# (unprivileged) computing-unit pods. It mounts under a host directory and the mount +# propagates into each CU pod via mountPropagation, so CU pods — which run untrusted user +# code — never need to be privileged. This is the ONLY privileged component of the mount +# feature, and it runs only trusted mount code (never user code). +apiVersion: apps/v1 +kind: DaemonSet +metadata: + name: {{ .Release.Name }}-mounter + namespace: {{ .Release.Namespace }} + labels: + app: {{ .Release.Name }}-mounter +spec: + selector: + matchLabels: + app: {{ .Release.Name }}-mounter + template: + metadata: + labels: + app: {{ .Release.Name }}-mounter + spec: + serviceAccountName: {{ .Values.mounter.serviceAccountName }} + tolerations: + - operator: "Exists" + containers: + - name: mounter + image: {{ .Values.texera.imageRegistry }}/{{ .Values.mounter.imageName }}:{{ .Values.texera.imageTag }} + imagePullPolicy: {{ .Values.mounter.imagePullPolicy }} + securityContext: + privileged: true + ports: + - name: mounter + containerPort: {{ .Values.mounter.port }} + hostPort: {{ .Values.mounter.port }} + env: + - name: MOUNTER_PORT + value: "{{ .Values.mounter.port }}" + - name: MOUNT_ROOT + value: "{{ .Values.mounter.hostMountRoot }}" + - name: POOL_NAMESPACE + value: "{{ .Values.workflowComputingUnitPool.namespace }}" + - name: CU_POD_NAME_PREFIX + value: "{{ .Values.workflowComputingUnitPool.podNamePrefix }}" + readinessProbe: + httpGet: + path: /healthz + port: mounter + initialDelaySeconds: 3 + periodSeconds: 10 + volumeMounts: + - name: texera-mounts + mountPath: {{ .Values.mounter.hostMountRoot }} + mountPropagation: Bidirectional + volumes: + - name: texera-mounts + hostPath: + path: {{ .Values.mounter.hostMountRoot }} + type: DirectoryOrCreate diff --git a/bin/k8s/templates/base/mounter/mounter-rbac.yaml b/bin/k8s/templates/base/mounter/mounter-rbac.yaml new file mode 100644 index 00000000000..0d35c5b41c0 --- /dev/null +++ b/bin/k8s/templates/base/mounter/mounter-rbac.yaml @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# The mounter watches computing-unit pods (in the pool namespace) so it can unmount a CU's +# mounts as soon as its pod is deleted. It needs only read access to pods. +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ .Values.mounter.serviceAccountName }} + namespace: {{ .Release.Namespace }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: {{ .Release.Name }}-mounter + namespace: {{ .Values.workflowComputingUnitPool.namespace }} +rules: + - apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: {{ .Release.Name }}-mounter-binding + namespace: {{ .Values.workflowComputingUnitPool.namespace }} +subjects: + - kind: ServiceAccount + name: {{ .Values.mounter.serviceAccountName }} + namespace: {{ .Release.Namespace }} +roleRef: + kind: Role + name: {{ .Release.Name }}-mounter + apiGroup: rbac.authorization.k8s.io diff --git a/bin/k8s/templates/base/workflow-computing-unit-manager/workflow-computing-unit-manager-deployment.yaml b/bin/k8s/templates/base/workflow-computing-unit-manager/workflow-computing-unit-manager-deployment.yaml index a9118450412..a41ed74ea00 100644 --- a/bin/k8s/templates/base/workflow-computing-unit-manager/workflow-computing-unit-manager-deployment.yaml +++ b/bin/k8s/templates/base/workflow-computing-unit-manager/workflow-computing-unit-manager-deployment.yaml @@ -64,6 +64,8 @@ spec: value: {{ .Values.workflowComputingUnitPool.namespace }} - name: KUBERNETES_COMPUTE_UNIT_SERVICE_NAME value: {{ .Values.workflowComputingUnitPool.name }}-svc + - name: KUBERNETES_COMPUTE_UNIT_POD_NAME_PREFIX + value: {{ .Values.workflowComputingUnitPool.podNamePrefix }} - name: KUBERNETES_IMAGE_NAME value: {{ .Values.texera.imageRegistry }}/{{ .Values.workflowComputingUnitPool.imageName }}:{{ .Values.texera.imageTag }} # TexeraDB Access diff --git a/bin/k8s/values.yaml b/bin/k8s/values.yaml index 4651bae6947..bb782993779 100644 --- a/bin/k8s/values.yaml +++ b/bin/k8s/values.yaml @@ -288,6 +288,9 @@ workflowComputingUnitPool: name: texera-workflow-computing-unit # Note: the namespace of the workflow computing unit pool might conflict when there are multiple texera deployments in the same cluster namespace: texera-workflow-computing-unit-pool + # Computing unit pods are named "-". The manager creates them under + # this name and the mounter reads a CU's id back out of it, so both are given this value. + podNamePrefix: computing-unit # Max number of resources allocated for computing units maxRequestedResources: cpu: 100 @@ -298,6 +301,16 @@ workflowComputingUnitPool: port: 8085 targetPort: 8085 +# Per-node privileged DaemonSet that performs the GeeseFS FUSE mount for computing-unit +# pods, so the CU pods stay unprivileged. port/hostMountRoot MUST match the values read by +# the computing-unit manager (kubernetes.mounter-port / kubernetes.mounter-host-root). +mounter: + imageName: texera-mounter + imagePullPolicy: IfNotPresent + port: 8100 + hostMountRoot: /var/lib/texera-mounts + serviceAccountName: texera-mounter-service-account + texeraEnvVars: - name: USER_SYS_ADMIN_USERNAME value: "texera" diff --git a/bin/mounter/mounter.py b/bin/mounter/mounter.py new file mode 100644 index 00000000000..88a6c8fc4f9 --- /dev/null +++ b/bin/mounter/mounter.py @@ -0,0 +1,442 @@ +#!/usr/bin/env python3 +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +texera-mounter: a per-node privileged service that performs GeeseFS FUSE mounts on behalf +of (unprivileged) computing-unit pods. + +A CU pod POSTs /mount with {cuid, repositoryName, commitHash, jwt, fileServiceBase}. The +mounter runs GeeseFS against file-service's JWT-authenticated S3 proxy (passing the pod's +JWT as the S3 access key) and mounts read-only under MOUNT_ROOT///. +That host directory is bind-mounted (mountPropagation: Bidirectional) into the mounter and +propagates back into the CU pod (mountPropagation: HostToContainer), so the CU pod sees +the mount without any privilege of its own. + +The mounter holds no LakeFS credentials; authorization stays entirely in file-service. +A background watcher unmounts a CU's directories as soon as its pod is deleted. +""" + +import json +import os +import shutil +import ssl +import subprocess +import threading +import time +import urllib.request +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import parse_qs, urlparse + +MOUNT_ROOT = os.environ.get("MOUNT_ROOT", "/var/lib/texera-mounts") +MOUNTER_PORT = int(os.environ.get("MOUNTER_PORT", "8100")) +POOL_NAMESPACE = os.environ.get("POOL_NAMESPACE", "texera-workflow-computing-unit-pool") +MOUNT_SECRET_PLACEHOLDER = "texera-jwt-mount" +MOUNT_TIMEOUT_S = 30 +# How long the API server keeps a watch open. Each expiry re-runs the reconcile, which is +# the safety net for events missed while disconnected and for orphans still busy earlier. +WATCH_TIMEOUT_S = 300 +WATCH_RETRY_S = 10 + +SA_DIR = "/var/run/secrets/kubernetes.io/serviceaccount" +# Computing-unit pods are named "-"; the helm chart gives this and the +# computing-unit manager's kubernetes.compute-unit-pod-name-prefix the same value. +CU_POD_NAME_PREFIX = os.environ.get("CU_POD_NAME_PREFIX", "computing-unit") +PROC_MOUNTS = "/proc/mounts" + + +def log(msg): + print(f"[mounter] {msg}", flush=True) + + +def _unescape(field): + """Decode the four characters /proc/mounts escapes octally.""" + for code, char in (("\\040", " "), ("\\011", "\t"), ("\\012", "\n"), ("\\134", "\\")): + field = field.replace(code, char) + return field + + +def mount_targets_under(path): + """Mount targets at or under `path`, deepest first. + + Read from /proc/mounts rather than using os.path.ismount(): when a FUSE server dies + (the mounter being restarted kills every GeeseFS it started) its mount entry survives + here, but stat()ing the mount point fails with ENOTCONN, which os.path.ismount() + reports as "not a mount point". Such a dead mount would then never be unmounted, and + its directory could never be removed. + """ + path = os.path.normpath(path) + prefix = path + "/" + targets = [] + try: + with open(PROC_MOUNTS) as mounts: + for line in mounts: + fields = line.split() + if len(fields) < 2: + continue + target = _unescape(fields[1]) + if target == path or target.startswith(prefix): + targets.append(target) + except OSError as e: + log(f"reading {PROC_MOUNTS} failed: {e}") + # Deepest first, so nested mounts are detached before their parents. + return sorted(targets, key=len, reverse=True) + + +def is_mounted(path): + """True if `path` itself is a mount point, alive or dead.""" + return os.path.normpath(path) in mount_targets_under(path) + + +def _responds(path): + """True if `path` can be stat()ed — false for a mount whose FUSE server is gone.""" + try: + os.stat(path) + return True + except OSError: + return False + + +def _remove_empty_dirs(path, cuid): + """Remove `path` and any parents left empty, stopping at the CU's own directory.""" + stop_at = os.path.normpath(os.path.join(MOUNT_ROOT, cuid)) + path = os.path.normpath(path) + while path.startswith(stop_at): + try: + os.rmdir(path) + except OSError: + return # not empty (another commit is mounted here) or already gone + path = os.path.dirname(path) + + +def ensure_shared_root(): + """Make MOUNT_ROOT a shared mount so mounts created under it propagate to peers.""" + os.makedirs(MOUNT_ROOT, exist_ok=True) + if not os.path.ismount(MOUNT_ROOT): + subprocess.run(["mount", "--bind", MOUNT_ROOT, MOUNT_ROOT], check=False) + subprocess.run(["mount", "--make-rshared", MOUNT_ROOT], check=False) + + +def do_mount(cuid, repo, commit, jwt, file_service_base): + """Idempotently mount repo:commit for cuid. Returns the mount target path.""" + if not cuid or not repo or not commit or not jwt or not file_service_base: + raise ValueError("cuid, repositoryName, commitHash, jwt and fileServiceBase are required") + + target = os.path.join(MOUNT_ROOT, cuid, repo, commit) + if is_mounted(target): + if _responds(target): + log(f"{repo}:{commit} already mounted for cu {cuid} at {target}") + return target + # The GeeseFS process backing this mount is gone — most likely the mounter was + # restarted, which kills every GeeseFS it started. The mount point survives but + # every access to it fails with ENOTCONN, so detach it and mount again. + log(f"{repo}:{commit} for cu {cuid} has a dead mount at {target}; remounting") + subprocess.run(["umount", "-l", target], capture_output=True, text=True) + + os.makedirs(target, exist_ok=True) + # allow_other: the mounter runs as root but the CU pod's UDF runs as a different + # (non-root) user, so the propagated FUSE mount must permit other users to access it. + cmd = [ + "geesefs", + "--endpoint", file_service_base, + "--memory-limit", "512", + "-o", "ro,allow_other", + f"{repo}:{commit}", + target, + ] + env = dict(os.environ) + env["AWS_ACCESS_KEY_ID"] = jwt + env["AWS_SECRET_ACCESS_KEY"] = MOUNT_SECRET_PLACEHOLDER + log(f"mounting {repo}:{commit} for cu {cuid} via: geesefs --endpoint {file_service_base} ... {target}") + result = subprocess.run(cmd, env=env, capture_output=True, text=True) + if result.returncode != 0: + # Nothing was mounted, so drop the directory just created for it rather than + # leaving an empty one behind until the next resync. A rejected mount (an + # unreadable repository, say) is a normal outcome, not a reason to litter. + _remove_empty_dirs(target, cuid) + raise RuntimeError( + f"geesefs exited {result.returncode}: {(result.stdout + result.stderr).strip()}" + ) + + # GeeseFS daemonizes after a successful mount; wait until the kernel reports it. + deadline = time.time() + MOUNT_TIMEOUT_S + while not is_mounted(target): + if time.time() > deadline: + raise RuntimeError(f"{repo}:{commit} did not appear as a mount at {target} in {MOUNT_TIMEOUT_S}s") + time.sleep(0.2) + log(f"mounted {repo}:{commit} for cu {cuid} at {target}") + return target + + +def list_mounts(cuid): + """List the datasets currently mounted for a CU as {repositoryName, commitHash, mountPath}. + + Derived entirely from /proc/mounts (via mount_targets_under) so the mounter stays + stateless: a mount at MOUNT_ROOT/// encodes its own identity. + """ + if not cuid: + raise ValueError("cuid is required") + cu_dir = os.path.normpath(os.path.join(MOUNT_ROOT, cuid)) + mounts = [] + for target in mount_targets_under(cu_dir): + if os.path.normpath(target) == cu_dir: + continue # the shared-root bind itself, never a dataset mount + parts = os.path.relpath(target, cu_dir).split(os.sep) + if len(parts) != 2: + continue # not a / mount point + repo, commit = parts + mounts.append({"repositoryName": repo, "commitHash": commit, "mountPath": target}) + return mounts + + +def do_unmount(cuid, repo, commit): + """Idempotently unmount a single dataset mounted for a CU (a user-initiated unmount, + as opposed to the watcher tearing down a whole departed CU).""" + if not cuid or not repo or not commit: + raise ValueError("cuid, repositoryName and commitHash are required") + + target = os.path.normpath(os.path.join(MOUNT_ROOT, cuid, repo, commit)) + if is_mounted(target): + # The mounter is root, so a plain lazy umount works (no setuid fusermount needed). + result = subprocess.run(["umount", "-l", target], capture_output=True, text=True) + if result.returncode != 0 and is_mounted(target): + raise RuntimeError( + f"umount {target} failed: {(result.stdout + result.stderr).strip()}" + ) + # Drop the now-empty directory (and any parents left empty up to the CU's own dir). + _remove_empty_dirs(target, cuid) + log(f"unmounted {repo}:{commit} for cu {cuid} at {target}") + + +class Handler(BaseHTTPRequestHandler): + def _send(self, code, obj): + body = json.dumps(obj).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + parsed = urlparse(self.path) + if parsed.path == "/healthz": + self._send(200, {"status": "ok"}) + elif parsed.path == "/mounts": + try: + cuid = parse_qs(parsed.query).get("cuid", [""])[0] + self._send(200, {"mounts": list_mounts(cuid)}) + except ValueError as e: + self._send(400, {"error": str(e)}) + except Exception as e: # noqa: BLE001 + log(f"listing mounts failed: {e}") + self._send(500, {"error": str(e)}) + else: + self._send(404, {"error": "not found"}) + + def do_DELETE(self): + # Params come in the query string (not a body): HttpURLConnection, which the + # computing-unit service uses to call this, cannot send a body with DELETE. + parsed = urlparse(self.path) + if parsed.path != "/mount": + self._send(404, {"error": "not found"}) + return + try: + params = parse_qs(parsed.query) + do_unmount( + params.get("cuid", [""])[0], + params.get("repositoryName", [""])[0], + params.get("commitHash", [""])[0], + ) + self._send(200, {"status": "ok"}) + except ValueError as e: + self._send(400, {"error": str(e)}) + except Exception as e: # noqa: BLE001 + log(f"unmount failed: {e}") + self._send(500, {"error": str(e)}) + + def do_POST(self): + if self.path != "/mount": + self._send(404, {"error": "not found"}) + return + try: + length = int(self.headers.get("Content-Length", "0")) + req = json.loads(self.rfile.read(length) or b"{}") + target = do_mount( + str(req.get("cuid", "")), + str(req.get("repositoryName", "")), + str(req.get("commitHash", "")), + str(req.get("jwt", "")), + str(req.get("fileServiceBase", "")), + ) + self._send(200, {"mountPath": target}) + except ValueError as e: + self._send(400, {"error": str(e)}) + except Exception as e: # noqa: BLE001 + log(f"mount failed: {e}") + self._send(500, {"error": str(e)}) + + def log_message(self, fmt, *args): # silence default per-request stderr logging + pass + + +# ---- watcher: unmount a CU's directories once its pod is deleted ---- + +def _k8s_open(path, timeout): + """GET against the in-cluster API server. Returns the response; the caller must close it.""" + with open(os.path.join(SA_DIR, "token")) as f: + token = f.read().strip() + host = os.environ.get("KUBERNETES_SERVICE_HOST", "kubernetes.default.svc") + port = os.environ.get("KUBERNETES_SERVICE_PORT", "443") + ctx = ssl.create_default_context(cafile=os.path.join(SA_DIR, "ca.crt")) + req = urllib.request.Request( + f"https://{host}:{port}{path}", headers={"Authorization": f"Bearer {token}"} + ) + return urllib.request.urlopen(req, timeout=timeout, context=ctx) + + +def _cuid_of(pod_name): + """The cuid a CU pod name belongs to, or None if it is not a CU pod.""" + prefix = CU_POD_NAME_PREFIX + "-" + if not pod_name.startswith(prefix): + return None + return pod_name[len(prefix):] or None + + +def _list_cu_pods(): + """(live cuids, resourceVersion) for the pool namespace, or (None, None) if unreachable.""" + try: + with _k8s_open(f"/api/v1/namespaces/{POOL_NAMESPACE}/pods", timeout=30) as response: + pods = json.loads(response.read()) + except Exception as e: # noqa: BLE001 + log(f"listing CU pods failed: {e}") + return None, None + live = { + cuid + for cuid in (_cuid_of(p.get("metadata", {}).get("name", "")) for p in pods.get("items", [])) + if cuid + } + return live, pods.get("metadata", {}).get("resourceVersion") + + +def clean_cu_dir(cuid, quiet=False): + """Unmount everything under a departed CU's directory and remove it. + + Returns True once the directory is gone. `quiet` suppresses the messages for an orphan + already reported, so a repeatedly retried one does not re-log on every resync. + """ + cu_dir = os.path.join(MOUNT_ROOT, cuid) + # Never treat the shared root itself as a CU directory: unmounting it would break + # propagation for every CU on this node. + if os.path.normpath(cu_dir) == os.path.normpath(MOUNT_ROOT) or not os.path.isdir(cu_dir): + return True + if not quiet: + log(f"cu {cuid} pod is gone; unmounting {cu_dir}") + + for target in mount_targets_under(cu_dir): + # The mounter is root, so a plain lazy umount works (no setuid fusermount needed). + result = subprocess.run(["umount", "-l", target], capture_output=True, text=True) + if result.returncode != 0 and not quiet: + log(f"cu {cuid}: umount -l {target} failed: {(result.stdout + result.stderr).strip()}") + + # A lazy umount only detaches once the last reference to the mount goes away, so a mount + # another namespace still holds can outlive this call. Removing the directory would then + # fail, so only remove it once nothing is mounted underneath and let the next resync retry + # the rest — the mounts are read-only, so an orphan lingering a while is harmless. + remaining = mount_targets_under(cu_dir) + if remaining: + if not quiet: + log(f"cu {cuid}: {len(remaining)} mount(s) still busy, leaving {cu_dir} for the next resync") + return False + + shutil.rmtree(cu_dir, ignore_errors=True) + if os.path.exists(cu_dir): + if not quiet: + log(f"cu {cuid}: could not remove {cu_dir}, leaving it for the next resync") + return False + log(f"cu {cuid}: unmounted and removed {cu_dir}") + return True + + +# cuids whose cleanup could not finish, retried on the next resync and logged only once. +_pending = set() + + +def reconcile(live_cuids): + """Clean up every mount directory with no live CU pod behind it.""" + global _pending + if not os.path.isdir(MOUNT_ROOT): + return + still_pending = set() + for cuid in os.listdir(MOUNT_ROOT): + if cuid in live_cuids or not os.path.isdir(os.path.join(MOUNT_ROOT, cuid)): + continue + if not clean_cu_dir(cuid, quiet=cuid in _pending): + still_pending.add(cuid) + _pending = still_pending + + +def _handle_event(line): + try: + event = json.loads(line) + except ValueError: + return + if event.get("type") != "DELETED": + return + cuid = _cuid_of(event.get("object", {}).get("metadata", {}).get("name", "")) + if cuid and not clean_cu_dir(cuid, quiet=cuid in _pending): + _pending.add(cuid) + + +def watch_loop(): + """List-then-watch CU pods, unmounting a CU's mounts as soon as its pod is deleted. + + The initial LIST reconciles whatever was missed while the mounter was down; the WATCH + then reacts to deletions immediately. The API server closes the watch every + WATCH_TIMEOUT_S, and the resulting re-LIST doubles as the safety net for events missed + across a disconnect and for orphans whose unmount could not complete earlier. + """ + if not os.path.exists(os.path.join(SA_DIR, "token")): + log("no service account token; not running in-cluster, pod watch disabled") + return + while True: + live, resource_version = _list_cu_pods() + if live is None: # API unreachable → keep every mount (fail safe) and retry + time.sleep(WATCH_RETRY_S) + continue + reconcile(live) + try: + with _k8s_open( + f"/api/v1/namespaces/{POOL_NAMESPACE}/pods?watch=1" + f"&resourceVersion={resource_version}&timeoutSeconds={WATCH_TIMEOUT_S}", + timeout=WATCH_TIMEOUT_S + 30, + ) as response: + for line in response: + _handle_event(line) + except Exception as e: # noqa: BLE001 + log(f"pod watch failed ({e}); resyncing") + time.sleep(WATCH_RETRY_S) + + +def main(): + ensure_shared_root() + threading.Thread(target=watch_loop, daemon=True).start() + log(f"listening on :{MOUNTER_PORT}, mount root {MOUNT_ROOT}, pool ns {POOL_NAMESPACE}") + ThreadingHTTPServer(("0.0.0.0", MOUNTER_PORT), Handler).serve_forever() + + +if __name__ == "__main__": + main() diff --git a/bin/mounter/tests/conftest.py b/bin/mounter/tests/conftest.py new file mode 100644 index 00000000000..e6525629032 --- /dev/null +++ b/bin/mounter/tests/conftest.py @@ -0,0 +1,111 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""pytest fixtures for the texera-mounter tests. + +The mounter is the entrypoint of its own image rather than a Python package, so we +load it with `importlib.util` the same way `bin/local-dev/tests` loads the TUI. + +The module is loaded fresh per test because the tests mutate module-level state +(`MOUNT_ROOT`, `PROC_MOUNTS`, the pending-cleanup set), and it is pointed at a fake +mount root and a fake /proc/mounts under `tmp_path` so nothing touches the host. Test +helpers are attached to the module object so tests can stay terse: `mounter.logs`, +`mounter.runs`, `mounter.mounts()` and `mounter.set_mounts()`.""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] +MOUNTER_PATH = REPO_ROOT / "bin" / "mounter" / "mounter.py" + + +@pytest.fixture +def mounter(tmp_path, monkeypatch): + spec = importlib.util.spec_from_file_location("texera_mounter", MOUNTER_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules["texera_mounter"] = module + spec.loader.exec_module(module) + + module.MOUNT_ROOT = str(tmp_path / "mounts") + module.PROC_MOUNTS = str(tmp_path / "proc_mounts") + os.makedirs(module.MOUNT_ROOT) + + module.logs = [] + module.log = module.logs.append + + def set_mounts(*targets): + """Rewrite the fake /proc/mounts to hold exactly these FUSE mount targets.""" + with open(module.PROC_MOUNTS, "w") as mounts: + # One unrelated entry, so tests also prove the filtering works. + mounts.write("/dev/sda1 / ext4 rw,relatime 0 0\n") + for target in targets: + escaped = target.replace("\\", "\\134").replace(" ", "\\040") + mounts.write(f"dataset-1 {escaped} fuse.geesefs ro,nosuid,allow_other 0 0\n") + + def mounts(): + """The FUSE targets currently in the fake /proc/mounts.""" + return module.mount_targets_under(module.MOUNT_ROOT) + + module.set_mounts = set_mounts + module.mounts = mounts + set_mounts() + + # Every subprocess call is recorded rather than run. `geesefs` and `umount` also + # update the fake /proc/mounts, so the mount-table logic is exercised for real. + module.runs = [] + module.umount_succeeds = True + module.geesefs_returncode = 0 + # Set False to model a GeeseFS that exits 0 without the mount ever appearing. + module.geesefs_mounts = True + + def fake_run(cmd, **kwargs): + module.runs.append((list(cmd), kwargs)) + returncode = 0 + live = module.mount_targets_under(module.MOUNT_ROOT) + if cmd[0] == "geesefs": + returncode = module.geesefs_returncode + if returncode == 0 and module.geesefs_mounts: + set_mounts(*live, cmd[-1]) + elif cmd[0] == "umount": + if module.umount_succeeds: + set_mounts(*[t for t in live if t != os.path.normpath(cmd[-1])]) + else: + returncode = 32 + return subprocess.CompletedProcess(cmd, returncode, stdout="", stderr="umount: target is busy") + + monkeypatch.setattr(module.subprocess, "run", fake_run) + return module + + +@pytest.fixture +def cu_dir(mounter): + """Create a CU's mount directory tree and return (cuid, cu_dir, mount target).""" + + def make(cuid, repo="dataset-1", commit="abc123"): + target = os.path.join(mounter.MOUNT_ROOT, cuid, repo, commit) + os.makedirs(target, exist_ok=True) + return os.path.join(mounter.MOUNT_ROOT, cuid), target + + return make diff --git a/bin/mounter/tests/test_mounter.py b/bin/mounter/tests/test_mounter.py new file mode 100644 index 00000000000..58f6e8dac1b --- /dev/null +++ b/bin/mounter/tests/test_mounter.py @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Unit tests for the texera-mounter (`bin/mounter/mounter.py`). + +Covers the logic that is easy to get wrong and expensive to debug in a cluster: how +mounts are discovered (from /proc/mounts, so a mount whose FUSE server died is still +seen), how a CU's mounts are torn down when its pod goes away, and how pod-watch +events are interpreted. The mounting itself is GeeseFS's job and is not re-tested +here — only the command we build for it and how we react to it failing.""" + +from __future__ import annotations + +import json +import os + +import pytest + + +# ─────────────────── mount_targets_under() / is_mounted() ─────────────────── + +def test_mount_targets_are_found_under_a_path_deepest_first(mounter, cu_dir): + _, target = cu_dir("7") + parent = os.path.dirname(target) + mounter.set_mounts(parent, target) + + found = mounter.mount_targets_under(os.path.join(mounter.MOUNT_ROOT, "7")) + + # Deepest first, so nested mounts are detached before the ones containing them. + assert found == [target, parent] + + +def test_mount_targets_exclude_unrelated_and_sibling_paths(mounter, cu_dir): + _, seven = cu_dir("7") + _, eight = cu_dir("8") + mounter.set_mounts(seven, eight, "/var/lib/something-else") + + assert mounter.mount_targets_under(os.path.join(mounter.MOUNT_ROOT, "7")) == [seven] + + +def test_mount_targets_do_not_match_a_path_that_is_only_a_string_prefix(mounter): + # "/mnt/7x" must not be treated as living under "/mnt/7". + seven = os.path.join(mounter.MOUNT_ROOT, "7") + mounter.set_mounts(seven + "x") + + assert mounter.mount_targets_under(seven) == [] + + +def test_mount_targets_decode_escaped_characters(mounter): + spaced = os.path.join(mounter.MOUNT_ROOT, "7", "data set", "abc") + mounter.set_mounts(spaced) + + assert mounter.mount_targets_under(mounter.MOUNT_ROOT) == [spaced] + + +def test_mount_targets_survive_an_unreadable_proc_mounts(mounter): + mounter.PROC_MOUNTS = os.path.join(mounter.MOUNT_ROOT, "does-not-exist") + + assert mounter.mount_targets_under(mounter.MOUNT_ROOT) == [] + assert any("failed" in line for line in mounter.logs) + + +def test_is_mounted_matches_the_path_itself_not_its_children(mounter, cu_dir): + _, target = cu_dir("7") + mounter.set_mounts(target) + + assert mounter.is_mounted(target) + assert not mounter.is_mounted(os.path.dirname(target)) + + +def test_a_dead_mount_is_still_reported_as_mounted(mounter, cu_dir): + """The regression this whole module exists for. + + Restarting the mounter kills every GeeseFS it started. The mount entry survives in + /proc/mounts but stat() on it fails, so os.path.ismount() — what this used to + use — reports False, and the mount was never cleaned up.""" + _, target = cu_dir("7") + mounter.set_mounts(target) + + assert not os.path.ismount(target) # what the buggy version asked + assert mounter.is_mounted(target) # what it should have asked + + +# ─────────────────── _cuid_of() ─────────────────── + +@pytest.mark.parametrize( + "pod_name, expected", + [ + ("computing-unit-17", "17"), + ("computing-unit-0", "0"), + ("computing-unit-", None), # no cuid at all + ("texera-file-service-abc", None), + ("", None), + ], +) +def test_cuid_is_parsed_only_from_computing_unit_pod_names(mounter, pod_name, expected): + assert mounter._cuid_of(pod_name) == expected + + +def test_the_pod_name_prefix_is_configurable(mounter, monkeypatch): + """The helm chart gives this and the CU manager the same prefix; honour it.""" + monkeypatch.setattr(mounter, "CU_POD_NAME_PREFIX", "other-deployment-cu") + + assert mounter._cuid_of("other-deployment-cu-17") == "17" + assert mounter._cuid_of("computing-unit-17") is None + + +# ─────────────────── do_mount() ─────────────────── + +@pytest.mark.parametrize("missing", ["cuid", "repo", "commit", "jwt", "base"]) +def test_do_mount_rejects_incomplete_requests(mounter, missing): + args = {"cuid": "7", "repo": "dataset-1", "commit": "abc", "jwt": "t", "base": "http://fs"} + args[missing] = "" + + with pytest.raises(ValueError): + mounter.do_mount(args["cuid"], args["repo"], args["commit"], args["jwt"], args["base"]) + + assert mounter.runs == [] + + +def test_do_mount_passes_the_jwt_as_the_s3_access_key(mounter): + mounter.do_mount("7", "dataset-1", "abc", "the-user-jwt", "http://file-service:9092") + + cmd, kwargs = mounter.runs[0] + assert cmd[0] == "geesefs" + assert "--endpoint" in cmd and "http://file-service:9092" in cmd + assert "dataset-1:abc" in cmd + # allow_other: the mounter is root but the UDF is not, so the propagated mount must + # be readable by another uid. ro: mounts are never writable. + assert "ro,allow_other" in cmd + # The pod's own JWT is the credential; no global LakeFS secret ever reaches the node. + assert kwargs["env"]["AWS_ACCESS_KEY_ID"] == "the-user-jwt" + + +def test_do_mount_is_idempotent_for_a_live_mount(mounter, cu_dir): + _, target = cu_dir("7", commit="abc") + mounter.set_mounts(target) + + assert mounter.do_mount("7", "dataset-1", "abc", "jwt", "http://fs") == target + assert mounter.runs == [] # no second GeeseFS for an already-mounted commit + + +def test_do_mount_replaces_a_dead_mount(mounter, cu_dir, monkeypatch): + _, target = cu_dir("7", commit="abc") + mounter.set_mounts(target) + monkeypatch.setattr(mounter, "_responds", lambda path: False) # FUSE server is gone + + assert mounter.do_mount("7", "dataset-1", "abc", "jwt", "http://fs") == target + + commands = [cmd[0] for cmd, _ in mounter.runs] + assert commands == ["umount", "geesefs"] # detached first, then mounted again + assert any("dead mount" in line for line in mounter.logs) + + +def test_do_mount_reports_what_geesefs_printed_when_it_fails(mounter): + mounter.geesefs_returncode = 1 + + with pytest.raises(RuntimeError, match="geesefs exited 1"): + mounter.do_mount("7", "dataset-1", "abc", "jwt", "http://fs") + + +def test_a_rejected_mount_leaves_no_empty_directory_behind(mounter): + """A mount the proxy refuses is routine; it should not litter until the next resync.""" + mounter.geesefs_returncode = 1 + + with pytest.raises(RuntimeError): + mounter.do_mount("7", "dataset-1", "abc", "jwt", "http://fs") + + assert not os.path.exists(os.path.join(mounter.MOUNT_ROOT, "7")) + + +def test_a_rejected_mount_keeps_a_sibling_commit_that_is_mounted(mounter, cu_dir): + _, live = cu_dir("7", commit="already-here") + mounter.set_mounts(live) + mounter.geesefs_returncode = 1 + + with pytest.raises(RuntimeError): + mounter.do_mount("7", "dataset-1", "new-commit", "jwt", "http://fs") + + assert not os.path.exists(os.path.join(mounter.MOUNT_ROOT, "7", "dataset-1", "new-commit")) + assert os.path.isdir(live) # the working mount is untouched + + +def test_do_mount_times_out_if_the_mount_never_appears(mounter, monkeypatch): + monkeypatch.setattr(mounter, "MOUNT_TIMEOUT_S", 0) + mounter.geesefs_mounts = False # exits 0 without the mount ever appearing + + with pytest.raises(RuntimeError, match="did not appear as a mount"): + mounter.do_mount("7", "dataset-1", "abc", "jwt", "http://fs") + + +# ─────────────────── clean_cu_dir() ─────────────────── + +def test_clean_unmounts_then_removes_the_directory(mounter, cu_dir): + directory, target = cu_dir("7") + mounter.set_mounts(target) + + assert mounter.clean_cu_dir("7") is True + assert [cmd for cmd, _ in mounter.runs] == [["umount", "-l", target]] + assert not os.path.exists(directory) + assert any("unmounted and removed" in line for line in mounter.logs) + + +def test_clean_unmounts_nested_mounts_deepest_first(mounter, cu_dir): + _, target = cu_dir("7") + parent = os.path.dirname(target) + mounter.set_mounts(parent, target) + + mounter.clean_cu_dir("7") + + assert [cmd[-1] for cmd, _ in mounter.runs] == [target, parent] + + +def test_clean_removes_a_directory_that_has_no_mounts(mounter, cu_dir): + directory, _ = cu_dir("7") + + assert mounter.clean_cu_dir("7") is True + assert mounter.runs == [] + assert not os.path.exists(directory) + + +def test_clean_keeps_a_directory_whose_mount_is_still_busy(mounter, cu_dir): + directory, target = cu_dir("7") + mounter.set_mounts(target) + mounter.umount_succeeds = False # lazy unmount has not completed yet + + assert mounter.clean_cu_dir("7") is False + # The directory must survive: removing it under a live mount is what the original + # code did, and it also made the removal fail on every cycle forever. + assert os.path.isdir(directory) + assert any("still busy" in line for line in mounter.logs) + + +def test_clean_is_silent_on_a_retry(mounter, cu_dir): + cu_dir("7") + mounter.set_mounts(os.path.join(mounter.MOUNT_ROOT, "7", "dataset-1", "abc123")) + mounter.umount_succeeds = False + + mounter.clean_cu_dir("7") + first_pass = len(mounter.logs) + mounter.clean_cu_dir("7", quiet=True) + + assert len(mounter.logs) == first_pass # a stuck orphan is reported once, not per cycle + + +def test_clean_never_touches_the_shared_root(mounter): + """A cuid of "" would resolve to MOUNT_ROOT, whose unmount would break every CU.""" + mounter.set_mounts(mounter.MOUNT_ROOT) + + assert mounter.clean_cu_dir("") is True + assert mounter.runs == [] + assert os.path.isdir(mounter.MOUNT_ROOT) + + +def test_clean_accepts_a_directory_that_is_already_gone(mounter): + assert mounter.clean_cu_dir("404") is True + + +# ─────────────────── reconcile() ─────────────────── + +def test_reconcile_removes_orphans_and_keeps_live_computing_units(mounter, cu_dir): + orphan, _ = cu_dir("7") + live, _ = cu_dir("8") + + mounter.reconcile({"8"}) + + assert not os.path.exists(orphan) + assert os.path.isdir(live) + + +def test_reconcile_retries_a_stuck_orphan_until_it_can_be_removed(mounter, cu_dir): + directory, target = cu_dir("7") + mounter.set_mounts(target) + mounter.umount_succeeds = False + + mounter.reconcile(set()) + assert mounter._pending == {"7"} + assert os.path.isdir(directory) + + mounter.umount_succeeds = True # the reference finally goes away + mounter.reconcile(set()) + + assert mounter._pending == set() + assert not os.path.exists(directory) + + +def test_reconcile_ignores_stray_files_in_the_mount_root(mounter): + stray = os.path.join(mounter.MOUNT_ROOT, "notes.txt") + open(stray, "w").close() + + mounter.reconcile(set()) + + assert os.path.exists(stray) + + +# ─────────────────── watch events ─────────────────── + +def _event(event_type, pod_name): + return json.dumps({"type": event_type, "object": {"metadata": {"name": pod_name}}}).encode() + + +def test_a_deleted_computing_unit_pod_is_cleaned_up_immediately(mounter, cu_dir): + directory, _ = cu_dir("7") + + mounter._handle_event(_event("DELETED", "computing-unit-7")) + + assert not os.path.exists(directory) + + +@pytest.mark.parametrize( + "line", + [ + _event("ADDED", "computing-unit-7"), + _event("MODIFIED", "computing-unit-7"), + _event("DELETED", "texera-file-service-xyz"), + _event("DELETED", "computing-unit-"), + b"{ not json\n", + b"\n", + ], + ids=["added", "modified", "other-pod", "no-cuid", "malformed", "blank"], +) +def test_irrelevant_watch_events_leave_every_mount_alone(mounter, cu_dir, line): + directory, _ = cu_dir("7") + + mounter._handle_event(line) + + assert os.path.isdir(directory) + assert mounter.runs == [] + + +def test_a_delete_that_cannot_finish_is_retried_on_the_next_resync(mounter, cu_dir): + directory, target = cu_dir("7") + mounter.set_mounts(target) + mounter.umount_succeeds = False + + mounter._handle_event(_event("DELETED", "computing-unit-7")) + + assert os.path.isdir(directory) + assert mounter._pending == {"7"} + + +# ─────────────────── _list_cu_pods() ─────────────────── + +def test_listing_returns_computing_unit_ids_and_the_resource_version(mounter, monkeypatch): + payload = { + "metadata": {"resourceVersion": "4242"}, + "items": [ + {"metadata": {"name": "computing-unit-7"}}, + {"metadata": {"name": "computing-unit-8"}}, + {"metadata": {"name": "some-other-pod"}}, + ], + } + monkeypatch.setattr(mounter, "_k8s_open", lambda path, timeout: _FakeResponse(payload)) + + assert mounter._list_cu_pods() == ({"7", "8"}, "4242") + + +def test_listing_failure_reports_no_answer_rather_than_an_empty_cluster(mounter, monkeypatch): + """A failed LIST must not read as "no CU pods exist" — that would unmount everything.""" + + def explode(path, timeout): + raise OSError("connection refused") + + monkeypatch.setattr(mounter, "_k8s_open", explode) + + live, resource_version = mounter._list_cu_pods() + + assert live is None and resource_version is None + assert any("listing CU pods failed" in line for line in mounter.logs) + + +class _FakeResponse: + def __init__(self, payload): + self._payload = json.dumps(payload).encode() + + def read(self): + return self._payload + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False diff --git a/common/config/src/main/resources/kubernetes.conf b/common/config/src/main/resources/kubernetes.conf index e85924e570c..04cc81624b6 100644 --- a/common/config/src/main/resources/kubernetes.conf +++ b/common/config/src/main/resources/kubernetes.conf @@ -29,6 +29,11 @@ kubernetes { compute-unit-service-name = "workflow-computing-unit-svc" compute-unit-service-name = ${?KUBERNETES_COMPUTE_UNIT_SERVICE_NAME} + # Computing-unit pods are named "-". The mounter derives a CU's id back + # out of the pod name, so both must be given the same value by the helm chart. + compute-unit-pod-name-prefix = "computing-unit" + compute-unit-pod-name-prefix = ${?KUBERNETES_COMPUTE_UNIT_POD_NAME_PREFIX} + image-name = "bobbai/texera-workflow-computing-unit:dev" image-name = ${?KUBERNETES_IMAGE_NAME} @@ -54,4 +59,12 @@ kubernetes { # GPU resource key used in Kubernetes (vendor-specific) computing-unit-gpu-resource-key = "nvidia.com/gpu" computing-unit-gpu-resource-key = ${?KUBERNETES_COMPUTING_UNIT_GPU_RESOURCE_KEY} + + # Per-node mounter (out-of-pod FUSE mounting). These must match the + # texera-mounter DaemonSet's hostPort and hostPath in the helm chart. + mounter-port = 8100 + mounter-port = ${?KUBERNETES_MOUNTER_PORT} + + mounter-host-root = "/var/lib/texera-mounts" + mounter-host-root = ${?KUBERNETES_MOUNTER_HOST_ROOT} } \ No newline at end of file diff --git a/common/config/src/main/scala/org/apache/texera/common/config/EnvironmentalVariable.scala b/common/config/src/main/scala/org/apache/texera/common/config/EnvironmentalVariable.scala index a335ddeff6c..2a6861d8774 100644 --- a/common/config/src/main/scala/org/apache/texera/common/config/EnvironmentalVariable.scala +++ b/common/config/src/main/scala/org/apache/texera/common/config/EnvironmentalVariable.scala @@ -45,6 +45,16 @@ object EnvironmentalVariable { val ENV_USER_JWT_TOKEN = "USER_JWT_TOKEN" val ENV_AUTH_JWT_SECRET = "AUTH_JWT_SECRET" + /** + * Node mounter (out-of-pod FUSE mounting) related vars, injected into the CU pod. + * The CU pod asks the per-node mounter (at NODE_IP:port) to mount, then reads the result + * under the in-pod mount root exposed via mount propagation. + */ + val ENV_NODE_IP = "NODE_IP" + val ENV_MOUNTER_PORT = "TEXERA_MOUNTER_PORT" + val ENV_CU_ID = "TEXERA_CU_ID" + val ENV_MOUNT_IN_POD_ROOT = "TEXERA_MOUNT_IN_POD_ROOT" + // JDBC val ENV_JDBC_URL = "STORAGE_JDBC_URL" val ENV_JDBC_USERNAME = "STORAGE_JDBC_USERNAME" diff --git a/common/config/src/main/scala/org/apache/texera/common/config/KubernetesConfig.scala b/common/config/src/main/scala/org/apache/texera/common/config/KubernetesConfig.scala index f6294767365..ec1e1bb0328 100644 --- a/common/config/src/main/scala/org/apache/texera/common/config/KubernetesConfig.scala +++ b/common/config/src/main/scala/org/apache/texera/common/config/KubernetesConfig.scala @@ -30,6 +30,7 @@ object KubernetesConfig { val computeUnitServiceName: String = conf.getString("kubernetes.compute-unit-service-name") val computeUnitPoolName: String = conf.getString("kubernetes.compute-unit-pool-name") val computeUnitPoolNamespace: String = conf.getString("kubernetes.compute-unit-pool-namespace") + val computeUnitPodNamePrefix: String = conf.getString("kubernetes.compute-unit-pod-name-prefix") val computeUnitImageName: String = conf.getString("kubernetes.image-name") val computingUnitImagePullPolicy: String = conf.getString("kubernetes.image-pull-policy") @@ -64,4 +65,8 @@ object KubernetesConfig { // GPU resource key used directly in Kubernetes resource specifications val gpuResourceKey: String = conf.getString("kubernetes.computing-unit-gpu-resource-key") + + // Per-node mounter that performs the FUSE mount outside the (unprivileged) CU pod. + val mounterPort: Int = conf.getInt("kubernetes.mounter-port") + val mounterHostRoot: String = conf.getString("kubernetes.mounter-host-root") } diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala index 4f1d391cb30..11adffc412f 100644 --- a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala @@ -22,7 +22,7 @@ package org.apache.texera.service.util import io.fabric8.kubernetes.api.model._ import io.fabric8.kubernetes.api.model.metrics.v1beta1.PodMetricsList import io.fabric8.kubernetes.client.KubernetesClientBuilder -import org.apache.texera.common.config.KubernetesConfig +import org.apache.texera.common.config.{EnvironmentalVariable, KubernetesConfig} import scala.jdk.CollectionConverters._ @@ -32,7 +32,7 @@ object KubernetesClient { private val client: io.fabric8.kubernetes.client.KubernetesClient = new KubernetesClientBuilder().build() private val namespace: String = KubernetesConfig.computeUnitPoolNamespace - private val podNamePrefix = "computing-unit" + private val podNamePrefix = KubernetesConfig.computeUnitPodNamePrefix def generatePodURI(cuid: Int): String = { s"${generatePodName(cuid)}.${KubernetesConfig.computeUnitServiceName}.$namespace.svc.cluster.local:${KubernetesConfig.computeUnitPortNumber}" @@ -92,16 +92,38 @@ object KubernetesClient { throw new Exception(s"Pod with cuid $cuid already exists") } - val envList = envVars - .map { - case (key, value) => - new EnvVarBuilder() - .withName(key) - .withValue(value.toString) - .build() - } - .toList - .asJava + val baseEnv = envVars.map { + case (key, value) => + new EnvVarBuilder().withName(key).withValue(value.toString).build() + }.toList + + // Env for the out-of-pod mounter: the node IP (downward API) plus this CU's id, + // the mounter port, and the in-pod root that the propagated mount appears under. + val inPodMountRoot = "/mnt/texera-mounts" + val mounterEnv = List( + new EnvVarBuilder() + .withName(EnvironmentalVariable.ENV_NODE_IP) + .withNewValueFrom() + .withNewFieldRef() + .withFieldPath("status.hostIP") + .endFieldRef() + .endValueFrom() + .build(), + new EnvVarBuilder() + .withName(EnvironmentalVariable.ENV_CU_ID) + .withValue(cuid.toString) + .build(), + new EnvVarBuilder() + .withName(EnvironmentalVariable.ENV_MOUNTER_PORT) + .withValue(KubernetesConfig.mounterPort.toString) + .build(), + new EnvVarBuilder() + .withName(EnvironmentalVariable.ENV_MOUNT_IN_POD_ROOT) + .withValue(inPodMountRoot) + .build() + ) + + val envList = (baseEnv ++ mounterEnv).asJava // Setup the resource requirements val resourceBuilder = new ResourceRequirementsBuilder() @@ -144,6 +166,16 @@ object KubernetesClient { .withEnv(envList) .withResources(resourceBuilder.build()) + // The FUSE mount is performed by the per-node texera-mounter (privileged), not here, + // so this pod stays UNPRIVILEGED. It only *receives* the mount via HostToContainer + // propagation from a host directory scoped to this CU id (see the hostPath volume below). + containerBuilder + .addNewVolumeMount() + .withName("texera-mounts") + .withMountPath(inPodMountRoot) + .withMountPropagation("HostToContainer") + .endVolumeMount() + // If shmSize requested, mount /dev/shm shmSize.foreach { _ => containerBuilder @@ -169,6 +201,17 @@ object KubernetesClient { .endVolume() } + // Per-CU host directory the mounter mounts into (DirectoryOrCreate so it exists + // before the mounter mounts). Scoped by cuid so a CU can only ever see its own mounts. + specBuilder + .addNewVolume() + .withName("texera-mounts") + .withNewHostPath() + .withPath(s"${KubernetesConfig.mounterHostRoot}/$cuid") + .withType("DirectoryOrCreate") + .endHostPath() + .endVolume() + val pod = specBuilder .withHostname(podName) .withSubdomain(KubernetesConfig.computeUnitServiceName) diff --git a/file-service/src/main/scala/org/apache/texera/service/FileService.scala b/file-service/src/main/scala/org/apache/texera/service/FileService.scala index d8e3ff122d5..ff112a0bf94 100644 --- a/file-service/src/main/scala/org/apache/texera/service/FileService.scala +++ b/file-service/src/main/scala/org/apache/texera/service/FileService.scala @@ -37,6 +37,7 @@ import org.apache.texera.service.resource.{ HealthCheckResource } import org.apache.texera.service.util.S3StorageClient +import org.apache.texera.service.util.S3ProxyServlet import org.apache.texera.service.util.LargeBinaryManager import org.apache.texera.service.util.StagedFileCleanupJob import org.eclipse.jetty.server.session.SessionHandler @@ -90,6 +91,12 @@ class FileService extends Application[FileServiceConfiguration] with LazyLogging environment.jersey.register(classOf[DatasetResource]) environment.jersey.register(classOf[DatasetAccessResource]) + // Register the read-only S3 proxy servlet for in-pod GeeseFS dataset mounts. GeeseFS + // authenticates with the pod's per-user JWT (carried as its S3 access key) and issues + // path-style requests at the root (//), while Jersey serves the REST API + // at /api/* (more specific, so it keeps taking precedence). + environment.servlets.addServlet("s3-mount-proxy", new S3ProxyServlet).addMapping("/*") + RoleAnnotationEnforcer.enforce(environment.jersey.getResourceConfig, "FileService") // Route request logs through SLF4J, controlled by TEXERA_SERVICE_LOG_LEVEL diff --git a/file-service/src/main/scala/org/apache/texera/service/util/S3ProxyServlet.scala b/file-service/src/main/scala/org/apache/texera/service/util/S3ProxyServlet.scala new file mode 100644 index 00000000000..85b8d7acd38 --- /dev/null +++ b/file-service/src/main/scala/org/apache/texera/service/util/S3ProxyServlet.scala @@ -0,0 +1,260 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.service.util + +import com.typesafe.scalalogging.LazyLogging +import jakarta.servlet.http.{HttpServlet, HttpServletRequest, HttpServletResponse} +import org.apache.texera.auth.JwtParser +import org.apache.texera.common.config.StorageConfig +import org.apache.texera.dao.SqlServer +import org.apache.texera.dao.SqlServer.withTransaction +import org.apache.texera.dao.jooq.generated.tables.daos.DatasetDao +import org.apache.texera.service.resource.DatasetAccessResource.userHasReadAccess +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials +import software.amazon.awssdk.auth.signer.AwsS3V4Signer +import software.amazon.awssdk.auth.signer.params.AwsS3V4SignerParams +import software.amazon.awssdk.http.{SdkHttpFullRequest, SdkHttpMethod} +import software.amazon.awssdk.regions.Region + +import java.net.{URI, URLDecoder} +import java.net.http.{HttpClient, HttpRequest, HttpResponse} +import java.time.Duration +import java.util.concurrent.ConcurrentHashMap +import scala.jdk.CollectionConverters._ +import scala.jdk.OptionConverters._ + +/** + * Read-only, JWT-authenticated, re-signing reverse proxy in front of the LakeFS S3 + * gateway. A computing-unit pod's GeeseFS mount talks to this servlet using the pod's + * own per-user JWT as the S3 credential: the JWT is passed to GeeseFS as + * `AWS_ACCESS_KEY_ID`, so it rides in the request's SigV4/SigV2 `Authorization` header. + * Reusing the JWT that is already present in the pod means no separate mount credential + * is ever issued, stored, or made multi-replica-consistent. The servlet: + * + * 1. reads the JWT back out of the incoming `Authorization` header (the JWT is the + * bearer capability; the pod-side S3 signature is not re-validated, and no LakeFS + * credentials ever leave this service), + * 2. verifies the JWT and checks that its user has read access to the requested + * repository (the S3 bucket), using the same `userHasReadAccess` gate as the + * dataset REST endpoints, and + * 3. re-signs the request with the global LakeFS credentials and forwards it to the + * LakeFS S3 gateway, streaming the response back. + * + * Because requests are forwarded verbatim, the proxy behaves identically to a direct + * GeeseFS -> LakeFS-gateway mount, just re-authenticated. GeeseFS mounts read-only, so + * only GET and HEAD are handled. + */ +class S3ProxyServlet extends HttpServlet with LazyLogging { + + // The LakeFS S3 gateway shares the LakeFS server address: the configured API endpoint + // with the trailing /api/v1 suffix removed. + private val gatewayEndpoint: URI = + URI.create(StorageConfig.lakefsEndpoint.stripSuffix("/").stripSuffix("/api/v1")) + + private val lakefsCredentials = + AwsBasicCredentials.create(StorageConfig.lakefsUsername, StorageConfig.lakefsPassword) + + // The S3-specific SigV4 signer adds and signs the x-amz-content-sha256 header (which + // S3 / the LakeFS gateway require) and disables path double-encoding — both essential + // for the signature to validate. The generic Aws4Signer omits x-amz-content-sha256. + private val signer = AwsS3V4Signer.create() + + private val httpClient: HttpClient = HttpClient + .newBuilder() + .followRedirects(HttpClient.Redirect.NEVER) + .connectTimeout(Duration.ofSeconds(10)) + .build() + + private val forwardedResponseHeaderPrefixes = + Seq("content-", "etag", "last-modified", "accept-ranges", "x-amz-") + + // Short-lived cache of (uid, repository) -> read-access decision. A mount issues many + // range reads for the same repository, so this avoids a DB round-trip per request. It + // is a pure optimization: each replica caches independently and a miss just re-queries, + // so unlike a shared session store it needs no cross-replica consistency. + private val AuthCacheTtlMs = 60000L + private case class CachedDecision(allowed: Boolean, expiresAtMs: Long) + private val authCache = new ConcurrentHashMap[String, CachedDecision]() + + override def doGet(req: HttpServletRequest, resp: HttpServletResponse): Unit = + proxy(req, resp, SdkHttpMethod.GET, streamBody = true) + + override def doHead(req: HttpServletRequest, resp: HttpServletResponse): Unit = + proxy(req, resp, SdkHttpMethod.HEAD, streamBody = false) + + private def proxy( + req: HttpServletRequest, + resp: HttpServletResponse, + method: SdkHttpMethod, + streamBody: Boolean + ): Unit = { + val user = S3ProxyServlet + .extractCredentialToken(req.getHeader("Authorization")) + .flatMap(token => JwtParser.parseToken(token).toScala) + if (user.isEmpty) { + // GeeseFS probes the bucket unauthenticated on mount, so this is expected noise. + resp.sendError(HttpServletResponse.SC_FORBIDDEN, "missing or invalid user token") + return + } + + val uid = user.get.getUid + val repositoryName = S3ProxyServlet.bucketFromUri(req.getRequestURI) + if (repositoryName.isEmpty || !authorizedToRead(uid, repositoryName)) { + logger.warn( + s"user $uid denied mount access to repository '$repositoryName' for ${req.getRequestURI}" + ) + resp.sendError(HttpServletResponse.SC_FORBIDDEN, "no read access to the requested repository") + return + } + + try { + writeResponse(forward(req, method), resp, streamBody) + } catch { + case e: Exception => + logger.error(s"error proxying ${req.getRequestURI} to LakeFS gateway", e) + resp.sendError(HttpServletResponse.SC_BAD_GATEWAY, "upstream error") + } + } + + /** + * True iff `uid` has read access to the dataset backing `repositoryName`, cached for a + * short window. Read access to a repository grants read to all of its commits, so no + * per-commit check is needed: a session addresses a single repository's data and any + * version the user may already read. + */ + private def authorizedToRead(uid: Integer, repositoryName: String): Boolean = { + val now = System.currentTimeMillis() + val cacheKey = s"$uid:$repositoryName" + val cached = authCache.get(cacheKey) + if (cached != null && cached.expiresAtMs > now) { + return cached.allowed + } + + val allowed = withTransaction(SqlServer.getInstance().createDSLContext()) { ctx => + val datasets = new DatasetDao(ctx.configuration()) + .fetchByRepositoryName(repositoryName) + .asScala + .toList + datasets.nonEmpty && userHasReadAccess(ctx, datasets.head.getDid, uid) + } + authCache.put(cacheKey, CachedDecision(allowed, now + AuthCacheTtlMs)) + allowed + } + + /** Re-sign the request with LakeFS credentials and send it to the LakeFS gateway. */ + private def forward( + req: HttpServletRequest, + method: SdkHttpMethod + ): HttpResponse[java.io.InputStream] = { + val query = Option(req.getQueryString).map("?" + _).getOrElse("") + val targetUri = URI.create( + gatewayEndpoint.getScheme + "://" + gatewayEndpoint.getAuthority + req.getRequestURI + query + ) + + val builder = SdkHttpFullRequest + .builder() + .method(method) + .uri(targetUri) + // Range must be part of the signed request so the gateway accepts it. + Option(req.getHeader("Range")).foreach(r => builder.putHeader("Range", r)) + + val signed = signer.sign( + builder.build(), + AwsS3V4SignerParams + .builder() + .awsCredentials(lakefsCredentials) + .signingName("s3") + .signingRegion(Region.US_EAST_1) + .build() + ) + + var outbound = HttpRequest + .newBuilder() + .uri(targetUri) + .method(method.name(), HttpRequest.BodyPublishers.noBody()) + // Forward every signed header except Host, which the HTTP client sets itself to the + // same value we signed (the gateway authority). + signed.headers().asScala.foreach { + case (name, values) => + if (!name.equalsIgnoreCase("Host")) { + values.asScala.foreach(v => outbound = outbound.header(name, v)) + } + } + + httpClient.send(outbound.build(), HttpResponse.BodyHandlers.ofInputStream()) + } + + private def writeResponse( + upstream: HttpResponse[java.io.InputStream], + resp: HttpServletResponse, + streamBody: Boolean + ): Unit = { + resp.setStatus(upstream.statusCode()) + upstream.headers().map().asScala.foreach { + case (name, values) => + val lower = name.toLowerCase + if (forwardedResponseHeaderPrefixes.exists(lower.startsWith)) { + values.asScala.foreach(v => resp.addHeader(name, v)) + } + } + + if (streamBody) { + val in = upstream.body() + try { + in.transferTo(resp.getOutputStream) + } finally { + in.close() + } + } else { + upstream.body().close() + } + } +} + +object S3ProxyServlet { + + /** + * Extract the credential token from an AWS `Authorization` header. GeeseFS carries the + * user JWT in the access-key-id position and signs with either SigV4 + * (`AWS4-HMAC-SHA256 Credential=//...`) or, against a plain-HTTP custom + * endpoint, SigV2 (`AWS :`); support both. A JWT is base64url with + * `.` separators, so it never contains the `/` or `:` these formats delimit on. The + * pod-side S3 signature itself is not re-validated — the JWT is the bearer capability — + * so only the token needs to be read out. + */ + private[util] def extractCredentialToken(authHeader: String): Option[String] = { + Option(authHeader).flatMap { h => + "Credential=([^/,\\s]+)/".r + .findFirstMatchIn(h) + .map(_.group(1)) // SigV4 + .orElse("^AWS ([^:\\s]+):".r.findFirstMatchIn(h.trim).map(_.group(1))) // SigV2 + } + } + + /** + * The repository (S3 bucket) a path-style request URI `//` targets: its + * first path segment, URL-decoded. Empty when the URI carries no bucket (root, or a + * service-level list-buckets), which is never authorized. + */ + private[util] def bucketFromUri(requestUri: String): String = { + val firstSegment = requestUri.stripPrefix("/").split("/", 2)(0) + if (firstSegment.isEmpty) "" else URLDecoder.decode(firstSegment, "UTF-8") + } +} diff --git a/file-service/src/test/scala/org/apache/texera/service/util/S3ProxyServletSpec.scala b/file-service/src/test/scala/org/apache/texera/service/util/S3ProxyServletSpec.scala new file mode 100644 index 00000000000..b7024778a6d --- /dev/null +++ b/file-service/src/test/scala/org/apache/texera/service/util/S3ProxyServletSpec.scala @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.service.util + +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +/** + * Unit tests for the pure request-parsing helpers of [[S3ProxyServlet]] — the credential + * token extraction (which must handle both the SigV4 and SigV2 `Authorization` formats + * GeeseFS emits, carrying the user JWT in the access-key position) and the path-style + * bucket/repository extraction that scopes each request. + */ +class S3ProxyServletSpec extends AnyFlatSpec with Matchers { + + // A representative JWT: base64url segments (A-Za-z0-9-_) joined by '.', so it contains + // none of the '/', ':' or whitespace that the two Authorization formats delimit on. + private val jwt = + "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ0ZXhlcmEiLCJ1c2VySWQiOjF9.abc-_DEF123" + + "extractCredentialToken" should "read the access key from a SigV4 Authorization header" in { + val header = + s"AWS4-HMAC-SHA256 Credential=$jwt/20260721/us-east-1/s3/aws4_request, " + + "SignedHeaders=host;x-amz-content-sha256;x-amz-date, Signature=deadbeef" + S3ProxyServlet.extractCredentialToken(header) shouldBe Some(jwt) + } + + it should "read the access key from a SigV2 Authorization header" in { + S3ProxyServlet.extractCredentialToken(s"AWS $jwt:c2lnbmF0dXJl") shouldBe Some(jwt) + } + + it should "handle a short (non-JWT) access key in both formats" in { + S3ProxyServlet.extractCredentialToken( + "AWS4-HMAC-SHA256 Credential=AKIAEXAMPLE/20260721/us-east-1/s3/aws4_request, " + + "SignedHeaders=host, Signature=abc" + ) shouldBe Some("AKIAEXAMPLE") + S3ProxyServlet.extractCredentialToken("AWS AKIAEXAMPLE:sig") shouldBe Some("AKIAEXAMPLE") + } + + it should "return None for a null header" in { + S3ProxyServlet.extractCredentialToken(null) shouldBe None + } + + it should "return None for a header in neither AWS format" in { + S3ProxyServlet.extractCredentialToken("Bearer some.jwt.token") shouldBe None + S3ProxyServlet.extractCredentialToken("") shouldBe None + S3ProxyServlet.extractCredentialToken("AWS4-HMAC-SHA256 SignedHeaders=host") shouldBe None + } + + "bucketFromUri" should "return the first path segment of a path-style object request" in { + S3ProxyServlet.bucketFromUri( + "/dataset-1/097b4111e0ac9f46/model-00001-of-00003.pt" + ) shouldBe "dataset-1" + } + + it should "return the bucket for a bucket-only request (with or without trailing slash)" in { + S3ProxyServlet.bucketFromUri("/dataset-1") shouldBe "dataset-1" + S3ProxyServlet.bucketFromUri("/dataset-1/") shouldBe "dataset-1" + } + + it should "URL-decode the bucket segment" in { + S3ProxyServlet.bucketFromUri("/my%20dataset/commit/f.txt") shouldBe "my dataset" + } + + it should "return empty for the root URI (no bucket to authorize)" in { + S3ProxyServlet.bucketFromUri("/") shouldBe "" + S3ProxyServlet.bucketFromUri("") shouldBe "" + } +} From c16ed4429da02c5112d0ecfc6da3ab7ae5ffede6 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 24 Jul 2026 14:04:57 -0700 Subject: [PATCH 2/3] feat(dataset-mount): drive dataset mounts from the platform Build on the mount infrastructure to let the platform and users mount datasets and use them in Python UDFs. - Engine mount client: DatasetMountManager asks the node mounter to mount a dataset for the computing unit; the locator flows through PhysicalOp and WorkerConfig, and the resolved local path is handed to the Python worker. - Per-computing-unit mount API: mount / list / unmount datasets on a CU (ComputingUnitManagingResource + MounterClient), FileResolver resolving dataset paths both ways. - "Mount datasets into computing unit" UI. - Python UDF dataset-variable bindings: bind each mounted dataset to a variable holding its local path. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01H27sHgS29u2JcLYPPR6Hdx --- .../architecture/managers/executor_manager.py | 11 ++ .../main/python/texera_run_python_worker.py | 9 +- .../pythonworker/PythonWorkflowWorker.scala | 19 +- .../scheduling/config/WorkerConfig.scala | 9 +- .../engine/common/DatasetMountManager.scala | 185 ++++++++++++++++++ .../src/test/python/test_run_python_worker.py | 24 +++ ...ythonWorkflowWorkerStartupConfigSpec.scala | 20 ++ .../amber/core/storage/FileResolver.scala | 120 +++++++++--- .../amber/core/workflow/PhysicalOp.scala | 10 +- .../udf/python/DatasetVariableMapping.scala | 42 ++++ .../udf/python/PythonUDFOpDescV2.scala | 46 ++++- .../ComputingUnitManagingResource.scala | 155 ++++++++++++++- .../texera/service/util/MounterClient.scala | 122 ++++++++++++ .../src/app/common/formly/formly-config.ts | 2 + ...orkflow-computing-unit-managing.service.ts | 45 +++++ .../computing-unit-mount-modal.component.html | 70 +++++++ .../computing-unit-mount-modal.component.scss | 74 +++++++ .../computing-unit-mount-modal.component.ts | 168 ++++++++++++++++ .../dataset-variables-editor.component.html | 82 ++++++++ .../dataset-variables-editor.component.scss | 59 ++++++ .../dataset-variables-editor.component.ts | 152 ++++++++++++++ .../computing-unit-selection.component.html | 10 + .../computing-unit-selection.component.ts | 15 ++ .../operator-property-edit-frame.component.ts | 12 ++ 24 files changed, 1425 insertions(+), 36 deletions(-) create mode 100644 amber/src/main/scala/org/apache/texera/amber/engine/common/DatasetMountManager.scala create mode 100644 common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMapping.scala create mode 100644 computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/MounterClient.scala create mode 100644 frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.html create mode 100644 frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.scss create mode 100644 frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.ts create mode 100644 frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.html create mode 100644 frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.scss create mode 100644 frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.ts diff --git a/amber/src/main/python/core/architecture/managers/executor_manager.py b/amber/src/main/python/core/architecture/managers/executor_manager.py index 7f2249476ad..1e432f0b462 100644 --- a/amber/src/main/python/core/architecture/managers/executor_manager.py +++ b/amber/src/main/python/core/architecture/managers/executor_manager.py @@ -19,6 +19,8 @@ import importlib import inspect import itertools +import json +import os import sys from cached_property import cached_property from fs.base import FS @@ -105,6 +107,15 @@ def load_executor_definition(self, code: str) -> type(Operator): executor_module = importlib.import_module(module_name) self.operator_module_name = module_name + # Expose each dataset mounted on this computing unit (set by + # texera_run_python_worker from the worker startup config, as a JSON object + # of {variableName: localMountPath}) to the UDF code as a module-level + # variable holding its local path, unless the UDF defines its own. + mounted_datasets = json.loads(os.environ.get("MOUNTED_DATASETS", "{}")) + for variable_name, mount_path in mounted_datasets.items(): + if not hasattr(executor_module, variable_name): + setattr(executor_module, variable_name, mount_path) + executors = list( filter(self.is_concrete_operator, executor_module.__dict__.values()) ) diff --git a/amber/src/main/python/texera_run_python_worker.py b/amber/src/main/python/texera_run_python_worker.py index 1e131d45739..ce0d8a8f5c7 100644 --- a/amber/src/main/python/texera_run_python_worker.py +++ b/amber/src/main/python/texera_run_python_worker.py @@ -17,6 +17,7 @@ import base64 import json +import os import sys from loguru import logger @@ -58,6 +59,7 @@ def init_loguru_logger(stream_log_level) -> None: { "workerId", "outputPort", + "mountedDatasets", "loggerLevel", "rPath", "icebergCatalogType", @@ -141,10 +143,13 @@ def main(raw_config: str) -> None: # Setting R_HOME environment variable for R-UDF usage r_path = config["rPath"] if r_path: - import os - os.environ["R_HOME"] = r_path + # Hand the mounted-dataset variable bindings (a JSON object of + # {variableName: localMountPath}) to ExecutorManager, which injects each one + # into the UDF module as a module-level variable holding its local path. + os.environ["MOUNTED_DATASETS"] = config["mountedDatasets"] + PythonWorker( worker_id=config["workerId"], host="localhost", diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala index 851420bc17b..c765a3e3d6e 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorker.scala @@ -35,6 +35,7 @@ import org.apache.texera.amber.engine.architecture.pythonworker.WorkerBatchInter } import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmbeddedControlMessage import org.apache.texera.amber.engine.architecture.scheduling.config.WorkerConfig +import org.apache.texera.amber.engine.common.DatasetMountManager import org.apache.texera.amber.engine.common.actormessage.{Backpressure, CreditUpdate} import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize import org.apache.texera.amber.engine.common.ambermessage._ @@ -84,13 +85,15 @@ object PythonWorkflowWorker { workerId: String, outputPort: String, rPath: String, - largeBinaryBaseUri: String + largeBinaryBaseUri: String, + mountedDatasets: String = "{}" ): Seq[(String, String)] = { val isPostgres = StorageConfig.icebergCatalogType == "postgres" val isRest = StorageConfig.icebergCatalogType == "rest" Seq( "workerId" -> workerId, "outputPort" -> outputPort, + "mountedDatasets" -> mountedDatasets, "loggerLevel" -> UdfConfig.pythonLogStreamHandlerLevel, "rPath" -> rPath, "icebergCatalogType" -> StorageConfig.icebergCatalogType, @@ -251,6 +254,17 @@ class PythonWorkflowWorker( val pythonBin: String = choosePythonBin() + // Ensure every bound dataset is mounted and resolve each to its in-pod path, keyed by + // the Python variable it will be exposed as. Serialized as a JSON object of + // {variableName: mountPath} so the startup config stays an all-string map. + val mountedDatasets: String = { + val variableToPath = workerConfig.mountedDatasets.map { + case (variableName, locator) => + variableName -> DatasetMountManager.ensureMounted(locator).toString + } + objectMapper.writeValueAsString(variableToPath) + } + // Pass startup configuration to the Python worker by name, as a single JSON // object, rather than by argv position. This way the two sides agree by key, // so adding/removing/reordering a field can no longer silently misassign @@ -259,7 +273,8 @@ class PythonWorkflowWorker( workerConfig.workerId.name, Integer.toString(pythonProxyServer.getPortNumber.get()), RENVPath, - workerConfig.largeBinaryBaseUri + workerConfig.largeBinaryBaseUri, + mountedDatasets ) pythonServerProcess = Process( diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala index da773d7a6b5..59faea4fc30 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/config/WorkerConfig.scala @@ -47,7 +47,8 @@ case object WorkerConfig { VirtualIdentityUtils.createWorkerIdentity(physicalOp.workflowId, physicalOp.id, idx), pveName = physicalOp.pveName, cuid = cuid, - largeBinaryBaseUri = LargeBinaryManager.baseUriForExecution(physicalOp.executionId.id) + largeBinaryBaseUri = LargeBinaryManager.baseUriForExecution(physicalOp.executionId.id), + mountedDatasets = physicalOp.mountedDatasets ) ) } @@ -59,5 +60,9 @@ case class WorkerConfig( cuid: Option[Int] = None, // Coordinator-named, execution-scoped base URI under which this worker's large binaries // live; create() appends a unique suffix. Empty when large binaries are unconfigured. - largeBinaryBaseUri: String = "" + largeBinaryBaseUri: String = "", + // datasets to bind as local-path variables in the Python worker: variable name -> + // dataset-version locator ":". Each is ensured mounted + // before the Python process starts; empty = no mounts. + mountedDatasets: Map[String, String] = Map.empty ) diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/common/DatasetMountManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/common/DatasetMountManager.scala new file mode 100644 index 00000000000..917860d7b2b --- /dev/null +++ b/amber/src/main/scala/org/apache/texera/amber/engine/common/DatasetMountManager.scala @@ -0,0 +1,185 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.amber.engine.common + +import com.fasterxml.jackson.databind.ObjectMapper +import com.typesafe.scalalogging.LazyLogging +import org.apache.texera.common.config.EnvironmentalVariable + +import java.net.{HttpURLConnection, URI, URL} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Path, Paths} +import scala.io.Source +import scala.util.Using + +/** + * Lazily FUSE-mounts dataset versions into the computing unit's local file system. + * + * The mount is NOT performed in this (user-accessible, unprivileged) pod. Instead the + * pod asks the per-node `texera-mounter` (a trusted, privileged DaemonSet) to run + * GeeseFS on its behalf; the resulting read-only mount is exposed back into this pod + * through Kubernetes mount propagation, under a host directory scoped to this CU id. + * A dataset version is addressed by a locator ":"; since a + * commit is immutable, a mount is created at most once per (repository, commit) for the + * lifetime of the pod and reused by every subsequent worker/execution. + * + * The mount is authorized with the pod's per-user JWT, not with global LakeFS + * credentials: the JWT is handed to the mounter, which passes it to GeeseFS as the S3 + * access key so it reaches file-service's JWT-authenticated S3 proxy exactly as a direct + * mount would. Neither this pod nor the mounter holds any global LakeFS credential. + */ +object DatasetMountManager extends LazyLogging { + + // Root, inside this pod, under which the agent's mounts appear via propagation. + private val inPodMountRoot: Path = + Paths.get(sys.env.getOrElse(EnvironmentalVariable.ENV_MOUNT_IN_POD_ROOT, "/mnt/texera-mounts")) + + private val mountTimeoutMs = 30000L + + private val jsonMapper = new ObjectMapper() + + private def env(name: String): String = sys.env.getOrElse(name, "").trim + + private def userJwtToken: String = env(EnvironmentalVariable.ENV_USER_JWT_TOKEN) + + /** + * Base URL of file-service as reachable from the node (scheme://authority), derived + * from the presigned-URL endpoint the pod already receives. The mounter's GeeseFS mount + * targets the S3 proxy hosted at its root. + */ + private def fileServiceBaseUrl: String = { + val presignEndpoint = sys.env.getOrElse( + EnvironmentalVariable.ENV_FILE_SERVICE_GET_PRESIGNED_URL_ENDPOINT, + "http://localhost:9092/api/dataset/presign-download" + ) + val uri = new URI(presignEndpoint) + s"${uri.getScheme}://${uri.getAuthority}" + } + + /** Base URL of the node-local mounter, reachable at the node IP + mounter port. */ + private def mounterBaseUrl: String = { + val nodeIp = env(EnvironmentalVariable.ENV_NODE_IP) + if (nodeIp.isEmpty) { + throw new RuntimeException( + s"No ${EnvironmentalVariable.ENV_NODE_IP} present in the computing unit; " + + "cannot reach the node mounter to mount a dataset." + ) + } + val port = sys.env.getOrElse(EnvironmentalVariable.ENV_MOUNTER_PORT, "8100").trim + s"http://$nodeIp:$port" + } + + /** + * Ensure the dataset version identified by the locator ":" + * is mounted, and return the local (in-pod) mount point. Thread-safe and idempotent. + */ + def ensureMounted(locator: String): Path = + synchronized { + val (repositoryName, commitHash) = locator.split(":", 2) match { + case Array(repo, commit) if repo.nonEmpty && commit.nonEmpty => (repo, commit) + case _ => + throw new IllegalArgumentException( + s"Invalid dataset mount locator '$locator'; expected :." + ) + } + + val mountPoint = inPodMountRoot.resolve(repositoryName).resolve(commitHash) + if (isMounted(mountPoint)) { + logger.info(s"Dataset $locator already mounted at $mountPoint") + return mountPoint + } + + val token = userJwtToken + if (token.isEmpty) { + throw new RuntimeException( + s"No ${EnvironmentalVariable.ENV_USER_JWT_TOKEN} present in the computing unit; " + + "cannot authorize a dataset mount without a user JWT." + ) + } + + requestMount(repositoryName, commitHash, token) + + // The mounter daemonizes GeeseFS; wait until the propagated mount appears in this pod. + val deadline = System.currentTimeMillis() + mountTimeoutMs + while (!isMounted(mountPoint)) { + if (System.currentTimeMillis() > deadline) { + throw new RuntimeException( + s"Dataset $locator did not appear as a mount at $mountPoint within ${mountTimeoutMs}ms." + ) + } + Thread.sleep(200) + } + logger.info(s"Dataset $locator mounted at $mountPoint") + mountPoint + } + + /** + * Ask the node mounter to mount `repositoryName`@`commitHash` for this CU. The mounter + * runs GeeseFS against file-service's S3 proxy with the pod's JWT and mounts into this + * CU's host directory, from where it propagates back into this pod. + */ + private def requestMount(repositoryName: String, commitHash: String, token: String): Unit = { + val payload = jsonMapper.createObjectNode() + payload.put("cuid", env(EnvironmentalVariable.ENV_CU_ID)) + payload.put("repositoryName", repositoryName) + payload.put("commitHash", commitHash) + payload.put("jwt", token) + payload.put("fileServiceBase", fileServiceBaseUrl) + val body = jsonMapper.writeValueAsBytes(payload) + + val url = s"$mounterBaseUrl/mount" + logger.info(s"Requesting mount of $repositoryName:$commitHash from node mounter at $url") + val connection = new URL(url).openConnection().asInstanceOf[HttpURLConnection] + connection.setRequestMethod("POST") + connection.setRequestProperty("Content-Type", "application/json") + connection.setDoOutput(true) + connection.setConnectTimeout(10000) + connection.setReadTimeout(mountTimeoutMs.toInt) + try { + Using(connection.getOutputStream)(_.write(body)) + val code = connection.getResponseCode + if (code != HttpURLConnection.HTTP_OK) { + val err = Option(connection.getErrorStream) + .map(s => new String(s.readAllBytes(), StandardCharsets.UTF_8)) + .getOrElse("") + throw new RuntimeException( + s"Node mounter failed to mount $repositoryName:$commitHash: HTTP $code $err" + ) + } + } finally { + connection.disconnect() + } + } + + private def isMounted(mountPoint: Path): Boolean = { + if (!Files.exists(mountPoint)) { + return false + } + val target = mountPoint.toAbsolutePath.toString + Using(Source.fromFile("/proc/mounts")) { source => + source + .getLines() + .exists(line => { + val fields = line.split(" ") + fields.length > 2 && fields(1) == target && fields(2).startsWith("fuse") + }) + }.getOrElse(false) + } +} diff --git a/amber/src/test/python/test_run_python_worker.py b/amber/src/test/python/test_run_python_worker.py index 7deec97d483..6c03c20c072 100644 --- a/amber/src/test/python/test_run_python_worker.py +++ b/amber/src/test/python/test_run_python_worker.py @@ -38,6 +38,7 @@ def _full_config() -> dict: return { "workerId": "worker-1", "outputPort": "5005", + "mountedDatasets": "{}", "loggerLevel": "INFO", "rPath": "", "icebergCatalogType": "postgres", @@ -133,6 +134,29 @@ def test_main_mapping_is_independent_of_key_order(): ) +def test_main_exposes_mounted_datasets_bindings_when_present(monkeypatch): + monkeypatch.delenv("MOUNTED_DATASETS", raising=False) + config = _full_config() + bindings = '{"A": "/tmp/texera-dataset-mounts/dataset-1/abc123"}' + config["mountedDatasets"] = bindings + storage_patch, worker_patch, _logger_patch = _patched_collaborators() + with storage_patch, worker_patch, _logger_patch: + import os + + entry.main(_encode(config)) + assert os.environ["MOUNTED_DATASETS"] == bindings + + +def test_main_exposes_empty_mounted_datasets_when_none_bound(monkeypatch): + monkeypatch.delenv("MOUNTED_DATASETS", raising=False) + storage_patch, worker_patch, _logger_patch = _patched_collaborators() + with storage_patch, worker_patch, _logger_patch: + import os + + entry.main(_encode(_full_config())) + assert os.environ["MOUNTED_DATASETS"] == "{}" + + def test_main_sets_r_home_when_r_path_present(monkeypatch): monkeypatch.delenv("R_HOME", raising=False) config = _full_config() diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerStartupConfigSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerStartupConfigSpec.scala index 87fceec483a..f3810361f37 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerStartupConfigSpec.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerStartupConfigSpec.scala @@ -60,6 +60,7 @@ class PythonWorkflowWorkerStartupConfigSpec extends AnyFlatSpec { private val expectedKeys = Set( "workerId", "outputPort", + "mountedDatasets", "loggerLevel", "rPath", "icebergCatalogType", @@ -92,6 +93,25 @@ class PythonWorkflowWorkerStartupConfigSpec extends AnyFlatSpec { assert(map("s3LargeBinariesBaseUri") == "s3://bucket/uri") } + it should "default the mounted datasets to an empty JSON object when the worker mounts nothing" in { + val map = + PythonWorkflowWorker.buildStartupConfig("worker-7", "6000", "/opt/R", "s3://bucket/uri").toMap + assert(map("mountedDatasets") == "{}") + } + + it should "carry the mounted dataset variable bindings when the worker mounts repositories" in { + val map = PythonWorkflowWorker + .buildStartupConfig( + "worker-7", + "6000", + "/opt/R", + "s3://bucket/uri", + """{"A":"/mnt/texera-mounts/r/c"}""" + ) + .toMap + assert(map("mountedDatasets") == """{"A":"/mnt/texera-mounts/r/c"}""") + } + it should "produce a config that round-trips through encodeStartupConfig" in { val encoded = PythonWorkflowWorker.encodeStartupConfig( PythonWorkflowWorker.buildStartupConfig("w", "1", "", "uri") diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/FileResolver.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/FileResolver.scala index c8a407df993..fc9788ebad8 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/FileResolver.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/storage/FileResolver.scala @@ -100,6 +100,48 @@ object FileResolver { Some((ownerEmail, datasetName, versionName, fileRelativePathSegments)) } + /** + * Fetches a dataset and one of its versions from the DB by owner email, dataset name, + * and version name. + * + * @param originalPath the caller's original path, used only for the error message + * @throws java.io.FileNotFoundException if the dataset or the version cannot be found + */ + private def fetchDatasetAndVersion( + ownerEmail: String, + datasetName: String, + versionName: String, + originalPath: String + ): (Dataset, DatasetVersion) = + withTransaction( + SqlServer + .getInstance() + .createDSLContext() + ) { ctx => + val dataset = ctx + .select(DATASET.fields: _*) + .from(DATASET) + .leftJoin(USER) + .on(USER.UID.eq(DATASET.OWNER_UID)) + .where(USER.EMAIL.eq(ownerEmail)) + .and(DATASET.NAME.eq(datasetName)) + .fetchOneInto(classOf[Dataset]) + + val datasetVersion = + if (dataset == null) null + else + ctx + .selectFrom(DATASET_VERSION) + .where(DATASET_VERSION.DID.eq(dataset.getDid)) + .and(DATASET_VERSION.NAME.eq(versionName)) + .fetchOneInto(classOf[DatasetVersion]) + + if (dataset == null || datasetVersion == null) { + throw new FileNotFoundException(s"Dataset file $originalPath not found.") + } + (dataset, datasetVersion) + } + /** * Attempts to resolve a given fileName to a URI. * @@ -123,33 +165,7 @@ object FileResolver { // fetch the dataset and version from DB to get dataset ID and version hash val (dataset, datasetVersion) = - withTransaction( - SqlServer - .getInstance() - .createDSLContext() - ) { ctx => - // fetch the dataset from DB - val dataset = ctx - .select(DATASET.fields: _*) - .from(DATASET) - .leftJoin(USER) - .on(USER.UID.eq(DATASET.OWNER_UID)) - .where(USER.EMAIL.eq(ownerEmail)) - .and(DATASET.NAME.eq(datasetName)) - .fetchOneInto(classOf[Dataset]) - - // fetch the dataset version from DB - val datasetVersion = ctx - .selectFrom(DATASET_VERSION) - .where(DATASET_VERSION.DID.eq(dataset.getDid)) - .and(DATASET_VERSION.NAME.eq(versionName)) - .fetchOneInto(classOf[DatasetVersion]) - - if (dataset == null || datasetVersion == null) { - throw new FileNotFoundException(s"Dataset file $fileName not found.") - } - (dataset, datasetVersion) - } + fetchDatasetAndVersion(ownerEmail, datasetName, versionName, fileName) // Convert each segment of fileRelativePath to an encoded String val encodedFileRelativePath = fileRelativePath @@ -178,6 +194,56 @@ object FileResolver { } } + /** + * Resolves a dataset version path to its LakeFS repository name and commit hash. + * Expected format: /ownerEmail/datasetName/versionName + * e.g. /bob@texera.com/twitterDataset/v1 + * + * @param datasetPath the dataset version path to resolve + * @throws java.io.FileNotFoundException if the dataset or version cannot be found + * @return (repositoryName, versionHash) of the dataset version + */ + def resolveDatasetVersion(datasetPath: String): (String, String) = { + val path = Paths.get(datasetPath) + val segments = (0 until path.getNameCount).map(path.getName(_).toString) + if (segments.length < 3) { + throw new FileNotFoundException( + s"Dataset version path $datasetPath is invalid; expected /ownerEmail/datasetName/versionName." + ) + } + val (dataset, datasetVersion) = + fetchDatasetAndVersion(segments(0), segments(1), segments(2), datasetPath) + (dataset.getRepositoryName, datasetVersion.getVersionHash) + } + + /** + * Reverse of [[resolveDatasetVersion]]: given a LakeFS repository name and commit hash, + * recover the human-readable dataset version path /ownerEmail/datasetName/versionName. + * Used to label datasets mounted on a computing unit (which the mounter only tracks by + * repository/commit). Returns None if no matching dataset version exists. + */ + def reverseResolveDatasetVersion( + repositoryName: String, + versionHash: String + ): Option[String] = + withTransaction( + SqlServer + .getInstance() + .createDSLContext() + ) { ctx => + val record = ctx + .select(USER.EMAIL, DATASET.NAME, DATASET_VERSION.NAME) + .from(DATASET) + .join(USER) + .on(USER.UID.eq(DATASET.OWNER_UID)) + .join(DATASET_VERSION) + .on(DATASET_VERSION.DID.eq(DATASET.DID)) + .where(DATASET.REPOSITORY_NAME.eq(repositoryName)) + .and(DATASET_VERSION.VERSION_HASH.eq(versionHash)) + .fetchOne() + Option(record).map(r => s"/${r.value1()}/${r.value2()}/${r.value3()}") + } + /** * Checks if a given file path has a valid scheme. * diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/PhysicalOp.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/PhysicalOp.scala index a1e93050665..8ee43ed0a1f 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/PhysicalOp.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/PhysicalOp.scala @@ -214,7 +214,11 @@ case class PhysicalOp( // hint for number of workers suggestedWorkerNum: Option[Int] = None, // name of the PVE to execute within - pveName: String = "" + pveName: String = "", + // datasets to expose to this operator's Python workers as local-path variables: + // variable name -> dataset-version locator ":". Each is + // ensured mounted before the worker starts and bound to the variable as its mount path. + mountedDatasets: Map[String, String] = Map.empty ) extends LazyLogging { // all the "dependee" links are also blocking @@ -393,6 +397,10 @@ case class PhysicalOp( this.copy(pveName = name) } + def withMountedDatasets(mountedDatasets: Map[String, String]): PhysicalOp = { + this.copy(mountedDatasets = mountedDatasets) + } + /** * creates a copy with an additional input link specified on an input port */ diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMapping.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMapping.scala new file mode 100644 index 00000000000..8a3335a8cda --- /dev/null +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMapping.scala @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.amber.operator.udf.python + +import com.fasterxml.jackson.annotation.{JsonProperty, JsonPropertyDescription} +import com.kjetland.jackson.jsonSchema.annotations.JsonSchemaTitle + +/** + * One row of the Python UDF's "Mounted dataset variables" property: it binds a dataset + * version already mounted on the computing unit to a Python variable. At runtime the + * variable holds the dataset's local (in-pod) mount path as a string. + */ +class DatasetVariableMapping { + @JsonProperty(required = true, defaultValue = "") + @JsonSchemaTitle("Variable name") + @JsonPropertyDescription( + "Name of the Python variable that will hold the dataset's local filesystem path" + ) + var variableName: String = "" + + @JsonProperty(required = true, defaultValue = "") + @JsonSchemaTitle("Mounted dataset") + @JsonPropertyDescription("A dataset version mounted on this computing unit") + var datasetPath: String = "" +} diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2.scala index 6739041a539..cf4647dccd0 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2.scala @@ -23,6 +23,7 @@ import com.fasterxml.jackson.annotation.{JsonProperty, JsonPropertyDescription} import com.google.common.base.Preconditions import com.kjetland.jackson.jsonSchema.annotations.JsonSchemaTitle import org.apache.texera.amber.core.executor.OpExecWithCode +import org.apache.texera.amber.core.storage.FileResolver import org.apache.texera.amber.core.tuple.{Attribute, Schema} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} import org.apache.texera.amber.core.workflow._ @@ -35,7 +36,12 @@ class PythonUDFOpDescV2 extends LogicalOp { @JsonProperty( required = true, defaultValue = - "# Choose from the following templates:\n" + + "# Datasets mounted on this computing unit and bound to variables in the\n" + + "# \"Mounted dataset variables\" property are available as local paths. For\n" + + "# example, if you bind a dataset to the variable A:\n" + + "# open(f\"{A}/file.csv\")\n" + + "# \n" + + "# Choose from the following templates:\n" + "# \n" + "# from pytexera import *\n" + "# \n" + @@ -91,6 +97,14 @@ class PythonUDFOpDescV2 extends LogicalOp { ) var outputColumns: List[Attribute] = List() + @JsonProperty() + @JsonSchemaTitle("Mounted dataset variables") + @JsonPropertyDescription( + "Bind datasets mounted on this computing unit to variables. In your code, each " + + "variable holds the local filesystem path to that dataset." + ) + var datasetVariables: List[DatasetVariableMapping] = List() + override def getPhysicalOp( workflowId: WorkflowIdentity, executionId: ExecutionIdentity @@ -156,6 +170,35 @@ class PythonUDFOpDescV2 extends LogicalOp { trimmed } + // Resolve each bound dataset version to a ":" locator, + // keyed by the Python variable it will be exposed as. + val variableBindings = Option(datasetVariables).getOrElse(List.empty).flatMap { mapping => + val variableName = Option(mapping.variableName).map(_.trim).getOrElse("") + val datasetPath = Option(mapping.datasetPath).map(_.trim).getOrElse("") + if (variableName.isEmpty && datasetPath.isEmpty) None // ignore fully blank rows + else { + if (!variableName.matches("[A-Za-z_][A-Za-z0-9_]*")) + throw new RuntimeException( + s"'$variableName' is not a valid Python variable name for a mounted dataset." + ) + if (datasetPath.isEmpty) + throw new RuntimeException( + s"No dataset selected for the mounted-dataset variable '$variableName'." + ) + val (repositoryName, versionHash) = FileResolver.resolveDatasetVersion(datasetPath) + Some(variableName -> s"$repositoryName:$versionHash") + } + } + val duplicateVariables = + variableBindings.map(_._1).groupBy(identity).collect { + case (name, xs) if xs.size > 1 => name + } + if (duplicateVariables.nonEmpty) + throw new RuntimeException( + s"Duplicate mounted-dataset variable name(s): ${duplicateVariables.mkString(", ")}" + ) + val mountedDatasets = variableBindings.toMap + physicalOp .withDerivePartition(_ => UnknownPartition()) .withInputPorts(operatorInfo.inputPorts) @@ -164,6 +207,7 @@ class PythonUDFOpDescV2 extends LogicalOp { .withIsOneToManyOp(true) .withPropagateSchema(SchemaPropagationFunc(propagateSchema)) .withPveName(pveName) + .withMountedDatasets(mountedDatasets) } override def operatorInfo: OperatorInfo = { diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala index aa02f73387e..b4c88e6978f 100644 --- a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala @@ -51,10 +51,12 @@ import org.apache.texera.dao.jooq.generated.tables.daos.{ import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowComputingUnit import org.apache.texera.service.resource.ComputingUnitManagingResource._ import org.apache.texera.service.resource.ComputingUnitState._ +import org.apache.texera.amber.core.storage.FileResolver import org.apache.texera.service.util.{ ComputingUnitManagingServiceException, InsufficientComputingUnitQuota, - KubernetesClient + KubernetesClient, + MounterClient } import org.jooq.{DSLContext, EnumType} import play.api.libs.json._ @@ -168,6 +170,29 @@ object ComputingUnitManagingResource { case class ComputingUnitTypesResponse( typeOptions: List[String] ) + + // A dataset mounted on a computing unit. `datasetPath` is the human-readable + // /ownerEmail/datasetName/versionName (empty if it could not be reverse-resolved); + // repositoryName/commitHash identify it to the mounter; mountPath is where it lives. + case class MountedDatasetInfo( + datasetPath: String, + repositoryName: String, + commitHash: String, + mountPath: String + ) + + case class DatasetMountParams(datasetPath: String) + + // Base URL (scheme://authority) of file-service as the mounter should reach it, derived + // from the presigned-URL endpoint this service is already configured with. The mounter's + // GeeseFS mount targets the JWT-authenticated S3 proxy hosted at this root. + private lazy val fileServiceBaseUrl: String = { + val endpoint = EnvironmentalVariable + .get(EnvironmentalVariable.ENV_FILE_SERVICE_GET_PRESIGNED_URL_ENDPOINT) + .getOrElse("http://localhost:9092/api/dataset/presign-download") + val uri = new java.net.URI(endpoint) + s"${uri.getScheme}://${uri.getAuthority}" + } } @Produces(Array(MediaType.APPLICATION_JSON)) @@ -765,4 +790,132 @@ class ComputingUnitManagingResource { val computingUnit = getComputingUnitByCuid(context, cuid.toInt) getComputingUnitResourceLimit(computingUnit) } + + // ── Dataset mounts ────────────────────────────────────────────────────── + // Users mount dataset versions onto a computing unit up front; the CU service acts as an + // authenticated proxy to that CU's node mounter. State lives only on the mounter (derived + // from the kernel's mount table), so there is nothing to persist here. + + private def requireMountAccess(cuid: Int, uid: Integer): Unit = { + if ( + !userOwnComputingUnit(context, cuid, uid) && + !ComputingUnitAccessResource.hasWriteAccess(cuid, uid) + ) { + throw new ForbiddenException( + "User does not have permission to manage mounts on this computing unit" + ) + } + } + + private def requireKubernetesUnit(cuid: Int): Unit = { + if (getComputingUnitByCuid(context, cuid).getType != WorkflowComputingUnitTypeEnum.kubernetes) { + throw new BadRequestException( + "Dataset mounting is only supported for Kubernetes computing units." + ) + } + } + + /** Node IP the CU pod is scheduled on, needed to reach that node's mounter. */ + private def mountNodeIp(cuid: Int): String = { + KubernetesClient + .getPodByName(KubernetesClient.generatePodName(cuid)) + .flatMap(pod => Option(pod.getStatus).flatMap(status => Option(status.getHostIP))) + .getOrElse( + throw new BadRequestException( + s"Computing unit $cuid is not running on a node yet; cannot manage mounts." + ) + ) + } + + private def resolveMountDatasetPath(datasetPath: String): (String, (String, String)) = { + val trimmed = Option(datasetPath).map(_.trim).getOrElse("") + if (trimmed.isEmpty) { + throw new BadRequestException("datasetPath is required") + } + (trimmed, FileResolver.resolveDatasetVersion(trimmed)) + } + + /** + * List the datasets currently mounted on the given computing unit. + */ + @GET + @RolesAllowed(Array("REGULAR", "ADMIN")) + @Produces(Array(MediaType.APPLICATION_JSON)) + @Path("/{cuid}/mounts") + def listMountedDatasets( + @PathParam("cuid") cuid: Integer, + @Auth user: SessionUser + ): List[MountedDatasetInfo] = { + requireMountAccess(cuid, user.getUid) + if (getComputingUnitByCuid(context, cuid).getType != WorkflowComputingUnitTypeEnum.kubernetes) { + return List.empty + } + val nodeIp = mountNodeIp(cuid) + MounterClient.listMounts(nodeIp, KubernetesConfig.mounterPort, cuid.toString).map { entry => + val datasetPath = + FileResolver + .reverseResolveDatasetVersion(entry.repositoryName, entry.commitHash) + .getOrElse("") + MountedDatasetInfo(datasetPath, entry.repositoryName, entry.commitHash, entry.mountPath) + } + } + + /** + * Mount a dataset version onto the given computing unit. + */ + @POST + @RolesAllowed(Array("REGULAR", "ADMIN")) + @Consumes(Array(MediaType.APPLICATION_JSON)) + @Produces(Array(MediaType.APPLICATION_JSON)) + @Path("/{cuid}/mounts") + def mountDataset( + @PathParam("cuid") cuid: Integer, + params: DatasetMountParams, + @Auth user: SessionUser + ): MountedDatasetInfo = { + requireMountAccess(cuid, user.getUid) + requireKubernetesUnit(cuid) + val (datasetPath, (repositoryName, commitHash)) = resolveMountDatasetPath(params.datasetPath) + val nodeIp = mountNodeIp(cuid) + // A short-lived token for the requesting user, forwarded to GeeseFS as the S3 access + // key; file-service verifies it and checks the user's read access to the dataset. + val userToken = JwtAuth.jwtToken(jwtClaims(user.user, TOKEN_EXPIRE_TIME_IN_MINUTES)) + val mountPath = MounterClient.mount( + nodeIp, + KubernetesConfig.mounterPort, + cuid.toString, + repositoryName, + commitHash, + userToken, + fileServiceBaseUrl + ) + MountedDatasetInfo(datasetPath, repositoryName, commitHash, mountPath) + } + + /** + * Unmount a dataset version from the given computing unit. + */ + @DELETE + @RolesAllowed(Array("REGULAR", "ADMIN")) + @Consumes(Array(MediaType.APPLICATION_JSON)) + @Produces(Array(MediaType.APPLICATION_JSON)) + @Path("/{cuid}/mounts") + def unmountDataset( + @PathParam("cuid") cuid: Integer, + params: DatasetMountParams, + @Auth user: SessionUser + ): Response = { + requireMountAccess(cuid, user.getUid) + requireKubernetesUnit(cuid) + val (_, (repositoryName, commitHash)) = resolveMountDatasetPath(params.datasetPath) + val nodeIp = mountNodeIp(cuid) + MounterClient.unmount( + nodeIp, + KubernetesConfig.mounterPort, + cuid.toString, + repositoryName, + commitHash + ) + Response.ok().build() + } } diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/MounterClient.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/MounterClient.scala new file mode 100644 index 00000000000..656b631e851 --- /dev/null +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/MounterClient.scala @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.service.util + +import play.api.libs.json._ + +import java.net.{HttpURLConnection, URL, URLEncoder} +import java.nio.charset.StandardCharsets +import scala.util.Using + +/** + * Thin HTTP client for the per-node `texera-mounter`. The computing-unit service uses it + * to mount/unmount/list datasets on a CU's node on the user's behalf: it resolves a + * dataset path to a repository/commit, finds the CU pod's node IP, and forwards the + * request here. The mounter holds no credentials; the JWT passed on mount is forwarded to + * GeeseFS as the S3 access key so authorization stays entirely in file-service. + */ +object MounterClient { + + case class MountEntry(repositoryName: String, commitHash: String, mountPath: String) + + private val connectTimeoutMs = 10000 + // The mounter waits up to 30s for a fresh mount to appear; give it a little headroom. + private val readTimeoutMs = 35000 + + private def baseUrl(nodeIp: String, port: Int): String = s"http://$nodeIp:$port" + + /** Mount `repositoryName:commitHash` for `cuid`; returns the mounter's host mount path. */ + def mount( + nodeIp: String, + port: Int, + cuid: String, + repositoryName: String, + commitHash: String, + jwt: String, + fileServiceBase: String + ): String = { + val body = Json.obj( + "cuid" -> cuid, + "repositoryName" -> repositoryName, + "commitHash" -> commitHash, + "jwt" -> jwt, + "fileServiceBase" -> fileServiceBase + ) + val resp = send("POST", s"${baseUrl(nodeIp, port)}/mount", Some(body)) + (resp \ "mountPath").asOpt[String].getOrElse("") + } + + /** Unmount a single dataset previously mounted for `cuid`. Idempotent on the mounter side. */ + def unmount( + nodeIp: String, + port: Int, + cuid: String, + repositoryName: String, + commitHash: String + ): Unit = { + val query = + s"cuid=${enc(cuid)}&repositoryName=${enc(repositoryName)}&commitHash=${enc(commitHash)}" + send("DELETE", s"${baseUrl(nodeIp, port)}/mount?$query", None) + } + + /** List the datasets currently mounted for `cuid` on this node. */ + def listMounts(nodeIp: String, port: Int, cuid: String): List[MountEntry] = { + val resp = send("GET", s"${baseUrl(nodeIp, port)}/mounts?cuid=${enc(cuid)}", None) + (resp \ "mounts").asOpt[List[JsValue]].getOrElse(Nil).map { m => + MountEntry( + (m \ "repositoryName").as[String], + (m \ "commitHash").as[String], + (m \ "mountPath").as[String] + ) + } + } + + private def enc(s: String): String = URLEncoder.encode(s, StandardCharsets.UTF_8) + + private def send(method: String, url: String, body: Option[JsObject]): JsValue = { + val connection = new URL(url).openConnection().asInstanceOf[HttpURLConnection] + connection.setRequestMethod(method) + connection.setConnectTimeout(connectTimeoutMs) + connection.setReadTimeout(readTimeoutMs) + body.foreach { _ => + connection.setRequestProperty("Content-Type", "application/json") + connection.setDoOutput(true) + } + try { + body.foreach(b => + Using(connection.getOutputStream)( + _.write(Json.stringify(b).getBytes(StandardCharsets.UTF_8)) + ) + ) + val code = connection.getResponseCode + val stream = + if (code >= 200 && code < 300) connection.getInputStream else connection.getErrorStream + val text = Option(stream) + .map(s => new String(s.readAllBytes(), StandardCharsets.UTF_8)) + .getOrElse("") + if (code < 200 || code >= 300) { + throw new RuntimeException(s"mounter $method $url failed: HTTP $code $text") + } + if (text.isEmpty) Json.obj() else Json.parse(text) + } finally { + connection.disconnect() + } + } +} diff --git a/frontend/src/app/common/formly/formly-config.ts b/frontend/src/app/common/formly/formly-config.ts index c4fc54fd77f..be3c2f1f938 100644 --- a/frontend/src/app/common/formly/formly-config.ts +++ b/frontend/src/app/common/formly/formly-config.ts @@ -29,6 +29,7 @@ import { CollabWrapperComponent } from "./collab-wrapper/collab-wrapper/collab-w import { FormlyRepeatDndComponent } from "./repeat-dnd/repeat-dnd.component"; import { UiUdfParametersComponent } from "../../workspace/component/ui-udf-parameters/ui-udf-parameters.component"; import { DatasetVersionSelectorComponent } from "../../workspace/component/dataset-version-selector/dataset-version-selector.component"; +import { DatasetVariablesEditorComponent } from "../../workspace/component/dataset-variables-editor/dataset-variables-editor.component"; import { HuggingFaceImageUploadComponent } from "../../workspace/component/hugging-face-image-upload/hugging-face-image-upload.component"; import { HuggingFaceComponent } from "../../workspace/component/hugging-face/hugging-face.component"; import { HuggingFaceAudioUploadComponent } from "../../workspace/component/hugging-face-audio-upload/hugging-face-audio-upload.component"; @@ -83,6 +84,7 @@ export const TEXERA_FORMLY_CONFIG = { { name: "codearea", component: CodeareaCustomTemplateComponent }, { name: "inputautocomplete", component: DatasetFileSelectorComponent, wrappers: ["form-field"] }, { name: "datasetversionselector", component: DatasetVersionSelectorComponent, wrappers: ["form-field"] }, + { name: "datasetvariables", component: DatasetVariablesEditorComponent, wrappers: ["form-field"] }, { name: "huggingface", component: HuggingFaceComponent, wrappers: ["form-field"] }, { name: "huggingface-audio-upload", component: HuggingFaceAudioUploadComponent, wrappers: ["form-field"] }, { name: "huggingface-image-upload", component: HuggingFaceImageUploadComponent, wrappers: ["form-field"] }, diff --git a/frontend/src/app/common/service/computing-unit/workflow-computing-unit/workflow-computing-unit-managing.service.ts b/frontend/src/app/common/service/computing-unit/workflow-computing-unit/workflow-computing-unit-managing.service.ts index 58768bb346d..ca567ce06c9 100644 --- a/frontend/src/app/common/service/computing-unit/workflow-computing-unit/workflow-computing-unit-managing.service.ts +++ b/frontend/src/app/common/service/computing-unit/workflow-computing-unit/workflow-computing-unit-managing.service.ts @@ -34,6 +34,16 @@ export const COMPUTING_UNIT_CREATE_URL = `${COMPUTING_UNIT_BASE_URL}/create`; export const COMPUTING_UNIT_LIST_URL = `${COMPUTING_UNIT_BASE_URL}`; export const COMPUTING_UNIT_TYPES_URL = `${COMPUTING_UNIT_BASE_URL}/types`; +/** A dataset version mounted on a computing unit. */ +export interface MountedDatasetInfo { + /** Human-readable /ownerEmail/datasetName/versionName ("" if it could not be resolved). */ + datasetPath: string; + repositoryName: string; + commitHash: string; + /** Local (in-pod) path where the dataset is mounted. */ + mountPath: string; +} + @Injectable({ providedIn: "root", }) @@ -191,4 +201,39 @@ export class WorkflowComputingUnitManagingService { {} ); } + + /** + * List the datasets currently mounted on a computing unit. + * @param cuid The ID of the computing unit. + */ + public listMountedDatasets(cuid: number): Observable { + return this.http.get( + `${AppSettings.getApiEndpoint()}/${COMPUTING_UNIT_BASE_URL}/${cuid}/mounts` + ); + } + + /** + * Mount a dataset version onto a computing unit. + * @param cuid The ID of the computing unit. + * @param datasetPath The /ownerEmail/datasetName/versionName path to mount. + */ + public mountDataset(cuid: number, datasetPath: string): Observable { + return this.http.post( + `${AppSettings.getApiEndpoint()}/${COMPUTING_UNIT_BASE_URL}/${cuid}/mounts`, + { datasetPath } + ); + } + + /** + * Unmount a dataset version from a computing unit. + * @param cuid The ID of the computing unit. + * @param datasetPath The /ownerEmail/datasetName/versionName path to unmount. + */ + public unmountDataset(cuid: number, datasetPath: string): Observable { + return this.http.request( + "delete", + `${AppSettings.getApiEndpoint()}/${COMPUTING_UNIT_BASE_URL}/${cuid}/mounts`, + { body: { datasetPath } } + ); + } } diff --git a/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.html b/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.html new file mode 100644 index 00000000000..82d96364a65 --- /dev/null +++ b/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.html @@ -0,0 +1,70 @@ + + +
+
+ Datasets you mount here are available to Python UDFs running on this computing unit. + +
+ + +
+
+ + {{ displayName(mount) }} + +
+
+ + + + +
diff --git a/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.scss b/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.scss new file mode 100644 index 00000000000..7c1485eb01f --- /dev/null +++ b/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.scss @@ -0,0 +1,74 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +.mount-modal { + display: flex; + flex-direction: column; + gap: 12px; +} + +.mount-modal-header { + display: flex; + align-items: center; + justify-content: space-between; + gap: 12px; +} + +.mount-modal-subtitle { + color: rgba(0, 0, 0, 0.55); + font-size: 12px; +} + +.mount-list { + display: flex; + flex-direction: column; + gap: 6px; +} + +.mount-row { + display: flex; + align-items: center; + gap: 8px; + padding: 6px 10px; + border: 1px solid #f0f0f0; + border-radius: 4px; +} + +.mount-row-icon { + color: #1890ff; +} + +.mount-row-name { + flex: 1; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + font-family: monospace; + font-size: 12px; +} + +.mount-row-remove { + color: #ff4d4f; + cursor: pointer; + + &.disabled { + color: rgba(0, 0, 0, 0.25); + cursor: not-allowed; + } +} diff --git a/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.ts b/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.ts new file mode 100644 index 00000000000..94f5d089d83 --- /dev/null +++ b/frontend/src/app/workspace/component/computing-unit-mount-modal/computing-unit-mount-modal.component.ts @@ -0,0 +1,168 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { Component, inject, OnInit } from "@angular/core"; +import { NZ_MODAL_DATA, NzModalService } from "ng-zorro-antd/modal"; +import { UntilDestroy, untilDestroyed } from "@ngneat/until-destroy"; +import { NgFor, NgIf } from "@angular/common"; +import { NzButtonComponent } from "ng-zorro-antd/button"; +import { NzWaveDirective } from "ng-zorro-antd/core/wave"; +import { ɵNzTransitionPatchDirective } from "ng-zorro-antd/core/transition-patch"; +import { NzIconDirective } from "ng-zorro-antd/icon"; +import { NzSpinComponent } from "ng-zorro-antd/spin"; +import { NzEmptyComponent } from "ng-zorro-antd/empty"; +import { NzTooltipDirective } from "ng-zorro-antd/tooltip"; +import { DatasetSelectionModalComponent } from "../dataset-selection-modal/dataset-selection-modal.component"; +import { + MountedDatasetInfo, + WorkflowComputingUnitManagingService, +} from "../../../common/service/computing-unit/workflow-computing-unit/workflow-computing-unit-managing.service"; +import { NotificationService } from "../../../common/service/notification/notification.service"; +import { extractErrorMessage } from "../../../common/util/error"; + +/** + * Modal for managing the datasets mounted on a single computing unit. Users add a dataset + * (picked with the shared dataset-selection modal), see everything currently mounted, and + * unmount individual datasets. Each action proxies through the computing-unit service to + * that CU's node mounter; state lives on the mounter, so the list is always re-fetched. + */ +@UntilDestroy() +@Component({ + templateUrl: "computing-unit-mount-modal.component.html", + styleUrls: ["computing-unit-mount-modal.component.scss"], + imports: [ + NgFor, + NgIf, + NzButtonComponent, + NzWaveDirective, + ɵNzTransitionPatchDirective, + NzIconDirective, + NzSpinComponent, + NzEmptyComponent, + NzTooltipDirective, + ], +}) +export class ComputingUnitMountModalComponent implements OnInit { + private readonly data = inject(NZ_MODAL_DATA) as { cuid: number; computingUnitName?: string }; + readonly cuid = this.data.cuid; + + mounts: MountedDatasetInfo[] = []; + loading = false; + mounting = false; + unmountingPath: string | null = null; + + constructor( + private modalService: NzModalService, + private computingUnitService: WorkflowComputingUnitManagingService, + private notificationService: NotificationService + ) {} + + ngOnInit(): void { + this.refresh(); + } + + refresh(): void { + this.loading = true; + this.computingUnitService + .listMountedDatasets(this.cuid) + .pipe(untilDestroyed(this)) + .subscribe({ + next: mounts => { + this.mounts = mounts; + this.loading = false; + }, + error: (err: unknown) => { + this.loading = false; + this.notificationService.error(extractErrorMessage(err)); + }, + }); + } + + onClickAddDataset(): void { + const modal = this.modalService.create({ + nzContent: DatasetSelectionModalComponent, + nzFooter: null, + nzData: { + fileMode: false, + selectedPath: null, + }, + nzBodyStyle: { + resize: "both", + overflow: "auto", + minHeight: "200px", + minWidth: "550px", + maxWidth: "90vw", + maxHeight: "80vh", + }, + nzWidth: "fit-content", + }); + + modal.afterClose.pipe(untilDestroyed(this)).subscribe((selectedPath: string | undefined) => { + if (selectedPath) { + this.mount(selectedPath); + } + }); + } + + private mount(datasetPath: string): void { + if (this.mounts.some(mount => mount.datasetPath === datasetPath)) { + this.notificationService.warning("That dataset is already mounted on this computing unit."); + return; + } + this.mounting = true; + this.computingUnitService + .mountDataset(this.cuid, datasetPath) + .pipe(untilDestroyed(this)) + .subscribe({ + next: () => { + this.mounting = false; + this.notificationService.success(`Mounted ${datasetPath}`); + this.refresh(); + }, + error: (err: unknown) => { + this.mounting = false; + this.notificationService.error(extractErrorMessage(err)); + }, + }); + } + + onClickUnmount(mount: MountedDatasetInfo): void { + if (!mount.datasetPath || this.unmountingPath) { + return; + } + this.unmountingPath = mount.datasetPath; + this.computingUnitService + .unmountDataset(this.cuid, mount.datasetPath) + .pipe(untilDestroyed(this)) + .subscribe({ + next: () => { + this.unmountingPath = null; + this.refresh(); + }, + error: (err: unknown) => { + this.unmountingPath = null; + this.notificationService.error(extractErrorMessage(err)); + }, + }); + } + + displayName(mount: MountedDatasetInfo): string { + return mount.datasetPath || `${mount.repositoryName}:${mount.commitHash}`; + } +} diff --git a/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.html b/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.html new file mode 100644 index 00000000000..4e0d366ad6f --- /dev/null +++ b/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.html @@ -0,0 +1,82 @@ + + +
+
+ Select a computing unit to mount datasets and bind them here. +
+
+ Dataset mounting is only available on Kubernetes computing units. +
+
+ No datasets are mounted on the active computing unit yet. Use "Mount datasets into computing unit" on the computing + unit to add one. +
+ +
+ + = + + + + +
+ + +
diff --git a/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.scss b/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.scss new file mode 100644 index 00000000000..4ff02cd19ad --- /dev/null +++ b/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.scss @@ -0,0 +1,59 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +.dataset-variables { + display: flex; + flex-direction: column; + gap: 6px; +} + +.dataset-variables-hint { + color: rgba(0, 0, 0, 0.45); + font-size: 12px; +} + +.dataset-variable-row { + display: flex; + align-items: center; + gap: 6px; +} + +.dataset-variable-name { + width: 30%; + min-width: 80px; + font-family: monospace; +} + +.dataset-variable-arrow { + color: rgba(0, 0, 0, 0.45); +} + +.dataset-variable-select { + flex: 1; + min-width: 0; +} + +.dataset-variable-remove { + color: #ff4d4f; + cursor: pointer; +} + +.dataset-variables-add { + align-self: flex-start; +} diff --git a/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.ts b/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.ts new file mode 100644 index 00000000000..4729c4259d4 --- /dev/null +++ b/frontend/src/app/workspace/component/dataset-variables-editor/dataset-variables-editor.component.ts @@ -0,0 +1,152 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { Component, OnInit } from "@angular/core"; +import { FieldType, FieldTypeConfig } from "@ngx-formly/core"; +import { UntilDestroy, untilDestroyed } from "@ngneat/until-destroy"; +import { NgFor, NgIf } from "@angular/common"; +import { FormsModule } from "@angular/forms"; +import { NzButtonComponent } from "ng-zorro-antd/button"; +import { NzWaveDirective } from "ng-zorro-antd/core/wave"; +import { ɵNzTransitionPatchDirective } from "ng-zorro-antd/core/transition-patch"; +import { NzInputDirective } from "ng-zorro-antd/input"; +import { NzIconDirective } from "ng-zorro-antd/icon"; +import { NzSelectComponent, NzOptionComponent } from "ng-zorro-antd/select"; +import { NzTooltipDirective } from "ng-zorro-antd/tooltip"; +import { ComputingUnitStatusService } from "../../../common/service/computing-unit/computing-unit-status/computing-unit-status.service"; +import { + MountedDatasetInfo, + WorkflowComputingUnitManagingService, +} from "../../../common/service/computing-unit/workflow-computing-unit/workflow-computing-unit-managing.service"; + +interface DatasetVariableRow { + variableName: string; + datasetPath: string; +} + +/** + * Property-editor widget for the Python UDF "Mounted dataset variables" property. It edits + * a list of {variableName, datasetPath} bindings: each row maps a Python variable to a + * dataset mounted on the active computing unit. The dataset dropdown is populated from the + * datasets currently mounted on that CU (see the "Mount datasets into computing unit" + * action). At runtime each variable holds the dataset's local mount path. + */ +@UntilDestroy() +@Component({ + templateUrl: "dataset-variables-editor.component.html", + styleUrls: ["dataset-variables-editor.component.scss"], + imports: [ + NgFor, + NgIf, + FormsModule, + NzButtonComponent, + NzWaveDirective, + ɵNzTransitionPatchDirective, + NzInputDirective, + NzIconDirective, + NzSelectComponent, + NzOptionComponent, + NzTooltipDirective, + ], +}) +export class DatasetVariablesEditorComponent extends FieldType implements OnInit { + rows: DatasetVariableRow[] = []; + mountedDatasetPaths: string[] = []; + loading = false; + activeCuid?: number; + activeCuIsKubernetes = false; + + constructor( + private computingUnitStatusService: ComputingUnitStatusService, + private computingUnitService: WorkflowComputingUnitManagingService + ) { + super(); + } + + ngOnInit(): void { + const value = this.formControl.value; + this.rows = Array.isArray(value) + ? value.map((row: Partial) => ({ + variableName: row?.variableName ?? "", + datasetPath: row?.datasetPath ?? "", + })) + : []; + + this.computingUnitStatusService + .getSelectedComputingUnit() + .pipe(untilDestroyed(this)) + .subscribe(unit => { + this.activeCuid = unit?.computingUnit?.cuid; + this.activeCuIsKubernetes = unit?.computingUnit?.type === "kubernetes"; + this.loadMountedDatasets(); + }); + } + + private loadMountedDatasets(): void { + this.mountedDatasetPaths = []; + if (this.activeCuid == null || !this.activeCuIsKubernetes) { + return; + } + this.loading = true; + this.computingUnitService + .listMountedDatasets(this.activeCuid) + .pipe(untilDestroyed(this)) + .subscribe({ + next: (mounts: MountedDatasetInfo[]) => { + this.mountedDatasetPaths = mounts.map(mount => mount.datasetPath).filter(path => !!path); + this.loading = false; + }, + error: () => { + this.loading = false; + }, + }); + } + + /** Options for a row's dropdown, always including the row's own saved value. */ + optionsForRow(row: DatasetVariableRow): string[] { + if (row.datasetPath && !this.mountedDatasetPaths.includes(row.datasetPath)) { + return [row.datasetPath, ...this.mountedDatasetPaths]; + } + return this.mountedDatasetPaths; + } + + addRow(): void { + this.rows = [...this.rows, { variableName: "", datasetPath: "" }]; + this.sync(); + } + + removeRow(index: number): void { + this.rows = this.rows.filter((_, rowIndex) => rowIndex !== index); + this.sync(); + } + + onRowChange(): void { + this.sync(); + } + + trackByIndex(index: number): number { + return index; + } + + private sync(): void { + this.formControl.setValue(this.rows.map(row => ({ ...row }))); + this.formControl.markAsDirty(); + this.formControl.markAsTouched(); + } +} diff --git a/frontend/src/app/workspace/component/power-button/computing-unit-selection.component.html b/frontend/src/app/workspace/component/power-button/computing-unit-selection.component.html index 747354a04ce..f39bc6bd49f 100644 --- a/frontend/src/app/workspace/component/power-button/computing-unit-selection.component.html +++ b/frontend/src/app/workspace/component/power-button/computing-unit-selection.component.html @@ -168,6 +168,16 @@ [nzTooltipTitle]="'Python Environment'" (click)="selectedComputingUnit = unit; showPVEmodalVisible(); $event.stopPropagation()"> + + Date: Fri, 24 Jul 2026 15:32:06 -0700 Subject: [PATCH 3/3] test(dataset-mount): unit-test the mount client, engine and UDF bindings Add coverage for the pure-logic paths of the platform integration: MounterClient (against a stub mounter HTTP server), DatasetMountManager locator/identity validation, the Python UDF datasetVariables validation, DatasetVariableMapping, and the executor-manager MOUNTED_DATASETS variable injection. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01H27sHgS29u2JcLYPPR6Hdx --- .../managers/test_executor_manager.py | 27 ++++ .../common/DatasetMountManagerSpec.scala | 44 ++++++ .../python/DatasetVariableMappingSpec.scala | 40 ++++++ .../udf/python/PythonUDFOpDescV2Spec.scala | 35 +++++ .../service/util/MounterClientSpec.scala | 130 ++++++++++++++++++ 5 files changed, 276 insertions(+) create mode 100644 amber/src/test/scala/org/apache/texera/amber/engine/common/DatasetMountManagerSpec.scala create mode 100644 common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMappingSpec.scala create mode 100644 computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/MounterClientSpec.scala diff --git a/amber/src/test/python/core/architecture/managers/test_executor_manager.py b/amber/src/test/python/core/architecture/managers/test_executor_manager.py index 1a21b106f30..b1b521fd9c0 100644 --- a/amber/src/test/python/core/architecture/managers/test_executor_manager.py +++ b/amber/src/test/python/core/architecture/managers/test_executor_manager.py @@ -180,6 +180,33 @@ def test_accept_python_language_source_operator(self, executor_manager): assert executor_manager.operator_module_name.startswith("udf-v") assert executor_manager.executor.is_source is True + def test_injects_mounted_dataset_variables_into_the_udf_module( + self, executor_manager, monkeypatch + ): + """Each dataset mounted on the CU (passed as the MOUNTED_DATASETS JSON) is exposed + to the UDF code as a module-level variable holding its local mount path.""" + monkeypatch.setenv( + "MOUNTED_DATASETS", + '{"A": "/mnt/texera-mounts/dataset-1/abc", ' + '"B": "/mnt/texera-mounts/dataset-2/def"}', + ) + executor_manager.initialize_executor( + code=SAMPLE_OPERATOR_CODE, is_source=False, language="python" + ) + module = sys.modules[executor_manager.operator_module_name] + assert module.A == "/mnt/texera-mounts/dataset-1/abc" + assert module.B == "/mnt/texera-mounts/dataset-2/def" + + def test_no_mounted_dataset_variables_when_env_absent( + self, executor_manager, monkeypatch + ): + """With no MOUNTED_DATASETS env set, nothing is injected and loading still works.""" + monkeypatch.delenv("MOUNTED_DATASETS", raising=False) + executor_manager.initialize_executor( + code=SAMPLE_OPERATOR_CODE, is_source=False, language="python" + ) + assert executor_manager.operator_module_name is not None + def test_reject_other_unsupported_languages(self, executor_manager): """Test that other arbitrary languages still work (no R-specific check).""" # Languages other than r-tuple and r-table should be allowed to pass diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/common/DatasetMountManagerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/common/DatasetMountManagerSpec.scala new file mode 100644 index 00000000000..093aab56f05 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/common/DatasetMountManagerSpec.scala @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.amber.engine.common + +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +class DatasetMountManagerSpec extends AnyFlatSpec with Matchers { + + "ensureMounted" should "reject a locator that is not :" in { + for (bad <- Seq("no-colon", "", ":", "repo:", ":commit")) { + withClue(s"locator '$bad': ") { + an[IllegalArgumentException] should be thrownBy DatasetMountManager.ensureMounted(bad) + } + } + } + + it should "not attempt a mount for a well-formed, unmounted locator without the pod's mount identity" in { + // The locator is valid and no such mount exists in this test process, so ensureMounted + // advances past parsing and the /proc/mounts check to the environment checks (user JWT / + // node IP), which must fail loudly rather than attempt an unauthorized/unroutable mount. + val ex = intercept[RuntimeException] { + DatasetMountManager.ensureMounted("dataset-does-not-exist:deadbeefdeadbeef") + } + ex.getMessage.toLowerCase should include("computing unit") + } +} diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMappingSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMappingSpec.scala new file mode 100644 index 00000000000..64378d7b266 --- /dev/null +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/DatasetVariableMappingSpec.scala @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.amber.operator.udf.python + +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +class DatasetVariableMappingSpec extends AnyFlatSpec with Matchers { + + "DatasetVariableMapping" should "default to empty variable and dataset" in { + val mapping = new DatasetVariableMapping() + mapping.variableName shouldBe "" + mapping.datasetPath shouldBe "" + } + + it should "hold the variable name and dataset path assigned to it" in { + val mapping = new DatasetVariableMapping() + mapping.variableName = "A" + mapping.datasetPath = "/bob@texera.com/twitterDataset/v1" + mapping.variableName shouldBe "A" + mapping.datasetPath shouldBe "/bob@texera.com/twitterDataset/v1" + } +} diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala index aea38193aa9..227778f3bf1 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala @@ -237,4 +237,39 @@ class PythonUDFOpDescV2Spec extends AnyFlatSpec with Matchers { case other => fail(s"expected OpExecWithCode, got $other") } } + + private def binding(name: String, path: String): DatasetVariableMapping = { + val m = new DatasetVariableMapping + m.variableName = name + m.datasetPath = path + m + } + + "PythonUDFOpDescV2 dataset-variable bindings" should + "carry no mounted datasets when nothing is bound" in { + val d = new PythonUDFOpDescV2 + d.code = "yield t" + d.getPhysicalOp(workflowId, executionId).mountedDatasets shouldBe empty + } + + it should "ignore fully blank dataset-variable rows" in { + val d = new PythonUDFOpDescV2 + d.code = "yield t" + d.datasetVariables = List(binding("", ""), binding(" ", " ")) + d.getPhysicalOp(workflowId, executionId).mountedDatasets shouldBe empty + } + + it should "reject a variable name that is not a valid Python identifier" in { + val d = new PythonUDFOpDescV2 + d.datasetVariables = List(binding("1bad", "/bob@texera.com/twitterDataset/v1")) + val ex = intercept[RuntimeException] { d.getPhysicalOp(workflowId, executionId) } + ex.getMessage should include("not a valid Python variable name") + } + + it should "reject a bound variable with no dataset selected" in { + val d = new PythonUDFOpDescV2 + d.datasetVariables = List(binding("A", " ")) + val ex = intercept[RuntimeException] { d.getPhysicalOp(workflowId, executionId) } + ex.getMessage should include("No dataset selected") + } } diff --git a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/MounterClientSpec.scala b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/MounterClientSpec.scala new file mode 100644 index 00000000000..907514a1d60 --- /dev/null +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/MounterClientSpec.scala @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.texera.service.util + +import com.sun.net.httpserver.{HttpExchange, HttpServer} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +import java.net.InetSocketAddress +import java.nio.charset.StandardCharsets +import scala.collection.mutable + +/** + * Tests [[MounterClient]] against a stub node-mounter HTTP server, verifying it builds the + * mount / unmount / list requests the mounter expects and parses the responses correctly. + */ +class MounterClientSpec extends AnyFlatSpec with Matchers with BeforeAndAfterAll { + + private var server: HttpServer = _ + private var port: Int = _ + // records the last request the stub received, per path + private val received = mutable.Map[String, (String, String, String)]() // path -> (method, query, body) + + private def body(ex: HttpExchange): String = + new String(ex.getRequestBody.readAllBytes(), StandardCharsets.UTF_8) + + private def reply(ex: HttpExchange, code: Int, json: String): Unit = { + val bytes = json.getBytes(StandardCharsets.UTF_8) + ex.sendResponseHeaders(code, bytes.length) + ex.getResponseBody.write(bytes) + ex.getResponseBody.close() + } + + override def beforeAll(): Unit = { + server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + server.createContext( + "/mount", + (ex: HttpExchange) => { + received("/mount") = (ex.getRequestMethod, Option(ex.getRequestURI.getQuery).getOrElse(""), body(ex)) + if (ex.getRequestMethod == "POST") + reply(ex, 200, """{"mountPath":"/var/lib/texera-mounts/7/dataset-1/abc123"}""") + else reply(ex, 200, """{"status":"ok"}""") // DELETE + } + ) + server.createContext( + "/mounts", + (ex: HttpExchange) => { + received("/mounts") = (ex.getRequestMethod, Option(ex.getRequestURI.getQuery).getOrElse(""), "") + reply( + ex, + 200, + """{"mounts":[{"repositoryName":"dataset-1","commitHash":"abc123","mountPath":"/var/lib/texera-mounts/7/dataset-1/abc123"}]}""" + ) + } + ) + server.start() + port = server.getAddress.getPort + + failing = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0) + failing.createContext("/", (ex: HttpExchange) => reply(ex, 500, """{"error":"boom"}""")) + failing.start() + failingPort = failing.getAddress.getPort + } + + private var failing: HttpServer = _ + private var failingPort: Int = _ + + override def afterAll(): Unit = { + if (server != null) server.stop(0) + if (failing != null) failing.stop(0) + } + + "mount" should "POST the mount request and return the mounter's mount path" in { + val path = + MounterClient.mount("127.0.0.1", port, "7", "dataset-1", "abc123", "the-jwt", "http://file-service:9092") + path shouldBe "/var/lib/texera-mounts/7/dataset-1/abc123" + val (method, _, requestBody) = received("/mount") + method shouldBe "POST" + requestBody should include(""""cuid":"7"""") + requestBody should include(""""repositoryName":"dataset-1"""") + requestBody should include(""""commitHash":"abc123"""") + requestBody should include(""""jwt":"the-jwt"""") + requestBody should include(""""fileServiceBase":"http://file-service:9092"""") + } + + "listMounts" should "GET with the cuid and parse the returned mounts" in { + val mounts = MounterClient.listMounts("127.0.0.1", port, "7") + mounts should have size 1 + mounts.head shouldBe MounterClient.MountEntry( + "dataset-1", + "abc123", + "/var/lib/texera-mounts/7/dataset-1/abc123" + ) + received("/mounts")._2 should include("cuid=7") + } + + "unmount" should "DELETE with the identifiers in the query string (no body)" in { + MounterClient.unmount("127.0.0.1", port, "7", "dataset-1", "abc123") + val (method, query, _) = received("/mount") + method shouldBe "DELETE" + query should include("cuid=7") + query should include("repositoryName=dataset-1") + query should include("commitHash=abc123") + } + + "a non-2xx mounter response" should "surface as a RuntimeException carrying the HTTP status" in { + val ex = intercept[RuntimeException] { + MounterClient.mount("127.0.0.1", failingPort, "7", "dataset-1", "abc123", "j", "http://f") + } + ex.getMessage should include("HTTP 500") + } +}