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

wandb/sdk/launch/sweeps/utils.py

@@ -7,7 +7,8 @@ import yaml
 
 import wandb
 from wandb import util
-from wandb.sdk.launch.utils import LaunchError
+from wandb.apis.public import Api as PublicApi
+from wandb.sdk.launch.errors import LaunchError
 
 DEFAULT_SWEEP_COMMAND: List[str] = [
     "${env}",
@@ -118,16 +119,18 @@ def load_launch_sweep_config(config: Optional[str]) -> Any:
     return parsed_config
 
 
-def construct_scheduler_entrypoint(
+def construct_scheduler_args(
     sweep_config: Dict[str, Any],
     queue: str,
     project: str,
-    num_workers: Union[str, int],
     author: Optional[str] = None,
-) -> Optional[List[str]]:
-    """Construct a sweep scheduler run spec.
+    sweep_type: Optional[str] = "wandb",
+    return_job: bool = False,
+) -> Union[List[str], Dict[str, str], None]:
+    """Construct sweep scheduler args.
 
-    logs error and returns None if misconfigured, otherwise returns entrypoint
+    logs error and returns None if misconfigured,
+    otherwise returns args as a dict if is_job else a list of strings.
     """
     job = sweep_config.get("job")
     image_uri = sweep_config.get("image_uri")
@@ -135,47 +138,53 @@ def construct_scheduler_entrypoint(
         wandb.termerror(
             "No 'job' nor 'image_uri' top-level key found in sweep config, exactly one is required for a launch-sweep"
         )
-        return []
+        return None
     elif job and image_uri:
         wandb.termerror(
             "Sweep config has both 'job' and 'image_uri' but a launch-sweep can use only one"
         )
-        return []
+        return None
 
-    if type(num_workers) is str:
-        if num_workers.isdigit():
-            num_workers = int(num_workers)
-        else:
-            wandb.termerror(
-                "'num_workers' must be an integer or a string that can be parsed as an integer"
-            )
-            return []
-
-    entrypoint = [
-        "wandb",
-        "scheduler",
-        "WANDB_SWEEP_ID",
+    # if scheduler is a job, return args as dict
+    if return_job:
+        args_dict: Dict[str, str] = {
+            "sweep_id": "WANDB_SWEEP_ID",
+            "queue": queue,
+            "project": project,
+        }
+        if job:
+            args_dict["job"] = job
+        elif image_uri:
+            args_dict["image_uri"] = image_uri
+
+        if author:
+            args_dict["author"] = author
+
+        return args_dict
+
+    # scheduler uses cli commands, pass args as param list
+    args = [
         "--queue",
         f"{queue!r}",
         "--project",
-        project,
-        "--num_workers",
-        f"{num_workers}",
+        f"{project!r}",
+        "--sweep_type",
+        f"{sweep_type}",
     ]
-
+    if author:
+        args += [
+            "--author",
+            f"{author!r}",
+        ]
     if job:
-        if ":" not in job:
-            wandb.termwarn("No alias specified for job, defaulting to 'latest'")
-            job += ":latest"
-
-        entrypoint += ["--job", job]
+        args += [
+            "--job",
+            f"{job!r}",
+        ]
     elif image_uri:
-        entrypoint += ["--image_uri", image_uri]
+        args += ["--image_uri", image_uri]
 
-    if author:
-        entrypoint += ["--author", author]
-
-    return entrypoint
+    return args
 
 
 def create_sweep_command(command: Optional[List] = None) -> List:
@@ -265,3 +274,41 @@ def make_launch_sweep_entrypoint(
         return None, macro_args
 
     return entry_point, macro_args
+
+
+def check_job_exists(public_api: PublicApi, job: Optional[str]) -> bool:
+    """Check if the job exists using the public api.
+
+    Returns: True if no job is passed, or if the job exists.
+    Returns: False if the job is misformatted or doesn't exist.
+    """
+    if not job:
+        return True
+
+    try:
+        public_api.job(job)
+    except Exception as e:
+        wandb.termerror(f"Failed to load job. {e}")
+        return False
+    return True
+
+
+def get_previous_args(
+    run_spec: Dict[str, Any]
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Parse through previous scheduler run_spec.
+
+    returns scheduler_args and settings.
+    """
+    scheduler_args = (
+        run_spec.get("overrides", {}).get("run_config", {}).get("scheduler", {})
+    )
+    # also pipe through top level resource setup
+    if run_spec.get("resource"):
+        scheduler_args["resource"] = run_spec["resource"]
+    if run_spec.get("resource_args"):
+        scheduler_args["resource_args"] = run_spec["resource_args"]
+
+    settings = run_spec.get("overrides", {}).get("run_config", {}).get("settings", {})
+
+    return scheduler_args, settings

wandb/sdk/launch/utils.py

@@ -13,7 +13,9 @@ import wandb
 import wandb.docker as docker
 from wandb import util
 from wandb.apis.internal import Api
-from wandb.errors import CommError, Error
+from wandb.errors import CommError
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.github_reference import GitHubReference
 from wandb.sdk.launch.wandb_reference import WandbReference
 
 from .builder.templates._wandb_bootstrap import (
@@ -26,31 +28,8 @@ FAILED_PACKAGES_REGEX = re.compile(
 )
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.apis.public import Artifact as PublicArtifact
-
-
-class LaunchError(Error):
-    """Raised when a known error occurs in wandb launch."""
-
-    pass
-
-
-class LaunchDockerError(Error):
-    """Raised when Docker daemon is not running."""
-
-    pass
-
-
-class ExecutionError(Error):
-    """Generic execution exception."""
-
-    pass
-
-
-class SweepError(Error):
-    """Raised when a known error occurs with wandb sweeps."""
-
-    pass
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 
 
 # TODO: this should be restricted to just Git repos and not S3 and stuff like that
@@ -69,7 +48,9 @@ _WANDB_LOCAL_DEV_URI_REGEX = re.compile(
     r"^https?://localhost"
 )  # for testing, not sure if we wanna keep this
 
-API_KEY_REGEX = r"WANDB_API_KEY=\w+"
+API_KEY_REGEX = r"WANDB_API_KEY=\w+(-\w+)?"
+
+MACRO_REGEX = re.compile(r"\$\{(\w+)\}")
 
 PROJECT_SYNCHRONOUS = "SYNCHRONOUS"
 
@@ -77,7 +58,6 @@ UNCATEGORIZED_PROJECT = "uncategorized"
 LAUNCH_CONFIG_FILE = "~/.config/wandb/launch-config.yaml"
 LAUNCH_DEFAULT_PROJECT = "model-registry"
 
-
 _logger = logging.getLogger(__name__)
 LOG_PREFIX = f"{click.style('launch:', fg='magenta')} "
 
@@ -164,6 +144,7 @@ def construct_launch_spec(
     run_id: Optional[str],
     repository: Optional[str],
     author: Optional[str],
+    sweep_id: Optional[str] = None,
 ) -> Dict[str, Any]:
     """Construct the launch specification from CLI arguments."""
     # override base config (if supplied) with supplied args
@@ -191,6 +172,8 @@ def construct_launch_spec(
         launch_spec["docker"] = {}
     if docker_image:
         launch_spec["docker"]["docker_image"] = docker_image
+    if sweep_id:  # all runs in a sweep have this set
+        launch_spec["sweep_id"] = sweep_id
 
     if "resource" not in launch_spec:
         launch_spec["resource"] = resource if resource else None
@@ -461,49 +444,16 @@ def _fetch_git_repo(dst_dir: str, uri: str, version: Optional[str]) -> str:
     """
     # We defer importing git until the last moment, because the import requires that the git
     # executable is available on the PATH, so we only want to fail if we actually need it.
-    import git  # type: ignore
 
     _logger.info("Fetching git repo")
-    repo = git.Repo.init(dst_dir)
-    origin = repo.create_remote("origin", uri)
-    refspec = _make_refspec_from_version(version)
-    origin.fetch(refspec=refspec, depth=1)
-
-    if version is not None:
-        try:
-            repo.git.checkout(version)
-        except git.exc.GitCommandError as e:
-            raise LaunchError(
-                f"Unable to checkout version '{version}' of git repo {uri}"
-                "- please ensure that the version exists in the repo. "
-                f"Error: {e}"
-            ) from e
-    else:
-        if getattr(repo, "references", None) is not None:
-            branches = [ref.name for ref in repo.references]
-        else:
-            branches = []
-        # Check if main is in origin, else set branch to master
-        if "main" in branches or "origin/main" in branches:
-            version = "main"
-        else:
-            version = "master"
-
-        try:
-            repo.create_head(version, origin.refs[version])
-            repo.heads[version].checkout()
-            wandb.termlog(
-                f"{LOG_PREFIX}No git branch passed, defaulted to branch: {version}"
-            )
-        except (AttributeError, IndexError) as e:
-            raise LaunchError(
-                f"Unable to checkout default version '{version}' of git repo {uri} "
-                "- to specify a git version use: --git-version \n"
-                f"Error: {e}"
-            ) from e
-
-    repo.submodule_update(init=True, recursive=True)
-    return version
+    ref = GitHubReference.parse(uri)
+    if ref is None:
+        raise LaunchError(f"Unable to parse git uri: {uri}")
+    if version:
+        ref.update_ref(version)
+    ref.fetch(dst_dir)
+    assert version is not None or ref.ref is not None or ref.default_branch is not None
+    return version or ref.ref or ref.default_branch  # type: ignore
 
 
 def merge_parameters(
@@ -522,9 +472,17 @@ def convert_jupyter_notebook_to_script(fname: str, project_dir: str) -> str:
     )
 
     _logger.info("Converting notebook to script")
-    new_name = fname.rstrip(".ipynb") + ".py"
+    new_name = fname.replace(".ipynb", ".py")
     with open(os.path.join(project_dir, fname)) as fh:
         nb = nbformat.reads(fh.read(), nbformat.NO_CONVERT)
+        for cell in nb.cells:
+            if cell.cell_type == "code":
+                source_lines = cell.source.split("\n")
+                modified_lines = []
+                for line in source_lines:
+                    if not line.startswith("!"):
+                        modified_lines.append(line)
+                cell.source = "\n".join(modified_lines)
 
     exporter = nbconvert.PythonExporter()
     source, meta = exporter.from_notebook_node(nb)
@@ -536,7 +494,7 @@ def convert_jupyter_notebook_to_script(fname: str, project_dir: str) -> str:
 
 def check_and_download_code_artifacts(
     entity: str, project: str, run_name: str, internal_api: Api, project_dir: str
-) -> Optional["PublicArtifact"]:
+) -> Optional["Artifact"]:
     _logger.info("Checking for code artifacts")
     public_api = wandb.PublicApi(
         overrides={"base_url": internal_api.settings("base_url")}
@@ -663,12 +621,23 @@ def make_name_dns_safe(name: str) -> str:
     return resp
 
 
-def warn_failed_packages_from_build_logs(log: str, image_uri: str) -> None:
+def warn_failed_packages_from_build_logs(
+    log: str, image_uri: str, api: Api, job_tracker: Optional["JobAndRunStatusTracker"]
+) -> None:
     match = FAILED_PACKAGES_REGEX.search(log)
     if match:
-        wandb.termwarn(
-            f"Failed to install the following packages: {match.group(1)} for image: {image_uri}. Will attempt to launch image without them."
-        )
+        _msg = f"Failed to install the following packages: {match.group(1)} for image: {image_uri}. Will attempt to launch image without them."
+        wandb.termwarn(_msg)
+        if job_tracker is not None:
+            res = job_tracker.saver.save_contents(
+                _msg, "failed-packages.log", "warning"
+            )
+            api.update_run_queue_item_warning(
+                job_tracker.run_queue_item_id,
+                "Some packages were not successfully installed during the build",
+                "build",
+                res,
+            )
 
 
 def docker_image_exists(docker_image: str, should_raise: bool = False) -> bool:
@@ -696,3 +665,48 @@ def pull_docker_image(docker_image: str) -> None:
         docker.run(["docker", "pull", docker_image])
     except docker.DockerError as e:
         raise LaunchError(f"Docker server returned error: {e}")
+
+
+def macro_sub(original: str, sub_dict: Dict[str, Optional[str]]) -> str:
+    """Substitute macros in a string.
+
+    Macros occur in the string in the ${macro} format. The macro names are
+    substituted with their values from the given dictionary. If a macro
+    is not found in the dictionary, it is left unchanged.
+
+    Args:
+        original: The string to substitute macros in.
+        sub_dict: A dictionary mapping macro names to their values.
+
+    Returns:
+        The string with the macros substituted.
+    """
+    return MACRO_REGEX.sub(
+        lambda match: str(sub_dict.get(match.group(1), match.group(0))), original
+    )
+
+
+def recursive_macro_sub(source: Any, sub_dict: Dict[str, Optional[str]]) -> Any:
+    """Recursively substitute macros in a parsed JSON or YAML blob.
+
+    Macros occur in strings at leaves of the blob in the ${macro} format.
+    The macro names are substituted with their values from the given dictionary.
+    If a macro is not found in the dictionary, it is left unchanged.
+
+    Arguments:
+        source: The JSON or YAML blob to substitute macros in.
+        sub_dict: A dictionary mapping macro names to their values.
+
+    Returns:
+        The blob with the macros substituted.
+    """
+    if isinstance(source, str):
+        return macro_sub(source, sub_dict)
+    elif isinstance(source, list):
+        return [recursive_macro_sub(item, sub_dict) for item in source]
+    elif isinstance(source, dict):
+        return {
+            key: recursive_macro_sub(value, sub_dict) for key, value in source.items()
+        }
+    else:
+        return source

wandb/sdk/lib/_settings_toposort_generated.py

@@ -12,6 +12,7 @@ else:
 
 _Setting = Literal[
     "_args",
+    "_aws_lambda",
     "_async_upload_concurrency_limit",
     "_cli_only_mode",
     "_colab",
@@ -22,6 +23,7 @@ _Setting = Literal[
     "_disable_service",
     "_disable_stats",
     "_disable_viewer",
+    "_disable_setproctitle",
     "_except_exit",
     "_executable",
     "_extra_http_headers",
@@ -47,6 +49,7 @@ _Setting = Literal[
     "_platform",
     "_python",
     "_runqueue_item_id",
+    "_require_nexus",
     "_save_requirements",
     "_service_transport",
     "_service_wait",
@@ -66,6 +69,7 @@ _Setting = Literal[
     "allow_val_change",
     "anonymous",
     "api_key",
+    "azure_account_url_to_access_key",
     "base_url",
     "code_dir",
     "config_paths",
@@ -90,6 +94,7 @@ _Setting = Literal[
     "ignore_globs",
     "init_timeout",
     "is_local",
+    "job_source",
     "label_disable",
     "launch",
     "launch_config_path",
@@ -162,11 +167,13 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "api_key",
     "base_url",
     "console",
+    "job_source",
     "mode",
     "problem",
     "project",
     "run_id",
     "start_method",
+    "_aws_lambda",
     "_colab",
     "_console",
     "_ipython",

wandb/sdk/lib/capped_dict.py

@@ -0,0 +1,26 @@
+import collections
+from typing import Any, Optional
+
+
+class CappedDict(collections.OrderedDict):
+    default_max_size = 50
+
+    def __init__(self, max_size: Optional[int] = None) -> None:
+        self.max_size = max_size or self.default_max_size
+        super().__init__()
+
+    def __setitem__(self, key: str, val: Any) -> None:
+        if key not in self:
+            max_size = self.max_size - 1
+            self._prune_dict(max_size)
+        super().__setitem__(key, val)
+
+    def update(self, **kwargs: Any) -> None:  # type: ignore[override]
+        super().update(**kwargs)
+        self._prune_dict(self.max_size)
+
+    def _prune_dict(self, max_size: int) -> None:
+        if len(self) >= max_size:
+            diff = len(self) - max_size
+            for k in list(self.keys())[:diff]:
+                del self[k]