docker_scheduler: add support for building a new image from a workspace (#347)

Summary:
This adds a new `build_workspace_image` method on the DockerScheduler. This takes in a base image and the fsspec workspace path and creates a new docker image with the workspace overlaid on the base image and returns it.

This is the first steps towards implementing #333.

When docker builds an image it uses a build context which is uploaded to it via a tarball. This builds the context by walking the fsspec path for the files and taring them.

## Dockerfile:

A generated Dockerfile is added which does a naive copy. If the workspace contains a `Dockerfile`, that Dockerfile is used instead of the generated one.

Dockerfile
```
FROM <base>
COPY . .
```

Pull Request resolved: #347

Test Plan:
pytest torchx/schedulers/tests/docker_scheduler_test.py
pyre

Reviewed By: kiukchung

Differential Revision: D32472559

Pulled By: d4l3k

fbshipit-source-id: 79cc5e01af777a586d3d723c690a859af648950a

Author: d4l3k

requirements.txt

@@ -3,3 +3,4 @@ docstring-parser==0.8.1
 pyyaml
 docker
 filelock
+fsspec

torchx/schedulers/docker_scheduler.py

@@ -5,13 +5,27 @@
 # LICENSE file in the root directory of this source tree.
 
 import fnmatch
+import io
 import logging
 import os.path
+import posixpath
+import tarfile
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Mapping, TYPE_CHECKING, Any, Dict, Iterable, Optional, List, Union
+from typing import (
+    Mapping,
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Optional,
+    List,
+    Union,
+    IO,
+)
 
+import fsspec
 import torchx
 import yaml
 from torchx.schedulers.api import (
@@ -390,6 +404,21 @@ def log_iter(
         else:
             return logs
 
+    def build_workspace_image(self, img: str, workspace: str) -> str:
+        """
+        build_workspace_image creates a new image with the files in workspace
+        overlaid on top of it.
+
+        Args:
+            img: a Docker image to use as a base
+            workspace: a fsspec path to a directory with contents to be overlaid
+
+        Returns:
+            The new Docker image ID.
+
+        """
+        return _build_container_from_workspace(self._client(), img, workspace)
+
 
 def _to_str(a: Union[str, bytes]) -> str:
     if isinstance(a, bytes):
@@ -403,3 +432,59 @@ def create_scheduler(session_name: str, **kwargs: Any) -> DockerScheduler:
     return DockerScheduler(
         session_name=session_name,
     )
+
+
+def _copy_to_tarfile(workspace: str, tf: tarfile.TarFile) -> None:
+    # TODO(d4l3k) implement docker ignore files
+
+    fs, path = fsspec.core.url_to_fs(workspace)
+    assert isinstance(path, str), "path must be str"
+
+    for dir, dirs, files in fs.walk(path, detail=True):
+        assert isinstance(dir, str), "path must be str"
+        relpath = posixpath.relpath(dir, path)
+        for file, info in files.items():
+            print(relpath, dir, file, info)
+            with fs.open(info["name"], "rb") as f:
+                tinfo = tarfile.TarInfo(posixpath.join(relpath, file))
+                tinfo.size = info["size"]
+                tf.addfile(tinfo, f)
+
+
+def _build_context(img: str, workspace: str) -> IO[bytes]:
+    # f is closed by parent, NamedTemporaryFile auto closes on GC
+    f = tempfile.NamedTemporaryFile(  # noqa P201
+        prefix="torchx-context",
+        suffix=".tar",
+    )
+    dockerfile = bytes(f"FROM {img}\nCOPY . .\n", encoding="utf-8")
+    with tarfile.open(fileobj=f, mode="w") as tf:
+        info = tarfile.TarInfo("Dockerfile")
+        info.size = len(dockerfile)
+        tf.addfile(info, io.BytesIO(dockerfile))
+
+        _copy_to_tarfile(workspace, tf)
+
+    f.seek(0)
+    return f
+
+
+def _build_container_from_workspace(
+    client: "DockerClient", img: str, workspace: str
+) -> str:
+    context = _build_context(img, workspace)
+
+    try:
+        image, logs = client.images.build(
+            fileobj=context,
+            custom_context=True,
+            pull=True,
+            rm=True,
+            labels={
+                LABEL_VERSION: torchx.__version__,
+            },
+        )
+    finally:
+        context.close()
+    print(image)
+    return image.id

torchx/schedulers/test/docker_scheduler_test.py

@@ -9,6 +9,7 @@
 from datetime import datetime, timedelta
 from unittest.mock import patch
 
+import fsspec
 from docker.types import DeviceRequest
 from torchx import specs
 from torchx.schedulers.api import Stream
@@ -347,3 +348,34 @@ def test_docker_submit_dist(self) -> None:
             self.assertEqual(
                 desc.roles_statuses[0].replicas[1].state, AppState.SUCCEEDED
             )
+
+        def test_docker_workspace(self) -> None:
+            fs = fsspec.filesystem("memory")
+            fs.mkdirs("test_workspace/bar", exist_ok=True)
+            with fs.open("test_workspace/bar/foo.sh", "w") as f:
+                f.write("exit 0")
+
+            img = self.scheduler.build_workspace_image(
+                "busybox",
+                "memory://test_workspace",
+            )
+
+            app = AppDef(
+                name="test-app",
+                roles=[
+                    Role(
+                        name="ping",
+                        image=img,
+                        entrypoint="sh",
+                        args=[
+                            "bar/foo.sh",
+                        ],
+                    ),
+                ],
+            )
+            app_id = self.scheduler.submit(app, {})
+            print(app_id)
+
+            desc = self.wait(app_id)
+            self.assertIsNotNone(desc)
+            self.assertEqual(AppState.SUCCEEDED, desc.state)