Merge pull request #2288 from transformerlab/fix/remote_launch_cwd

Fix remote working directory by storing in a file

Author: dadmobile

api/localprovider_pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
   "soundfile==0.13.1",
   "tensorboardX==2.6.2.2",
   "timm==1.0.15",
-  "transformerlab==0.1.42",
+  "transformerlab==0.1.43",
   "transformerlab-inference==0.2.52",
   "transformers==4.57.1",
   "wandb==0.23.1",

api/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
   "python-dotenv==1.2.2",
   "python-multipart==0.0.32",
   "sqlalchemy[asyncio]==2.0.50",
-  "transformerlab==0.1.42",
+  "transformerlab==0.1.43",
   "transformerlab-inference==0.2.52",
   "uvicorn==0.35.0",
   "watchfiles==1.2.0",

api/test/test_remote_workdir_cd.py

@@ -0,0 +1,77 @@
+"""Tests for prepend_remote_workdir_cd.
+
+Regression coverage for the `//main.py` quirk: on remote providers the run command
+executes from the provider's default cwd (e.g. `/` for Lambda cloud-init), not the
+directory where lab.copy_file_mounts() synced the task files. prepend_remote_workdir_cd
+prefixes a `cd` into the recorded workdir so a bare `python main.py` resolves.
+"""
+
+from transformerlab.services.compute_provider.launch_credentials import WORKDIR_FILE_PATH
+from transformerlab.services.compute_provider.launch_template import prepend_remote_workdir_cd
+from transformerlab.shared.models.models import ProviderType
+
+EXPECTED_PREFIX = f'cd "$(cat {WORKDIR_FILE_PATH} 2>/dev/null || echo "$HOME")" && '
+
+
+def test_prefixes_cd_for_remote_provider_with_file_mounts():
+    for provider_type in (
+        ProviderType.LAMBDA.value,
+        ProviderType.RUNPOD.value,
+        ProviderType.SKYPILOT.value,
+        ProviderType.VASTAI.value,
+        ProviderType.DSTACK.value,
+    ):
+        out = prepend_remote_workdir_cd(
+            "python main.py",
+            task_id="task-123",
+            file_mounts=True,
+            provider_type=provider_type,
+        )
+        assert out == EXPECTED_PREFIX + "python main.py", provider_type
+
+
+def test_local_provider_is_not_prefixed():
+    # Local already forces cwd == $HOME == the file-drop dir.
+    out = prepend_remote_workdir_cd(
+        "python main.py",
+        task_id="task-123",
+        file_mounts=True,
+        provider_type=ProviderType.LOCAL.value,
+    )
+    assert out == "python main.py"
+
+
+def test_no_prefix_when_file_mounts_not_enabled():
+    # No synced files (inlined run, or dict-form file_mounts handled natively):
+    # copy_file_mounts isn't injected, so there's no workdir file and nothing to cd into.
+    for file_mounts in (False, None, {"/remote": "/local"}):
+        out = prepend_remote_workdir_cd(
+            "python main.py",
+            task_id="task-123",
+            file_mounts=file_mounts,
+            provider_type=ProviderType.LAMBDA.value,
+        )
+        assert out == "python main.py", file_mounts
+
+
+def test_no_prefix_without_task_id():
+    out = prepend_remote_workdir_cd(
+        "python main.py",
+        task_id=None,
+        file_mounts=True,
+        provider_type=ProviderType.LAMBDA.value,
+    )
+    assert out == "python main.py"
+
+
+def test_preserves_shell_operators_in_command():
+    # The whole command is later shlex.quoted as one payload to tfl-remote-trap, so
+    # operators must be carried through verbatim after the cd prefix.
+    cmd = "python main.py --epochs 3 && echo done"
+    out = prepend_remote_workdir_cd(
+        cmd,
+        task_id="task-123",
+        file_mounts=True,
+        provider_type=ProviderType.RUNPOD.value,
+    )
+    assert out == EXPECTED_PREFIX + cmd

api/transformerlab/services/compute_provider/launch_credentials.py

@@ -9,6 +9,11 @@
 # lab.init() not required; copy_file_mounts uses _TFL_JOB_ID, _TFL_EXPERIMENT_ID / TFL_EXPERIMENT_ID, and job_data
 COPY_FILE_MOUNTS_SETUP = 'python -c "from lab import lab; lab.copy_file_mounts()"'
 
+# File that lab.copy_file_mounts() writes on the remote box, recording the absolute
+# directory where task files were synced. The remote run command cd's into it so a bare
+# `python main.py` resolves. Keep in sync with lab-sdk's copy_file_mounts().
+WORKDIR_FILE_PATH = "/tmp/.tfl_task_workdir"
+
 # RunPod (and similar) use /workspace as a writable persistent path
 RUNPOD_AWS_CREDENTIALS_DIR = "/workspace/.aws"
 

api/transformerlab/services/compute_provider/launch_template.py

@@ -16,6 +16,7 @@
 from transformerlab.services.compute_provider.launch_credentials import (
     COPY_FILE_MOUNTS_SETUP,
     RUNPOD_AWS_CREDENTIALS_DIR,
+    WORKDIR_FILE_PATH,
     generate_aws_credentials_setup,
     generate_azure_credentials_setup,
     generate_gcp_credentials_setup,
@@ -55,6 +56,36 @@
 from werkzeug.utils import secure_filename
 
 
+def prepend_remote_workdir_cd(
+    command: str,
+    *,
+    task_id: Optional[str],
+    file_mounts: Any,
+    provider_type: str,
+) -> str:
+    """Prefix a remote run command with a `cd` into the synced task-file directory.
+
+    When task files are synced to a remote box via lab.copy_file_mounts() (any non-Local
+    provider with file_mounts enabled), the files land in the directory copy_file_mounts
+    resolved at runtime (~/sky_workdir or $HOME) and recorded in the WORKDIR_FILE_PATH
+    file. The remote run command, however, executes from the provider's default cwd —
+    `/` for Lambda cloud-init, the image WORKDIR for RunPod — so a bare `python main.py`
+    can't find the synced file (CPython reports it as `//main.py` when cwd is `/`). cd into
+    the recorded workdir first so the user's run command sees its files, matching the Local
+    provider's cwd == $HOME invariant.
+
+    We read the workdir file rather than hardcoding $HOME because the run shell's $HOME may
+    not match expanduser("~") used during sync (e.g. Lambda cloud-init runs as root with
+    HOME unset, so bash $HOME is empty while expanduser("~") resolves to /root). The `||
+    echo "$HOME"` is a best-effort fallback for the case where the file is somehow absent.
+
+    Local is excluded because it already forces cwd == $HOME == the file-drop dir.
+    """
+    if task_id and file_mounts is True and provider_type != ProviderType.LOCAL.value:
+        return f'cd "$(cat {WORKDIR_FILE_PATH} 2>/dev/null || echo "$HOME")" && {command}'
+    return command
+
+
 async def launch_template_on_provider(
     provider_id: str,
     request: ProviderTemplateLaunchRequest,
@@ -669,6 +700,15 @@ async def launch_template_on_provider(
             post_hook=str(post_setup_hook) if post_setup_hook is not None else None,
         )
 
+    # On remote providers, cd into the directory where lab.copy_file_mounts() synced the
+    # task files so a bare `python main.py` resolves (see prepend_remote_workdir_cd).
+    command_with_hooks = prepend_remote_workdir_cd(
+        command_with_hooks,
+        task_id=request.task_id,
+        file_mounts=request.file_mounts,
+        provider_type=provider.type,
+    )
+
     # Wrap the user command with tfl-remote-trap so we can track live_status in job_data.
     # This uses the tfl-remote-trap helper from the transformerlab SDK, which:
     #   - sets job_data.live_status="started" when execution begins

lab-sdk/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab"
-version = "0.1.42"
+version = "0.1.43"
 description = "Python SDK for Transformer Lab"
 readme = "README.md"
 requires-python = ">=3.10"

lab-sdk/src/lab/lab_facade.py

@@ -376,6 +376,20 @@ async def _copy_file_mounts_async(self, job_id: str) -> None:
         else:
             dest_dir = home_dir
 
+        # Record the resolved destination so the remote run command can `cd` into the
+        # exact directory where these task files land. The launch template prepends
+        # `cd "$(cat /tmp/.tfl_task_workdir ...)"` to the run command on remote
+        # providers. We write the absolute path here rather than letting the run shell
+        # recompute it because the run shell's $HOME may not match expanduser("~") used
+        # during sync — e.g. Lambda's cloud-init runs as root with HOME unset, so bash
+        # `$HOME` is empty while expanduser("~") resolves to /root. (Keep this path in
+        # sync with launch_template.py's cd prefix.)
+        try:
+            with open("/tmp/.tfl_task_workdir", "w") as workdir_file:
+                workdir_file.write(dest_dir)
+        except Exception:
+            logger.debug("Failed to write task workdir file", exc_info=True)
+
         # Determine if we're working with a remote URI (e.g., s3://)
         is_remote = storage.is_remote_path(task_dir)
         protocol_prefix = ""

lab-sdk/tests/test_lab_facade.py

@@ -1291,6 +1291,11 @@ async def test_copy_file_mounts_without_lab_init(tmp_path, monkeypatch):
     with open(dest) as f:
         assert f.read() == "hi"
 
+    # The workdir file records the absolute dir where files landed, so the remote
+    # run command can `cd` into it (no ~/sky_workdir here, so dest is $HOME).
+    with open("/tmp/.tfl_task_workdir") as f:
+        assert f.read() == str(home)
+
 
 def test_download_registry_model(tmp_path, monkeypatch):
     _fresh(monkeypatch)