wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -1,10 +1,13 @@
 """Implementation of Elastic Container Registry class for wandb launch."""
 import base64
 import logging
+import re
 from typing import Dict, Tuple
 
+import yaml
+
 from wandb.sdk.launch.environment.aws_environment import AwsEnvironment
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 from wandb.util import get_module
 
 from .abstract import AbstractRegistry
@@ -52,7 +55,7 @@ class ElasticContainerRegistry(AbstractRegistry):
     @classmethod
     def from_config(  # type: ignore[override]
         cls,
-        config: Dict,
+        config: Dict[str, str],
         environment: AwsEnvironment,
         verify: bool = True,
     ) -> "ElasticContainerRegistry":
@@ -70,10 +73,41 @@ class ElasticContainerRegistry(AbstractRegistry):
                 f"Could not create ElasticContainerRegistry from config. Expected type 'ecr' "
                 f"but got '{config.get('type')}'."
             )
-        repository = config.get("repository")
-        if not repository:
+        if ("uri" in config) == ("repository" in config):
+            raise LaunchError(
+                "Could not create ElasticContainerRegistry from config. Either 'uri' or "
+                f"'repository' is required. The config received was:\n{yaml.dump(config)}."
+            )
+        if "repository" in config:
+            repository = config.get("repository")
+        else:
+            match = re.match(
+                r"^(?P<account>.*)\.dkr\.ecr\.(?P<region>.*)\.amazonaws\.com/(?P<repository>.*)/?$",
+                config["uri"],
+            )
+            if not match:
+                raise LaunchError(
+                    f"Could not create ElasticContainerRegistry from config. The uri "
+                    f"{config.get('uri')} is invalid."
+                )
+            repository = match.group("repository")
+            if match.group("region") != environment.region:
+                raise LaunchError(
+                    f"Could not create ElasticContainerRegistry from config. The uri "
+                    f"{config.get('uri')} is in region {match.group('region')} but the "
+                    f"environment is in region {environment.region}."
+                )
+            if match.group("account") != environment._account:
+                raise LaunchError(
+                    f"Could not create ElasticContainerRegistry from config. The uri "
+                    f"{config.get('uri')} is in account {match.group('account')} but the "
+                    f"account being used is {environment._account}."
+                )
+        if not isinstance(repository, str):
+            # This is for mypy. We should never get here.
             raise LaunchError(
-                "Could not create ElasticContainerRegistry from config. 'repository' is required."
+                f"Could not create ElasticContainerRegistry from config. The repository "
+                f"{repository} is invalid: repository should be a string."
             )
         return cls(repository, environment)
 

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -3,8 +3,10 @@ import logging
 import re
 from typing import Tuple
 
+import yaml
+
 from wandb.sdk.launch.environment.gcp_environment import GcpEnvironment
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 from wandb.util import get_module
 
 from .abstract import AbstractRegistry
@@ -99,20 +101,57 @@ class GoogleArtifactRegistry(AbstractRegistry):
         Arguments:
             config: A dictionary containing the following keys:
                 repository: The repository name.
-                image_name: The image 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."
+        if "uri" in config:
+            if "repository" in config or "image-name" in config:
+                raise LaunchError(
+                    "The Google Artifact Registry must be specified with either "
+                    "the uri key or the repository and image-name keys, but not both. "
+                    f"The provided config is:\n{yaml.dump(config)}."
+                )
+            match = re.match(
+                r"^(?P<region>[\w-]+)-docker\.pkg\.dev/(?P<project>[\w-]+)/(?P<repository>[\w-]+)/(?P<image_name>[\w-]+)$",
+                config["uri"],
             )
-        image_name = config.get("image_name")
-        if not image_name:
-            raise LaunchError("The image name must be specified.")
+            if not match:
+                raise LaunchError(
+                    f"The Google Artifact Registry uri {config['uri']} is invalid. "
+                    "Please provide a uri of the form "
+                    "REGION-docker.pkg.dev/PROJECT/REPOSITORY/IMAGE_NAME."
+                )
+            else:
+                repository = match.group("repository")
+                image_name = match.group("image_name")
+                if match.group("region") != environment.region:
+                    raise LaunchError(
+                        f"The Google Artifact Registry uri {config['uri']} does not "
+                        f"match the configured region {environment.region}."
+                    )
+                if match.group("project") != environment.project:
+                    raise LaunchError(
+                        f"The Google Artifact Registry uri {config['uri']} does not "
+                        f"match the configured project {environment.project}."
+                    )
+        else:
+            repository = config.get("repository")
+            if not repository:
+                raise LaunchError(
+                    "The Google Artifact Registry repository must be specified "
+                    "by setting the either the uri or  repository key of your "
+                    f"registry config. The provided config is:\n{yaml.dump(config)}."
+                )
+            image_name = config.get("image-name")
+            if not image_name:
+                raise LaunchError(
+                    "The Google Artifact Registry repository must be specified "
+                    "by setting the either the uri or  repository key of your "
+                    f"registry config. The provided config is:\n{yaml.dump(config)}."
+                )
         return cls(repository, image_name, environment, verify=verify)
 
     def verify(self) -> None:
@@ -182,22 +221,25 @@ class GoogleArtifactRegistry(AbstractRegistry):
         _logger.info(
             f"Checking if image {image_uri} exists. In Google Artifact Registry {self.uri}."
         )
+        repo_uri, tag = 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(name=image_uri)
+        parent = f"projects/{self.environment.project}/locations/{self.environment.region}/repositories/{self.repository}"
+        client = google.cloud.artifactregistry.ArtifactRegistryClient(
+            credentials=credentials
+        )
+        try:
+            for image in client.list_docker_images(request={"parent": parent}):
+                if tag in image.tags:
+                    return True
+        except google.api_core.exceptions.NotFound as e:
+            raise LaunchError(
+                f"The Google Artifact Registry repository {self.repository} does not exist."
+            ) from e
         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

@@ -2,7 +2,8 @@
 import logging
 from typing import Tuple
 
-from wandb.sdk.launch.utils import LaunchError, docker_image_exists
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.utils import docker_image_exists
 
 from ..environment.abstract import AbstractEnvironment
 from .abstract import AbstractRegistry

wandb/sdk/launch/runner/abstract.py

@@ -1,9 +1,14 @@
+"""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 subprocess
 import sys
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from dockerpycreds.utils import find_executable  # type: ignore
 
@@ -17,6 +22,9 @@ from .._project_spec import LaunchProject
 
 _logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
+
 
 if sys.version_info >= (3, 8):
     from typing import Literal
@@ -24,7 +32,14 @@ else:
     from typing_extensions import Literal
 
 State = Literal[
-    "unknown", "starting", "running", "failed", "finished", "stopping", "stopped"
+    "unknown",
+    "starting",
+    "running",
+    "failed",
+    "finished",
+    "stopping",
+    "stopped",
+    "preempted",
 ]
 
 
@@ -56,6 +71,11 @@ class AbstractRun(ABC):
     def status(self) -> Status:
         return self._status
 
+    @abstractmethod
+    def get_logs(self) -> Optional[str]:
+        """Return the logs associated with the run."""
+        pass
+
     def _run_cmd(
         self, cmd: List[str], output_only: Optional[bool] = False
     ) -> Optional[Union["subprocess.Popen[bytes]", bytes]]:
@@ -101,7 +121,7 @@ class AbstractRun(ABC):
 
     @property
     @abstractmethod
-    def id(self) -> str:
+    def id(self) -> Optional[str]:
         pass
 
 
@@ -147,6 +167,7 @@ class AbstractRunner(ABC):
         self,
         launch_project: LaunchProject,
         builder: AbstractBuilder,
+        job_tracker: Optional["JobAndRunStatusTracker"] = None,
     ) -> Optional[AbstractRun]:
         """Submit an LaunchProject to be run.