feat(dataset-mount): drive dataset mounts from the platform

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())
         )

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",

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(

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 "<repositoryName>:<commitHash>". Each is ensured mounted
+    // before the Python process starts; empty = no mounts.
+    mountedDatasets: Map[String, String] = Map.empty
 )

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 "<repositoryName>:<commitHash>"; 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 "<repositoryName>:<commitHash>"
+    * 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 <repositoryName>:<commitHash>."
+          )
+      }
+
+      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)
+  }
+}

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

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()