feat: Add Volume class to docker client and --user flag to cli (#241)

#### Description of changes

- Add volume class to docker_client and docker_cleanup
- Add user flag to docker container run and serve pipelines
- Add tests to explore volume permission and user id switching

#### Testing done
Unit Testing

---------

Co-authored-by: Alessandro Angioi <alessandro.angioi@simulation.science>

Author: angela-ko

tesseract_core/sdk/cli.py

@@ -510,6 +510,15 @@ def serve(
             ),
         ),
     ] = None,
+    user: Annotated[
+        str | None,
+        typer.Option(
+            "--user",
+            help=(
+                "User to run the Tesseracts as e.g. '1000' or '1000:1000' (uid:gid)."
+            ),
+        ),
+    ] = None,
 ) -> None:
     """Serve one or more Tesseract images.
 
@@ -555,6 +564,7 @@ def serve(
             num_workers,
             no_compose,
             service_names_list,
+            user,
         )
     except RuntimeError as ex:
         raise UserError(
@@ -830,6 +840,13 @@ def run_container(
             ),
         ),
     ] = None,
+    user: Annotated[
+        str | None,
+        typer.Option(
+            "--user",
+            help=("User to run the Tesseract as e.g. '1000' or '1000:1000' (uid:gid)."),
+        ),
+    ] = None,
 ) -> None:
     """Execute a command in a Tesseract.
 
@@ -873,7 +890,7 @@ def run_container(
 
     try:
         result_out, result_err = engine.run_tesseract(
-            tesseract_image, cmd, args, volumes=volume, gpus=gpus
+            tesseract_image, cmd, args, volumes=volume, gpus=gpus, user=user
         )
 
     except ImageNotFound as e:

tesseract_core/sdk/docker_client.py

@@ -512,13 +512,15 @@ def run(
         ports: dict | None = None,
         stdout: bool = True,
         stderr: bool = False,
+        user: str | None = None,
     ) -> Container | tuple[bytes, bytes] | bytes:
         """Run a command in a container from an image.
 
         Params:
             image: The image name or id to run the command in.
             command: The command to run in the container.
             volumes: A dict of volumes to mount in the container.
+            user: String of user information to run command as in the format "uid:(optional)gid".
             device_requests: A list of device requests for the container.
             detach: If True, run the container in detached mode. Detach must be set to
                     True if we wish to retrieve the container id of the running container,
@@ -556,6 +558,9 @@ def run(
                 )
             optional_args.extend(volume_args)
 
+        if user:
+            optional_args.extend(["-u", user])
+
         if device_requests:
             gpus_str = ",".join(device_requests)
             optional_args.extend(["--gpus", f'"device={gpus_str}"'])
@@ -774,6 +779,116 @@ def _update_projects(include_stopped: bool = False) -> dict[str, list_[str]]:
         return project_container_map
 
 
+@dataclass
+class Volume:
+    """Volume class to wrap Docker volumes."""
+
+    name: str
+    attrs: dict
+
+    @classmethod
+    def from_dict(cls, json_dict: dict) -> "Volume":
+        """Create an Image object from a json dictionary.
+
+        Params:
+            json_dict: The json dictionary to create the object from.
+
+        Returns:
+            The created volume object.
+        """
+        return cls(
+            name=json_dict.get("Name", None),
+            attrs=json_dict,
+        )
+
+    def remove(self, force: bool = False) -> None:
+        """Remove a Docker volume.
+
+        Params:
+            force: If True, force the removal of the volume.
+        """
+        docker = _get_executable("docker")
+        try:
+            _ = subprocess.run(
+                [*docker, "volume", "rm", "--force" if force else "", self.name],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+        except subprocess.CalledProcessError as ex:
+            raise NotFound(f"Error removing volume {self.name}: {ex}") from ex
+
+
+class Volumes:
+    """Volume class to wrap Docker volumes."""
+
+    @staticmethod
+    def create(name: str) -> Volume:
+        """Create a Docker volume.
+
+        Params:
+            name: The name of the volume to create.
+
+        Returns:
+            The created volume object.
+        """
+        docker = _get_executable("docker")
+        try:
+            _ = subprocess.run(
+                [*docker, "volume", "create", name],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            return Volumes.get(name)
+        except subprocess.CalledProcessError as ex:
+            raise NotFound(f"Error creating volume {name}: {ex}") from ex
+
+    @staticmethod
+    def get(name: str) -> Volume:
+        """Get a Docker volume.
+
+        Params:
+            name: The name of the volume to get.
+
+        Returns:
+            The volume object.
+        """
+        docker = _get_executable("docker")
+        try:
+            result = subprocess.run(
+                [*docker, "volume", "inspect", name],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            json_dict = json.loads(result.stdout)
+        except subprocess.CalledProcessError as ex:
+            raise NotFound(f"Volume {name} not found: {ex}") from ex
+        if not json_dict:
+            raise NotFound(f"Volume {name} not found.")
+        return Volume.from_dict(json_dict[0])
+
+    @staticmethod
+    def list() -> list[str]:
+        """List all Docker volumes.
+
+        Returns:
+            List of volume names.
+        """
+        docker = _get_executable("docker")
+        try:
+            result = subprocess.run(
+                [*docker, "volume", "ls", "-q"],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+        except subprocess.CalledProcessError as ex:
+            raise APIError(f"Error listing volumes: {ex}") from ex
+        return result.stdout.strip().split("\n")
+
+
 class DockerException(Exception):
     """Base class for Docker CLI exceptions."""
 
@@ -847,6 +962,7 @@ def __init__(self) -> None:
         self.containers = Containers()
         self.images = Images()
         self.compose = Compose()
+        self.volumes = Volumes()
 
     @staticmethod
     def info() -> tuple:

tesseract_core/sdk/engine.py

@@ -549,6 +549,7 @@ def serve(
     num_workers: int = 1,
     no_compose: bool = False,
     service_names: list[str] | None = None,
+    user: str | None = None,
 ) -> str:
     """Serve one or more Tesseract images.
 
@@ -566,6 +567,7 @@ def serve(
         num_workers: number of workers to use for serving the Tesseracts.
         no_compose: if True, do not use Docker Compose to serve the Tesseracts.
         service_names: list of service names under which to expose each Tesseract container on the shared network.
+        user: user to run the Tesseracts as, e.g. '1000' or '1000:1000' (uid:gid).
 
     Returns:
         A string representing the Tesseract project ID.
@@ -642,6 +644,7 @@ def serve(
             ports=port_mappings,
             detach=True,
             volumes=volumes,
+            user=user,
         )
         # wait for server to start
         timeout = 30
@@ -671,6 +674,7 @@ def serve(
         gpus,
         num_workers,
         debug=debug,
+        user=user,
     )
     compose_fname = f"docker-compose-{_id_generator()}.yml"
 
@@ -696,6 +700,7 @@ def _create_docker_compose_template(
     gpus: list[str] | None = None,
     num_workers: int = 1,
     debug: bool = False,
+    user: str | None = None,
 ) -> str:
     """Create Docker Compose template."""
     services = []
@@ -744,6 +749,7 @@ def _create_docker_compose_template(
     for i, image_id in enumerate(image_ids):
         service = {
             "name": service_names[i],
+            "user": user,
             "image": image_id,
             "port": f"{ports[i]}:8000",
             "volumes": volumes,
@@ -756,8 +762,28 @@ def _create_docker_compose_template(
         }
 
         services.append(service)
+
+    docker_volumes = {}  # Dictionary of volume names mapped to whether or not they already exist
+    if volumes:
+        for volume in volumes:
+            source = volume.split(":")[0]
+            # Check if source exists to determine if specified volume is a docker volume
+            if not Path(source).exists():
+                # Check if volume exists
+                if not docker_client.volumes.get(source):
+                    if "/" not in source:
+                        docker_volumes[source] = False
+                    else:
+                        raise ValueError(
+                            f"Volume/Path {source} does not already exist, "
+                            "and new volume cannot be created due to '/' in name."
+                        )
+                else:
+                    # Docker volume is external
+                    docker_volumes[source] = True
+
     template = ENV.get_template("docker-compose.yml")
-    return template.render(services=services)
+    return template.render(services=services, docker_volumes=docker_volumes)
 
 
 def _id_generator(
@@ -817,6 +843,7 @@ def run_tesseract(
     volumes: list[str] | None = None,
     gpus: list[int | str] | None = None,
     ports: dict[str, str] | None = None,
+    user: str | None = None,
 ) -> tuple[str, str]:
     """Start a Tesseract and execute a given command.
 
@@ -828,6 +855,7 @@ def run_tesseract(
         gpus: list of GPUs, as indices or names, to passthrough the container.
         ports: dictionary of ports to bind to the host. Key is the host port,
             value is the container port.
+        user: user to run the Tesseract as, e.g. '1000' or '1000:1000' (uid:gid).
 
     Returns:
         Tuple with the stdout and stderr of the Tesseract.
@@ -896,6 +924,7 @@ def run_tesseract(
         detach=False,
         remove=True,
         stderr=True,
+        user=user,
     )
     stdout = stdout.decode("utf-8")
     stderr = stderr.decode("utf-8")

tesseract_core/sdk/templates/docker-compose.yml

@@ -4,6 +4,9 @@ services:
     image: {{ service.image }}
     restart: unless-stopped
     command: ["serve", "--host", "0.0.0.0", "--num-workers", "{{ service.num_workers }}"]
+    {%- if service.user %}
+    user: "{{ service.user }}"
+    {%- endif %}
     ports:
       - {{ service.port }}
     {% if service.environment.TESSERACT_DEBUG == "1" %}
@@ -39,5 +42,15 @@ services:
     {% endif %}
 {% endfor %}
 
+{%- if docker_volumes %}
+volumes:
+{%- for name, external in docker_volumes.items() %}
+  {{ name }}:
+    {%- if external %}
+    external: true
+    {% endif %}
+{% endfor %}
+{% endif %}
+
 networks:
   multi-tesseract-network:

tests/conftest.py

@@ -207,6 +207,20 @@ def docker_client():
     return docker_client_module.CLIDockerClient()
 
 
+@pytest.fixture
+def docker_volume(docker_client):
+    # Create the Docker volume
+    volume = docker_client.volumes.create(name="docker_client_test_volume")
+    try:
+        yield volume
+    finally:
+        try:
+            volume.remove()
+        except Exception:
+            # already removed
+            pass
+
+
 @pytest.fixture(scope="module")
 def docker_cleanup_module(docker_client, request):
     """Clean up all tesseracts created by the tests after the module exits."""
@@ -222,7 +236,7 @@ def docker_cleanup(docker_client, request):
 def _docker_cleanup(docker_client, request):
     """Clean up all tesseracts created by the tests."""
     # Shared object to track what objects need to be cleaned up in each test
-    context = {"images": [], "project_ids": [], "containers": []}
+    context = {"images": [], "project_ids": [], "containers": [], "volumes": []}
 
     def pprint_exc(e: BaseException) -> str:
         """Pretty print exception."""
@@ -268,6 +282,18 @@ def cleanup_func():
             except Exception as e:
                 failures.append(f"Failed to remove image {image}: {pprint_exc(e)}")
 
+        # Remove volumes
+        for volume in context["volumes"]:
+            try:
+                if isinstance(volume, str):
+                    volume_obj = docker_client.volumes.get(volume)
+                else:
+                    volume_obj = volume
+
+                volume_obj.remove(force=True)
+            except Exception as e:
+                failures.append(f"Failed to remove volume {volume}: {pprint_exc(e)}")
+
         if failures:
             raise RuntimeError(
                 "Failed to clean up some Docker objects during test teardown:\n"