wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/launch/loader.py

@@ -0,0 +1,249 @@
+"""Utilities for the agent."""
+
+from typing import Any, Dict, Optional
+
+import wandb
+from wandb.apis.internal import Api
+from wandb.docker import is_docker_installed
+from wandb.sdk.launch.errors import LaunchError
+
+from .builder.abstract import AbstractBuilder
+from .environment.abstract import AbstractEnvironment
+from .registry.abstract import AbstractRegistry
+from .runner.abstract import AbstractRunner
+
+WANDB_RUNNERS = {
+    "local-container",
+    "local-process",
+    "kubernetes",
+    "vertex",
+    "sagemaker",
+}
+
+
+def environment_from_config(config: Optional[Dict[str, Any]]) -> AbstractEnvironment:
+    """Create an environment from a config.
+
+    This helper function is used to create an environment from a config. The
+    config should have a "type" key that specifies the type of environment to
+    create. The remaining keys are passed to the environment's from_config
+    method. If the config is None or empty, a LocalEnvironment is returned.
+
+    Arguments:
+        config (Dict[str, Any]): The config.
+
+    Returns:
+        Environment: The environment constructed.
+    """
+    if not config:
+        from .environment.local_environment import LocalEnvironment
+
+        return LocalEnvironment()  # This is the default, dummy environment.
+    env_type = config.get("type")
+    if not env_type:
+        raise LaunchError(
+            "Could not create environment from config. Environment type not specified!"
+        )
+    if env_type == "local":
+        from .environment.local_environment import LocalEnvironment
+
+        return LocalEnvironment.from_config(config)
+    if env_type == "aws":
+        from .environment.aws_environment import AwsEnvironment
+
+        return AwsEnvironment.from_config(config)
+    if env_type == "gcp":
+        from .environment.gcp_environment import GcpEnvironment
+
+        return GcpEnvironment.from_config(config)
+    if env_type == "azure":
+        from .environment.azure_environment import AzureEnvironment
+
+        return AzureEnvironment.from_config(config)
+    raise LaunchError(
+        f"Could not create environment from config. Invalid type: {env_type}"
+    )
+
+
+def registry_from_config(
+    config: Optional[Dict[str, Any]], environment: AbstractEnvironment
+) -> AbstractRegistry:
+    """Create a registry from a config.
+
+    This helper function is used to create a registry from a config. The
+    config should have a "type" key that specifies the type of registry to
+    create. The remaining keys are passed to the registry's from_config
+    method. If the config is None or empty, a LocalRegistry is returned.
+
+    Arguments:
+        config (Dict[str, Any]): The registry config.
+        environment (Environment): The environment of the registry.
+
+    Returns:
+        The registry if config is not None, otherwise None.
+
+    Raises:
+        LaunchError: If the registry is not configured correctly.
+    """
+    if not config:
+        from .registry.local_registry import LocalRegistry
+
+        return LocalRegistry()  # This is the default, dummy registry.
+
+    wandb.termwarn(
+        "The `registry` block of the launch agent config is being deprecated. "
+        "Please specify an image repository URI under the `builder.destination` "
+        "key of your launch agent config. See "
+        "https://docs.wandb.ai/guides/launch/setup-agent-advanced#agent-configuration "
+        "for more information."
+    )
+
+    registry_type = config.get("type")
+    if registry_type is None or registry_type == "local":
+        from .registry.local_registry import LocalRegistry
+
+        return LocalRegistry()  # This is the default, dummy registry.
+    if registry_type == "ecr":
+        from .registry.elastic_container_registry import ElasticContainerRegistry
+
+        return ElasticContainerRegistry.from_config(config)
+    if registry_type == "gcr":
+        from .registry.google_artifact_registry import GoogleArtifactRegistry
+
+        return GoogleArtifactRegistry.from_config(config)
+    if registry_type == "acr":
+        from .registry.azure_container_registry import AzureContainerRegistry
+
+        return AzureContainerRegistry.from_config(config)
+    raise LaunchError(
+        f"Could not create registry from config. Invalid registry type: {registry_type}"
+    )
+
+
+def builder_from_config(
+    config: Optional[Dict[str, Any]],
+    environment: AbstractEnvironment,
+    registry: AbstractRegistry,
+) -> AbstractBuilder:
+    """Create a builder from a config.
+
+    This helper function is used to create a builder from a config. The
+    config should have a "type" key that specifies the type of builder to import
+    and create. The remaining keys are passed to the builder's from_config
+    method. If the config is None or empty, a default builder is returned.
+
+    The default builder will be a DockerBuilder if we find a working docker cli
+    on the system, otherwise it will be a NoOpBuilder.
+
+    Arguments:
+        config (Dict[str, Any]): The builder config.
+        registry (Registry): The registry of the builder.
+
+    Returns:
+        The builder.
+
+    Raises:
+        LaunchError: If the builder is not configured correctly.
+    """
+    if not config:
+        if is_docker_installed():
+            from .builder.docker_builder import DockerBuilder
+
+            return DockerBuilder.from_config(
+                {}, environment, registry
+            )  # This is the default builder.
+
+        from .builder.noop import NoOpBuilder
+
+        return NoOpBuilder.from_config({}, environment, registry)
+
+    builder_type = config.get("type")
+    if builder_type is None:
+        raise LaunchError(
+            "Could not create builder from config. Builder type not specified"
+        )
+    if builder_type == "docker":
+        from .builder.docker_builder import DockerBuilder
+
+        return DockerBuilder.from_config(config, environment, registry)
+    if builder_type == "kaniko":
+        from .builder.kaniko_builder import KanikoBuilder
+
+        return KanikoBuilder.from_config(config, environment, registry)
+    if builder_type == "noop":
+        from .builder.noop import NoOpBuilder
+
+        return NoOpBuilder.from_config(config, environment, registry)
+    raise LaunchError(
+        f"Could not create builder from config. Invalid builder type: {builder_type}"
+    )
+
+
+def runner_from_config(
+    runner_name: str,
+    api: Api,
+    runner_config: Dict[str, Any],
+    environment: AbstractEnvironment,
+    registry: AbstractRegistry,
+) -> AbstractRunner:
+    """Create a runner from a config.
+
+    This helper function is used to create a runner from a config. The
+    config should have a "type" key that specifies the type of runner to import
+    and create. The remaining keys are passed to the runner's from_config
+    method. If the config is None or empty, a LocalContainerRunner is returned.
+
+    Arguments:
+        runner_name (str): The name of the backend.
+        api (Api): The API.
+        runner_config (Dict[str, Any]): The backend config.
+
+    Returns:
+        The runner.
+
+    Raises:
+        LaunchError: If the runner is not configured correctly.
+    """
+    if not runner_name or runner_name in ["local-container", "local"]:
+        from .runner.local_container import LocalContainerRunner
+
+        return LocalContainerRunner(api, runner_config, environment, registry)
+    if runner_name == "local-process":
+        from .runner.local_process import LocalProcessRunner
+
+        return LocalProcessRunner(api, runner_config)
+    if runner_name == "sagemaker":
+        from .environment.aws_environment import AwsEnvironment
+
+        if not isinstance(environment, AwsEnvironment):
+            try:
+                environment = AwsEnvironment.from_default()
+            except LaunchError as e:
+                raise LaunchError(
+                    "Could not create Sagemaker runner. "
+                    "Environment must be an instance of AwsEnvironment."
+                ) from e
+        from .runner.sagemaker_runner import SageMakerRunner
+
+        return SageMakerRunner(api, runner_config, environment, registry)
+    if runner_name in ["vertex", "gcp-vertex"]:
+        from .environment.gcp_environment import GcpEnvironment
+
+        if not isinstance(environment, GcpEnvironment):
+            try:
+                environment = GcpEnvironment.from_default()
+            except LaunchError as e:
+                raise LaunchError(
+                    "Could not create Vertex runner. "
+                    "Environment must be an instance of GcpEnvironment."
+                ) from e
+        from .runner.vertex_runner import VertexRunner
+
+        return VertexRunner(api, runner_config, environment, registry)
+    if runner_name == "kubernetes":
+        from .runner.kubernetes_runner import KubernetesRunner
+
+        return KubernetesRunner(api, runner_config, environment, registry)
+    raise LaunchError(
+        f"Could not create runner from config. Invalid runner name: {runner_name}"
+    )

wandb/sdk/launch/registry/abstract.py

@@ -0,0 +1,48 @@
+"""Abstract base class for registries."""
+
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+
+class AbstractRegistry(ABC):
+    """Abstract base class for registries."""
+
+    uri: str
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry.
+
+        Returns:
+            (str, str): The username and password.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get_repo_uri(self) -> str:
+        """Get the URI for a repository.
+
+        Returns:
+            str: The URI.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if an image exists in the registry.
+
+        Arguments:
+            image_uri (str): The URI of the image.
+
+        Returns:
+            bool: True if the image exists.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(
+        cls,
+        config: dict,
+    ) -> "AbstractRegistry":
+        """Create a registry from a config."""
+        raise NotImplementedError

wandb/sdk/launch/registry/anon.py

@@ -0,0 +1,29 @@
+from typing import Tuple
+
+from wandb.docker import is_docker_installed
+from wandb.sdk.launch.utils import docker_image_exists
+
+from .abstract import AbstractRegistry
+
+
+class AnonynmousRegistry(AbstractRegistry):
+    def __init__(self, uri: str) -> None:
+        """Initialize the registry."""
+        self.uri = uri
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry."""
+        raise NotImplementedError("Anonymous registry does not require authentication")
+
+    async def get_repo_uri(self) -> str:
+        return self.uri
+
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if an image exists in the registry."""
+        if not is_docker_installed():
+            return False
+        return docker_image_exists(image_uri)
+
+    @classmethod
+    def from_config(cls, config: dict) -> "AbstractRegistry":
+        return cls(uri=config["uri"])

wandb/sdk/launch/registry/azure_container_registry.py

@@ -0,0 +1,124 @@
+"""Implementation of AzureContainerRegistry class."""
+
+import re
+from typing import TYPE_CHECKING, Optional, Tuple
+
+from wandb.sdk.launch.environment.azure_environment import AzureEnvironment
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.utils import AZURE_CONTAINER_REGISTRY_URI_REGEX
+from wandb.util import get_module
+
+from .abstract import AbstractRegistry
+
+if TYPE_CHECKING:
+    from azure.containerregistry import ContainerRegistryClient  # type: ignore
+    from azure.core.exceptions import ResourceNotFoundError  # type: ignore
+
+
+ContainerRegistryClient = get_module(  # noqa: F811
+    "azure.containerregistry",
+    required="The azure-containerregistry package is required to use launch with Azure. Please install it with `pip install azure-containerregistry`.",
+).ContainerRegistryClient
+
+ResourceNotFoundError = get_module(  # noqa: F811
+    "azure.core.exceptions",
+    required="The azure-core package is required to use launch with Azure. Please install it with `pip install azure-core`.",
+).ResourceNotFoundError
+
+
+class AzureContainerRegistry(AbstractRegistry):
+    """Helper for accessing Azure Container Registry resources."""
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        registry_name: Optional[str] = None,
+        repo_name: Optional[str] = None,
+    ):
+        """Initialize an AzureContainerRegistry."""
+        if uri is not None:
+            self.uri = uri
+            if any(x is not None for x in (registry_name, repo_name)):
+                raise LaunchError(
+                    "Please specify either a registry name and repo name or a registry URI."
+                )
+            if self.uri.startswith("https://"):
+                self.uri = self.uri[len("https://") :]
+            match = AZURE_CONTAINER_REGISTRY_URI_REGEX.match(self.uri)
+            if match is None:
+                raise LaunchError(
+                    f"Unable to parse Azure Container Registry URI: {self.uri}"
+                )
+            self.registry_name = match.group(1)
+            self.repo_name = match.group(2)
+        else:
+            if any(x is None for x in (registry_name, repo_name)):
+                raise LaunchError(
+                    "Please specify both a registry name and repo name or a registry URI."
+                )
+            self.registry_name = registry_name
+            self.repo_name = repo_name
+            self.uri = f"{self.registry_name}.azurecr.io/{self.repo_name}"
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict,
+    ) -> "AzureContainerRegistry":
+        """Create an AzureContainerRegistry from a config dict.
+
+        Args:
+            config (dict): The config dict.
+            environment (AbstractEnvironment): The environment to use.
+            verify (bool, optional): Whether to verify the registry. Defaults to True.
+
+        Returns:
+            AzureContainerRegistry: The registry.
+
+        Raises:
+            LaunchError: If the config is invalid.
+        """
+        uri = config.get("uri")
+        if uri is None:
+            raise LaunchError(
+                "Please specify a registry name to use under the registry.uri."
+            )
+        return cls(
+            uri=uri,
+        )
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get username and password for container registry."""
+        raise NotImplementedError
+
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if image exists in container registry.
+
+        Args:
+            image_uri (str): Image URI to check.
+
+        Returns:
+            bool: True if image exists, False otherwise.
+        """
+        match = re.match(AZURE_CONTAINER_REGISTRY_URI_REGEX, image_uri)
+        if match is None:
+            raise LaunchError(
+                f"Unable to parse Azure Container Registry URI: {image_uri}"
+            )
+        registry = match.group(1)
+        repository = match.group(2)
+        tag = match.group(3)
+        credential = AzureEnvironment.get_credentials()
+        client = ContainerRegistryClient(f"https://{registry}.azurecr.io", credential)
+        try:
+            client.get_manifest_properties(repository, tag)
+            return True
+        except ResourceNotFoundError:
+            return False
+        except Exception as e:
+            raise LaunchError(
+                f"Unable to check if image exists in Azure Container Registry: {e}"
+            ) from e
+
+    async def get_repo_uri(self) -> str:
+        return f"{self.registry_name}.azurecr.io/{self.repo_name}"

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -0,0 +1,192 @@
+"""Implementation of Elastic Container Registry class for wandb launch."""
+
+import base64
+import logging
+from typing import Dict, Optional, Tuple
+
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.registry.abstract import AbstractRegistry
+from wandb.sdk.launch.utils import (
+    ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
+    event_loop_thread_exec,
+)
+from wandb.util import get_module
+
+_logger = logging.getLogger(__name__)
+
+botocore = get_module(  # noqa: F811
+    "botocore",
+    required="The boto3 package is required to use launch with AWS. Please install it with `pip install wandb[launch]`.",
+)
+boto3 = get_module(  # noqa: F811
+    "boto3",
+    required="The boto3 package is required to use launch with AWS. Please install it with `pip install wandb[launch]`.",
+)
+
+
+class ElasticContainerRegistry(AbstractRegistry):
+    """Elastic Container Registry class."""
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        account_id: Optional[str] = None,
+        region: Optional[str] = None,
+        repo_name: Optional[str] = None,
+    ) -> None:
+        """Initialize the Elastic Container Registry.
+
+        Arguments:
+            uri: The uri of the repository.
+            account_id: The AWS account id.
+            region: The AWS region of the container registry.
+            repository: The name of the repository.
+
+        Raises:
+            LaunchError: If there is an error initializing the Elastic Container Registry helper.
+        """
+        if uri:
+            self.uri = uri
+            if any([account_id, region, repo_name]):
+                raise LaunchError(
+                    "Could not create ElasticContainerRegistry from config. Either 'uri' or "
+                    "'account_id', 'region', and 'repo_name' are required."
+                )
+            match = ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(
+                self.uri,
+            )
+            if not match:
+                raise LaunchError(
+                    f"Could not create ElasticContainerRegistry from config. The uri "
+                    f"{self.uri} is invalid."
+                )
+            self.account_id = match.group("account")
+            self.region = match.group("region")
+            self.repo_name = match.group("repository")
+        else:
+            if not all([account_id, region, repo_name]):
+                raise LaunchError(
+                    "Could not create ElasticContainerRegistry from config. Either 'uri' or "
+                    "'account_id', 'region', and 'repo_name' are required."
+                )
+            self.account_id = account_id
+            self.region = region
+            self.repo_name = repo_name
+            self.uri = f"{self.account_id}.dkr.ecr.{self.region}.amazonaws.com/{self.repo_name}"
+        if self.account_id is None:
+            raise LaunchError(
+                "Could not create ElasticContainerRegistry from config. Either 'uri' or "
+                "'account_id' is required."
+            )
+        if self.region is None:
+            raise LaunchError(
+                "Could not create ElasticContainerRegistry from config. Either 'uri' or "
+                "'region' is required."
+            )
+        if self.repo_name is None:
+            raise LaunchError(
+                "Could not create ElasticContainerRegistry from config. Either 'uri' or "
+                "'repository' is required."
+            )
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, str],
+    ) -> "ElasticContainerRegistry":
+        """Create an Elastic Container Registry from a config.
+
+        Arguments:
+            config (dict): The config.
+
+        Returns:
+            ElasticContainerRegistry: The Elastic Container Registry.
+        """
+        # TODO: Replace this with pydantic.
+        acceptable_keys = {
+            "uri",
+            "type",
+            "account_id",
+            "region",
+            "repo_name",
+        }
+        unsupported_keys = set(config.keys()) - acceptable_keys
+        if unsupported_keys:
+            raise LaunchError(
+                f"The Elastic Container Registry config contains unsupported keys: "
+                f"{unsupported_keys}. Please remove these keys. The acceptable "
+                f"keys are: {acceptable_keys}."
+            )
+        return cls(
+            uri=config.get("uri"),
+            account_id=config.get("account_id"),
+            region=config.get("region"),
+            repo_name=config.get("repository"),
+        )
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry.
+
+        Returns:
+            (str, str): The username and password.
+
+        Raises:
+            RegistryError: If there is an error getting the username and password.
+        """
+        _logger.debug("Getting username and password for Elastic Container Registry.")
+        try:
+            session = boto3.Session(region_name=self.region)
+            client = await event_loop_thread_exec(session.client)("ecr")
+            response = await event_loop_thread_exec(client.get_authorization_token)()
+            username, password = base64.standard_b64decode(
+                response["authorizationData"][0]["authorizationToken"]
+            ).split(b":")
+            return username.decode("utf-8"), password.decode("utf-8")
+
+        except botocore.exceptions.ClientError as e:
+            code = e.response["Error"]["Code"]
+            msg = e.response["Error"]["Message"]
+            # TODO: Log the code and the message here?
+            raise LaunchError(f"Error getting username and password: {code} {msg}")
+
+    async def get_repo_uri(self) -> str:
+        """Get the uri of the repository.
+
+        Returns:
+            str: The uri of the repository.
+        """
+        return f"{self.account_id}.dkr.ecr.{self.region}.amazonaws.com/{self.repo_name}"
+
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if the image tag exists.
+
+        Arguments:
+            image_uri (str): The full image_uri.
+
+        Returns:
+            bool: True if the image tag exists.
+        """
+        if ":" not in image_uri:
+            tag = image_uri
+        else:
+            uri, tag = image_uri.split(":")
+            repo_uri = await self.get_repo_uri()
+            if uri != repo_uri:
+                raise LaunchError(
+                    f"Image uri {image_uri} does not match Elastic Container Registry uri {repo_uri}."
+                )
+        _logger.debug(f"Checking if image tag {tag} exists in repository {self.uri}")
+        try:
+            session = boto3.Session(region_name=self.region)
+            client = await event_loop_thread_exec(session.client)("ecr")
+            response = await event_loop_thread_exec(client.describe_images)(
+                repositoryName=self.repo_name, imageIds=[{"imageTag": tag}]
+            )
+            return len(response["imageDetails"]) > 0
+
+        except botocore.exceptions.ClientError as e:
+            code = e.response["Error"]["Code"]
+            if code == "ImageNotFoundException":
+                return False
+            msg = e.response["Error"]["Message"]
+            raise LaunchError(f"Error checking if image tag exists: {code} {msg}")