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

wandb/sdk/launch/builder/kaniko_builder.py

@@ -1,13 +1,11 @@
 import base64
 import json
+import logging
 import tarfile
 import tempfile
 import time
 from typing import Optional
 
-import kubernetes  # type: ignore
-from kubernetes import client
-
 import wandb
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
@@ -16,7 +14,7 @@ from wandb.sdk.launch.registry.elastic_container_registry import (
     ElasticContainerRegistry,
 )
 from wandb.sdk.launch.registry.google_artifact_registry import GoogleArtifactRegistry
-from wandb.sdk.launch.utils import LaunchError
+from wandb.util import get_module
 
 from .._project_spec import (
     EntryPoint,
@@ -24,8 +22,28 @@ from .._project_spec import (
     create_metadata_file,
     get_entry_point_command,
 )
-from ..utils import LOG_PREFIX, get_kube_context_and_api_client, sanitize_wandb_api_key
-from .build import _create_docker_build_ctx, generate_dockerfile
+from ..utils import (
+    LOG_PREFIX,
+    LaunchError,
+    get_kube_context_and_api_client,
+    sanitize_wandb_api_key,
+    warn_failed_packages_from_build_logs,
+)
+from .build import (
+    _create_docker_build_ctx,
+    generate_dockerfile,
+    image_tag_from_dockerfile_and_source,
+)
+
+get_module(
+    "kubernetes",
+    required="Kaniko builder requires the kubernetes package. Please install it with `pip install wandb[launch]`.",
+)
+
+import kubernetes  # type: ignore # noqa: E402
+from kubernetes import client  # noqa: E402
+
+_logger = logging.getLogger(__name__)
 
 _DEFAULT_BUILD_TIMEOUT_SECS = 1800  # 30 minute build timeout
 
@@ -155,7 +173,7 @@ class KanikoBuilder(AbstractBuilder):
         pass
 
     def _create_docker_ecr_config_map(
-        self, corev1_client: client.CoreV1Api, repository: str
+        self, job_name: str, corev1_client: client.CoreV1Api, repository: str
     ) -> None:
         if self.registry is None:
             raise LaunchError("No registry specified for Kaniko build.")
@@ -165,7 +183,7 @@ class KanikoBuilder(AbstractBuilder):
             api_version="v1",
             kind="ConfigMap",
             metadata=client.V1ObjectMeta(
-                name="docker-config",
+                name=f"docker-config-{job_name}",
                 namespace="wandb",
             ),
             data={
@@ -177,8 +195,11 @@ class KanikoBuilder(AbstractBuilder):
         )
         corev1_client.create_namespaced_config_map("wandb", ecr_config_map)
 
-    def _delete_docker_ecr_config_map(self, client: client.CoreV1Api) -> None:
-        client.delete_namespaced_config_map("docker-config", "wandb")
+    def _delete_docker_ecr_config_map(
+        self, job_name: str, client: client.CoreV1Api
+    ) -> None:
+        if self.secret_name:
+            client.delete_namespaced_config_map(f"docker-config-{job_name}", "wandb")
 
     def _upload_build_context(self, run_id: str, context_path: str) -> str:
         # creat a tar archive of the build context and upload it to s3
@@ -197,20 +218,28 @@ class KanikoBuilder(AbstractBuilder):
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
     ) -> str:
+        # TODO: this should probably throw an error if the registry is a local registry
         if not self.registry:
             raise LaunchError("No registry specified for Kaniko build.")
+        # kaniko builder doesn't seem to work with a custom user id, need more investigation
+        dockerfile_str = generate_dockerfile(
+            launch_project, entrypoint, launch_project.resource, "kaniko"
+        )
+        image_tag = image_tag_from_dockerfile_and_source(launch_project, dockerfile_str)
         repo_uri = self.registry.get_repo_uri()
-        image_uri = repo_uri + ":" + launch_project.image_tag
+        image_uri = repo_uri + ":" + image_tag
+
+        if not launch_project.build_required() and self.registry.check_image_exists(
+            image_uri
+        ):
+            return image_uri
+
+        _logger.info(f"Building image {image_uri}...")
 
         entry_cmd = " ".join(
             get_entry_point_command(entrypoint, launch_project.override_args)
         )
 
-        # kaniko builder doesn't seem to work with a custom user id, need more investigation
-        dockerfile_str = generate_dockerfile(
-            launch_project, entrypoint, launch_project.resource, self.type
-        )
-
         create_metadata_file(
             launch_project,
             image_uri,
@@ -240,7 +269,8 @@ class KanikoBuilder(AbstractBuilder):
 
         try:
             # core_v1.create_namespaced_config_map("wandb", dockerfile_config_map)
-            self._create_docker_ecr_config_map(core_v1, repo_uri)
+            if self.secret_name:
+                self._create_docker_ecr_config_map(build_job_name, core_v1, repo_uri)
             batch_v1.create_namespaced_job("wandb", build_job)
 
             # wait for double the job deadline since it might take time to schedule
@@ -248,6 +278,13 @@ class KanikoBuilder(AbstractBuilder):
                 batch_v1, build_job_name, 3 * _DEFAULT_BUILD_TIMEOUT_SECS
             ):
                 raise Exception(f"Failed to build image in kaniko for job {run_id}")
+            try:
+                logs = batch_v1.read_namespaced_job_log(build_job_name, "wandb")
+                warn_failed_packages_from_build_logs(logs, image_uri)
+            except Exception as e:
+                wandb.termwarn(
+                    f"{LOG_PREFIX}Failed to get logs for kaniko job {build_job_name}: {e}"
+                )
         except Exception as e:
             wandb.termerror(
                 f"{LOG_PREFIX}Exception when creating Kubernetes resources: {e}\n"
@@ -258,7 +295,8 @@ class KanikoBuilder(AbstractBuilder):
             try:
                 # should we clean up the s3 build contexts? can set bucket level policy to auto deletion
                 # core_v1.delete_namespaced_config_map(config_map_name, "wandb")
-                self._delete_docker_ecr_config_map(core_v1)
+                if self.secret_name:
+                    self._delete_docker_ecr_config_map(build_job_name, core_v1)
                 batch_v1.delete_namespaced_job(build_job_name, "wandb")
             except Exception as e:
                 raise LaunchError(f"Exception during Kubernetes resource clean up {e}")
@@ -273,25 +311,34 @@ class KanikoBuilder(AbstractBuilder):
         build_context_path: str,
     ) -> "client.V1Job":
         env = []
-        volume_mounts = [
-            client.V1VolumeMount(name="docker-config", mount_path="/kaniko/.docker/"),
-        ]
-        volumes = [
-            client.V1Volume(
-                name="docker-config",
-                config_map=client.V1ConfigMapVolumeSource(
-                    name="docker-config",
-                ),
-            ),
-        ]
+        volume_mounts = []
+        volumes = []
         if bool(self.secret_name) != bool(self.secret_key):
             raise LaunchError(
                 "Both secret_name and secret_key or neither must be specified "
                 "for kaniko build. You provided only one of them."
             )
+        if isinstance(self.registry, ElasticContainerRegistry):
+            env += [
+                client.V1EnvVar(
+                    name="AWS_REGION",
+                    value=self.registry.environment.region,
+                )
+            ]
         if self.secret_name and self.secret_key:
-            # TODO: We should validate that the secret exists and has the key
-            # before creating the job. Or when we create the builder.
+            volumes += [
+                client.V1Volume(
+                    name="docker-config",
+                    config_map=client.V1ConfigMapVolumeSource(
+                        name=f"docker-config-{job_name}",
+                    ),
+                ),
+            ]
+            volume_mounts += [
+                client.V1VolumeMount(
+                    name="docker-config", mount_path="/kaniko/.docker/"
+                ),
+            ]
             # TODO: I don't like conditioning on the registry type here. As a
             # future change I want the registry and environment classes to provide
             # a list of environment variables and volume mounts that need to be
@@ -303,12 +350,6 @@ class KanikoBuilder(AbstractBuilder):
             if isinstance(self.registry, ElasticContainerRegistry):
                 mount_path = "/root/.aws"
                 key = "credentials"
-                env += [
-                    client.V1EnvVar(
-                        name="AWS_REGION",
-                        value=self.registry.environment.region,
-                    )
-                ]
             elif isinstance(self.registry, GoogleArtifactRegistry):
                 mount_path = "/kaniko/.config/gcloud"
                 key = "config.json"

wandb/sdk/launch/builder/templates/_wandb_bootstrap.py

@@ -1,10 +1,13 @@
 import json
 import multiprocessing
 import os
+import re
 import subprocess
 import sys
 from typing import List, Optional, Set
 
+FAILED_PACKAGES_PREFIX = "ERROR: Failed to install: "
+FAILED_PACKAGES_POSTFIX = ". During automated build process."
 CORES = multiprocessing.cpu_count()
 ONLY_INCLUDE = {x for x in os.getenv("WANDB_ONLY_INCLUDE", "").split(",") if x != ""}
 OPTS = []
@@ -21,7 +24,10 @@ else:
 
 
 def install_deps(
-    deps: List[str], failed: Optional[Set[str]] = None
+    deps: List[str],
+    failed: Optional[Set[str]] = None,
+    extra_index: Optional[str] = None,
+    opts: Optional[List[str]] = None,
 ) -> Optional[Set[str]]:
     """Install pip dependencies.
 
@@ -35,33 +41,45 @@ def install_deps(
     try:
         # Include only uri if @ is present
         clean_deps = [d.split("@")[-1].strip() if "@" in d else d for d in deps]
-
+        index_args = ["--extra-index-url", extra_index] if extra_index else []
         print("installing {}...".format(", ".join(clean_deps)))
+        opts = opts or []
+        args = ["pip", "install"] + opts + clean_deps + index_args
         sys.stdout.flush()
-        subprocess.check_output(
-            ["pip", "install"] + OPTS + clean_deps, stderr=subprocess.STDOUT
-        )
-        if failed is not None and len(failed) > 0:
-            sys.stderr.write(
-                "ERROR: Unable to install: {}".format(", ".join(clean_deps))
-            )
-            sys.stderr.flush()
+        subprocess.check_output(args, stderr=subprocess.STDOUT)
         return failed
     except subprocess.CalledProcessError as e:
         if failed is None:
             failed = set()
         num_failed = len(failed)
-        for line in e.output.decode("utf8"):
+        for line in e.output.decode("utf8").splitlines():
             if line.startswith("ERROR:"):
-                failed.add(line.split(" ")[-1])
-        if len(failed) > num_failed:
-            return install_deps(list(set(clean_deps) - failed), failed)
+                clean_dep = find_package_in_error_string(clean_deps, line)
+                if clean_dep is not None:
+                    if clean_dep in deps:
+                        failed.add(clean_dep)
+                    else:
+                        for d in deps:
+                            if clean_dep in d:
+                                failed.add(d.replace(" ", ""))
+                                break
+        if len(set(clean_deps) - failed) == 0:
+            return failed
+        elif len(failed) > num_failed:
+            return install_deps(
+                list(set(clean_deps) - failed),
+                failed,
+                extra_index=extra_index,
+                opts=opts,
+            )
         else:
             return failed
 
 
 def main() -> None:
     """Install deps in requirements.frozen.txt."""
+    extra_index = None
+    torch_reqs = []
     if os.path.exists("requirements.frozen.txt"):
         with open("requirements.frozen.txt") as f:
             print("Installing frozen dependencies...")
@@ -72,28 +90,60 @@ def main() -> None:
                     # can't pip install wandb==0.*.*.dev1 through pip. Lets just install wandb for now
                     if req.startswith("wandb==") and "dev1" in req:
                         req = "wandb"
-                    reqs.append(req.strip().replace(" ", ""))
+                    match = re.match(
+                        r"torch(vision|audio)?==\d+\.\d+\.\d+(\+(?:cu[\d]{2,3})|(?:cpu))?",
+                        req,
+                    )
+                    if match:
+                        variant = match.group(2)
+                        if variant:
+                            extra_index = (
+                                f"https://download.pytorch.org/whl/{variant[1:]}"
+                            )
+                        torch_reqs.append(req.strip().replace(" ", ""))
+                    else:
+                        reqs.append(req.strip().replace(" ", ""))
                 else:
                     print(f"Ignoring requirement: {req} from frozen requirements")
                 if len(reqs) >= CORES:
-                    deps_failed = install_deps(reqs)
+                    deps_failed = install_deps(reqs, opts=OPTS)
                     reqs = []
                     if deps_failed is not None:
                         failed = failed.union(deps_failed)
             if len(reqs) > 0:
-                deps_failed = install_deps(reqs)
+                deps_failed = install_deps(reqs, opts=OPTS)
                 if deps_failed is not None:
                     failed = failed.union(deps_failed)
             with open("_wandb_bootstrap_errors.json", "w") as f:
                 f.write(json.dumps({"pip": list(failed)}))
             if len(failed) > 0:
                 sys.stderr.write(
-                    "ERROR: Failed to install: {}".format(",".join(failed))
+                    FAILED_PACKAGES_PREFIX + ",".join(failed) + FAILED_PACKAGES_POSTFIX
                 )
                 sys.stderr.flush()
+        install_deps(torch_reqs, extra_index=extra_index)
     else:
         print("No frozen requirements found")
 
 
+# hacky way to get the name of the requirement that failed
+# attempt last word which is the name of the package often
+# fall back to checking all words in the line for the package name
+def find_package_in_error_string(deps: List[str], line: str) -> Optional[str]:
+    # if the last word in the error string is in the list of deps, return it
+    last_word = line.split(" ")[-1]
+    if last_word in deps:
+        return last_word
+    # if the last word is not in the list of deps, check all words
+    # TODO: this could report the wrong package if the error string
+    # contains a reference to another package in the deps
+    # before the package that failed to install
+    for word in line.split(" "):
+        if word in deps:
+            return word
+    # if we can't find the package, return None
+    return None
+
+
 if __name__ == "__main__":
     main()

wandb/sdk/launch/environment/aws_environment.py

@@ -70,6 +70,10 @@ class AwsEnvironment(AbstractEnvironment):
             session = boto3.Session()
             region = region or session.region_name
             credentials = session.get_credentials()
+            if not credentials:
+                raise LaunchError(
+                    "Could not create AWS environment from default environment. Please verify that your AWS credentials are configured correctly."
+                )
             access_key = credentials.access_key
             secret_key = credentials.secret_key
             session_token = credentials.token

wandb/sdk/launch/launch.py

@@ -55,7 +55,7 @@ def resolve_agent_config(  # noqa: C901
         "api_key": api.api_key,
         "base_url": api.settings("base_url"),
         "registry": {},
-        "build": {},
+        "builder": {},
         "runner": {},
     }
     user_set_project = False
@@ -151,7 +151,6 @@ def _run(
     resource_args: Optional[Dict[str, Any]],
     launch_config: Optional[Dict[str, Any]],
     synchronous: Optional[bool],
-    cuda: Optional[bool],
     api: Api,
     run_id: Optional[str],
     repository: Optional[str],
@@ -171,7 +170,6 @@ def _run(
         parameters,
         resource_args,
         launch_config,
-        cuda,
         run_id,
         repository,
     )
@@ -217,7 +215,6 @@ def run(
     docker_image: Optional[str] = None,
     config: Optional[Dict[str, Any]] = None,
     synchronous: Optional[bool] = True,
-    cuda: Optional[bool] = None,
     run_id: Optional[str] = None,
     repository: Optional[str] = None,
 ) -> AbstractRun:
@@ -247,7 +244,6 @@ def run(
         asynchronous runs launched via this method will be terminated. If
         ``synchronous`` is True and the run fails, the current process will
         error out as well.
-    cuda: Whether to build a CUDA-enabled docker image or not
     run_id: ID for the run (To ultimately replace the :name: field)
     repository: string name of repository path for remote registry
 
@@ -290,7 +286,6 @@ def run(
         resource_args=resource_args,
         launch_config=config,
         synchronous=synchronous,
-        cuda=cuda,
         api=api,
         run_id=run_id,
         repository=repository,

wandb/sdk/launch/launch_add.py

@@ -44,7 +44,6 @@ def launch_add(
     params: Optional[Dict[str, Any]] = None,
     project_queue: Optional[str] = None,
     resource_args: Optional[Dict[str, Any]] = None,
-    cuda: Optional[bool] = None,
     run_id: Optional[str] = None,
     build: Optional[bool] = False,
     repository: Optional[str] = None,
@@ -69,7 +68,6 @@ def launch_add(
         the parameters used to run the original run.
     resource_args: Resource related arguments for launching runs onto a remote backend.
         Will be stored on the constructed launch config under ``resource_args``.
-    cuda: Whether to build a CUDA-enabled docker image or not
     run_id: optional string indicating the id of the launched run
     build: optional flag defaulting to false, requires queue to be set
         if build, an image is created, creates a job artifact, pushes a reference
@@ -116,7 +114,6 @@ def launch_add(
         params,
         project_queue,
         resource_args,
-        cuda,
         run_id=run_id,
         build=build,
         repository=repository,
@@ -139,7 +136,6 @@ def _launch_add(
     params: Optional[Dict[str, Any]],
     project_queue: Optional[str],
     resource_args: Optional[Dict[str, Any]] = None,
-    cuda: Optional[bool] = None,
     run_id: Optional[str] = None,
     build: Optional[bool] = False,
     repository: Optional[str] = None,
@@ -158,7 +154,6 @@ def _launch_add(
         params,
         resource_args,
         config,
-        cuda,
         run_id,
         repository,
     )

wandb/sdk/launch/registry/abstract.py

@@ -33,6 +33,18 @@ class AbstractRegistry(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    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(

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -87,7 +87,6 @@ class ElasticContainerRegistry(AbstractRegistry):
         try:
             session = self.environment.get_session()
             client = session.client("ecr")
-            response = client.describe_registry()
             response = client.describe_repositories(repositoryNames=[self.repo_name])
             self.uri = response["repositories"][0]["repositoryUri"].split("/")[0]
 
@@ -131,3 +130,34 @@ class ElasticContainerRegistry(AbstractRegistry):
             str: The uri of the repository.
         """
         return self.uri + "/" + self.repo_name
+
+    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.
+        """
+        uri, tag = image_uri.split(":")
+        if uri != self.get_repo_uri():
+            raise LaunchError(
+                f"Image uri {image_uri} does not match Elastic Container Registry uri {self.get_repo_uri()}."
+            )
+
+        _logger.debug("Checking if image tag exists.")
+        try:
+            session = self.environment.get_session()
+            client = session.client("ecr")
+            response = 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}")

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -169,3 +169,35 @@ class GoogleArtifactRegistry(AbstractRegistry):
             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

@@ -1,11 +1,14 @@
 """Local registry implementation."""
+import logging
 from typing import Tuple
 
-from wandb.sdk.launch.utils import LaunchError
+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.
@@ -46,3 +49,14 @@ class LocalRegistry(AbstractRegistry):
         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
 
@@ -143,19 +142,6 @@ 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,