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

wandb/sdk/launch/registry/azure_container_registry.py

@@ -0,0 +1,132 @@
+"""Implementation of AzureContainerRegistry class."""
+import re
+from typing import TYPE_CHECKING, Tuple
+
+from wandb.util import get_module
+
+from ..environment.abstract import AbstractEnvironment
+from ..environment.azure_environment import AzureEnvironment
+from ..errors import LaunchError
+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,
+        environment: AzureEnvironment,
+        uri: str,
+        verify: bool = True,
+    ):
+        """Initialize an AzureContainerRegistry."""
+        self.environment = environment
+        self.uri = uri
+        if verify:
+            self.verify()
+
+    @classmethod
+    def from_config(
+        cls, config: dict, environment: AbstractEnvironment, verify: bool = True
+    ) -> "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.
+        """
+        if not isinstance(environment, AzureEnvironment):
+            raise LaunchError(
+                "AzureContainerRegistry requires an AzureEnvironment to be passed in."
+            )
+        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,
+            environment=environment,
+            verify=verify,
+        )
+
+    def get_username_password(self) -> Tuple[str, str]:
+        """Get username and password for container registry."""
+        raise NotImplementedError
+
+    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.
+        """
+        credential = self.environment.get_credentials()
+        registry, repository, tag = self.parse_azurecr_uri(image_uri)
+        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
+
+    def get_repo_uri(self) -> str:
+        return self.uri
+
+    def verify(self) -> None:
+        try:
+            _ = self.registry_name
+        except Exception as e:
+            raise LaunchError(f"Unable to verify Azure Container Registry: {e}") from e
+
+    @property
+    def registry_name(self) -> str:
+        """Get registry name."""
+        return self.parse_azurecr_uri(self.uri)[0]
+
+    @staticmethod
+    def parse_azurecr_uri(uri: str) -> Tuple[str, str, str]:
+        """Parse an Azure Container Registry URI.
+
+        Args:
+            uri (str): URI to parse.
+
+        Returns:
+            Tuple[str, str, str]: Tuple of registry name, repository name, and tag.
+
+        Raises:
+            LaunchError: If unable to parse URI.
+        """
+        regex = r"(?:https://)([\w]+)\.azurecr\.io/([\w\-]+):?(.*)"
+        match = re.match(regex, uri)
+        if match is None:
+            raise LaunchError(f"Unable to parse Azure Container Registry URI: {uri}")
+        return match.group(1), match.group(2), match.group(3)

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -1,8 +1,11 @@
 """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.errors import LaunchError
 from wandb.util import get_module
@@ -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,6 +3,8 @@ import logging
 import re
 from typing import Tuple
 
+import yaml
+
 from wandb.sdk.launch.environment.gcp_environment import GcpEnvironment
 from wandb.sdk.launch.errors import LaunchError
 from wandb.util import get_module
@@ -105,14 +107,51 @@ class GoogleArtifactRegistry(AbstractRegistry):
         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:

wandb/sdk/launch/runner/abstract.py

@@ -8,7 +8,7 @@ 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
 
@@ -22,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
@@ -29,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",
 ]
 
 
@@ -61,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]]:
@@ -106,7 +121,7 @@ class AbstractRun(ABC):
 
     @property
     @abstractmethod
-    def id(self) -> str:
+    def id(self) -> Optional[str]:
         pass
 
 
@@ -152,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.
 

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -4,13 +4,15 @@ import base64
 import json
 import logging
 import time
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import wandb
 from wandb.apis.internal import Api
+from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 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.azure_container_registry import AzureContainerRegistry
 from wandb.sdk.launch.registry.local_registry import LocalRegistry
 from wandb.sdk.launch.runner.abstract import State, Status
 from wandb.util import get_module
@@ -104,6 +106,22 @@ class KubernetesSubmittedRun(AbstractRun):
         """Return the run id."""
         return self.name
 
+    def get_logs(self) -> Optional[str]:
+        try:
+            logs = self.core_api.read_namespaced_pod_log(
+                name=self.pod_names[0], namespace=self.namespace
+            )
+            if logs:
+                return str(logs)
+            else:
+                wandb.termwarn(
+                    f"Retrieved no logs for kubernetes pod(s): {self.pod_names}"
+                )
+            return None
+        except Exception as e:
+            wandb.termerror(f"{LOG_PREFIX}Failed to get pod logs: {e}")
+            return None
+
     def get_job(self) -> "V1Job":
         """Return the job object."""
         return self.batch_api.read_namespaced_job(
@@ -128,14 +146,22 @@ class KubernetesSubmittedRun(AbstractRun):
 
     def get_status(self) -> Status:
         """Return the run status."""
-        job_response = self.batch_api.read_namespaced_job_status(
-            name=self.name, namespace=self.namespace
-        )
-        status = job_response.status
+        try:
+            job_response = self.batch_api.read_namespaced_job_status(
+                name=self.name, namespace=self.namespace
+            )
+            status = job_response.status
+
+            pod = self.core_api.read_namespaced_pod(
+                name=self.pod_names[0], namespace=self.namespace
+            )
+        except ApiException as e:
+            if "(404)" not in str(e):
+                raise
+            # 404 = Pod/job not reachable
+            wandb.termlog(f"{LOG_PREFIX}Job or pod disconnected for job: {self.name}")
+            return Status("preempted")
 
-        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:
@@ -154,7 +180,13 @@ class KubernetesSubmittedRun(AbstractRun):
         if status.succeeded == 1:
             return_status = Status("finished")
         elif status.failed is not None and status.failed >= 1:
-            return_status = Status("failed")
+            if status.conditions[0].reason == "BackoffLimitExceeded":
+                wandb.termlog(
+                    f"{LOG_PREFIX}Job or pod disconnected for job: {self.name}"
+                )
+                return_status = Status("preempted")
+            else:
+                return_status = Status("failed")
         elif status.active == 1:
             return Status("running")
         elif status.conditions is not None and status.conditions[0].type == "Suspended":
@@ -261,6 +293,23 @@ class CrdSubmittedRun(AbstractRun):
         """Get the name of the custom object."""
         return self.name
 
+    def get_logs(self) -> Optional[str]:
+        """Get logs for custom object."""
+        # TODO: test more carefully once we release multi-node support
+        logs: Dict[str, Optional[str]] = {}
+        try:
+            for pod_name in self.pod_names:
+                logs[pod_name] = self.core_api.read_namespaced_pod_log(
+                    name=pod_name, namespace=self.namespace
+                )
+        except ApiException as e:
+            wandb.termwarn(f"Failed to get logs for {self.name}: {str(e)}")
+            return None
+        if not logs:
+            return None
+        logs_as_array = [f"Pod {pod_name}:\n{log}" for pod_name, log in logs.items()]
+        return "\n".join(logs_as_array)
+
     def get_status(self) -> Status:
         """Get status of custom object."""
         try:
@@ -403,6 +452,7 @@ class KubernetesRunner(AbstractRunner):
         builder: Optional[AbstractBuilder],
         namespace: str,
         core_api: "CoreV1Api",
+        job_tracker: Optional[JobAndRunStatusTracker],
     ) -> Tuple[Dict[str, Any], Optional["V1Secret"]]:
         """Apply our default values, return job dict and secret.
 
@@ -454,9 +504,10 @@ class KubernetesRunner(AbstractRunner):
                     "Invalid specification of multiple containers. See https://docs.wandb.ai/guides/launch for guidance on submitting jobs."
                 )
             # dont specify run id if user provided image, could have multiple runs
-            containers[0]["image"] = launch_project.docker_image
+            image_uri = launch_project.docker_image
+            containers[0]["image"] = image_uri
+            launch_project.fill_macros(image_uri)
             # TODO: handle secret pulling image from registry
-            launch_project.fill_macros(launch_project.docker_image)
         elif not any(["image" in cont for cont in containers]):
             if len(containers) > 1:
                 raise LaunchError(
@@ -464,7 +515,8 @@ class KubernetesRunner(AbstractRunner):
                 )
             assert entry_point is not None
             assert builder is not None
-            image_uri = builder.build_image(launch_project, entry_point)
+            image_uri = builder.build_image(launch_project, entry_point, job_tracker)
+            image_uri = image_uri.replace("https://", "")
             launch_project.fill_macros(image_uri)
             # in the non instance case we need to make an imagePullSecret
             # so the new job can pull the image
@@ -510,6 +562,7 @@ class KubernetesRunner(AbstractRunner):
         self,
         launch_project: LaunchProject,
         builder: AbstractBuilder,
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> Optional[AbstractRun]:  # noqa: C901
         """Execute a launch project on Kubernetes.
 
@@ -545,7 +598,7 @@ class KubernetesRunner(AbstractRunner):
                 image_uri = launch_project.docker_image
             else:
                 assert entrypoint is not None
-                image_uri = builder.build_image(launch_project, entrypoint)
+                image_uri = builder.build_image(launch_project, entrypoint, job_tracker)
             launch_project.fill_macros(image_uri)
             env_vars = get_env_vars_dict(launch_project, self._api)
             # Crawl the resource args and add our env vars to the containers.
@@ -607,11 +660,7 @@ class KubernetesRunner(AbstractRunner):
         namespace = self.get_namespace(resource_args, context)
 
         job, secret = self._inject_defaults(
-            resource_args,
-            launch_project,
-            builder,
-            namespace,
-            core_api,
+            resource_args, launch_project, builder, namespace, core_api, job_tracker
         )
 
         msg = "Creating Kubernetes job"
@@ -681,7 +730,9 @@ def maybe_create_imagepull_secret(
         A secret if one was created, otherwise None.
     """
     secret = None
-    if isinstance(registry, LocalRegistry):
+    if isinstance(registry, LocalRegistry) or isinstance(
+        registry, AzureContainerRegistry
+    ):
         # Secret not required
         return None
     uname, token = registry.get_username_password()
@@ -715,30 +766,37 @@ def add_wandb_env(root: Union[dict, list], env_vars: Dict[str, str]) -> None:
     Recursively walks the spec and injects the environment variables into
     every container spec. Containers are identified by the "containers" key.
 
+    This function treats the WANDB_RUN_ID and WANDB_GROUP_ID environment variables
+    specially. If they are present in the spec, they will be overwritten. If a setting
+    for WANDB_RUN_ID is provided in env_vars, then that environment variable will only be
+    set in the first container modified by this function.
+
     Arguments:
         root: The spec to modify.
         env_vars: The environment variables to inject.
 
     Returns: None.
     """
-    if isinstance(root, dict):
-        for k, v in root.items():
-            if k == "containers":
-                if isinstance(v, list):
-                    for cont in v:
-                        env = cont.get("env", [])
-                        env.extend(
-                            [
-                                {"name": key, "value": value}
-                                for key, value in env_vars.items()
-                            ]
-                        )
-                        cont["env"] = env
-            elif isinstance(v, (dict, list)):
-                add_wandb_env(v, env_vars)
-    elif isinstance(root, list):
-        for item in root:
-            add_wandb_env(item, env_vars)
+
+    def yield_containers(root: Any) -> Iterator[dict]:
+        if isinstance(root, dict):
+            for k, v in root.items():
+                if k == "containers":
+                    if isinstance(v, list):
+                        yield from v
+                elif isinstance(v, (dict, list)):
+                    yield from yield_containers(v)
+        elif isinstance(root, list):
+            for item in root:
+                yield from yield_containers(item)
+
+    for cont in yield_containers(root):
+        env = cont.setdefault("env", [])
+        env.extend([{"name": key, "value": value} for key, value in env_vars.items()])
+        cont["env"] = env
+        # After we have set WANDB_RUN_ID once, we don't want to set it again
+        if "WANDB_RUN_ID" in env_vars:
+            env_vars.pop("WANDB_RUN_ID")
 
 
 def add_label_to_pods(
@@ -757,16 +815,22 @@ def add_label_to_pods(
 
     Returns: None.
     """
-    if isinstance(manifest, list):
-        for item in manifest:
-            add_label_to_pods(item, label_key, label_value)
-    elif isinstance(manifest, dict):
-        if "spec" in manifest and "containers" in manifest["spec"]:
-            metadata = manifest.setdefault("metadata", {})
-            labels = metadata.setdefault("labels", {})
-            labels[label_key] = label_value
-        for value in manifest.values():
-            add_label_to_pods(value, label_key, label_value)
+
+    def yield_pods(manifest: Any) -> Iterator[dict]:
+        if isinstance(manifest, list):
+            for item in manifest:
+                yield from yield_pods(item)
+        elif isinstance(manifest, dict):
+            if "spec" in manifest and "containers" in manifest["spec"]:
+                yield manifest
+            for value in manifest.values():
+                if isinstance(value, (dict, list)):
+                    yield from yield_pods(value)
+
+    for pod in yield_pods(manifest):
+        metadata = pod.setdefault("metadata", {})
+        labels = metadata.setdefault("labels", {})
+        labels[label_key] = label_value
 
 
 def add_entrypoint_args_overrides(manifest: Union[dict, list], overrides: dict) -> None: