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

wandb/sdk/launch/builder/kaniko_builder.py

@@ -0,0 +1,419 @@
+import base64
+import json
+import logging
+import tarfile
+import tempfile
+import time
+from typing import Optional
+
+import wandb
+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.elastic_container_registry import (
+    ElasticContainerRegistry,
+)
+from wandb.sdk.launch.registry.google_artifact_registry import GoogleArtifactRegistry
+from wandb.util import get_module
+
+from .._project_spec import (
+    EntryPoint,
+    LaunchProject,
+    create_metadata_file,
+    get_entry_point_command,
+)
+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
+
+
+def _wait_for_completion(
+    batch_client: client.BatchV1Api, job_name: str, deadline_secs: Optional[int] = None
+) -> bool:
+    start_time = time.time()
+    while True:
+        job = batch_client.read_namespaced_job_status(job_name, "wandb")
+        if job.status.succeeded is not None and job.status.succeeded >= 1:
+            return True
+        elif job.status.failed is not None and job.status.failed >= 1:
+            wandb.termerror(f"{LOG_PREFIX}Build job {job.status.failed} failed {job}")
+            return False
+        wandb.termlog(f"{LOG_PREFIX}Waiting for build job to complete...")
+        if deadline_secs is not None and time.time() - start_time > deadline_secs:
+            return False
+
+        time.sleep(5)
+
+
+class KanikoBuilder(AbstractBuilder):
+    """Builds a docker image for a project using Kaniko."""
+
+    type = "kaniko"
+
+    build_job_name: str
+    build_context_store: str
+    secret_name: Optional[str]
+    secret_key: Optional[str]
+
+    def __init__(
+        self,
+        environment: AbstractEnvironment,
+        registry: AbstractRegistry,
+        build_job_name: str = "wandb-launch-container-build",
+        build_context_store: str = "",
+        secret_name: str = "",
+        secret_key: str = "",
+        verify: bool = True,
+    ):
+        """Initialize a KanikoBuilder.
+
+        Arguments:
+            environment (AbstractEnvironment): The environment to use.
+            registry (AbstractRegistry): The registry to use.
+            build_job_name (str, optional): The name of the build job.
+            build_context_store (str, optional): The name of the build context store.
+            secret_name (str, optional): The name of the secret to use for the registry.
+            secret_key (str, optional): The key of the secret to use for the registry.
+            verify (bool, optional): Whether to verify the functionality of the builder.
+                Defaults to True.
+        """
+        if build_context_store is None:
+            raise LaunchError(
+                "You are required to specify an external build "
+                "context store for Kaniko builds. Please specify a  storage url "
+                "in the 'build-context-store' field of your builder config."
+            )
+        self.environment = environment
+        self.registry = registry
+        self.build_job_name = build_job_name
+        self.build_context_store = build_context_store.rstrip("/")
+        self.secret_name = secret_name
+        self.secret_key = secret_key
+        if verify:
+            self.verify()
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict,
+        environment: AbstractEnvironment,
+        registry: AbstractRegistry,
+        verify: bool = True,
+        login: bool = True,
+    ) -> "AbstractBuilder":
+        """Create a KanikoBuilder from a config dict.
+
+        Arguments:
+            config: A dict containing the builder config. Must contain a "type" key
+                with value "kaniko".
+            environment: The environment to use for the build.
+            registry: The registry to use for the build.
+            verify: Whether to verify the builder config.
+
+        Returns:
+            A KanikoBuilder instance.
+        """
+        if config.get("type") != "kaniko":
+            raise LaunchError(
+                "Builder config must include 'type':'kaniko' to create a KanikoBuilder."
+            )
+        build_context_store = config.get("build-context-store")
+        if build_context_store is None:
+            raise LaunchError(
+                "You are required to specify an external build "
+                "context store for Kaniko builds. Please specify a  "
+                "storage url in the 'build_context_store' field of your builder config."
+            )
+        build_job_name = config.get("build-job-name", "wandb-launch-container-build")
+        secret_name = config.get("secret-name", "")
+        secret_key = config.get("secret-key", "")
+        return cls(
+            environment,
+            registry,
+            build_context_store=build_context_store,
+            build_job_name=build_job_name,
+            secret_name=secret_name,
+            secret_key=secret_key,
+            verify=verify,
+        )
+
+    def verify(self) -> None:
+        """Verify that the builder config is valid.
+
+        Raises:
+            LaunchError: If the builder config is invalid.
+        """
+        if self.environment is None:
+            raise LaunchError("No environment specified for Kaniko build.")
+        self.environment.verify_storage_uri(self.build_context_store)
+
+    def login(self) -> None:
+        """Login to the registry."""
+        pass
+
+    def _create_docker_ecr_config_map(
+        self, job_name: str, corev1_client: client.CoreV1Api, repository: str
+    ) -> None:
+        if self.registry is None:
+            raise LaunchError("No registry specified for Kaniko build.")
+        username, password = self.registry.get_username_password()
+        encoded = base64.b64encode(f"{username}:{password}".encode()).decode("utf-8")
+        ecr_config_map = client.V1ConfigMap(
+            api_version="v1",
+            kind="ConfigMap",
+            metadata=client.V1ObjectMeta(
+                name=f"docker-config-{job_name}",
+                namespace="wandb",
+            ),
+            data={
+                "config.json": json.dumps(
+                    {"auths": {f"{self.registry.get_repo_uri()}": {"auth": encoded}}}
+                )
+            },
+            immutable=True,
+        )
+        corev1_client.create_namespaced_config_map("wandb", ecr_config_map)
+
+    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
+        context_file = tempfile.NamedTemporaryFile(delete=False)
+        with tarfile.TarFile.open(fileobj=context_file, mode="w:gz") as context_tgz:
+            context_tgz.add(context_path, arcname=".")
+        context_file.close()
+        destination = f"{self.build_context_store}/{run_id}.tgz"
+        if self.environment is None:
+            raise LaunchError("No environment specified for Kaniko build.")
+        self.environment.upload_file(context_file.name, destination)
+        return destination
+
+    def build_image(
+        self,
+        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 + ":" + 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)
+        )
+
+        create_metadata_file(
+            launch_project,
+            image_uri,
+            sanitize_wandb_api_key(entry_cmd),
+            sanitize_wandb_api_key(dockerfile_str),
+        )
+        context_path = _create_docker_build_ctx(launch_project, dockerfile_str)
+        run_id = launch_project.run_id
+
+        _, api_client = get_kube_context_and_api_client(
+            kubernetes, launch_project.resource_args
+        )
+        build_job_name = f"{self.build_job_name}-{run_id}"
+
+        build_context = self._upload_build_context(run_id, context_path)
+        build_job = self._create_kaniko_job(
+            build_job_name,
+            repo_uri,
+            image_uri,
+            build_context,
+        )
+        wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}")
+
+        # TODO: use same client as kuberentes.py
+        batch_v1 = client.BatchV1Api(api_client)
+        core_v1 = client.CoreV1Api(api_client)
+
+        try:
+            # core_v1.create_namespaced_config_map("wandb", dockerfile_config_map)
+            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
+            if not _wait_for_completion(
+                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"
+            )
+            raise e
+        finally:
+            wandb.termlog(f"{LOG_PREFIX}Cleaning up resources")
+            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")
+                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}")
+
+        return image_uri
+
+    def _create_kaniko_job(
+        self,
+        job_name: str,
+        repository: str,
+        image_tag: str,
+        build_context_path: str,
+    ) -> "client.V1Job":
+        env = []
+        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:
+            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
+            # added to the job. The environment class provides credentials for
+            # build context access, and the registry class provides credentials
+            # for pushing the image. This way we can have separate secrets for
+            # each and support build contexts and registries that require
+            # different credentials.
+            if isinstance(self.registry, ElasticContainerRegistry):
+                mount_path = "/root/.aws"
+                key = "credentials"
+            elif isinstance(self.registry, GoogleArtifactRegistry):
+                mount_path = "/kaniko/.config/gcloud"
+                key = "config.json"
+                env += [
+                    client.V1EnvVar(
+                        name="GOOGLE_APPLICATION_CREDENTIALS",
+                        value="/kaniko/.config/gcloud/config.json",
+                    )
+                ]
+            else:
+                raise LaunchError(
+                    f"Registry type {type(self.registry)} not supported by kaniko"
+                )
+            volume_mounts += [
+                client.V1VolumeMount(
+                    name=self.secret_name,
+                    mount_path=mount_path,
+                    read_only=True,
+                )
+            ]
+            volumes += [
+                client.V1Volume(
+                    name=self.secret_name,
+                    secret=client.V1SecretVolumeSource(
+                        secret_name=self.secret_name,
+                        items=[client.V1KeyToPath(key=self.secret_key, path=key)],
+                    ),
+                )
+            ]
+
+        args = [
+            f"--context={build_context_path}",
+            "--dockerfile=Dockerfile.wandb-autogenerated",
+            f"--destination={image_tag}",
+            "--cache=true",
+            f"--cache-repo={repository}",
+            "--snapshotMode=redo",
+            "--compressed-caching=false",
+        ]
+        container = client.V1Container(
+            name="wandb-container-build",
+            image="gcr.io/kaniko-project/executor:v1.8.0",
+            args=args,
+            volume_mounts=volume_mounts,
+            env=env if env else None,
+        )
+        # Create and configure a spec section
+        template = client.V1PodTemplateSpec(
+            metadata=client.V1ObjectMeta(labels={"wandb": "launch"}),
+            spec=client.V1PodSpec(
+                restart_policy="Never",
+                active_deadline_seconds=_DEFAULT_BUILD_TIMEOUT_SECS,
+                containers=[container],
+                volumes=volumes,
+            ),
+        )
+        # Create the specification of job
+        spec = client.V1JobSpec(template=template, backoff_limit=1)
+        job = client.V1Job(
+            api_version="batch/v1",
+            kind="Job",
+            metadata=client.V1ObjectMeta(
+                name=job_name, namespace="wandb", labels={"wandb": "launch"}
+            ),
+            spec=spec,
+        )
+        return job

wandb/sdk/launch/builder/noop.py

@@ -1,32 +1,51 @@
-from typing import Any, Dict, Optional
+"""NoOp builder implementation."""
+from typing import Any, Dict
 
-from wandb.errors import LaunchError
 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.utils import LaunchError
 
 from .._project_spec import EntryPoint, LaunchProject
 
 
 class NoOpBuilder(AbstractBuilder):
+    """NoOp builder."""
 
     type = "noop"
 
-    def __init__(self, builder_config: Dict[str, Any]) -> None:
-        self.builder_config = builder_config
+    def __init__(
+        self,
+        builder_config: Dict[str, Any],
+        environment: AbstractEnvironment,
+        registry: AbstractRegistry,
+    ) -> None:
+        """Initialize a NoOpBuilder."""
+        pass
+
+    @classmethod
+    def from_config(
+        cls,
+        config: dict,
+        environment: AbstractEnvironment,
+        registry: AbstractRegistry,
+        verify: bool = True,
+    ) -> "AbstractBuilder":
+        """Create a noop builder from a config."""
+        return cls(config, environment, registry)
+
+    def verify(self) -> None:
+        """Verify the builder."""
+        raise LaunchError("Attempted to verify noop builder.")
 
     def build_image(
         self,
         launch_project: LaunchProject,
-        registry: Optional[str],
         entrypoint: EntryPoint,
     ) -> str:
-        """Build the image for the given project.
-
-        Args:
-            launch_project: The project to build.
-            build_ctx_path: The path to the build context.
+        """Build the image.
 
-        Returns:
-            The image name.
+        For this we raise a launch error since it can't build.
         """
         raise LaunchError(
             "Attempted build with noop builder. Specify a builder in your launch config at ~/.config/wandb/launch-config.yaml"

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,9 +24,12 @@ 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
+    """Install pip dependencies.
 
     Arguments:
         deps {List[str]} -- List of dependencies to install
@@ -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"""
+    """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/abstract.py

@@ -0,0 +1,28 @@
+"""Abstract base class for environments."""
+from abc import ABC, abstractmethod
+
+
+class AbstractEnvironment(ABC):
+    """Abstract base class for environments."""
+
+    region: str
+
+    @abstractmethod
+    def verify(self) -> None:
+        """Verify that the environment is configured correctly."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def upload_file(self, source: str, destination: str) -> None:
+        """Upload a file from the local filesystem to storage in the environment."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def upload_dir(self, source: str, destination: str) -> None:
+        """Upload the contents of a directory from the local filesystem to the environment."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def verify_storage_uri(self, uri: str) -> None:
+        """Verify that the storage URI is configured correctly."""
+        raise NotImplementedError