From 1c634089a5d6349734eb27473c8a020f55c7a9ca Mon Sep 17 00:00:00 2001 From: Sergey Serebryakov Date: Fri, 4 Aug 2023 14:39:52 -0700 Subject: [PATCH] Adding support for `docker:` protocol for DockerImage.from_string method. (#333) This method correctly processes image names that start with `docker:[/]*` protocol specification. --- .../mlcube_singularity/singularity_client.py | 5 ++++- .../tests/test_singularity_client.py | 10 +++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/runners/mlcube_singularity/mlcube_singularity/singularity_client.py b/runners/mlcube_singularity/mlcube_singularity/singularity_client.py index 03a3953..542d15b 100644 --- a/runners/mlcube_singularity/mlcube_singularity/singularity_client.py +++ b/runners/mlcube_singularity/mlcube_singularity/singularity_client.py @@ -1,4 +1,3 @@ -import json import logging import typing as t from enum import Enum @@ -77,6 +76,10 @@ def from_string(cls, name: str) -> "DockerImage": Returns: DockerImage instance with parsed components. """ + # Remove protocol if present + if name.startswith("docker:"): + name = name[7:].lstrip("/") + # Split into parts that are separated by "/". parts: t.List[str] = name.strip().split("/") diff --git a/runners/mlcube_singularity/mlcube_singularity/tests/test_singularity_client.py b/runners/mlcube_singularity/mlcube_singularity/tests/test_singularity_client.py index a351222..92dfdea 100644 --- a/runners/mlcube_singularity/mlcube_singularity/tests/test_singularity_client.py +++ b/runners/mlcube_singularity/mlcube_singularity/tests/test_singularity_client.py @@ -2,7 +2,8 @@ from unittest import TestCase import semver -from mlcube_singularity.singularity_client import Client, DockerImage, Runtime, Version +from mlcube_singularity.singularity_client import (Client, DockerImage, + Runtime, Version) class TestSingularityRunner(TestCase): @@ -165,13 +166,16 @@ def test_docker_image_from_string(self) -> None: ) def test_docker_image_to_string(self) -> None: - names = [ + names: t.List[t.Union[t.Tuple[str, str], str]] = [ "LOCATION-docker.pkg.dev/PROJECT-ID/REPOSITORY/IMAGE", "LOCATION-docker.pkg.dev/PROJECT-ID/REPOSITORY/IMAGE:TAG", "LOCATION-docker.pkg.dev/PROJECT-ID/REPOSITORY/IMAGE@IMG-DIGEST", "USERNAME/REPOSITORY:TAG", "mlcommons/hello_world:0.0.1", + ("docker:mlcommons/hello_world:0.0.1", "mlcommons/hello_world:0.0.1"), + ("docker:///mlcommons/hello_world:0.0.1", "mlcommons/hello_world:0.0.1"), "mlcommons/hello_world@sha256:9b77d4cb97f8dcf14ac137bf65185fc8980578", ] for name in names: - self.assertEqual(str(DockerImage.from_string(name)), name) + name_in, name_out = (name, name) if isinstance(name, str) else (name[0], name[1]) + self.assertEqual(str(DockerImage.from_string(name_in)), name_out)