wandb 0.16.6__py3-none-any.whl → 0.17.0__py3-none-any.whl

wandb/sdk/launch/agent/config.py

@@ -80,17 +80,7 @@ class RegistryConfig(BaseModel):
     @validator("uri")  # type: ignore
     @classmethod
     def validate_uri(cls, uri: str) -> str:
-        for regex in [
-            GCP_ARTIFACT_REGISTRY_URI_REGEX,
-            AZURE_CONTAINER_REGISTRY_URI_REGEX,
-            ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
-        ]:
-            if regex.match(uri):
-                return uri
-        raise ValueError(
-            "Invalid uri. URI must be a repository URI for an "
-            "ECR, ACR, or GCP Artifact Registry."
-        )
+        return validate_registry_uri(uri)
 
 
 class EnvironmentConfig(BaseModel):
@@ -186,6 +176,14 @@ class BuilderConfig(BaseModel):
         """Right now there are no required fields for docker builds."""
         return values
 
+    @validator("destination")  # type: ignore
+    @classmethod
+    def validate_destination(cls, destination: Optional[str]) -> Optional[str]:
+        """Validate that the destination is a valid container registry URI."""
+        if destination is None:
+            return None
+        return validate_registry_uri(destination)
+
 
 class AgentConfig(BaseModel):
     """Configuration for the Launch agent."""
@@ -194,9 +192,6 @@ class AgentConfig(BaseModel):
         default=[],
         description="The queues to use for this agent.",
     )
-    project: Optional[str] = Field(
-        description="The W&B project to use for this agent.",
-    )
     entity: Optional[str] = Field(
         description="The W&B entity to use for this agent.",
     )
@@ -236,3 +231,66 @@ class AgentConfig(BaseModel):
 
     class Config:
         extra = "forbid"
+
+
+def validate_registry_uri(uri: str) -> str:
+    """Validate that the registry URI is a valid container registry URI.
+
+    The URI should resolve to an image name in a container registry. The recognized
+    formats are for ECR, ACR, and GCP Artifact Registry. If the URI does not match
+    any of these formats, a warning is printed indicating the registry type is not
+    recognized and the agent can't guarantee that images can be pushed.
+
+    If the format is recognized but does not resolve to an image name, an
+    error is raised. For example, if the URI is an ECR URI but does not include
+    an image name or includes a tag as well as an image name, an error is raised.
+    """
+    tag_msg = (
+        "Destination for built images may not include a tag, but the URI provided "
+        "includes the suffix '{tag}'. Please remove the tag and try again. The agent "
+        "will automatically tag each image with a unique hash of the source code."
+    )
+    if uri.startswith("https://"):
+        uri = uri[8:]
+
+    match = GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("image_name"):
+            raise ValueError(
+                "An image name must be specified in the URI for a GCP Artifact Registry. "
+                "Please provide a uri with the format "
+                "'https://<region>-docker.pkg.dev/<project>/<repository>/<image>'."
+            )
+        return uri
+
+    match = AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("repository"):
+            raise ValueError(
+                "A repository name must be specified in the URI for an "
+                "Azure Container Registry. Please provide a uri with the format "
+                "'https://<registry-name>.azurecr.io/<repository>'."
+            )
+        return uri
+
+    match = ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("repository"):
+            raise ValueError(
+                "A repository name must be specified in the URI for an "
+                "Elastic Container Registry. Please provide a uri with the format "
+                "'https://<account-id>.dkr.ecr.<region>.amazonaws.com/<repository>'."
+            )
+        return uri
+
+    wandb.termwarn(
+        f"Unable to recognize registry type in URI {uri}. You are responsible "
+        "for ensuring the agent can push images to this registry."
+    )
+    return uri

wandb/sdk/launch/builder/abstract.py

@@ -1,4 +1,5 @@
 """Abstract plugin class defining the interface needed to build container images for W&B Launch."""
+
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
@@ -6,6 +7,12 @@ from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
 from .._project_spec import EntryPoint, LaunchProject
+from ..registry.anon import AnonynmousRegistry
+from ..utils import (
+    AZURE_CONTAINER_REGISTRY_URI_REGEX,
+    ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
+    GCP_ARTIFACT_REGISTRY_URI_REGEX,
+)
 
 if TYPE_CHECKING:
     from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
@@ -34,7 +41,7 @@ class AbstractBuilder(ABC):
             verify: Whether to verify the functionality of the builder.
 
         Raises:
-            LaunchError: If the builder cannot be intialized or verified.
+            LaunchError: If the builder cannot be initialized or verified.
         """
         raise NotImplementedError
 
@@ -86,3 +93,64 @@ class AbstractBuilder(ABC):
             LaunchError: If the builder cannot be used to build images.
         """
         raise NotImplementedError
+
+
+def registry_from_uri(uri: str) -> AbstractRegistry:
+    """Create a registry helper object from a uri.
+
+    This function parses the URI and determines which supported registry it
+    belongs to. It then creates a registry helper object for that registry.
+    The supported remote registry types are:
+    - Azure Container Registry
+    - Google Container Registry
+    - AWS Elastic Container Registry
+
+    The format of the URI is as follows:
+    - Azure Container Registry: <registry-name>.azurecr.io/<repo-name>/<image-name>
+    - Google Container Registry: <location>-docker.pkg.dev/<project-id>/<repo-name>/<image-name>
+    - AWS Elastic Container Registry: <account-id>.dkr.ecr.<region>.amazonaws.com/<repo-name>/<image-name>
+
+    Our classification of the registry is based on the domain name. For example,
+    if the uri contains `.azurecr.io`, we classify it as an Azure
+    Container Registry. If the uri contains `.dkr.ecr`, we classify
+    it as an AWS Elastic Container Registry. If the uri contains
+    `-docker.pkg.dev`, we classify it as a Google Artifact Registry.
+
+    This function will attempt to load the approriate cloud helpers for the
+
+    `https://` prefix is optional for all of the above.
+
+    Arguments:
+        uri: The uri to create a registry from.
+
+    Returns:
+        The registry.
+
+    Raises:
+        LaunchError: If the registry helper cannot be loaded for the given URI.
+    """
+    if uri.startswith("https://"):
+        uri = uri[len("https://") :]
+
+    if AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.azure_container_registry import (
+            AzureContainerRegistry,
+        )
+
+        return AzureContainerRegistry(uri=uri)
+
+    elif GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.google_artifact_registry import (
+            GoogleArtifactRegistry,
+        )
+
+        return GoogleArtifactRegistry(uri=uri)
+
+    elif ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.elastic_container_registry import (
+            ElasticContainerRegistry,
+        )
+
+        return ElasticContainerRegistry(uri=uri)
+
+    return AnonynmousRegistry(uri=uri)