wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -0,0 +1,219 @@
+"""Implementation of Google Artifact Registry for wandb launch."""
+
+import logging
+from typing import Optional, Tuple
+
+import google.auth  # type: ignore
+import google.cloud.artifactregistry  # type: ignore
+
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.utils import (
+    GCP_ARTIFACT_REGISTRY_URI_REGEX,
+    event_loop_thread_exec,
+)
+from wandb.util import get_module
+
+from .abstract import AbstractRegistry
+
+_logger = logging.getLogger(__name__)
+
+google = get_module(
+    "google",
+    required="The google package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
+)
+google.auth = get_module(
+    "google.auth",
+    required="The google-auth package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
+)
+
+google.cloud.artifactregistry = get_module(
+    "google.cloud.artifactregistry",
+    required="The google-cloud-artifactregistry package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
+)
+
+
+class GoogleArtifactRegistry(AbstractRegistry):
+    """Google Artifact Registry helper for interacting with the registry.
+
+    This helper should be constructed from either a uri or a repository,
+    project, and optional image-name. If constructed from a uri, the uri
+    must be of the form REGION-docker.pkg.dev/PROJECT/REPOSITORY/[IMAGE_NAME],
+    with an optional https:// preceding.
+    """
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        repository: Optional[str] = None,
+        image_name: Optional[str] = None,
+        project: Optional[str] = None,
+        region: Optional[str] = None,
+    ) -> None:
+        """Initialize the Google Artifact Registry.
+
+        Either uri or repository and image_name must be provided. Project and
+        region are optional, and will be inferred from the uri if provided, or
+        from the default credentials if not.
+
+        Arguments:
+            uri (optional): The uri of the repository.
+            repository (optional): The repository name.
+            image_name (optional): The image name.
+            project (optional): The GCP project name.
+            region (optional): The GCP region name.
+
+        Raises:
+            LaunchError: If verify is True and the container registry or its
+                environment have not been properly configured. Or if the environment
+                is not an instance of GcpEnvironment.
+        """
+        _logger.info(
+            f"Initializing Google Artifact Registry with repository {repository} "
+            f"and image name {image_name}"
+        )
+
+        if uri is not None:
+            self.uri = uri
+            # Raise an error if any other kwargs were provided in addition to uri.
+            if any([repository, image_name, project, region]):
+                raise LaunchError(
+                    "The Google Artifact Registry must be specified with either "
+                    "the uri key or the repository, image-name, project and region "
+                    "keys, but not both."
+                )
+            match = GCP_ARTIFACT_REGISTRY_URI_REGEX.match(self.uri)
+            if not match:
+                raise LaunchError(
+                    f"The Google Artifact Registry uri {self.uri} is invalid. "
+                    "Please provide a uri of the form "
+                    "REGION-docker.pkg.dev/PROJECT/REPOSITORY/IMAGE_NAME."
+                )
+            self.project = match.group("project")
+            self.region = match.group("region")
+            self.repository = match.group("repository")
+            self.image_name = match.group("image_name")
+        else:
+            if any(x is None for x in (repository, region, image_name)):
+                raise LaunchError(
+                    "The Google Artifact Registry must be specified with either "
+                    "the uri key or the repository, image-name, project and region "
+                    "keys."
+                )
+            self.project = project
+            self.region = region
+            self.repository = repository
+            self.image_name = image_name
+            self.uri = f"{self.region}-docker.pkg.dev/{self.project}/{self.repository}/{self.image_name}"
+
+        _missing_kwarg_msg = (
+            "The Google Artifact Registry is missing the {} kwarg. "
+            "Please specify it by name or as part of the uri argument."
+        )
+        if not self.region:
+            raise LaunchError(_missing_kwarg_msg.format("region"))
+        if not self.repository:
+            raise LaunchError(_missing_kwarg_msg.format("repository"))
+        if not self.image_name:
+            raise LaunchError(_missing_kwarg_msg.format("image-name"))
+        # Try to load default project from the default credentials.
+        self.credentials, project = google.auth.default()
+        self.project = self.project or project
+        self.credentials.refresh(google.auth.transport.requests.Request())
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict,
+    ) -> "GoogleArtifactRegistry":
+        """Create a Google Artifact Registry from a config.
+
+        Arguments:
+            config: A dictionary containing the following keys:
+                repository: The repository name.
+                image-name: The image name.
+            environment: A GcpEnvironment configured for access to this registry.
+
+        Returns:
+            A GoogleArtifactRegistry.
+        """
+        # TODO: Replace this with pydantic.
+        acceptable_keys = {
+            "uri",
+            "type",
+            "repository",
+            "image-name",
+            "region",
+            "project",
+        }
+        unacceptable_keys = set(config.keys()) - acceptable_keys
+        if unacceptable_keys:
+            raise LaunchError(
+                f"The Google Artifact Registry config contains unacceptable keys: "
+                f"{unacceptable_keys}. Please remove these keys. The acceptable "
+                f"keys are: {acceptable_keys}."
+            )
+        return cls(
+            uri=config.get("uri"),
+            repository=config.get("repository"),
+            image_name=config.get("image-name"),
+            project=config.get("project"),
+            region=config.get("region"),
+        )
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry.
+
+        Returns:
+            A tuple of the username and password.
+        """
+        if not self.credentials.token:
+            self.credentials.refresh(google.auth.transport.requests.Request())
+        return "oauth2accesstoken", self.credentials.token
+
+    async def get_repo_uri(self) -> str:
+        """Get the URI for the given repository.
+
+        Arguments:
+            repo_name: The repository name.
+
+        Returns:
+            The repository URI.
+        """
+        return (
+            f"{self.region}-docker.pkg.dev/"
+            f"{self.project}/{self.repository}/{self.image_name}"
+        )
+
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if the image exists.
+
+        Arguments:
+            image_uri: The image URI.
+
+        Returns:
+            True if the image exists, False otherwise.
+        """
+        _logger.info(f"Checking if image {image_uri} exists")
+        repo_uri, tag = image_uri.split(":")
+        self_repo_uri = await self.get_repo_uri()
+        if repo_uri != self_repo_uri:
+            raise LaunchError(
+                f"The image {image_uri} does not match to the image uri "
+                f"repository {self.uri}."
+            )
+        parent = f"projects/{self.project}/locations/{self.region}/repositories/{self.repository}"
+        artifact_registry_client = event_loop_thread_exec(
+            google.cloud.artifactregistry.ArtifactRegistryClient
+        )
+        client = await artifact_registry_client(credentials=self.credentials)
+        list_images = event_loop_thread_exec(client.list_docker_images)
+        try:
+            for image in await list_images(request={"parent": parent}):
+                if tag in image.tags:
+                    return True
+        except google.api_core.exceptions.NotFound as e:  # type: ignore[attr-defined]
+            raise LaunchError(
+                f"The Google Artifact Registry repository {self.repository} "
+                f"does not exist. Please create it or modify your registry configuration."
+            ) from e
+        return False

wandb/sdk/launch/registry/local_registry.py

@@ -0,0 +1,65 @@
+"""Local registry implementation."""
+
+import logging
+from typing import Tuple
+
+from wandb.docker import is_docker_installed
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.utils import docker_image_exists
+
+from .abstract import AbstractRegistry
+
+_logger = logging.getLogger(__name__)
+
+
+class LocalRegistry(AbstractRegistry):
+    """A local registry.
+
+    This is a dummy registry that is used when no registry is configured.
+    """
+
+    def __init__(self) -> None:
+        """Initialize a local registry."""
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict,
+    ) -> "LocalRegistry":
+        """Create a local registry from a config.
+
+        Arguments:
+            config (dict): The config. This is ignored.
+            environment (AbstractEnvironment): The environment. This is ignored.
+
+        Returns:
+            LocalRegistry: The local registry.
+        """
+        return cls()
+
+    async def verify(self) -> None:
+        """Verify the local registry by doing nothing."""
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password of the local registry."""
+        raise LaunchError("Attempted to get username and password for LocalRegistry.")
+
+    async def get_repo_uri(self) -> str:
+        """Get the uri of the local registry.
+
+        Returns: An empty string.
+        """
+        return ""
+
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if an image exists in the local registry.
+
+        Arguments:
+            image_uri (str): The uri of the image.
+
+        Returns:
+            bool: True.
+        """
+        if is_docker_installed():
+            return docker_image_exists(image_uri)
+        return False

wandb/sdk/launch/runner/abstract.py

@@ -0,0 +1,185 @@
+"""Implementation of the abstract runner class.
+
+This class defines the interface that the W&B launch runner uses to manage the lifecycle
+of runs launched in different environments (e.g. runs launched locally or in a cluster).
+"""
+
+import logging
+import os
+import shutil
+import subprocess
+import sys
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import wandb
+from wandb.apis.internal import Api
+from wandb.sdk.lib import runid
+
+from .._project_spec import LaunchProject
+
+_logger = logging.getLogger(__name__)
+
+
+State = Literal[
+    "unknown",
+    "starting",
+    "running",
+    "failed",
+    "finished",
+    "stopping",
+    "stopped",
+    "preempted",
+]
+
+
+class Status:
+    def __init__(self, state: "State" = "unknown", messages: List[str] = None):  # type: ignore
+        self.state = state
+        self.messages = messages or []
+
+    def __repr__(self) -> "State":
+        return self.state
+
+    def __str__(self) -> str:
+        return self.state
+
+    def __eq__(self, __value: object) -> bool:
+        if isinstance(__value, Status):
+            return self.state == __value.state
+        else:
+            return self.state == __value
+
+    def __hash__(self) -> int:
+        return hash(self.state)
+
+
+class AbstractRun(ABC):
+    """Wrapper around a W&B launch run.
+
+    A launched run is a subprocess running an entry point
+    command, that exposes methods for waiting on and cancelling the run.
+    This class defines the interface that the W&B launch runner uses to manage the lifecycle
+    of runs launched in different environments (e.g. runs launched locally or in a cluster).
+    ``AbstractRun`` is not thread-safe. That is, concurrent calls to wait() / cancel()
+    from multiple threads may inadvertently kill resources (e.g. local processes) unrelated to the
+    run.
+    """
+
+    def __init__(self) -> None:
+        self._status = Status()
+
+    @property
+    def status(self) -> Status:
+        return self._status
+
+    @abstractmethod
+    async def get_logs(self) -> Optional[str]:
+        """Return the logs associated with the run."""
+
+    def _run_cmd(
+        self, cmd: List[str], output_only: Optional[bool] = False
+    ) -> Optional[Union["subprocess.Popen[bytes]", bytes]]:
+        """Run the command and returns a popen object or the stdout of the command.
+
+        Arguments:
+        cmd: The command to run
+        output_only: If true just return the stdout bytes
+        """
+        try:
+            env = os.environ
+            popen = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE)
+            if output_only:
+                popen.wait()
+                if popen.stdout is not None:
+                    return popen.stdout.read()
+            return popen
+        except subprocess.CalledProcessError as e:
+            wandb.termerror(f"Command failed: {e}")
+            return None
+
+    @abstractmethod
+    async def wait(self) -> bool:
+        """Wait for the run to finish, returning True if the run succeeded and false otherwise.
+
+        Note that in some cases, we may wait until the remote job completes rather than until the W&B run completes.
+        """
+
+    @abstractmethod
+    async def get_status(self) -> Status:
+        """Get status of the run."""
+
+    @abstractmethod
+    async def cancel(self) -> None:
+        """Cancel the run (interrupts the command subprocess, cancels the run, etc).
+
+        Cancels the run and waits for it to terminate. The W&B run status may not be
+        set correctly upon run cancellation.
+        """
+
+    @property
+    @abstractmethod
+    def id(self) -> Optional[str]:
+        pass
+
+
+class AbstractRunner(ABC):
+    """Abstract plugin class defining the interface needed to execute W&B Launches.
+
+    You can define subclasses of ``AbstractRunner`` and expose them as third-party
+    plugins to enable running W&B projects against custom execution backends
+    (e.g. to run projects against your team's in-house cluster or job scheduler).
+    """
+
+    _type: str
+
+    def __init__(
+        self,
+        api: Api,
+        backend_config: Dict[str, Any],
+    ) -> None:
+        self._api = api
+        self.backend_config = backend_config
+        self._cwd = os.getcwd()
+        self._namespace = runid.generate_id()
+
+    def find_executable(
+        self,
+        cmd: str,
+    ) -> Union[str, None]:
+        """Cross platform utility for checking if a program is available."""
+        return shutil.which(cmd)
+
+    @property
+    def api_key(self) -> Any:
+        return self._api.api_key
+
+    def verify(self) -> bool:
+        """This is called on first boot to verify the needed commands, and permissions are available.
+
+        For now just call `wandb.termerror` and `sys.exit(1)`
+        """
+        if self._api.api_key is None:
+            wandb.termerror(
+                "Couldn't find W&B api key, run wandb login or set WANDB_API_KEY"
+            )
+            sys.exit(1)
+        return True
+
+    @abstractmethod
+    async def run(
+        self,
+        launch_project: LaunchProject,
+        image_uri: str,
+    ) -> Optional[AbstractRun]:
+        """Submit an LaunchProject to be run.
+
+        Returns a SubmittedRun object to track the execution
+        Arguments:
+        launch_project: Object of _project_spec.LaunchProject class representing a wandb launch project
+
+        Returns:
+            A :py:class:`wandb.sdk.launch.runners.SubmittedRun`. This function is expected to run
+            the project asynchronously, i.e. it should trigger project execution and then
+            immediately return a `SubmittedRun` to track execution status.
+        """