feat: Add Python Virtual Environment Support: Installing User Defined Packages

amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveManager.scala

@@ -32,7 +32,8 @@ import org.apache.texera.amber.config.PythonUtils
   * for each Computing Unit
   *
   * It supports:
-  * - Creating and initializing isolated Python environments
+  * - Creating and initializing isolated Python environments (with system packages)
+  * - Installing user defined packages
   * - Streaming pip output logs back to the caller
   *
   * Each PVE is stored under:
@@ -41,6 +42,11 @@ import org.apache.texera.amber.config.PythonUtils
 
 object PveManager {
 
+  case class PvePackageResponse(
+      pveName: String,
+      userPackages: Seq[String]
+  )
+
   private val VenvRoot: Path = Paths.get("/tmp/texera-pve/venvs")
 
   private def cuidDir(cuid: Int, pveName: String): Path = {
@@ -61,9 +67,63 @@ object PveManager {
       "PIP_NO_INPUT" -> "1"
     )
 
-  def getSystemPackages(): Seq[String] = {
-    val python = PythonUtils.getPythonExecutable
-    Process(Seq(python, "-m", "pip", "freeze")).!!.split("\n").map(_.trim).filter(_.nonEmpty).toSeq
+  private def readPackageFile(path: Path): Seq[String] = {
+    if (Files.exists(path)) {
+      Files
+        .readAllLines(path)
+        .asScala
+        .map(_.trim)
+        .filter(_.nonEmpty)
+        .toSeq
+    } else {
+      Seq()
+    }
+  }
+
+  private def getSystemPath(isLocal: Boolean): Path = {
+    Paths.get(
+      if (isLocal) "amber/system-requirements-lock.txt"
+      else "/tmp/system-requirements-lock.txt"
+    )
+  }
+
+  def getSystemPackages(isLocal: Boolean): Seq[String] = {
+    if (!Files.exists(getSystemPath(isLocal))) {
+      Seq()
+    } else {
+      Files
+        .readAllLines(getSystemPath(isLocal))
+        .asScala
+        .map(_.trim)
+        .filter(line => line.nonEmpty && !line.startsWith("#"))
+        .toSeq
+    }
+  }
+
+  private def runPipInstall(
+      python: String,
+      args: Seq[String],
+      queue: BlockingQueue[String]
+  ): Int = {
+    Process(
+      Seq(
+        python,
+        "-u",
+        "-m",
+        "pip",
+        "install",
+        "--progress-bar",
+        "off",
+        "--no-input"
+      ) ++ args,
+      None,
+      pipEnv.toSeq: _*
+    ).!(
+      ProcessLogger(
+        out => queue.put(s"[pip] $out"),
+        err => queue.put(s"[pip][ERR] $err")
+      )
+    )
   }
 
   /**
@@ -85,23 +145,18 @@ object PveManager {
     queue.put(s"[PVE] Creating new PVE for cuid: $cuid with name: $pveName")
 
     // NOTE: These paths are derived from computing-unit-master.dockerfile.
-    // If requirements.txt or operator-requirements.txt locations change, update these paths.
+    // If requirements.txt location changes, update these paths.
     val requirementsPath =
       if (isLocal) Paths.get("amber", "requirements.txt")
       else Paths.get("/tmp", "requirements.txt")
 
-    val operatorRequirementsPath =
-      if (isLocal) Paths.get("amber", "operator-requirements.txt")
-      else Paths.get("/tmp", "operator-requirements.txt")
-
-    if (!Files.exists(requirementsPath) || !Files.exists(operatorRequirementsPath)) {
+    if (!Files.exists(requirementsPath)) {
       queue.put(s"[PVE][ERR] System requirements not found")
       return
     }
 
     val venvDirPath = pveDir(cuid, pveName).toAbsolutePath
     val python = pythonBinPath(cuid, pveName).toAbsolutePath.toString
-    val envVars = pipEnv
 
     val createVenvPython = PythonUtils.getPythonExecutable
 
@@ -121,43 +176,17 @@ object PveManager {
       return
     }
 
-    if (!Files.exists(requirementsPath)) {
-      queue.put(s"[PVE][ERR] requirements.txt not found at ${requirementsPath.toAbsolutePath}")
-      return
-    }
-
-    if (!Files.exists(operatorRequirementsPath)) {
-      queue.put(
-        s"[PVE][ERR] operator-requirements.txt not found at ${operatorRequirementsPath.toAbsolutePath}"
-      )
-      return
-    }
-
     queue.put(
-      s"[PVE] Installing requirements from ${requirementsPath.toAbsolutePath} and ${operatorRequirementsPath.toAbsolutePath}"
+      s"[PVE] Installing requirements from ${requirementsPath.toAbsolutePath}"
     )
 
-    val installReqCode = Process(
+    val installReqCode = runPipInstall(
+      python,
       Seq(
-        python,
-        "-u",
-        "-m",
-        "pip",
-        "install",
-        "--progress-bar",
-        "off",
         "-r",
-        requirementsPath.toString,
-        "-r",
-        operatorRequirementsPath.toString
+        requirementsPath.toString
       ),
-      None,
-      envVars.toSeq: _*
-    ).!(
-      ProcessLogger(
-        out => queue.put(s"[pip] $out"),
-        err => queue.put(s"[pip][ERR] $err")
-      )
+      queue
     )
 
     queue.put(s"[PVE] requirements install finished with exit code $installReqCode")
@@ -170,7 +199,8 @@ object PveManager {
     queue.put(s"[PVE] Created new environment for cuid = $cuid")
   }
 
-  def getEnvironments(cuid: Int): List[String] = {
+  // returns list of PVE names and corresponding user packages for a given CU
+  def getEnvironments(cuid: Int): List[PvePackageResponse] = {
 
     val cuPath = VenvRoot.resolve(cuid.toString)
 
@@ -185,7 +215,17 @@ object PveManager {
         .iterator()
         .asScala
         .filter(path => Files.isDirectory(path))
-        .map(path => path.getFileName.toString)
+        .map { path =>
+          val pveName = path.getFileName.toString
+          val metadataPath = path.resolve("user-packages.txt")
+
+          val userPackages = readPackageFile(metadataPath)
+
+          PvePackageResponse(
+            pveName = pveName,
+            userPackages = userPackages
+          )
+        }
         .toList
     } finally {
       stream.close()
@@ -212,4 +252,93 @@ object PveManager {
       stream.close()
     }
   }
+
+  /**
+    * Installs user requested Python packages into the PVE.
+    *
+    * 1. Executes pip install for each package
+    * 2. Prevents conflicts with system dependencies.
+    * 3. Updates user metadata file
+    * 4. Streams logs back via queue
+    */
+  def installUserPackages(
+      packages: List[String],
+      cuid: Int,
+      queue: BlockingQueue[String],
+      pveName: String,
+      isLocal: Boolean
+  ): Unit = {
+
+    val python = pythonBinPath(cuid, pveName).toAbsolutePath.toString
+
+    if (!Files.exists(Paths.get(python))) {
+      queue.put(s"[PVE][ERR] Python executable not found for PVE: $python")
+      return
+    }
+
+    val metadataPath = cuidDir(cuid, pveName).resolve("user-packages.txt")
+
+    var installedPackages = readPackageFile(metadataPath).toSet
+
+    val systemPackages =
+      if (Files.exists(getSystemPath(isLocal))) {
+        Files
+          .readAllLines(getSystemPath(isLocal))
+          .asScala
+          .map(_.trim)
+          .filter(line => line.nonEmpty && !line.startsWith("#"))
+          .map(line => line.split("==")(0).trim.toLowerCase)
+          .toSet
+      } else {
+        Set[String]()
+      }
+
+    packages.foreach { pkg =>
+      val trimmedPkg = pkg.trim
+
+      if (trimmedPkg.nonEmpty) {
+
+        val userPackageName = trimmedPkg.split("==")(0).trim.toLowerCase
+
+        if (systemPackages.contains(userPackageName)) {
+          queue.put(
+            s"[PVE][ERR] $trimmedPkg is a system package and cannot be installed or modified by the user."
+          )
+          return
+        }
+
+        queue.put(s"[PVE] Installing package: $trimmedPkg")
+
+        val code = runPipInstall(
+          python,
+          Seq(
+            "--constraint", // check against system-requirements-lock
+            getSystemPath(isLocal).toString,
+            trimmedPkg
+          ),
+          queue
+        )
+
+        queue.put(s"[pip] install($trimmedPkg) finished with exit code $code")
+
+        if (code != 0) {
+          queue.put(s"[PVE][ERR] Failed to install package: $trimmedPkg")
+          return
+        }
+
+        installedPackages = installedPackages + trimmedPkg
+
+        Files.write(
+          metadataPath,
+          installedPackages.toSeq.sorted.asJava
+        )
+      }
+    }
+
+    queue.put("[PVE] Final user package list:")
+
+    installedPackages.toSeq.sorted.foreach { pkg =>
+      queue.put(s"[user-package] $pkg")
+    }
+  }
 }

amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveResource.scala

@@ -28,24 +28,32 @@ import java.util
 @Consumes(Array(MediaType.APPLICATION_JSON))
 class PveResource {
   // --------------------------------------------------
-  // Get installed packages
+  // Get system packages
   // --------------------------------------------------
   @GET
   @Path("/system")
   @Produces(Array(MediaType.APPLICATION_JSON))
   def getSystemPackages: util.Map[String, util.List[String]] = {
     try {
-      val systemPkgs = PveManager.getSystemPackages().toList.asJava
+
+      // TODO: Support Kubernetes environment handling
+      val isLocal = true
+
+      val systemPkgs =
+        PveManager.getSystemPackages(isLocal).toList.asJava
+
       Map("system" -> systemPkgs).asJava
     } catch {
       case e: Exception =>
         e.printStackTrace()
-        throw new InternalServerErrorException("Failed to get system packages.")
+        throw new InternalServerErrorException(
+          "Failed to get system packages."
+        )
     }
   }
 
   // --------------------------------------------------
-  // Fetch PVEs
+  // Fetch PVEs and Installed User Packages
   // --------------------------------------------------
   @GET
   @Path("/pves")
@@ -54,9 +62,10 @@ class PveResource {
     try {
       PveManager
         .getEnvironments(cuid)
-        .map { pveName =>
+        .map { pve =>
           Map(
-            "pveName" -> pveName.asInstanceOf[Object]
+            "pveName" -> pve.pveName.asInstanceOf[Object],
+            "userPackages" -> pve.userPackages.asJava.asInstanceOf[Object]
           ).asJava
         }
         .asJava

amber/src/main/scala/org/apache/texera/web/resource/pythonvirtualenvironment/PveWebsocketResource.scala

@@ -26,9 +26,9 @@ import scala.concurrent.Future
 import scala.concurrent.ExecutionContext.Implicits.global
 
 /**
-  * WebSocket endpoint for PVE creation that streams pip installation logs
-  * to the frontend in real time. The environment setup runs asynchronously,
-  * and output is pushed to the client until completion.
+  *  WebSocket endpoint for PVE creation and user package installation that streams
+  *  pip installation logs  to the frontend in real time. The environment setup runs
+  *  asynchronously, and output is pushed to the client until completion.
   */
 
 @ServerEndpoint("/wsapi/pve")
@@ -42,12 +42,33 @@ class PveWebsocketResource {
     val cuid = params.get("cuid").get(0).toInt
     val pveName = params.get("pveName").get(0)
     val isLocal = params.get("isLocal").get(0).toBoolean
+    val action = params.getOrDefault("action", java.util.List.of("create")).get(0)
 
     val queue = new LinkedBlockingQueue[String]()
 
     Future {
       try {
-        PveManager.createNewPve(cuid, queue, pveName, isLocal)
+        action match {
+          case "create" =>
+            PveManager.createNewPve(cuid, queue, pveName, isLocal)
+
+          case "install" =>
+            val packages =
+              params
+                .getOrDefault("packages", java.util.List.of("[]"))
+                .get(0)
+                .stripPrefix("[")
+                .stripSuffix("]")
+                .split(",")
+                .toList
+                .map(_.replace("\"", "").trim)
+                .filter(_.nonEmpty)
+
+            PveManager.installUserPackages(packages, cuid, queue, pveName, isLocal)
+
+          case _ =>
+            queue.put(s"[ERR] Unknown action: $action")
+        }
       } catch {
         case e: Exception =>
           queue.put(s"[ERR] ${e.getMessage}")
@@ -61,7 +82,6 @@ class PveWebsocketResource {
 
       while (!done && session.isOpen) {
         val line = queue.take()
-
         session.getBasicRemote.sendText(line)
 
         if (line == "__DONE__") {