wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -0,0 +1,203 @@
+"""Implementation of Google Artifact Registry for wandb launch."""
+import logging
+import re
+from typing import Tuple
+
+from wandb.sdk.launch.environment.gcp_environment import GcpEnvironment
+from wandb.sdk.launch.utils import LaunchError
+from wandb.util import get_module
+
+from .abstract import AbstractRegistry
+
+google = get_module(
+    "google",
+    required="Google Cloud Platform support requires the google package. Please"
+    " install it with `pip install wandb[launch]`.",
+)
+
+google.cloud.artifactregistry = get_module(
+    "google.cloud.artifactregistry",
+    required="Google Cloud Platform support requires the google-cloud-artifact-registry package. "
+    "Please install it with `pip install wandb[launch]`.",
+)
+
+google.auth.credentials = get_module(
+    "google.auth.credentials",
+    required="Google Cloud Platform support requires google-auth. "
+    "Please install it with `pip install wandb[launch]`.",
+)
+
+_logger = logging.getLogger(__name__)
+
+
+class GoogleArtifactRegistry(AbstractRegistry):
+    """Google Artifact Registry.
+
+    Attributes:
+        repository: The repository name.
+        environment: A GcpEnvironment configured for access to this registry.
+    """
+
+    repository: str
+    image_name: str
+    environment: GcpEnvironment
+
+    def __init__(
+        self,
+        repository: str,
+        image_name: str,
+        environment: GcpEnvironment,
+        verify: bool = True,
+    ) -> None:
+        """Initialize the Google Artifact Registry.
+
+        Arguments:
+            repository: The repository name.
+            image_name: The image name.
+            environment: A GcpEnvironment configured for access to this registry.
+            verify: Whether to verify the credentials, region, and project.
+
+        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}"
+        )
+        self.repository = repository
+        self.image_name = image_name
+        if not re.match(r"^\w[\w.-]+$", image_name):
+            raise LaunchError(
+                f"The image name {image_name} is invalid. The image name must "
+                "consist of alphanumeric characters and underscores."
+            )
+        self.environment = environment
+        if verify:
+            self.verify()
+
+    @property
+    def uri(self) -> str:
+        """The uri of the registry."""
+        return f"{self.environment.region}-docker.pkg.dev/{self.environment.project}/{self.repository}/{self.image_name}"
+
+    @uri.setter
+    def uri(self, uri: str) -> None:
+        """Set the uri of the registry."""
+        raise LaunchError("The uri of the Google Artifact Registry cannot be set.")
+
+    @classmethod
+    def from_config(  # type: ignore[override]
+        cls,
+        config: dict,
+        environment: GcpEnvironment,
+        verify: bool = True,
+    ) -> "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.
+        """
+        repository = config.get("repository")
+        if not repository:
+            raise LaunchError(
+                "The Google Artifact Registry repository must be specified."
+            )
+        image_name = config.get("image_name")
+        if not image_name:
+            raise LaunchError("The image name must be specified.")
+        return cls(repository, image_name, environment, verify=verify)
+
+    def verify(self) -> None:
+        """Verify the registry is properly configured.
+
+        Raises:
+            LaunchError: If the registry is not properly configured.
+        """
+        credentials = self.environment.get_credentials()
+        parent = (
+            f"projects/{self.environment.project}/locations/{self.environment.region}"
+        )
+        # We need to list the repositories to verify that the repository exists.
+        request = google.cloud.artifactregistry.ListRepositoriesRequest(parent=parent)
+        client = google.cloud.artifactregistry.ArtifactRegistryClient(
+            credentials=credentials
+        )
+        try:
+            response = client.list_repositories(request=request)
+        except google.api_core.exceptions.PermissionDenied:
+            raise LaunchError(
+                "The provided credentials do not have permission to access the "
+                f"Google Artifact Registry repository {self.repository}."
+            )
+        # Look for self.repository in the list of responses.
+        for repo in response:
+            if repo.name.endswith(self.repository):
+                break
+        # If we didn't find the repository, raise an error.
+        else:
+            raise LaunchError(
+                f"The Google Artifact Registry repository {self.repository} does not exist."
+            )
+
+    def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry.
+
+        Returns:
+            A tuple of the username and password.
+        """
+        credentials = self.environment.get_credentials()
+        return "oauth2accesstoken", credentials.token
+
+    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.environment.region}-docker.pkg.dev/"
+            f"{self.environment.project}/{self.repository}/{self.image_name}"
+        )
+
+    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. In Google Artifact Registry {self.uri}."
+        )
+
+        return False
+        # TODO: Test GCP Artifact Registry image exists to get working
+        # repo_uri, _ = image_uri.split(":")
+        # if repo_uri != self.get_repo_uri():
+        #     raise LaunchError(
+        #         f"The image {image_uri} does not belong to the Google Artifact "
+        #         f"Repository {self.get_repo_uri()}."
+        #     )
+        # credentials = self.environment.get_credentials()
+        # request = google.cloud.artifactregistry.GetTagRequest(parent=image_uri)
+        # client = google.cloud.artifactregistry.ArtifactRegistryClient(
+        #     credentials=credentials
+        # )
+        # try:
+        #     client.get_tag(request=request)
+        #     return True
+        # except google.api_core.exceptions.NotFound:
+        #     return False

wandb/sdk/launch/registry/local_registry.py

@@ -0,0 +1,62 @@
+"""Local registry implementation."""
+import logging
+from typing import Tuple
+
+from wandb.sdk.launch.utils import LaunchError, docker_image_exists
+
+from ..environment.abstract import AbstractEnvironment
+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."""
+        pass
+
+    @classmethod
+    def from_config(
+        cls, config: dict, environment: "AbstractEnvironment", verify: bool = True
+    ) -> "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()
+
+    def verify(self) -> None:
+        """Verify the local registry by doing nothing."""
+        pass
+
+    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.")
+
+    def get_repo_uri(self) -> str:
+        """Get the uri of the local registry.
+
+        Returns: An empty string.
+        """
+        return ""
+
+    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.
+        """
+        return docker_image_exists(image_uri)

wandb/sdk/launch/runner/abstract.py

@@ -10,7 +10,6 @@ from dockerpycreds.utils import find_executable  # type: ignore
 import wandb
 from wandb import Settings
 from wandb.apis.internal import Api
-from wandb.errors import CommError
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.lib import runid
 
@@ -60,7 +59,7 @@ class AbstractRun(ABC):
     def _run_cmd(
         self, cmd: List[str], output_only: Optional[bool] = False
     ) -> Optional[Union["subprocess.Popen[bytes]", bytes]]:
-        """Runs the command and returns a popen object or the stdout of the command.
+        """Run the command and returns a popen object or the stdout of the command.
 
         Arguments:
         cmd: The command to run
@@ -143,25 +142,11 @@ class AbstractRunner(ABC):
             sys.exit(1)
         return True
 
-    def ack_run_queue_item(self, launch_project: LaunchProject) -> bool:
-        if self.backend_config.get("runQueueItemId"):
-            try:
-                self._api.ack_run_queue_item(
-                    self.backend_config["runQueueItemId"], launch_project.run_id
-                )
-            except CommError:
-                wandb.termerror(
-                    "Error acking run queue item. Item lease may have ended or another process may have acked it."
-                )
-                return False
-        return True
-
     @abstractmethod
     def run(
         self,
         launch_project: LaunchProject,
         builder: AbstractBuilder,
-        registry_config: Dict[str, Any],
     ) -> Optional[AbstractRun]:
         """Submit an LaunchProject to be run.
 

wandb/sdk/launch/runner/{kubernetes.py → kubernetes_runner.py}

@@ -1,17 +1,15 @@
 import base64
 import json
+import logging
 import time
 from typing import Any, Dict, List, Optional
 
-from kubernetes import client  # type: ignore
-from kubernetes.client.api.batch_v1_api import BatchV1Api  # type: ignore
-from kubernetes.client.api.core_v1_api import CoreV1Api  # type: ignore
-from kubernetes.client.models.v1_job import V1Job  # type: ignore
-from kubernetes.client.models.v1_secret import V1Secret  # type: ignore
-
 import wandb
-from wandb.errors import LaunchError
+from wandb.apis.internal import Api
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
+from wandb.sdk.launch.environment.abstract import AbstractEnvironment
+from wandb.sdk.launch.registry.abstract import AbstractRegistry
+from wandb.sdk.launch.registry.local_registry import LocalRegistry
 from wandb.util import get_module, load_json_yaml_dict
 
 from .._project_spec import LaunchProject, get_entry_point_command
@@ -19,15 +17,30 @@ from ..builder.build import get_env_vars_dict
 from ..utils import (
     LOG_PREFIX,
     PROJECT_SYNCHRONOUS,
+    LaunchError,
     get_kube_context_and_api_client,
     make_name_dns_safe,
 )
 from .abstract import AbstractRun, AbstractRunner, Status
 
+get_module(
+    "kubernetes",
+    required="Kubernetes runner requires the kubernetes package. Please install it with `pip install wandb[launch]`.",
+)
+
+from kubernetes import client  # type: ignore # noqa: E402
+from kubernetes.client.api.batch_v1_api import BatchV1Api  # type: ignore # noqa: E402
+from kubernetes.client.api.core_v1_api import CoreV1Api  # type: ignore # noqa: E402
+from kubernetes.client.models.v1_job import V1Job  # type: ignore # noqa: E402
+from kubernetes.client.models.v1_secret import V1Secret  # type: ignore # noqa: E402
+
 TIMEOUT = 5
 MAX_KUBERNETES_RETRIES = (
     60  # default 10 second loop time on the agent, this is 10 minutes
 )
+FAIL_MESSAGE_INTERVAL = 60
+
+_logger = logging.getLogger(__name__)
 
 
 class KubernetesSubmittedRun(AbstractRun):
@@ -76,20 +89,23 @@ class KubernetesSubmittedRun(AbstractRun):
             name=self.name, namespace=self.namespace
         )
         status = job_response.status
-        try:
-            self.core_api.read_namespaced_pod_log(
-                name=self.pod_names[0], namespace=self.namespace
-            )
-        except Exception as e:
-            if self._fail_count == 1:
+
+        pod = self.core_api.read_namespaced_pod(
+            name=self.pod_names[0], namespace=self.namespace
+        )
+        if pod.status.phase in ["Pending", "Unknown"]:
+            now = time.time()
+            if self._fail_count == 0:
+                self._fail_first_msg_time = now
+                self._fail_last_msg_time = 0.0
+            self._fail_count += 1
+            if now - self._fail_last_msg_time > FAIL_MESSAGE_INTERVAL:
                 wandb.termlog(
-                    f"{LOG_PREFIX}Failed to get pod status for job: {self.name}. Will wait up to 10 minutes for job to start."
+                    f"{LOG_PREFIX}Pod has not started yet for job: {self.name}. Will wait up to {round(10 - (now - self._fail_first_msg_time)/60)} minutes."
                 )
-            self._fail_count += 1
+                self._fail_last_msg_time = now
             if self._fail_count > MAX_KUBERNETES_RETRIES:
-                raise LaunchError(
-                    f"Failed to start job {self.name}, because of error {str(e)}"
-                )
+                raise LaunchError(f"Failed to start job {self.name}")
         # todo: we only handle the 1 pod case. see https://kubernetes.io/docs/concepts/workloads/controllers/job/#parallel-jobs for multipod handling
         return_status = None
         if status.succeeded == 1:
@@ -145,6 +161,12 @@ class KubernetesSubmittedRun(AbstractRun):
 
 
 class KubernetesRunner(AbstractRunner):
+    def __init__(
+        self, api: Api, backend_config: Dict[str, Any], environment: AbstractEnvironment
+    ) -> None:
+        super().__init__(api, backend_config)
+        self.environment = environment
+
     def populate_job_spec(
         self, job_spec: Dict[str, Any], resource_args: Dict[str, Any]
     ) -> None:
@@ -180,7 +202,6 @@ class KubernetesRunner(AbstractRunner):
     def populate_container_resources(
         self, containers: List[Dict[str, Any]], resource_args: Dict[str, Any]
     ) -> None:
-
         if resource_args.get("container_name"):
             if len(containers) > 1:
                 raise LaunchError(
@@ -247,9 +268,7 @@ class KubernetesRunner(AbstractRunner):
         )
         return pod_names
 
-    def get_namespace(
-        self, resource_args: Dict[str, Any]
-    ) -> Optional[str]:  # noqa: C901
+    def get_namespace(self, resource_args: Dict[str, Any]) -> Optional[str]:
         return self.backend_config.get("runner", {}).get(
             "namespace"
         ) or resource_args.get("namespace")
@@ -257,18 +276,19 @@ class KubernetesRunner(AbstractRunner):
     def run(
         self,
         launch_project: LaunchProject,
-        builder: AbstractBuilder,
-        registry_config: Dict[str, Any],
+        builder: Optional[AbstractBuilder],
     ) -> Optional[AbstractRun]:  # noqa: C901
         kubernetes = get_module(  # noqa: F811
-            "kubernetes", "KubernetesRunner requires kubernetes to be installed"
+            "kubernetes",
+            required="Kubernetes runner requires the kubernetes package. Please"
+            " install it with `pip install wandb[launch]`.",
         )
-
         resource_args = launch_project.resource_args.get("kubernetes", {})
         if not resource_args:
             wandb.termlog(
                 f"{LOG_PREFIX}Note: no resource args specified. Add a Kubernetes yaml spec or other options in a json file with --resource-args <json>."
             )
+        _logger.info(f"Running Kubernetes job with resource args: {resource_args}")
         context, api_client = get_kube_context_and_api_client(kubernetes, resource_args)
 
         batch_api = kubernetes.client.BatchV1Api(api_client)
@@ -332,35 +352,25 @@ class KubernetesRunner(AbstractRunner):
                     "Multiple container configurations should be specified in a yaml file supplied via job_spec."
                 )
             # dont specify run id if user provided image, could have multiple runs
-            env_vars.pop("WANDB_RUN_ID")
             containers[0]["image"] = launch_project.docker_image
             image_uri = launch_project.docker_image
             # TODO: handle secret pulling image from registry
-        elif any(["image" in cont for cont in containers]):
-            # user specified image configurations via kubernetes yaml, could have multiple images
-            # dont specify run id if user provided image, could have multiple runs
-            env_vars.pop("WANDB_RUN_ID")
-            # TODO: handle secret pulling image from registries?
-        else:
+        elif not any(["image" in cont for cont in containers]):
             if len(containers) > 1:
                 raise LaunchError(
                     "Launch only builds one container at a time. Multiple container configurations should be pre-built and specified in a yaml file supplied via job_spec."
                 )
-            given_reg = resource_args.get("registry", "")
-            repository: Optional[str] = (
-                given_reg if given_reg != "" else registry_config.get("url")
-            )
-            if repository is None:
-                # allow local registry usage for eg local clusters but throw a warning
-                wandb.termwarn(
-                    f"{LOG_PREFIX}Warning: No Docker repository specified. Image will be hosted on local registry, which may not be accessible to your training cluster."
-                )
             assert entry_point is not None
-            image_uri = builder.build_image(launch_project, repository, entry_point)
+            assert builder is not None
+            image_uri = builder.build_image(launch_project, entry_point)
             # in the non instance case we need to make an imagePullSecret
             # so the new job can pull the image
+            if not builder.registry:
+                raise LaunchError(
+                    "No registry specified. Please specify a registry in your wandb/settings file or pass a registry to the builder."
+                )
             secret = maybe_create_imagepull_secret(
-                core_api, registry_config, launch_project.run_id, namespace
+                core_api, builder.registry, launch_project.run_id, namespace
             )
 
             containers[0]["image"] = image_uri
@@ -370,6 +380,9 @@ class KubernetesRunner(AbstractRunner):
         kubernetes_style_env_vars = [
             {"name": k, "value": v} for k, v in env_vars.items()
         ]
+        _logger.info(
+            f"Using environment variables: {given_env_vars + kubernetes_style_env_vars}"
+        )
         for cont in containers:
             cont["env"] = given_env_vars + kubernetes_style_env_vars
         pod_spec["containers"] = containers
@@ -385,9 +398,7 @@ class KubernetesRunner(AbstractRunner):
         job_dict["metadata"] = job_metadata
         job_dict["status"] = job_status
 
-        if not self.ack_run_queue_item(launch_project):
-            return None
-
+        _logger.info(f"Creating Kubernetes job from: {job_dict}")
         job_response = kubernetes.utils.create_from_yaml(
             api_client, yaml_objects=[job_dict], namespace=namespace
         )[0][
@@ -409,57 +420,34 @@ class KubernetesRunner(AbstractRunner):
 
 def maybe_create_imagepull_secret(
     core_api: "CoreV1Api",
-    registry_config: Dict[str, Any],
+    registry: AbstractRegistry,
     run_id: str,
     namespace: str,
 ) -> Optional["V1Secret"]:
     secret = None
-    ecr_provider = registry_config.get("ecr-provider", "").lower()
-    if (
-        ecr_provider
-        and ecr_provider == "aws"
-        and registry_config.get("url") is not None
-        and registry_config.get("credentials") is not None
-    ):
-        boto3 = get_module(
-            "boto3", "AWS ECR requires boto3,  install with pip install wandb[launch]"
-        )
-        ecr_client = boto3.client("ecr")
-        try:
-            encoded_token = ecr_client.get_authorization_token()["authorizationData"][
-                0
-            ]["authorizationToken"]
-            decoded_token = base64.b64decode(encoded_token.encode()).decode()
-            uname, token = decoded_token.split(":")
-        except Exception as e:
-            raise LaunchError(f"Could not get authorization token for ECR, error: {e}")
-        creds_info = {
-            "auths": {
-                registry_config.get("url"): {
-                    "username": uname,
-                    "password": token,
-                    # need an email but the use is deprecated
-                    "email": "deprecated@wandblaunch.com",
-                    "auth": encoded_token,
-                }
+    if isinstance(registry, LocalRegistry):
+        # Secret not required
+        return None
+    uname, token = registry.get_username_password()
+    creds_info = {
+        "auths": {
+            registry.uri: {
+                "auth": base64.b64encode(f"{uname}:{token}".encode()).decode(),
+                # need an email but the use is deprecated
+                "email": "deprecated@wandblaunch.com",
             }
         }
-        secret_data = {
-            ".dockerconfigjson": base64.b64encode(
-                json.dumps(creds_info).encode()
-            ).decode()
-        }
-        secret = client.V1Secret(
-            data=secret_data,
-            metadata=client.V1ObjectMeta(name=f"regcred-{run_id}", namespace=namespace),
-            kind="Secret",
-            type="kubernetes.io/dockerconfigjson",
-        )
-        try:
-            core_api.create_namespaced_secret(namespace, secret)
-        except Exception as e:
-            raise LaunchError(f"Exception when creating Kubernetes secret: {str(e)}\n")
-    # TODO: support other ecr providers
-    elif ecr_provider and ecr_provider != "aws":
-        raise LaunchError(f"Registry provider not supported: {ecr_provider}")
-    return secret
+    }
+    secret_data = {
+        ".dockerconfigjson": base64.b64encode(json.dumps(creds_info).encode()).decode()
+    }
+    secret = client.V1Secret(
+        data=secret_data,
+        metadata=client.V1ObjectMeta(name=f"regcred-{run_id}", namespace=namespace),
+        kind="Secret",
+        type="kubernetes.io/dockerconfigjson",
+    )
+    try:
+        return core_api.create_namespaced_secret(namespace, secret)
+    except Exception as e:
+        raise LaunchError(f"Exception when creating Kubernetes secret: {str(e)}\n")