wandb 0.16.3__py3-none-any.whl → 0.16.5__py3-none-any.whl

wandb/sdk/launch/_launch_add.py

@@ -6,7 +6,6 @@ import wandb
 import wandb.apis.public as public
 from wandb.apis.internal import Api
 from wandb.errors import CommError
-from wandb.sdk.launch._project_spec import create_project_from_spec
 from wandb.sdk.launch.builder.build import build_image_from_project
 from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.utils import (
@@ -16,6 +15,8 @@ from wandb.sdk.launch.utils import (
     validate_launch_spec_source,
 )
 
+from ._project_spec import LaunchProject
+
 
 def push_to_queue(
     api: Api,
@@ -106,34 +107,32 @@ def launch_add(
     """
     api = Api()
 
-    return asyncio.run(
-        _launch_add(
-            api,
-            uri,
-            job,
-            config,
-            template_variables,
-            project,
-            entity,
-            queue_name,
-            resource,
-            entry_point,
-            name,
-            version,
-            docker_image,
-            project_queue,
-            resource_args,
-            run_id=run_id,
-            build=build,
-            repository=repository,
-            sweep_id=sweep_id,
-            author=author,
-            priority=priority,
-        )
+    return _launch_add(
+        api,
+        uri,
+        job,
+        config,
+        template_variables,
+        project,
+        entity,
+        queue_name,
+        resource,
+        entry_point,
+        name,
+        version,
+        docker_image,
+        project_queue,
+        resource_args,
+        run_id=run_id,
+        build=build,
+        repository=repository,
+        sweep_id=sweep_id,
+        author=author,
+        priority=priority,
     )
 
 
-async def _launch_add(
+def _launch_add(
     api: Api,
     uri: Optional[str],
     job: Optional[str],
@@ -185,9 +184,9 @@ async def _launch_add(
             wandb.termwarn("Build doesn't support setting a job. Overwriting job.")
             launch_spec["job"] = None
 
-        launch_project = create_project_from_spec(launch_spec, api)
-        docker_image_uri = await build_image_from_project(
-            launch_project, api, config or {}
+        launch_project = LaunchProject.from_spec(launch_spec, api)
+        docker_image_uri = asyncio.run(
+            build_image_from_project(launch_project, api, config or {})
         )
         run = wandb.run or wandb.init(
             project=launch_spec["project"],

wandb/sdk/launch/_project_spec.py

@@ -3,7 +3,6 @@
 Arguments can come from a launch spec or call to wandb launch.
 """
 import enum
-import json
 import logging
 import os
 import tempfile
@@ -25,7 +24,6 @@ if TYPE_CHECKING:
 
 _logger = logging.getLogger(__name__)
 
-DEFAULT_LAUNCH_METADATA_PATH = "launch_metadata.json"
 
 # need to make user root for sagemaker, so users have access to /opt/ml directories
 # that let users create artifacts and access input data
@@ -46,7 +44,22 @@ class EntrypointDefaults(List[str]):
 
 
 class LaunchProject:
-    """A launch project specification."""
+    """A launch project specification.
+
+    The LaunchProject is initialized from a raw launch spec an internal API
+    object. The project encapsulates logic for taking a launch spec and converting
+    it into the executable code.
+
+    The LaunchProject needs to ultimately produce a full container spec for
+    execution in docker, k8s, sagemaker, or vertex. This container spec includes:
+    - container image uri
+    - environment variables for configuring wandb etc.
+    - entrypoint command and arguments
+    - additional arguments specific to the target resource (e.g. instance type, node selector)
+
+    This class is stateful and certain methods can only be called after
+    `LaunchProject.fetch_and_validate_project()` has been called.
+    """
 
     def __init__(
         self,
@@ -122,7 +135,7 @@ class LaunchProject:
         if override_entrypoint:
             _logger.info("Adding override entry point")
             self.override_entrypoint = EntryPoint(
-                name=self._get_entrypoint_file(override_entrypoint),
+                name=_get_entrypoint_file(override_entrypoint),
                 command=override_entrypoint,
             )
 
@@ -160,22 +173,42 @@ class LaunchProject:
             self.source = LaunchSource.LOCAL
             self.project_dir = self.uri
 
-        self.aux_dir = tempfile.mkdtemp()
+    def __repr__(self) -> str:
+        """String representation of LaunchProject."""
+        if self.source == LaunchSource.JOB:
+            return f"{self.job}"
+        return f"{self.uri}"
 
-    @property
-    def base_image(self) -> str:
-        """Returns {PROJECT}_base:{PYTHON_VERSION}."""
-        # TODO: this should likely be source_project when we have it...
+    @classmethod
+    def from_spec(cls, launch_spec: Dict[str, Any], api: Api) -> "LaunchProject":
+        """Constructs a LaunchProject instance using a launch spec.
 
-        # don't make up a separate base image name if user provides a docker image
-        if self.docker_image is not None:
-            return self.docker_image
+        Arguments:
+            launch_spec: Dictionary representation of launch spec
+            api: Instance of wandb.apis.internal Api
 
-        python_version = (self.python_version or "3").replace("+", "dev")
-        generated_name = "{}_base:{}".format(
-            self.target_project.replace(" ", "-"), python_version
+        Returns:
+            An initialized `LaunchProject` object
+        """
+        name: Optional[str] = None
+        if launch_spec.get("name"):
+            name = launch_spec["name"]
+        return LaunchProject(
+            launch_spec.get("uri"),
+            launch_spec.get("job"),
+            api,
+            launch_spec,
+            launch_spec["entity"],
+            launch_spec["project"],
+            name,
+            launch_spec.get("docker", {}),
+            launch_spec.get("git", {}),
+            launch_spec.get("overrides", {}),
+            launch_spec.get("resource", None),
+            launch_spec.get("resource_args", {}),
+            launch_spec.get("run_id", None),
+            launch_spec.get("sweep_id", {}),
         )
-        return self._base_image or generated_name
 
     @property
     def image_name(self) -> str:
@@ -215,15 +248,6 @@ class LaunchProject:
     def run_queue_item_id(self, value: str) -> None:
         self._run_queue_item_id = value
 
-    def _get_entrypoint_file(self, entrypoint: List[str]) -> Optional[str]:
-        if not entrypoint:
-            return None
-        if entrypoint[0].endswith(".py") or entrypoint[0].endswith(".sh"):
-            return entrypoint[0]
-        if len(entrypoint) < 2:
-            return None
-        return entrypoint[1]
-
     def fill_macros(self, image: str) -> Dict[str, Any]:
         """Substitute values for macros in resource arguments.
 
@@ -277,10 +301,25 @@ class LaunchProject:
 
     @property
     def docker_image(self) -> Optional[str]:
+        """Returns the Docker image associated with this LaunchProject.
+
+        This will only be set if an image_uri is being run outside a job.
+
+        Returns:
+            Optional[str]: The Docker image or None if not specified.
+        """
         return self._docker_image
 
     @docker_image.setter
     def docker_image(self, value: str) -> None:
+        """Sets the Docker image for the project.
+
+        Args:
+            value (str): The Docker image to set.
+
+        Returns:
+            None
+        """
         self._docker_image = value
         self._ensure_not_docker_image_and_local_process()
 
@@ -305,24 +344,39 @@ class LaunchProject:
         self._entry_point = new_entrypoint
         return new_entrypoint
 
-    def _ensure_not_docker_image_and_local_process(self) -> None:
-        if self.docker_image is not None and self.resource == "local-process":
-            raise LaunchError(
-                "Cannot specify docker image with local-process resource runner"
-            )
+    def fetch_and_validate_project(self) -> None:
+        """Fetches a project into a local directory, adds the config values to the directory, and validates the first entrypoint for the project.
 
-    def _fetch_job(self) -> None:
-        public_api = wandb.apis.public.Api()
-        job_dir = tempfile.mkdtemp()
-        try:
-            job = public_api.job(self.job, path=job_dir)
-        except CommError as e:
-            msg = e.message
-            raise LaunchError(
-                f"Error accessing job {self.job}: {msg} on {public_api.settings.get('base_url')}"
-            )
-        job.configure_launch_project(self)
-        self._job_artifact = job._job_artifact
+        Arguments:
+            launch_project: LaunchProject to fetch and validate.
+            api: Instance of wandb.apis.internal Api
+
+        Returns:
+            A validated `LaunchProject` object.
+
+        """
+        if self.source == LaunchSource.DOCKER:
+            return
+        if self.source == LaunchSource.LOCAL:
+            if not self._entry_point:
+                wandb.termlog(
+                    f"{LOG_PREFIX}Entry point for repo not specified, defaulting to `python main.py`"
+                )
+                self.set_entry_point(EntrypointDefaults.PYTHON)
+        elif self.source == LaunchSource.JOB:
+            self._fetch_job()
+        else:
+            self._fetch_project_local(internal_api=self.api)
+
+        assert self.project_dir is not None
+        # this prioritizes pip, and we don't support any cases where both are present conda projects when uploaded to
+        # wandb become pip projects via requirements.frozen.txt, wandb doesn't preserve conda envs
+        if os.path.exists(
+            os.path.join(self.project_dir, "requirements.txt")
+        ) or os.path.exists(os.path.join(self.project_dir, "requirements.frozen.txt")):
+            self.deps_type = "pip"
+        elif os.path.exists(os.path.join(self.project_dir, "environment.yml")):
+            self.deps_type = "conda"
 
     def get_image_source_string(self) -> str:
         """Returns a unique string identifying the source of an image."""
@@ -349,6 +403,35 @@ class LaunchProject:
         else:
             raise LaunchError("Unknown source type when determing image source string")
 
+    def _ensure_not_docker_image_and_local_process(self) -> None:
+        """Ensure that docker image is not specified with local-process resource runner.
+
+        Raises:
+            LaunchError: If docker image is specified with local-process resource runner.
+        """
+        if self.docker_image is not None and self.resource == "local-process":
+            raise LaunchError(
+                "Cannot specify docker image with local-process resource runner"
+            )
+
+    def _fetch_job(self) -> None:
+        """Fetches the job details from the public API and configures the launch project.
+
+        Raises:
+            LaunchError: If there is an error accessing the job.
+        """
+        public_api = wandb.apis.public.Api()
+        job_dir = tempfile.mkdtemp()
+        try:
+            job = public_api.job(self.job, path=job_dir)
+        except CommError as e:
+            msg = e.message
+            raise LaunchError(
+                f"Error accessing job {self.job}: {msg} on {public_api.settings.get('base_url')}"
+            )
+        job.configure_launch_project(self)
+        self._job_artifact = job._job_artifact
+
     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
@@ -453,6 +536,24 @@ class LaunchProject:
                 self.git_version = branch_name
 
 
+def _get_entrypoint_file(entrypoint: List[str]) -> Optional[str]:
+    """Get the entrypoint file from the given command.
+
+    Args:
+        entrypoint (List[str]): List of command and arguments.
+
+    Returns:
+        Optional[str]: The entrypoint file if found, otherwise None.
+    """
+    if not entrypoint:
+        return None
+    if entrypoint[0].endswith(".py") or entrypoint[0].endswith(".sh"):
+        return entrypoint[0]
+    if len(entrypoint) < 2:
+        return None
+    return entrypoint[1]
+
+
 class EntryPoint:
     """An entry point into a wandb launch specification."""
 
@@ -467,6 +568,11 @@ class EntryPoint:
             return ret + user_parameters
         return ret
 
+    def update_entrypoint_path(self, new_path: str) -> None:
+        """Updates the entrypoint path to a new path."""
+        if len(self.command) == 2 and self.command[0] in ["python", "bash"]:
+            self.command[1] = new_path
+
 
 def get_entry_point_command(
     entry_point: Optional["EntryPoint"], parameters: List[str]
@@ -483,97 +589,3 @@ def get_entry_point_command(
     if entry_point is None:
         return []
     return entry_point.compute_command(parameters)
-
-
-def create_project_from_spec(launch_spec: Dict[str, Any], api: Api) -> LaunchProject:
-    """Constructs a LaunchProject instance using a launch spec.
-
-    Arguments:
-    launch_spec: Dictionary representation of launch spec
-    api: Instance of wandb.apis.internal Api
-
-    Returns:
-        An initialized `LaunchProject` object
-    """
-    name: Optional[str] = None
-    if launch_spec.get("name"):
-        name = launch_spec["name"]
-    return LaunchProject(
-        launch_spec.get("uri"),
-        launch_spec.get("job"),
-        api,
-        launch_spec,
-        launch_spec["entity"],
-        launch_spec["project"],
-        name,
-        launch_spec.get("docker", {}),
-        launch_spec.get("git", {}),
-        launch_spec.get("overrides", {}),
-        launch_spec.get("resource", None),
-        launch_spec.get("resource_args", {}),
-        launch_spec.get("run_id", None),
-        launch_spec.get("sweep_id", {}),
-    )
-
-
-def fetch_and_validate_project(
-    launch_project: LaunchProject, api: Api
-) -> LaunchProject:
-    """Fetches a project into a local directory, adds the config values to the directory, and validates the first entrypoint for the project.
-
-    Arguments:
-    launch_project: LaunchProject to fetch and validate.
-    api: Instance of wandb.apis.internal Api
-
-    Returns:
-        A validated `LaunchProject` object.
-
-    """
-    if launch_project.source == LaunchSource.DOCKER:
-        return launch_project
-    if launch_project.source == LaunchSource.LOCAL:
-        if not launch_project._entry_point:
-            wandb.termlog(
-                f"{LOG_PREFIX}Entry point for repo not specified, defaulting to `python main.py`"
-            )
-            launch_project.set_entry_point(EntrypointDefaults.PYTHON)
-    elif launch_project.source == LaunchSource.JOB:
-        launch_project._fetch_job()
-    else:
-        launch_project._fetch_project_local(internal_api=api)
-
-    assert launch_project.project_dir is not None
-    # this prioritizes pip, and we don't support any cases where both are present conda projects when uploaded to
-    # wandb become pip projects via requirements.frozen.txt, wandb doesn't preserve conda envs
-    if os.path.exists(
-        os.path.join(launch_project.project_dir, "requirements.txt")
-    ) or os.path.exists(
-        os.path.join(launch_project.project_dir, "requirements.frozen.txt")
-    ):
-        launch_project.deps_type = "pip"
-    elif os.path.exists(os.path.join(launch_project.project_dir, "environment.yml")):
-        launch_project.deps_type = "conda"
-
-    return launch_project
-
-
-def create_metadata_file(
-    launch_project: LaunchProject,
-    image_uri: str,
-    sanitized_entrypoint_str: str,
-    sanitized_dockerfile_contents: str,
-) -> None:
-    assert launch_project.project_dir is not None
-    with open(
-        os.path.join(launch_project.project_dir, DEFAULT_LAUNCH_METADATA_PATH),
-        "w",
-    ) as f:
-        json.dump(
-            {
-                **launch_project.launch_spec,
-                "image_uri": image_uri,
-                "command": sanitized_entrypoint_str,
-                "dockerfile_contents": sanitized_dockerfile_contents,
-            },
-            f,
-        )

wandb/sdk/launch/agent/agent.py

@@ -20,11 +20,7 @@ from wandb.sdk.launch.sweeps.scheduler import Scheduler
 from wandb.sdk.lib import runid
 
 from .. import loader
-from .._project_spec import (
-    LaunchProject,
-    create_project_from_spec,
-    fetch_and_validate_project,
-)
+from .._project_spec import LaunchProject
 from ..builder.build import construct_agent_configs
 from ..errors import LaunchDockerError, LaunchError
 from ..utils import (
@@ -630,7 +626,7 @@ class LaunchAgent:
         thread_id: int,
         job_tracker: JobAndRunStatusTracker,
     ) -> None:
-        project = create_project_from_spec(launch_spec, api)
+        project = LaunchProject.from_spec(launch_spec, api)
         self._set_queue_and_rqi_in_project(project, job, job_tracker.queue)
         ack = event_loop_thread_exec(api.ack_run_queue_item)
         await ack(job["runQueueItemId"], project.run_id)
@@ -639,7 +635,7 @@ class LaunchAgent:
 
         job_tracker.update_run_info(project)
         _logger.info("Fetching and validating project...")
-        project = fetch_and_validate_project(project, api)
+        project.fetch_and_validate_project()
         _logger.info("Fetching resource...")
         resource = launch_spec.get("resource") or "local-container"
         backend_config: Dict[str, Any] = {

wandb/sdk/launch/agent/config.py

@@ -128,22 +128,6 @@ class BuilderConfig(BaseModel):
         "the image will be pushed to the registry.",
     )
 
-    @validator("destination")  # type: ignore
-    @classmethod
-    def validate_destination(cls, destination: str) -> str:
-        """Validate that the destination is a valid container registry URI."""
-        for regex in [
-            GCP_ARTIFACT_REGISTRY_URI_REGEX,
-            AZURE_CONTAINER_REGISTRY_URI_REGEX,
-            ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
-        ]:
-            if regex.match(destination):
-                return destination
-        raise ValueError(
-            "Invalid destination. Destination must be a repository URI for an "
-            "ECR, ACR, or GCP Artifact Registry."
-        )
-
     platform: Optional[TargetPlatform] = Field(
         None,
         description="The platform to use for the built image. If not provided, "
@@ -196,17 +180,6 @@ class BuilderConfig(BaseModel):
             "S3 bucket, GCS bucket, or Azure blob."
         )
 
-    @root_validator(pre=True)  # type: ignore
-    @classmethod
-    def validate_kaniko(cls, values: dict) -> dict:
-        """Validate that kaniko is configured correctly."""
-        if values.get("type") == BuilderType.kaniko:
-            if values.get("build-context-store") is None:
-                raise ValueError(
-                    "builder.build-context-store is required if builder.type is set to kaniko."
-                )
-        return values
-
     @root_validator(pre=True)  # type: ignore
     @classmethod
     def validate_docker(cls, values: dict) -> dict:

wandb/sdk/launch/builder/build.py

@@ -2,13 +2,13 @@ import hashlib
 import json
 import logging
 import os
+import pathlib
 import shlex
 import shutil
 import sys
 import tempfile
 from typing import Any, Dict, List, Optional, Tuple
 
-import pkg_resources
 import yaml
 from dockerpycreds.utils import find_executable  # type: ignore
 from six.moves import shlex_quote
@@ -22,15 +22,12 @@ from wandb.sdk.launch.loader import (
     environment_from_config,
     registry_from_config,
 )
+from wandb.util import get_module
 
-from .._project_spec import (
-    EntryPoint,
-    EntrypointDefaults,
-    LaunchProject,
-    fetch_and_validate_project,
-)
+from .._project_spec import EntryPoint, EntrypointDefaults, LaunchProject
 from ..errors import ExecutionError, LaunchError
 from ..registry.abstract import AbstractRegistry
+from ..registry.anon import AnonynmousRegistry
 from ..utils import (
     AZURE_CONTAINER_REGISTRY_URI_REGEX,
     ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
@@ -105,8 +102,7 @@ def registry_from_uri(uri: str) -> AbstractRegistry:
 
         return ElasticContainerRegistry(uri=uri)
 
-    else:
-        raise LaunchError(f"Unsupported registry URI: {uri}. Unable to load helper.")
+    return AnonynmousRegistry(uri=uri)
 
 
 async def validate_docker_installation() -> None:
@@ -315,7 +311,6 @@ def get_env_vars_dict(
     _inject_wandb_config_env_vars(
         launch_project.override_config, env_vars, max_env_length
     )
-    # env_vars["WANDB_CONFIG"] = json.dumps(launch_project.override_config)
     artifacts = {}
     # if we're spinning up a launch process from a job
     # we should tell the run to use that artifact
@@ -340,27 +335,56 @@ def get_requirements_section(launch_project: LaunchProject, builder_type: str) -
         buildx_installed = False
     if launch_project.deps_type == "pip":
         requirements_files = []
-        if launch_project.project_dir is not None and os.path.exists(
-            os.path.join(launch_project.project_dir, "requirements.txt")
-        ):
+        deps_install_line = None
+        assert launch_project.project_dir is not None
+        base_path = pathlib.Path(launch_project.project_dir)
+        # If there is a requirements.txt at root of build context, use that.
+        if (base_path / "requirements.txt").exists():
             requirements_files += ["src/requirements.txt"]
-            pip_install_line = "pip install -r requirements.txt"
-        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
-            requirements_files += ["src/requirements.frozen.txt", "_wandb_bootstrap.py"]
-            pip_install_line = (
+            deps_install_line = "pip install -r requirements.txt"
+        # Elif there is pyproject.toml at build context, convert the dependencies
+        # section to a requirements.txt and use that.
+        elif (base_path / "pyproject.toml").exists():
+            tomli = get_module("tomli")
+            if tomli is None:
+                wandb.termwarn(
+                    "pyproject.toml found but tomli could not be loaded. To "
+                    "install dependencies from pyproject.toml please run "
+                    "`pip install tomli` and try again."
+                )
+            else:
+                # First try to read deps from standard pyproject format.
+                with open(base_path / "pyproject.toml", "rb") as f:
+                    contents = tomli.load(f)
+                project_deps = [
+                    str(d) for d in contents.get("project", {}).get("dependencies", [])
+                ]
+                if project_deps:
+                    with open(base_path / "requirements.txt", "w") as f:
+                        f.write("\n".join(project_deps))
+                    requirements_files += ["src/requirements.txt"]
+                    deps_install_line = "pip install -r requirements.txt"
+        # Else use frozen requirements from wandb run.
+        if not deps_install_line and (base_path / "requirements.frozen.txt").exists():
+            requirements_files += [
+                "src/requirements.frozen.txt",
+                "_wandb_bootstrap.py",
+            ]
+            deps_install_line = (
                 _parse_existing_requirements(launch_project)
                 + "python _wandb_bootstrap.py"
             )
+
+        if not deps_install_line:
+            raise LaunchError(f"No dependency sources found for {launch_project}")
+
         if buildx_installed:
             prefix = "RUN --mount=type=cache,mode=0777,target=/root/.cache/pip"
 
         requirements_line = PIP_TEMPLATE.format(
             buildx_optional_prefix=prefix,
             requirements_files=" ".join(requirements_files),
-            pip_install=pip_install_line,
+            pip_install=deps_install_line,
         )
     elif launch_project.deps_type == "conda":
         if buildx_installed:
@@ -446,13 +470,9 @@ def generate_dockerfile(
     return dockerfile_contents
 
 
-def construct_gcp_registry_uri(
-    gcp_repo: str, gcp_project: str, gcp_registry: str
-) -> str:
-    return "/".join([gcp_registry, gcp_project, gcp_repo])
-
-
 def _parse_existing_requirements(launch_project: LaunchProject) -> str:
+    import pkg_resources
+
     requirements_line = ""
     assert launch_project.project_dir is not None
     base_requirements = os.path.join(launch_project.project_dir, "requirements.txt")
@@ -506,6 +526,12 @@ def _create_docker_build_ctx(
                 dirs_exist_ok=True,
                 ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"),
             )
+            # TODO: remove this once we make things more explicit for users
+            if entrypoint_dir:
+                new_path = os.path.basename(entrypoint.name)
+                entrypoint = launch_project.get_single_entry_point()
+                if entrypoint is not None:
+                    entrypoint.update_entrypoint_path(new_path)
             return directory
 
     dst_path = os.path.join(directory, "src")
@@ -613,7 +639,7 @@ async def build_image_from_project(
     if not builder:
         raise LaunchError("Unable to build image. No builder found.")
 
-    launch_project = fetch_and_validate_project(launch_project, api)
+    launch_project.fetch_and_validate_project()
 
     entry_point: EntryPoint = launch_project.get_single_entry_point() or EntryPoint(
         name=EntrypointDefaults.PYTHON[-1],