diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 0d55832e65..216abecb99 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -90,6 +90,13 @@ def __post_init__(self): self._is_force_push = os.environ.get(FLYTE_FORCE_PUSH_IMAGE_SPEC, False) # False by default if self.registry: self.registry = self.registry.lower() + if not validate_container_registry_name(self.registry): + raise ValueError( + f"Invalid container registry name: '{self.registry}'.\n Expected formats:\n" + f"- 'localhost:30000' (for local registries)\n" + f"- 'ghcr.io/username' (for GitHub Container Registry)\n" + f"- 'docker.io/username' (for docker hub)\n" + ) # If not set, help the user set this option as well, to support the older default behavior where existence # of the source root implied that copying of files was needed. @@ -407,3 +414,12 @@ def _get_builder(cls, builder: str) -> ImageSpecBuilder: f" Please upgrade envd to v0.3.39+." ) return cls._REGISTRY[builder][0] + + +def validate_container_registry_name(name: str) -> bool: + """Validate Docker container registry name.""" + # Define the regular expression for the registry name + registry_pattern = r"^(localhost:\d{1,5}|([a-z\d\._-]+)(:\d{1,5})?)(/[\w\.-]+)*$" + + # Use regex to validate the given name + return bool(re.match(registry_pattern, name)) diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 2694d028e5..6a102292ed 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -215,3 +215,29 @@ def test_update_image_spec_copy_handling(): update_image_spec_copy_handling(image_spec, ss) assert image_spec.source_copy_mode is None assert image_spec.source_root is None + + +def test_registry_name(): + invalid_registry_names = [ + "invalid:port:50000", + "ghcr.io/flyteorg:latest", + "flyteorg:latest" + ] + for invalid_registry_name in invalid_registry_names: + with pytest.raises(ValueError, match="Invalid container registry name"): + ImageSpec(registry=invalid_registry_name) + + valid_registry_names = [ + "localhost:30000", + "localhost:30000/flyte", + "192.168.1.1:30000", + "192.168.1.1:30000/myimage", + "ghcr.io/flyteorg", + "my.registry.com/myimage", + "my.registry.com:5000/myimage", + "myregistry:5000/myimage", + "us-west1-docker.pkg.dev/example.com/my-project/my-repo" + "flyteorg", + ] + for valid_registry_name in valid_registry_names: + ImageSpec(registry=valid_registry_name)