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

wandb/sdk/launch/agent/run_queue_item_file_saver.py

@@ -0,0 +1,45 @@
+"""Implementation of the run queue item file saver class."""
+
+import os
+import sys
+from typing import List, Optional, Union
+
+import wandb
+from wandb.sdk.lib import RunDisabled
+from wandb.sdk.wandb_run import Run
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+FileSubtypes = Literal["warning", "error"]
+
+
+class RunQueueItemFileSaver:
+    def __init__(
+        self, agent_run: Optional[Union[Run, RunDisabled]], run_queue_item_id: str
+    ):
+        self.run_queue_item_id = run_queue_item_id
+        self.run = agent_run
+
+    def save_contents(
+        self, contents: str, fname: str, file_sub_type: FileSubtypes
+    ) -> Optional[List[str]]:
+        if not isinstance(self.run, Run):
+            wandb.termwarn("Not saving file contents because agent has no run")
+            return None
+        root_dir = self.run._settings.files_dir
+        saved_run_path = os.path.join(self.run_queue_item_id, file_sub_type, fname)
+        local_path = os.path.join(root_dir, saved_run_path)
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        with open(local_path, "w") as f:
+            f.write(contents)
+        res = self.run.save(local_path, base_path=root_dir, policy="now")
+        if isinstance(res, list):
+            return [saved_run_path]
+        else:
+            wandb.termwarn(
+                f"Failed to save files for run queue item: {self.run_queue_item_id}"
+            )
+            return None

wandb/sdk/launch/builder/abstract.py

@@ -1,12 +1,15 @@
 """Abstract plugin class defining the interface needed to build container images for W&B Launch."""
 from abc import ABC, abstractmethod
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
 from .._project_spec import EntryPoint, LaunchProject
 
+if TYPE_CHECKING:
+    from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
+
 
 class AbstractBuilder(ABC):
     """Abstract plugin class defining the interface needed to build container images for W&B Launch."""
@@ -63,6 +66,7 @@ class AbstractBuilder(ABC):
         self,
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
+        job_tracker: Optional["JobAndRunStatusTracker"] = None,
     ) -> str:
         """Build the image for the given project.
 

wandb/sdk/launch/builder/build.py

@@ -28,13 +28,8 @@ from .._project_spec import (
     LaunchProject,
     fetch_and_validate_project,
 )
-from ..utils import (
-    LAUNCH_CONFIG_FILE,
-    LOG_PREFIX,
-    ExecutionError,
-    LaunchError,
-    resolve_build_and_registry_config,
-)
+from ..errors import ExecutionError, LaunchError
+from ..utils import LAUNCH_CONFIG_FILE, LOG_PREFIX, resolve_build_and_registry_config
 from .abstract import AbstractBuilder
 
 _logger = logging.getLogger(__name__)
@@ -43,8 +38,6 @@ _logger = logging.getLogger(__name__)
 _GENERATED_DOCKERFILE_NAME = "Dockerfile.wandb-autogenerated"
 DEFAULT_ENTRYPOINT = "_wandb_default_entrypoint"
 
-DEFAULT_CUDA_VERSION = "10.0"
-
 
 def validate_docker_installation() -> None:
     """Verify if Docker is installed on host machine."""
@@ -108,8 +101,12 @@ FROM {py_base_image} as base
 """
 
 # this goes into base_setup in TEMPLATE
-CUDA_SETUP_TEMPLATE = """
-FROM {cuda_base_image} as base
+ACCELERATOR_SETUP_TEMPLATE = """
+FROM {accelerator_base_image} as base
+
+# make non-interactive so build doesn't block on questions
+ENV DEBIAN_FRONTEND=noninteractive
+
 # TODO: once NVIDIA their linux repository keys for all docker images
 RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/$(cat /etc/os-release | grep ^ID= |  cut -d "=" -f2 )$(cat /etc/os-release | grep ^VERSION_ID= |  cut -d "=" -f2 | sed -e 's/[\".]//g' )/$(uname -i)/3bf863cc.pub
 RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/$(cat /etc/os-release | grep ^ID= |  cut -d "=" -f2 )$(cat /etc/os-release | grep ^VERSION_ID= |  cut -d "=" -f2 | sed -e 's/[\".]//g' )/$(uname -i)/7fa2af80.pub
@@ -189,12 +186,14 @@ def get_base_setup(
 ) -> str:
     """Fill in the Dockerfile templates for stage 2 of build.
 
-    CPU version is built on python, GPU version is built on nvidia:cuda.
+    CPU version is built on python, Accelerator version is built on user provided.
     """
     python_base_image = f"python:{py_version}-buster"
-    if launch_project.cuda_base_image:
-        _logger.info(f"Using cuda base image: {launch_project.cuda_base_image}")
-        # cuda image doesn't come with python tooling
+    if launch_project.accelerator_base_image:
+        _logger.info(
+            f"Using accelerator base image: {launch_project.accelerator_base_image}"
+        )
+        # accelerator base images doesn't come with python tooling
         if py_major == "2":
             python_packages = [
                 f"python{py_version}",
@@ -209,8 +208,8 @@ def get_base_setup(
                 "python3-pip",
                 "python3-setuptools",
             ]
-        base_setup = CUDA_SETUP_TEMPLATE.format(
-            cuda_base_image=launch_project.cuda_base_image,
+        base_setup = ACCELERATOR_SETUP_TEMPLATE.format(
+            accelerator_base_image=launch_project.accelerator_base_image,
             python_packages=" \\\n".join(python_packages),
             py_version=py_version,
         )
@@ -246,6 +245,10 @@ def get_env_vars_dict(launch_project: LaunchProject, api: Api) -> Dict[str, str]
         env_vars["WANDB_NAME"] = launch_project.name
     if "author" in launch_project.launch_spec and not override_api_key:
         env_vars["WANDB_USERNAME"] = launch_project.launch_spec["author"]
+    if launch_project.sweep_id:
+        env_vars["WANDB_SWEEP_ID"] = launch_project.sweep_id
+    if launch_project.launch_spec.get("_resume_count"):
+        env_vars["WANDB_RESUME"] = "must"
 
     # TODO: handle env vars > 32760 characters
     env_vars["WANDB_CONFIG"] = json.dumps(launch_project.override_config)
@@ -288,7 +291,7 @@ def get_requirements_section(launch_project: LaunchProject, builder_type: str) -
         ):
             requirements_files += ["src/requirements.txt"]
             pip_install_line = "pip install -r requirements.txt"
-        if launch_project.project_dir is not None and os.path.exists(
+        elif launch_project.project_dir is not None and os.path.exists(
             os.path.join(launch_project.project_dir, "requirements.frozen.txt")
         ):
             # if we have frozen requirements stored, copy those over and have them take precedence

wandb/sdk/launch/builder/docker_builder.py

@@ -1,10 +1,11 @@
 """Implementation of the docker builder."""
 import logging
 import os
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import wandb
 import wandb.docker as docker
+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
@@ -15,11 +16,10 @@ from .._project_spec import (
     create_metadata_file,
     get_entry_point_command,
 )
+from ..errors import LaunchDockerError, LaunchError
 from ..registry.local_registry import LocalRegistry
 from ..utils import (
     LOG_PREFIX,
-    LaunchDockerError,
-    LaunchError,
     sanitize_wandb_api_key,
     warn_failed_packages_from_build_logs,
 )
@@ -112,6 +112,7 @@ class DockerBuilder(AbstractBuilder):
         self,
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         """Build the image for the given project.
 
@@ -160,9 +161,14 @@ class DockerBuilder(AbstractBuilder):
                 context_path=build_ctx_path,
                 platform=self.config.get("platform"),
             )
-            warn_failed_packages_from_build_logs(output, image_uri)
+
+            warn_failed_packages_from_build_logs(
+                output, image_uri, launch_project.api, job_tracker
+            )
 
         except docker.DockerError as e:
+            if job_tracker:
+                job_tracker.set_err_stage("build")
             raise LaunchDockerError(f"Error communicating with docker client: {e}")
 
         try:

wandb/sdk/launch/builder/kaniko_builder.py

@@ -1,15 +1,19 @@
 import base64
 import json
 import logging
+import os
 import tarfile
 import tempfile
 import time
 from typing import Optional
 
 import wandb
+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.environment.azure_environment import AzureEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
+from wandb.sdk.launch.registry.azure_container_registry import AzureContainerRegistry
 from wandb.sdk.launch.registry.elastic_container_registry import (
     ElasticContainerRegistry,
 )
@@ -22,9 +26,9 @@ from .._project_spec import (
     create_metadata_file,
     get_entry_point_command,
 )
+from ..errors import LaunchError
 from ..utils import (
     LOG_PREFIX,
-    LaunchError,
     get_kube_context_and_api_client,
     sanitize_wandb_api_key,
     warn_failed_packages_from_build_logs,
@@ -47,13 +51,21 @@ _logger = logging.getLogger(__name__)
 
 _DEFAULT_BUILD_TIMEOUT_SECS = 1800  # 30 minute build timeout
 
+SERVICE_ACCOUNT_NAME = os.environ.get("WANDB_LAUNCH_SERVICE_ACCOUNT_NAME", "default")
+
+if os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/namespace"):
+    with open("/var/run/secrets/kubernetes.io/serviceaccount/namespace") as f:
+        NAMESPACE = f.read().strip()
+else:
+    NAMESPACE = "wandb"
+
 
 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")
+        job = batch_client.read_namespaced_job_status(job_name, NAMESPACE)
         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:
@@ -75,6 +87,7 @@ class KanikoBuilder(AbstractBuilder):
     build_context_store: str
     secret_name: Optional[str]
     secret_key: Optional[str]
+    image: str
 
     def __init__(
         self,
@@ -84,6 +97,7 @@ class KanikoBuilder(AbstractBuilder):
         build_context_store: str = "",
         secret_name: str = "",
         secret_key: str = "",
+        image: str = "gcr.io/kaniko-project/executor:v1.11.0",
         verify: bool = True,
     ):
         """Initialize a KanikoBuilder.
@@ -110,6 +124,7 @@ class KanikoBuilder(AbstractBuilder):
         self.build_context_store = build_context_store.rstrip("/")
         self.secret_name = secret_name
         self.secret_key = secret_key
+        self.image = image
         if verify:
             self.verify()
 
@@ -148,6 +163,7 @@ class KanikoBuilder(AbstractBuilder):
         build_job_name = config.get("build-job-name", "wandb-launch-container-build")
         secret_name = config.get("secret-name", "")
         secret_key = config.get("secret-key", "")
+        image = config.get("kaniko-image", "gcr.io/kaniko-project/executor:v1.11.0")
         return cls(
             environment,
             registry,
@@ -155,6 +171,7 @@ class KanikoBuilder(AbstractBuilder):
             build_job_name=build_job_name,
             secret_name=secret_name,
             secret_key=secret_key,
+            image=image,
             verify=verify,
         )
 
@@ -184,7 +201,7 @@ class KanikoBuilder(AbstractBuilder):
             kind="ConfigMap",
             metadata=client.V1ObjectMeta(
                 name=f"docker-config-{job_name}",
-                namespace="wandb",
+                namespace=NAMESPACE,
             ),
             data={
                 "config.json": json.dumps(
@@ -193,13 +210,13 @@ class KanikoBuilder(AbstractBuilder):
             },
             immutable=True,
         )
-        corev1_client.create_namespaced_config_map("wandb", ecr_config_map)
+        corev1_client.create_namespaced_config_map(NAMESPACE, 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")
+            client.delete_namespaced_config_map(f"docker-config-{job_name}", NAMESPACE)
 
     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
@@ -217,6 +234,7 @@ class KanikoBuilder(AbstractBuilder):
         self,
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         # TODO: this should probably throw an error if the registry is a local registry
         if not self.registry:
@@ -252,35 +270,52 @@ class KanikoBuilder(AbstractBuilder):
         _, api_client = get_kube_context_and_api_client(
             kubernetes, launch_project.resource_args
         )
+        # TODO: use same client as kuberentes_runner.py
+        batch_v1 = client.BatchV1Api(api_client)
+        core_v1 = client.CoreV1Api(api_client)
+
         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,
+            build_job_name, repo_uri, image_uri, build_context, core_v1
         )
         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:
+            if isinstance(self.registry, AzureContainerRegistry):
+                dockerfile_config_map = client.V1ConfigMap(
+                    metadata=client.V1ObjectMeta(
+                        name=f"docker-config-{build_job_name}"
+                    ),
+                    data={
+                        "config.json": json.dumps(
+                            {
+                                "credHelpers": {
+                                    f"{self.registry.registry_name}.azurecr.io": "acr-env"
+                                }
+                            }
+                        )
+                    },
+                )
+                core_v1.create_namespaced_config_map("wandb", dockerfile_config_map)
             # 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)
+            batch_v1.create_namespaced_job(NAMESPACE, 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
             ):
+                if job_tracker:
+                    job_tracker.set_err_stage("build")
                 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)
+                logs = batch_v1.read_namespaced_job_log(build_job_name, NAMESPACE)
+                warn_failed_packages_from_build_logs(
+                    logs, image_uri, launch_project.api, job_tracker
+                )
             except Exception as e:
                 wandb.termwarn(
                     f"{LOG_PREFIX}Failed to get logs for kaniko job {build_job_name}: {e}"
@@ -295,9 +330,13 @@ 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")
+                if isinstance(self.registry, AzureContainerRegistry):
+                    core_v1.delete_namespaced_config_map(
+                        f"docker-config-{build_job_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")
+                batch_v1.delete_namespaced_job(build_job_name, NAMESPACE)
             except Exception as e:
                 raise LaunchError(f"Exception during Kubernetes resource clean up {e}")
 
@@ -309,6 +348,7 @@ class KanikoBuilder(AbstractBuilder):
         repository: str,
         image_tag: str,
         build_context_path: str,
+        core_client: client.CoreV1Api,
     ) -> "client.V1Job":
         env = []
         volume_mounts = []
@@ -325,6 +365,33 @@ class KanikoBuilder(AbstractBuilder):
                     value=self.registry.environment.region,
                 )
             ]
+        # TODO: Refactor all of this environment/registry
+        # specific stuff into methods of those classes.
+        if isinstance(self.environment, AzureEnvironment):
+            # Use the core api to check if the secret exists
+            try:
+                core_client.read_namespaced_secret(
+                    "azure-storage-access-key",
+                    "wandb",
+                )
+            except Exception as e:
+                raise LaunchError(
+                    "Secret azure-storage-access-key does not exist in "
+                    "namespace wandb. Please create it with the key password "
+                    "set to your azure storage access key."
+                ) from e
+            env += [
+                client.V1EnvVar(
+                    name="AZURE_STORAGE_ACCESS_KEY",
+                    value_from=client.V1EnvVarSource(
+                        secret_key_ref=client.V1SecretKeySelector(
+                            name="azure-storage-access-key",
+                            key="password",
+                        )
+                    ),
+                )
+            ]
+
         if self.secret_name and self.secret_key:
             volumes += [
                 client.V1Volume(
@@ -379,31 +446,54 @@ class KanikoBuilder(AbstractBuilder):
                     ),
                 )
             ]
-
+        if isinstance(self.registry, AzureContainerRegistry):
+            # ADd the docker config map
+            volume_mounts += [
+                client.V1VolumeMount(
+                    name="docker-config", mount_path="/kaniko/.docker/"
+                ),
+            ]
+            volumes += [
+                client.V1Volume(
+                    name="docker-config",
+                    config_map=client.V1ConfigMapVolumeSource(
+                        name=f"docker-config-{job_name}",
+                    ),
+                ),
+            ]
+        # Kaniko doesn't want https:// at the begining of the image tag.
+        destination = image_tag
+        if destination.startswith("https://"):
+            destination = destination.replace("https://", "")
         args = [
             f"--context={build_context_path}",
             "--dockerfile=Dockerfile.wandb-autogenerated",
-            f"--destination={image_tag}",
+            f"--destination={destination}",
             "--cache=true",
-            f"--cache-repo={repository}",
+            f"--cache-repo={repository.replace('https://', '')}",
             "--snapshotMode=redo",
             "--compressed-caching=false",
         ]
         container = client.V1Container(
             name="wandb-container-build",
-            image="gcr.io/kaniko-project/executor:v1.8.0",
+            image=self.image,
             args=args,
             volume_mounts=volume_mounts,
             env=env if env else None,
         )
         # Create and configure a spec section
+        labels = {"wandb": "launch"}
+        # This annotation is required to enable azure workload identity.
+        if isinstance(self.registry, AzureContainerRegistry):
+            labels["azure.workload.identity/use"] = "true"
         template = client.V1PodTemplateSpec(
-            metadata=client.V1ObjectMeta(labels={"wandb": "launch"}),
+            metadata=client.V1ObjectMeta(labels=labels),
             spec=client.V1PodSpec(
                 restart_policy="Never",
                 active_deadline_seconds=_DEFAULT_BUILD_TIMEOUT_SECS,
                 containers=[container],
                 volumes=volumes,
+                service_account_name=SERVICE_ACCOUNT_NAME,
             ),
         )
         # Create the specification of job
@@ -412,7 +502,7 @@ class KanikoBuilder(AbstractBuilder):
             api_version="batch/v1",
             kind="Job",
             metadata=client.V1ObjectMeta(
-                name=job_name, namespace="wandb", labels={"wandb": "launch"}
+                name=job_name, namespace=NAMESPACE, labels={"wandb": "launch"}
             ),
             spec=spec,
         )

wandb/sdk/launch/builder/noop.py

@@ -1,12 +1,13 @@
 """NoOp builder implementation."""
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
+from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
-from wandb.sdk.launch.utils import LaunchError
 
 from .._project_spec import EntryPoint, LaunchProject
+from ..agent.job_status_tracker import JobAndRunStatusTracker
 
 
 class NoOpBuilder(AbstractBuilder):
@@ -21,7 +22,8 @@ class NoOpBuilder(AbstractBuilder):
         registry: AbstractRegistry,
     ) -> None:
         """Initialize a NoOpBuilder."""
-        pass
+        self.environment = environment
+        self.registry = registry
 
     @classmethod
     def from_config(
@@ -42,6 +44,7 @@ class NoOpBuilder(AbstractBuilder):
         self,
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         """Build the image.
 

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

@@ -1,5 +1,4 @@
 import json
-import multiprocessing
 import os
 import re
 import subprocess
@@ -8,7 +7,6 @@ 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 = []
 # If the builder doesn't support buildx no need to use the cache
@@ -52,8 +50,16 @@ def install_deps(
         if failed is None:
             failed = set()
         num_failed = len(failed)
+        current_pkg = None
         for line in e.output.decode("utf8").splitlines():
-            if line.startswith("ERROR:"):
+            # Since the name of the package might not be on the same line as
+            # the error msg, keep track of the currently installing package
+            current_pkg = get_current_package(line, clean_deps, current_pkg)
+
+            if "error: subprocess-exited-with-error" in line:
+                if current_pkg is not None:
+                    failed.add(current_pkg)
+            elif line.startswith("ERROR:"):
                 clean_dep = find_package_in_error_string(clean_deps, line)
                 if clean_dep is not None:
                     if clean_dep in deps:
@@ -84,7 +90,6 @@ def main() -> None:
         with open("requirements.frozen.txt") as f:
             print("Installing frozen dependencies...")
             reqs = []
-            failed: Set[str] = set()
             for req in f:
                 if (
                     len(ONLY_INCLUDE) == 0
@@ -109,15 +114,7 @@ def main() -> None:
                         reqs.append(req.strip().replace(" ", ""))
                 else:
                     print(f"Ignoring requirement: {req} from frozen requirements")
-                if len(reqs) >= CORES:
-                    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, opts=OPTS)
-                if deps_failed is not None:
-                    failed = failed.union(deps_failed)
+            failed = install_deps(reqs, opts=OPTS) or set()
             with open("_wandb_bootstrap_errors.json", "w") as f:
                 f.write(json.dumps({"pip": list(failed)}))
             if len(failed) > 0:
@@ -130,6 +127,41 @@ def main() -> None:
         print("No frozen requirements found")
 
 
+def add_version_to_package_name(deps: List[str], package: str) -> Optional[str]:
+    """Add the associated version to a package name.
+
+    For example: `my-package` -> `my-package==1.0.0`
+    """
+    for dep in deps:
+        if dep.split("==")[0] == package:
+            return dep
+    return None
+
+
+def get_current_package(
+    line: str, deps: List[str], current_pkg: Optional[str]
+) -> Optional[str]:
+    """Tries to pull a package name from the line.
+
+    Used to keep track of what the currently-installing package is,
+    in case an error message isn't on the same line as the package
+    """
+    # "Collecting my-package==1.0.0"
+    if line.startswith("Collecting"):
+        return line.split(" ")[1]
+    # "Building wheel for my-package (pyproject.toml): finished with status 'error'"
+    elif line.strip().startswith("Building wheel") and line.strip().endswith(
+        "finished with status 'error'"
+    ):
+        return add_version_to_package_name(deps, line.strip().split(" ")[3])
+    # "Running setup.py install for my-package: finished with status 'error'"
+    elif line.strip().startswith("Running setup.py install") and line.strip().endswith(
+        "finished with status 'error'"
+    ):
+        return add_version_to_package_name(deps, line.strip().split(" ")[4][:-1])
+    return current_pkg
+
+
 # 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
@@ -143,7 +175,7 @@ def find_package_in_error_string(deps: List[str], line: str) -> Optional[str]:
     # 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:
+        if word.strip(",") in deps:
             return word
     # if we can't find the package, return None
     return None