feat: add venv parameter for Python virtualenv activation

Allows instances to run in a pre-existing virtualenv instead of the
worker's Python environment. Users are responsible for creating and
maintaining the venv on shared filesystem.

Changes:
- Add `venv` field to InstanceSubmissionRequest, Instance, DesiredInstance
- Add `venv` column to instances table in SQLite
- Worker activates venv before command: `source {venv}/bin/activate && {cmd}`
- Server validates: must be absolute path, no '..' allowed
- CLI: `--venv /path/to/venv` flag
- Python API: `venv` parameter in submit()

Example:
  pylet submit "python train.py" --venv /home/user/torch-env

Author: future-xy

pylet/_sync_api.py

@@ -84,6 +84,8 @@ def submit(
     exclusive: bool = True,
     labels: Optional[Dict[str, str]] = None,
     env: Optional[Dict[str, str]] = None,
+    # Venv support
+    venv: Optional[str] = None,
 ) -> Instance:
     """
     Submit a new instance.
@@ -99,6 +101,7 @@ def submit(
         exclusive: If False, GPUs don't block allocation pool (default True)
         labels: Custom metadata dict
         env: Environment variables to set
+        venv: Path to pre-existing virtualenv (must be absolute path)
 
     Returns:
         Instance handle
@@ -120,6 +123,13 @@ def submit(
             exclusive=False,
             labels={"type": "sllm-store"},
         )
+
+        # Venv example
+        instance = pylet.submit(
+            "python train.py",
+            venv="/home/user/my-venv",
+            gpu=1,
+        )
     """
     client = _get_client()
     address = _get_head_address()
@@ -145,6 +155,8 @@ def submit(
         submission["target_worker"] = target_worker
     if gpu_indices is not None:
         submission["gpu_indices"] = gpu_indices
+    if venv is not None:
+        submission["venv"] = venv
 
     # Submit instance
     response = client.post(f"{address}/instances", json=submission)

pylet/aio/__init__.py

@@ -13,7 +13,7 @@ async def main():
 """
 
 import shlex
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import httpx
 
@@ -97,16 +97,30 @@ async def submit(
     gpu: int = 0,
     cpu: int = 1,
     memory: int = 512,
+    # SLLM support parameters
+    target_worker: Optional[str] = None,
+    gpu_indices: Optional[List[int]] = None,
+    exclusive: bool = True,
+    labels: Optional[Dict[str, str]] = None,
+    env: Optional[Dict[str, str]] = None,
+    # Venv support
+    venv: Optional[str] = None,
 ) -> Instance:
     """
     Submit a new instance.
 
     Args:
         command: Shell command string, or list of args (auto shell-escaped)
         name: Optional instance name for service discovery
-        gpu: GPU units required (default 0)
+        gpu: GPU units required (default 0, ignored if gpu_indices specified)
         cpu: CPU cores required (default 1)
         memory: Memory in MB required (default 512)
+        target_worker: Place on specific worker node
+        gpu_indices: Request specific physical GPU indices
+        exclusive: If False, GPUs don't block allocation pool (default True)
+        labels: Custom metadata dict
+        env: Environment variables to set
+        venv: Path to pre-existing virtualenv (must be absolute path)
 
     Returns:
         Instance handle
@@ -122,19 +136,28 @@ async def submit(
     if isinstance(command, list):
         command = shlex.join(command)
 
-    # Submit instance
-    response = await client.post(
-        f"{address}/instances",
-        json={
-            "command": command,
-            "resource_requirements": {
-                "cpu_cores": cpu,
-                "gpu_units": gpu,
-                "memory_mb": memory,
-            },
-            "name": name,
+    # Build submission request
+    submission = {
+        "command": command,
+        "resource_requirements": {
+            "cpu_cores": cpu,
+            "gpu_units": gpu,
+            "memory_mb": memory,
         },
-    )
+        "name": name,
+        "exclusive": exclusive,
+        "labels": labels or {},
+        "env": env or {},
+    }
+    if target_worker is not None:
+        submission["target_worker"] = target_worker
+    if gpu_indices is not None:
+        submission["gpu_indices"] = gpu_indices
+    if venv is not None:
+        submission["venv"] = venv
+
+    # Submit instance
+    response = await client.post(f"{address}/instances", json=submission)
     response.raise_for_status()
 
     instance_id = response.json()["instance_id"]

pylet/cli.py

@@ -83,8 +83,10 @@ def start(head, cpu_cores, gpu_units, memory_mb):
 @click.option("--exclusive/--no-exclusive", default=True, help="GPU exclusivity mode.")
 @click.option("--label", multiple=True, help="Labels (key=value).")
 @click.option("--env", "env_vars", multiple=True, help="Env vars (key=value).")
+# Venv support
+@click.option("--venv", default=None, help="Path to pre-existing virtualenv (must be absolute).")
 def submit(command, config, cpu_cores, gpu_units, memory_mb, name,
-           target_worker, gpu_indices, exclusive, label, env_vars):
+           target_worker, gpu_indices, exclusive, label, env_vars, venv):
     """Submit a new instance to the PyLet cluster.
 
     Precedence (highest wins): CLI args > Config file > Defaults
@@ -192,6 +194,7 @@ async def submit_instance():
                 exclusive=exclusive,
                 labels=labels if labels else None,
                 env=env if env else None,
+                venv=venv,
             )
             click.echo(f"Instance submitted with ID: {instance_id}")
         except Exception as e:

pylet/client.py

@@ -31,6 +31,8 @@ async def submit_instance(
         exclusive: bool = True,
         labels: Optional[Dict[str, str]] = None,
         env: Optional[Dict[str, str]] = None,
+        # Venv support
+        venv: Optional[str] = None,
     ) -> str:
         """Submit a new instance for execution."""
         submission_data: Dict[str, Any] = {
@@ -46,6 +48,8 @@ async def submit_instance(
             submission_data["target_worker"] = target_worker
         if gpu_indices is not None:
             submission_data["gpu_indices"] = gpu_indices
+        if venv is not None:
+            submission_data["venv"] = venv
 
         response = await self.client.post(
             f"{self.api_server_url}/instances", json=submission_data

pylet/controller.py

@@ -214,6 +214,7 @@ async def _get_desired_instances(self, worker_id: str) -> List[DesiredInstance]:
                     command=row["command"],
                     gpu_indices=gpu_indices,
                     env=env,
+                    venv=row.get("venv"),
                     expected_status=row["status"],
                 )
             )
@@ -360,6 +361,8 @@ async def submit_instance(
         exclusive: bool = True,
         labels: Optional[Dict[str, str]] = None,
         env: Optional[Dict[str, str]] = None,
+        # Venv support
+        venv: Optional[str] = None,
     ) -> str:
         """Submit a new instance for execution."""
         async with self.lock:
@@ -391,6 +394,7 @@ async def submit_instance(
                 env=env or {},
                 target_worker=target_worker,
                 gpu_indices=gpu_indices,
+                venv=venv,
             )
 
             logger.info(

pylet/db.py

@@ -86,6 +86,8 @@ async def _create_tables(self) -> None:
                 env TEXT,                     -- JSON object
                 target_worker TEXT,           -- Placement constraint
                 requested_gpu_indices TEXT,   -- JSON array
+                -- Venv support
+                venv TEXT,                    -- Absolute path to virtualenv
                 -- Execution
                 exit_code INTEGER,
                 stdout_log TEXT,
@@ -317,16 +319,18 @@ async def insert_instance(
         env: Optional[Dict[str, str]] = None,
         target_worker: Optional[str] = None,
         gpu_indices: Optional[List[int]] = None,
+        # Venv support
+        venv: Optional[str] = None,
     ) -> None:
         """Insert a new instance."""
         async with self.transaction():
             await self._conn.execute(
                 """
                 INSERT INTO instances (
                     id, name, command, status, cpu_cores, gpu_units, memory_mb,
-                    exclusive, labels, env, target_worker, requested_gpu_indices
+                    exclusive, labels, env, target_worker, requested_gpu_indices, venv
                 )
-                VALUES (?, ?, ?, 'PENDING', ?, ?, ?, ?, ?, ?, ?, ?)
+                VALUES (?, ?, ?, 'PENDING', ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """,
                 (
                     instance_id, name, command, cpu_cores, gpu_units, memory_mb,
@@ -335,6 +339,7 @@ async def insert_instance(
                     json.dumps(env) if env else None,
                     target_worker,
                     json.dumps(gpu_indices) if gpu_indices else None,
+                    venv,
                 ),
             )
 

pylet/schemas.py

@@ -108,6 +108,8 @@ class Instance(BaseModel):
     exclusive: bool = True
     labels: Dict[str, str] = Field(default_factory=dict)
     env: Dict[str, str] = Field(default_factory=dict)
+    # Venv support
+    venv: Optional[str] = None  # Path to pre-existing virtualenv (must be absolute path)
 
 
 def get_display_status(status: InstanceStatus, cancellation_requested_at: Optional[datetime]) -> str:
@@ -146,6 +148,8 @@ class InstanceSubmissionRequest(BaseModel):
     exclusive: bool = True                    # GPU exclusivity mode
     labels: Dict[str, str] = Field(default_factory=dict)  # Custom metadata
     env: Dict[str, str] = Field(default_factory=dict)     # Environment variables
+    # Venv support
+    venv: Optional[str] = None  # Path to pre-existing virtualenv (must be absolute path)
 
 
 class WorkerRegistrationRequest(BaseModel):
@@ -187,6 +191,7 @@ class DesiredInstance(BaseModel):
     command: str
     gpu_indices: List[int] = []
     env: Dict[str, str] = {}
+    venv: Optional[str] = None  # Path to pre-existing virtualenv
     expected_status: str = "ASSIGNED"  # ASSIGNED means "please start"
 
 

pylet/server.py

@@ -46,6 +46,8 @@ class InstanceSubmissionRequest(BaseModel):
     exclusive: bool = True
     labels: Dict[str, str] = Field(default_factory=dict)
     env: Dict[str, str] = Field(default_factory=dict)
+    # Venv support
+    venv: Optional[str] = None  # Path to pre-existing virtualenv (must be absolute path)
 
 
 class WorkerRegistrationRequest(BaseModel):
@@ -158,6 +160,8 @@ async def _build_instance_response(inst: Dict) -> Dict:
         "labels": _parse_json_field(inst.get("labels"), {}),
         "env": _parse_json_field(inst.get("env"), {}),
         "target_worker": inst.get("target_worker"),
+        # Venv support
+        "venv": inst.get("venv"),
     }
 
 
@@ -227,6 +231,14 @@ async def _validate_submission(request: InstanceSubmissionRequest) -> None:
     if conflicts:
         logger.warning(f"User env vars {conflicts} will be overridden by Pylet")
 
+    # 5. Validate venv path if specified
+    if request.venv is not None:
+        if not request.venv.startswith("/"):
+            raise ValueError("venv must be an absolute path")
+        if ".." in request.venv:
+            raise ValueError("venv path cannot contain '..'")
+        # Note: We don't check if the path exists - that happens at runtime on the worker
+
 
 @app.post("/instances")
 async def submit_instance(request: InstanceSubmissionRequest):
@@ -244,6 +256,7 @@ async def submit_instance(request: InstanceSubmissionRequest):
             exclusive=request.exclusive,
             labels=request.labels,
             env=request.env,
+            venv=request.venv,
         )
         logger.info(
             f"Client submitted instance {instance_id} ('{request.name}') "

pylet/worker.py

@@ -418,9 +418,19 @@ async def _start_instance(self, inst: DesiredInstance) -> None:
         # Uses trap "" PIPE to let task continue if sidecar crashes
         log_dir = str(config.LOG_DIR)
         escaped_cmd = inst.command.replace("'", "'\"'\"'")  # Escape single quotes
+
+        # Activate venv if specified (before user command, not sidecar)
+        if inst.venv:
+            # Use double quotes inside the single-quoted bash command to handle spaces
+            # Escape any double quotes or backslashes in the path
+            safe_venv = inst.venv.replace("\\", "\\\\").replace('"', '\\"')
+            venv_activate = f'source "{safe_venv}/bin/activate" && '
+        else:
+            venv_activate = ""
+
         shell_cmd = (
             f"/bin/bash -c 'set -o pipefail; trap \"\" PIPE; "
-            f"({escaped_cmd}) 2>&1 | "
+            f"({venv_activate}{escaped_cmd}) 2>&1 | "
             f"python3 -m pylet.log_sidecar {shlex.quote(log_dir)} "
             f"{shlex.quote(inst.instance_id)}'"
         )