fix(connector): clean up the local code bundle copy

Without the async LRU cache the connector rebuilds a tarball on every remote execution, and each one lands in a random local path that nobody deletes. Solved this by using a temporary directory that is removed once the upload finishes.

Signed-off-by: Jeff Chung <sh1001309@gmail.com>

Author: jeffoodchain

src/flyte/connectors/_connector.py

@@ -1,6 +1,8 @@
 import asyncio
 import json
 import os
+import pathlib
+import tempfile
 import typing
 from abc import ABC, abstractmethod
 from dataclasses import asdict, dataclass
@@ -248,24 +250,12 @@ async def execute(self, **kwargs) -> Any:
             if not storage.is_remote(tctx.raw_data_path.path):
                 return await TaskTemplate.execute(self, **kwargs)
             else:
-                local_code_bundle = await build_code_bundle(
-                    from_dir=cfg.root_dir,
-                    dryrun=True,
-                )
-                if local_code_bundle.tgz is None:
-                    raise RuntimeError("no tgz found in code bundle")
-                remote_code_path = await storage.put(
-                    local_code_bundle.tgz, prefix + "/code_bundle/" + os.path.basename(local_code_bundle.tgz)
-                )
+                code_bundle = await _build_and_upload_code_bundle(cfg.root_dir, prefix)
                 sc = SerializationContext(
                     project=tctx.action.project,
                     domain=tctx.action.domain,
                     org=tctx.action.org,
-                    code_bundle=CodeBundle(
-                        tgz=remote_code_path,
-                        computed_version=local_code_bundle.computed_version,
-                        destination="/opt/flyte/",
-                    ),
+                    code_bundle=code_bundle,
                     version=tctx.version,
                     image_cache=await build_images.aio(task.parent_env()) if task.parent_env else None,
                     root_dir=cfg.root_dir,
@@ -331,6 +321,25 @@ async def execute(self, **kwargs) -> Any:
         return tuple(resource.outputs.values())
 
 
+async def _build_and_upload_code_bundle(from_dir: pathlib.Path, prefix: str) -> CodeBundle:
+    with tempfile.TemporaryDirectory(prefix="flyte-code-bundle-") as tmp_dir:
+        local_code_bundle = await build_code_bundle(
+            from_dir=from_dir,
+            dryrun=True,
+            copy_bundle_to=pathlib.Path(tmp_dir),
+        )
+        if local_code_bundle.tgz is None:
+            raise RuntimeError("no tgz found in code bundle")
+        remote_code_path = await storage.put(
+            local_code_bundle.tgz, prefix + "/code_bundle/" + os.path.basename(local_code_bundle.tgz)
+        )
+    return CodeBundle(
+        tgz=remote_code_path,
+        computed_version=local_code_bundle.computed_version,
+        destination="/opt/flyte/",
+    )
+
+
 async def get_resource_proto(resource: Resource) -> connector_pb2.Resource:
     if resource.outputs:
         interface = NativeInterface.from_types(inputs={}, outputs={k: type(v) for k, v in resource.outputs.items()})

tests/flyte/connector/test_connector_executor.py

@@ -0,0 +1,78 @@
+import pathlib
+
+import pytest
+
+from flyte.connectors import _connector
+from flyte.models import CodeBundle
+
+
+@pytest.mark.asyncio
+async def test_upload_code_bundle_cleans_up_local_copy(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    bundle_dirs: list[pathlib.Path] = []
+
+    async def build_code_bundle(
+        *, from_dir: pathlib.Path, dryrun: bool, copy_bundle_to: pathlib.Path | None = None
+    ) -> CodeBundle:
+        assert from_dir == tmp_path
+        assert dryrun is True
+        assert copy_bundle_to is not None
+        bundle_dirs.append(copy_bundle_to)
+        local_bundle = copy_bundle_to / "bundle.tar.gz"
+        local_bundle.write_bytes(b"bundle")
+        return CodeBundle(tgz=str(local_bundle), computed_version="v1")
+
+    async def put(from_path: str, to_path: str) -> str:
+        local_path = pathlib.Path(from_path)
+        assert bundle_dirs
+        assert local_path.parent == bundle_dirs[0]
+        assert local_path.exists()
+        assert to_path == "s3://bucket/prefix/code_bundle/bundle.tar.gz"
+        return "s3://bucket/uploaded/bundle.tar.gz"
+
+    monkeypatch.setattr(_connector, "build_code_bundle", build_code_bundle)
+    monkeypatch.setattr(_connector.storage, "put", put)
+
+    result = await _connector._build_and_upload_code_bundle(tmp_path, "s3://bucket/prefix")
+
+    assert result == CodeBundle(
+        tgz="s3://bucket/uploaded/bundle.tar.gz",
+        computed_version="v1",
+        destination="/opt/flyte/",
+    )
+    assert len(bundle_dirs) == 1
+    assert not bundle_dirs[0].exists()
+
+
+@pytest.mark.asyncio
+async def test_upload_code_bundle_cleans_up_local_copy_when_upload_fails(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    bundle_dirs: list[pathlib.Path] = []
+
+    async def build_code_bundle(
+        *, from_dir: pathlib.Path, dryrun: bool, copy_bundle_to: pathlib.Path | None = None
+    ) -> CodeBundle:
+        assert from_dir == tmp_path
+        assert dryrun is True
+        assert copy_bundle_to is not None
+        bundle_dirs.append(copy_bundle_to)
+        local_bundle = copy_bundle_to / "bundle.tar.gz"
+        local_bundle.write_bytes(b"bundle")
+        return CodeBundle(tgz=str(local_bundle), computed_version="v1")
+
+    async def put(from_path: str, to_path: str) -> str:
+        local_path = pathlib.Path(from_path)
+        assert bundle_dirs
+        assert local_path.parent == bundle_dirs[0]
+        assert local_path.exists()
+        assert to_path == "s3://bucket/prefix/code_bundle/bundle.tar.gz"
+        raise OSError("upload failed")
+
+    monkeypatch.setattr(_connector, "build_code_bundle", build_code_bundle)
+    monkeypatch.setattr(_connector.storage, "put", put)
+
+    with pytest.raises(OSError, match="upload failed"):
+        await _connector._build_and_upload_code_bundle(tmp_path, "s3://bucket/prefix")
+
+    assert len(bundle_dirs) == 1
+    assert not bundle_dirs[0].exists()