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

wandb/sdk/launch/_project_spec.py

@@ -2,7 +2,6 @@
 
 Arguments can come from a launch spec or call to wandb launch.
 """
-import binascii
 import enum
 import json
 import logging
@@ -60,7 +59,6 @@ class LaunchProject:
         overrides: Dict[str, Any],
         resource: str,
         resource_args: Dict[str, Any],
-        cuda: Optional[bool],
         run_id: Optional[str],
     ):
         if uri is not None and utils.is_bare_wandb_uri(uri):
@@ -76,10 +74,16 @@ class LaunchProject:
         self.target_entity = target_entity
         self.target_project = target_project.lower()
         self.name = name  # TODO: replace with run_id
+        # the builder key can be passed in through the resource args
+        # but these resource_args are then passed to the appropriate
+        # runner, so we need to pop the builder key out
+        resource_args_build = resource_args.get(resource, {}).pop("builder", {})
         self.resource = resource
         self.resource_args = resource_args
         self.python_version: Optional[str] = launch_spec.get("python_version")
-        self.cuda_version: Optional[str] = launch_spec.get("cuda_version")
+        self.cuda_base_image: Optional[str] = resource_args_build.get("cuda", {}).get(
+            "base_image"
+        )
         self._base_image: Optional[str] = launch_spec.get("base_image")
         self.docker_image: Optional[str] = docker_config.get(
             "docker_image"
@@ -96,11 +100,8 @@ class LaunchProject:
         self.override_artifacts: Dict[str, Any] = overrides.get("artifacts", {})
         self.override_entrypoint: Optional[EntryPoint] = None
         self.deps_type: Optional[str] = None
-        self.cuda = cuda
         self._runtime: Optional[str] = None
         self.run_id = run_id or generate_id()
-        self._image_tag: str = self._initialize_image_job_tag() or self.run_id
-        wandb.termlog(f"{LOG_PREFIX}Launch project using image tag {self._image_tag}")
         self._entry_points: Dict[
             str, EntryPoint
         ] = {}  # todo: keep multiple entrypoint support?
@@ -140,8 +141,6 @@ class LaunchProject:
                 )
             self.source = LaunchSource.LOCAL
             self.project_dir = self.uri
-        if launch_spec.get("resource_args"):
-            self.resource_args = launch_spec["resource_args"]
 
         self.aux_dir = tempfile.mkdtemp()
         self.clear_parameter_run_config_collisions()
@@ -175,24 +174,15 @@ class LaunchProject:
             assert self.job is not None
             return wandb.util.make_docker_image_name_safe(self.job.split(":")[0])
 
-    def _initialize_image_job_tag(self) -> Optional[str]:
-        if self.job is not None:
-            job_name, alias = self.job.split(":")
-            # Alias is used to differentiate images between jobs of the same sequence
-            _image_tag = f"{alias}-{job_name}"
-            _logger.debug(f"{LOG_PREFIX}Setting image tag {_image_tag}")
-            return wandb.util.make_docker_image_name_safe(_image_tag)
-        return None
-
-    @property
-    def image_uri(self) -> str:
-        if self.docker_image:
-            return self.docker_image
-        return f"{self.image_name}:{self.image_tag}"
-
-    @property
-    def image_tag(self) -> str:
-        return self._image_tag[:IMAGE_TAG_MAX_LENGTH]
+    def build_required(self) -> bool:
+        """Checks the source to see if a build is required."""
+        # since the image tag for images built from jobs
+        # is based on the job version index, which is immutable
+        # we don't need to build the image for a job if that tag
+        # already exists
+        if self.source != LaunchSource.JOB:
+            return True
+        return False
 
     @property
     def docker_image(self) -> Optional[str]:
@@ -243,10 +233,37 @@ class LaunchProject:
         try:
             job = public_api.job(self.job, path=job_dir)
         except CommError:
-            raise LaunchError(f"Job {self.job} not found")
+            raise LaunchError(
+                f"Job {self.job} not found. Jobs have the format: <entity>/<project>/<name>:<alias>"
+            )
         job.configure_launch_project(self)
         self._job_artifact = job._job_artifact
 
+    def get_image_source_string(self) -> str:
+        """Returns a unique string identifying the source of an image."""
+        if self.source == LaunchSource.LOCAL:
+            # TODO: more correct to get a hash of local uri contents
+            assert isinstance(self.uri, str)
+            return self.uri
+        elif self.source == LaunchSource.JOB:
+            assert self._job_artifact is not None
+            return f"{self._job_artifact.name}:v{self._job_artifact.version}"
+        elif self.source == LaunchSource.GIT:
+            assert isinstance(self.uri, str)
+            ret = self.uri
+            if self.git_version:
+                ret += self.git_version
+            return ret
+        elif self.source == LaunchSource.WANDB:
+            assert isinstance(self.uri, str)
+            return self.uri
+        elif self.source == LaunchSource.DOCKER:
+            assert isinstance(self.docker_image, str)
+            _logger.debug("")
+            return self.docker_image
+        else:
+            raise LaunchError("Unknown source type when determing image source string")
+
     def _fetch_project_local(self, internal_api: Api) -> None:
         """Fetch a project (either wandb run or git repo) into a local directory, returning the path to the local project directory."""
         # these asserts are all guaranteed to pass, but are required by mypy
@@ -263,24 +280,6 @@ class LaunchProject:
             )
             program_name = run_info.get("codePath") or run_info["program"]
 
-            if run_info.get("cudaVersion"):
-                original_cuda_version = ".".join(run_info["cudaVersion"].split(".")[:2])
-
-                if self.cuda is None:
-                    # only set cuda on by default if cuda is None (unspecified), not False (user specifically requested cpu image)
-                    wandb.termlog(
-                        f"{LOG_PREFIX}Original wandb run {source_run_name} was run with cuda version {original_cuda_version}. Enabling cuda builds by default; to build on a CPU-only image, run again with --cuda=False"
-                    )
-                    self.cuda_version = original_cuda_version
-                    self.cuda = True
-                if (
-                    self.cuda
-                    and self.cuda_version
-                    and self.cuda_version != original_cuda_version
-                ):
-                    wandb.termlog(
-                        f"{LOG_PREFIX}Specified cuda version {self.cuda_version} differs from original cuda version {original_cuda_version}. Running with specified version {self.cuda_version}"
-                    )
             self.python_version = run_info.get("python", "3")
             downloaded_code_artifact = utils.check_and_download_code_artifacts(
                 source_entity,
@@ -289,11 +288,7 @@ class LaunchProject:
                 internal_api,
                 self.project_dir,
             )
-            if downloaded_code_artifact:
-                self._image_tag = binascii.hexlify(
-                    downloaded_code_artifact.digest.encode()
-                ).decode()
-            else:
+            if not downloaded_code_artifact:
                 if not run_info["git"]:
                     raise LaunchError(
                         "Reproducing a run requires either an associated git repo or a code artifact logged with `run.log_code()`"
@@ -308,12 +303,8 @@ class LaunchProject:
                 patch = utils.fetch_project_diff(
                     source_entity, source_project, source_run_name, internal_api
                 )
-                tag_string = run_info["git"]["remote"] + run_info["git"]["commit"]
                 if patch:
                     utils.apply_patch(patch, self.project_dir)
-                    tag_string += patch
-
-                self._image_tag = binascii.hexlify(tag_string.encode()).decode()
 
                 # For cases where the entry point wasn't checked into git
                 if not os.path.exists(os.path.join(self.project_dir, program_name)):
@@ -450,7 +441,6 @@ def create_project_from_spec(launch_spec: Dict[str, Any], api: Api) -> LaunchPro
         launch_spec.get("overrides", {}),
         launch_spec.get("resource", None),
         launch_spec.get("resource_args", {}),
-        launch_spec.get("cuda", None),
         launch_spec.get("run_id", None),
     )
 

wandb/sdk/launch/agent/agent.py

@@ -13,6 +13,8 @@ from typing import Any, Dict, List, Optional, Union
 import wandb
 import wandb.util as util
 from wandb.apis.internal import Api
+from wandb.errors import CommError
+from wandb.sdk.launch._project_spec import LaunchProject
 from wandb.sdk.launch.runner.local_container import LocalSubmittedRun
 from wandb.sdk.launch.sweeps import SCHEDULER_URI
 from wandb.sdk.lib import runid
@@ -21,7 +23,13 @@ from .. import loader
 from .._project_spec import create_project_from_spec, fetch_and_validate_project
 from ..builder.build import construct_builder_args
 from ..runner.abstract import AbstractRun
-from ..utils import LAUNCH_DEFAULT_PROJECT, LOG_PREFIX, PROJECT_SYNCHRONOUS, LaunchError
+from ..utils import (
+    LAUNCH_DEFAULT_PROJECT,
+    LOG_PREFIX,
+    PROJECT_SYNCHRONOUS,
+    LaunchDockerError,
+    LaunchError,
+)
 
 AGENT_POLLING_INTERVAL = 10
 ACTIVE_SWEEP_POLLING_INTERVAL = 1  # more frequent when we know we have jobs
@@ -37,6 +45,9 @@ _logger = logging.getLogger(__name__)
 
 @dataclass
 class JobAndRunStatus:
+    run_queue_item_id: str
+    run_id: Optional[str] = None
+    project: Optional[str] = None
     run: Optional[AbstractRun] = None
     failed_to_start: bool = False
     completed: bool = False
@@ -46,6 +57,10 @@ class JobAndRunStatus:
     def job_completed(self) -> bool:
         return self.completed or self.failed_to_start
 
+    def update_run_info(self, launch_project: LaunchProject) -> None:
+        self.run_id = launch_project.run_id
+        self.project = launch_project.target_project
+
 
 def _convert_access(access: str) -> str:
     """Convert access string to a value accepted by wandb."""
@@ -90,7 +105,20 @@ def _job_is_scheduler(run_spec: Dict[str, Any]) -> bool:
     if not run_spec:
         _logger.debug("Recieved runSpec in _job_is_scheduler that was empty")
 
-    return run_spec.get("uri") == SCHEDULER_URI
+    if run_spec.get("uri") != SCHEDULER_URI:
+        return False
+
+    if run_spec.get("resource") == "local-process":
+        # If a scheduler is a local-process (100%), also
+        #    confirm command is in format: [wandb scheduler <sweep>]
+        cmd = run_spec.get("overrides", {}).get("entry_point", [])
+        if len(cmd) < 3:
+            return False
+
+        if cmd[:2] != ["wandb", "scheduler"]:
+            return False
+
+    return True
 
 
 class LaunchAgent:
@@ -119,7 +147,6 @@ class LaunchAgent:
         self._max_schedulers = _max_from_config(config, "max_schedulers")
         self._pool = ThreadPool(
             processes=int(min(MAX_THREADS, self._max_jobs + self._max_schedulers)),
-            # initializer=init_pool_processes,
             initargs=(self._jobs, self._jobs_lock),
         )
         self.default_config: Dict[str, Any] = config
@@ -128,6 +155,10 @@ class LaunchAgent:
         self.gorilla_supports_agents = (
             self._api.launch_agent_introspection() is not None
         )
+        self._gorilla_supports_fail_run_queue_items = (
+            self._api.fail_run_queue_item_introspection()
+        )
+
         self._queues = config.get("queues", ["default"])
         create_response = self._api.create_launch_agent(
             self._entity,
@@ -137,6 +168,14 @@ class LaunchAgent:
         )
         self._id = create_response["launchAgentId"]
         self._name = ""  # hacky: want to display this to the user but we don't get it back from gql until polling starts. fix later
+        if self._api.entity_is_team(self._entity):
+            wandb.termwarn(
+                f"{LOG_PREFIX}Agent is running on team entity ({self._entity}). Members of this team will be able to run code on this device."
+            )
+
+    def fail_run_queue_item(self, run_queue_item_id: str) -> None:
+        if self._gorilla_supports_fail_run_queue_items:
+            self._api.fail_run_queue_item(run_queue_item_id)
 
     @property
     def thread_ids(self) -> List[int]:
@@ -214,9 +253,28 @@ class LaunchAgent:
 
     def finish_thread_id(self, thread_id: int) -> None:
         """Removes the job from our list for now."""
+        job_and_run_status = self._jobs[thread_id]
+        if not job_and_run_status.run_id or not job_and_run_status.project:
+            self.fail_run_queue_item(job_and_run_status.run_queue_item_id)
+        else:
+            run_info = None
+            # sweep runs exist but have no info before they are started
+            # so run_info returned will be None
+            # normal runs just throw a comm error
+            try:
+                run_info = self._api.get_run_info(
+                    self._entity, job_and_run_status.project, job_and_run_status.run_id
+                )
+
+            except CommError:
+                pass
+            if run_info is None:
+                self.fail_run_queue_item(job_and_run_status.run_queue_item_id)
+
         # TODO:  keep logs or something for the finished jobs
         with self._jobs_lock:
             del self._jobs[thread_id]
+
         # update status back to polling if no jobs are running
         if len(self.thread_ids) == 0:
             self.update_status(AGENT_POLLING)
@@ -295,16 +353,12 @@ class LaunchAgent:
 
                             try:
                                 self.run_job(job)
-                            except Exception:
+                            except Exception as e:
                                 wandb.termerror(
                                     f"{LOG_PREFIX}Error running job: {traceback.format_exc()}"
                                 )
-                                try:
-                                    self._api.ack_run_queue_item(job["runQueueItemId"])
-                                except Exception:
-                                    _logger.error(
-                                        f"{LOG_PREFIX}Error acking job when job errored: {traceback.format_exc()}"
-                                    )
+                                util.sentry_exc(e)
+                                self.fail_run_queue_item(job["runQueueItemId"])
 
                 for thread_id in self.thread_ids:
                     self._update_finished(thread_id)
@@ -340,11 +394,20 @@ class LaunchAgent:
         default_config: Dict[str, Any],
         api: Api,
     ) -> None:
+        thread_id = threading.current_thread().ident
+        assert thread_id is not None
         try:
-            self._thread_run_job(launch_spec, job, default_config, api)
-        except Exception:
+            self._thread_run_job(launch_spec, job, default_config, api, thread_id)
+        except LaunchDockerError as e:
+            wandb.termerror(
+                f"{LOG_PREFIX}agent {self._name} encountered an issue while starting Docker, see above output for details."
+            )
+            self.finish_thread_id(thread_id)
+            util.sentry_exc(e)
+        except Exception as e:
             wandb.termerror(f"{LOG_PREFIX}Error running job: {traceback.format_exc()}")
-            api.ack_run_queue_item(job["runQueueItemId"])
+            self.finish_thread_id(thread_id)
+            util.sentry_exc(e)
 
     def _thread_run_job(
         self,
@@ -352,13 +415,13 @@ class LaunchAgent:
         job: Dict[str, Any],
         default_config: Dict[str, Any],
         api: Api,
+        thread_id: int,
     ) -> None:
-        thread_id = threading.current_thread().ident
-        assert thread_id is not None
-        job_tracker = JobAndRunStatus()
+        job_tracker = JobAndRunStatus(job["runQueueItemId"])
         with self._jobs_lock:
             self._jobs[thread_id] = job_tracker
         project = create_project_from_spec(launch_spec, api)
+        job_tracker.update_run_info(project)
         _logger.info("Fetching and validating project...")
         project = fetch_and_validate_project(project, api)
         _logger.info("Fetching resource...")
@@ -366,8 +429,6 @@ class LaunchAgent:
         backend_config: Dict[str, Any] = {
             PROJECT_SYNCHRONOUS: False,  # agent always runs async
         }
-
-        backend_config["runQueueItemId"] = job["runQueueItemId"]
         _logger.info("Loading backend")
         override_build_config = launch_spec.get("builder")
 
@@ -382,6 +443,7 @@ class LaunchAgent:
         builder = loader.builder_from_config(build_config, environment, registry)
         backend = loader.runner_from_config(resource, api, backend_config, environment)
         _logger.info("Backend loaded...")
+        api.ack_run_queue_item(job["runQueueItemId"], project.run_id)
         run = backend.run(project, builder)
 
         if _job_is_scheduler(launch_spec):

wandb/sdk/launch/builder/build.py

@@ -1,3 +1,4 @@
+import hashlib
 import json
 import logging
 import os
@@ -21,7 +22,6 @@ from wandb.sdk.launch.loader import (
     registry_from_config,
 )
 
-from ...lib.git import GitRepo
 from .._project_spec import (
     EntryPoint,
     EntrypointDefaults,
@@ -132,6 +132,7 @@ PIP_TEMPLATE = """
 RUN python -m venv /env
 # make sure we install into the env
 ENV PATH="/env/bin:$PATH"
+
 COPY {requirements_files} ./
 {buildx_optional_prefix} {pip_install}
 """
@@ -192,8 +193,8 @@ def get_base_setup(
     CPU version is built on python, GPU version is built on nvidia:cuda.
     """
     python_base_image = f"python:{py_version}-buster"
-    if launch_project.cuda:
-        cuda_version = launch_project.cuda_version or DEFAULT_CUDA_VERSION
+    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 py_major == "2":
             python_packages = [
@@ -210,7 +211,7 @@ def get_base_setup(
                 "python3-setuptools",
             ]
         base_setup = CUDA_SETUP_TEMPLATE.format(
-            cuda_base_image=f"nvidia/cuda:{cuda_version}-runtime",
+            cuda_base_image=launch_project.cuda_base_image,
             python_packages=" \\\n".join(python_packages),
             py_version=py_version,
         )
@@ -390,57 +391,6 @@ def generate_dockerfile(
     return dockerfile_contents
 
 
-_inspected_images = {}
-
-
-def docker_image_exists(docker_image: str, should_raise: bool = False) -> bool:
-    """Check if a specific image is already available.
-
-    Optionally raises an exception if the image is not found.
-    """
-    _logger.info("Checking if base image exists...")
-    try:
-        data = docker.run(["docker", "image", "inspect", docker_image])
-        # always true, since return stderr defaults to false
-        assert isinstance(data, str)
-        parsed = json.loads(data)[0]
-        _inspected_images[docker_image] = parsed
-        return True
-    except (docker.DockerError, ValueError) as e:
-        if should_raise:
-            raise e
-        _logger.info("Base image not found. Generating new base image")
-        return False
-
-
-def docker_image_inspect(docker_image: str) -> Dict[str, Any]:
-    """Get the parsed json result of docker inspect image_name."""
-    if _inspected_images.get(docker_image) is None:
-        docker_image_exists(docker_image, True)
-    return _inspected_images.get(docker_image, {})
-
-
-def pull_docker_image(docker_image: str) -> None:
-    """Pull the requested docker image."""
-    if docker_image_exists(docker_image):
-        # don't pull images if they exist already, eg if they are local images
-        return
-    try:
-        docker.run(["docker", "pull", docker_image])
-    except docker.DockerError as e:
-        raise LaunchError(f"Docker server returned error: {e}")
-
-
-def construct_gcp_image_uri(
-    launch_project: LaunchProject,
-    gcp_repo: str,
-    gcp_project: str,
-    gcp_registry: str,
-) -> str:
-    base_uri = launch_project.image_uri
-    return "/".join([gcp_registry, gcp_project, gcp_repo, base_uri])
-
-
 def construct_gcp_registry_uri(
     gcp_repo: str, gcp_project: str, gcp_registry: str
 ) -> str:
@@ -474,24 +424,6 @@ def _parse_existing_requirements(launch_project: LaunchProject) -> str:
     return requirements_line
 
 
-def _get_docker_image_uri(name: Optional[str], work_dir: str, image_id: str) -> str:
-    """Create a Docker image URI for a project.
-
-    The resulting URI is based on the git hash of the specified working directory.
-    :param name: The URI of the Docker repository with which to tag the image. The
-                           repository URI is used as the prefix of the image URI.
-    :param work_dir: Path to the working directory in which to search for a git commit hash.
-    """
-    name = name.replace(" ", "-") if name else "wandb-launch"
-    # Optionally include first 7 digits of git SHA in tag name, if available.
-
-    git_commit = GitRepo(work_dir).last_commit
-    version_string = (
-        ":" + str(git_commit[:7]) + image_id if git_commit else ":" + image_id
-    )
-    return name + version_string
-
-
 def _create_docker_build_ctx(
     launch_project: LaunchProject,
     dockerfile_contents: str,
@@ -537,7 +469,7 @@ def construct_builder_args(
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     registry_config = None
     if launch_config is not None:
-        build_config = launch_config.get("build")
+        build_config = launch_config.get("builder")
         registry_config = launch_config.get("registry")
 
     default_launch_config = None
@@ -625,3 +557,13 @@ def build_image_from_project(
         raise LaunchError("Error building image uri")
     else:
         return image_uri
+
+
+def image_tag_from_dockerfile_and_source(
+    launch_project: LaunchProject, dockerfile_contents: str
+) -> str:
+    """Hashes the source and dockerfile contents into a unique tag."""
+    image_source_string = launch_project.get_image_source_string()
+    unique_id_string = image_source_string + dockerfile_contents
+    image_tag = hashlib.sha256(unique_id_string.encode("utf-8")).hexdigest()[:8]
+    return image_tag

wandb/sdk/launch/builder/docker_builder.py

@@ -16,10 +16,17 @@ from .._project_spec import (
     get_entry_point_command,
 )
 from ..registry.local_registry import LocalRegistry
-from ..utils import LOG_PREFIX, LaunchError, sanitize_wandb_api_key
+from ..utils import (
+    LOG_PREFIX,
+    LaunchDockerError,
+    LaunchError,
+    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,
     validate_docker_installation,
 )
 
@@ -110,15 +117,32 @@ class DockerBuilder(AbstractBuilder):
             launch_project (LaunchProject): The project to build.
             entrypoint (EntryPoint): The entrypoint to use.
         """
+        dockerfile_str = generate_dockerfile(
+            launch_project, entrypoint, launch_project.resource, "docker"
+        )
+
+        image_tag = image_tag_from_dockerfile_and_source(launch_project, dockerfile_str)
+
         repository = None if not self.registry else self.registry.get_repo_uri()
+        # if repo is set, use the repo name as the image name
         if repository:
-            image_uri = f"{repository}:{launch_project.image_tag}"
+            image_uri = f"{repository}:{image_tag}"
+        # otherwise, base the image name off of the source
+        # which the launch_project checks in image_name
         else:
-            image_uri = launch_project.image_uri
-        entry_cmd = get_entry_point_command(entrypoint, launch_project.override_args)
-        dockerfile_str = generate_dockerfile(
-            launch_project, entrypoint, launch_project.resource, "docker"
+            image_uri = f"{launch_project.image_name}:{image_tag}"
+
+        if not launch_project.build_required() and self.registry.check_image_exists(
+            image_uri
+        ):
+            return image_uri
+
+        _logger.info(
+            f"image {image_uri} does not already exist in repository, building."
         )
+
+        entry_cmd = get_entry_point_command(entrypoint, launch_project.override_args)
+
         create_metadata_file(
             launch_project,
             image_uri,
@@ -128,9 +152,13 @@ class DockerBuilder(AbstractBuilder):
         build_ctx_path = _create_docker_build_ctx(launch_project, dockerfile_str)
         dockerfile = os.path.join(build_ctx_path, _GENERATED_DOCKERFILE_NAME)
         try:
-            docker.build(tags=[image_uri], file=dockerfile, context_path=build_ctx_path)
+            output = docker.build(
+                tags=[image_uri], file=dockerfile, context_path=build_ctx_path
+            )
+            warn_failed_packages_from_build_logs(output, image_uri)
+
         except docker.DockerError as e:
-            raise LaunchError(f"Error communicating with docker client: {e}")
+            raise LaunchDockerError(f"Error communicating with docker client: {e}")
 
         try:
             os.remove(build_ctx_path)