wandb 0.17.0rc1__py3-none-any.whl → 0.17.1__py3-none-any.whl

wandb/sdk/internal/job_builder.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Un
 import wandb
 from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.data_types._dtypes import TypeRegistry
+from wandb.sdk.internal.internal_api import Api
 from wandb.sdk.lib.filenames import DIFF_FNAME, METADATA_FNAME, REQUIREMENTS_FNAME
 from wandb.util import make_artifact_name_safe
 
@@ -23,7 +24,7 @@ else:
 _logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from wandb.proto.wandb_internal_pb2 import ArtifactRecord, UseArtifactRecord
+    from wandb.proto.wandb_internal_pb2 import ArtifactRecord
 
 FROZEN_REQUIREMENTS_FNAME = "requirements.frozen.txt"
 JOB_FNAME = "wandb-job.json"
@@ -32,6 +33,43 @@ JOB_ARTIFACT_TYPE = "job"
 LOG_LEVEL = Literal["log", "warn", "error"]
 
 
+class Version:
+    def __init__(self, major: int, minor: int, patch: int):
+        self._major = major
+        self._minor = minor
+        self._patch = patch
+
+    def __repr__(self) -> str:
+        return f"{self._major}.{self._minor}.{self._patch}"
+
+    def __lt__(self, other: "Version") -> bool:
+        if self._major < other._major:
+            return True
+        elif self._major == other._major:
+            if self._minor < other._minor:
+                return True
+            elif self._minor == other._minor:
+                if self._patch < other._patch:
+                    return True
+        return False
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return NotImplemented
+        return (
+            self._major == other._major
+            and self._minor == other._minor
+            and self._patch == other._patch
+        )
+
+
+# Minimum supported wandb version for keys in the source dict of wandb-job.json
+SOURCE_KEYS_MIN_SUPPORTED_VERSION = {
+    "dockerfile": Version(0, 17, 0),
+    "build_context": Version(0, 17, 0),
+}
+
+
 class GitInfo(TypedDict):
     remote: str
     commit: str
@@ -41,12 +79,16 @@ class GitSourceDict(TypedDict):
     git: GitInfo
     entrypoint: List[str]
     notebook: bool
+    build_context: Optional[str]
+    dockerfile: Optional[str]
 
 
 class ArtifactSourceDict(TypedDict):
     artifact: str
     entrypoint: List[str]
     notebook: bool
+    build_context: Optional[str]
+    dockerfile: Optional[str]
 
 
 class ImageSourceDict(TypedDict):
@@ -60,12 +102,6 @@ class JobSourceDict(TypedDict, total=False):
     input_types: Dict[str, Any]
     output_types: Dict[str, Any]
     runtime: Optional[str]
-    _partial: Optional[str]  # flag to indicate incomplete job
-
-
-class PartialJobSourceDict(TypedDict):
-    job_name: str
-    job_source_info: JobSourceDict
 
 
 class ArtifactInfoForJob(TypedDict):
@@ -73,6 +109,19 @@ class ArtifactInfoForJob(TypedDict):
     name: str
 
 
+def get_min_supported_for_source_dict(
+    source: Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict],
+) -> Optional[Version]:
+    """Get the minimum supported wandb version the source dict of wandb-job.json."""
+    min_seen = None
+    for key in source:
+        new_ver = SOURCE_KEYS_MIN_SUPPORTED_VERSION.get(key)
+        if new_ver:
+            if min_seen is None or new_ver < min_seen:
+                min_seen = new_ver
+    return min_seen
+
+
 class JobArtifact(Artifact):
     def __init__(self, name: str, *args: Any, **kwargs: Any):
         super().__init__(name, "placeholder", *args, **kwargs)
@@ -87,7 +136,7 @@ class JobBuilder:
     _summary: Optional[Dict[str, Any]]
     _logged_code_artifact: Optional[ArtifactInfoForJob]
     _disable: bool
-    _partial_source: Optional[PartialJobSourceDict]
+    _partial_source_id: Optional[str]  # Partial job source artifact id.
     _aliases: List[str]
     _job_seq_id: Optional[str]
     _job_version_alias: Optional[str]
@@ -104,13 +153,14 @@ class JobBuilder:
         self._job_seq_id = None
         self._job_version_alias = None
         self._disable = settings.disable_job_creation
-        self._partial_source = None
+        self._partial_source_id = None
         self._aliases = []
         self._source_type: Optional[Literal["repo", "artifact", "image"]] = (
             settings.job_source  # type: ignore[assignment]
         )
         self._is_notebook_run = self._get_is_notebook_run()
         self._verbose = verbose
+        self._partial = False
 
     def set_config(self, config: Dict[str, Any]) -> None:
         self._config = config
@@ -126,6 +176,17 @@ class JobBuilder:
     def disable(self, val: bool) -> None:
         self._disable = val
 
+    @property
+    def input_types(self) -> Dict[str, Any]:
+        return TypeRegistry.type_of(self._config).to_json()
+
+    @property
+    def output_types(self) -> Dict[str, Any]:
+        return TypeRegistry.type_of(self._summary).to_json()
+
+    def set_partial_source_id(self, source_id: str) -> None:
+        self._partial_source_id = source_id
+
     def _handle_server_artifact(
         self, res: Optional[Dict], artifact: "ArtifactRecord"
     ) -> None:
@@ -197,6 +258,8 @@ class JobBuilder:
             "git": {"remote": remote, "commit": commit},
             "entrypoint": entrypoint,
             "notebook": self._is_notebook_run,
+            "build_context": metadata.get("build_context"),
+            "dockerfile": metadata.get("dockerfile"),
         }
         name = self._make_job_name(f"{remote}_{program_relpath}")
 
@@ -247,6 +310,8 @@ class JobBuilder:
             "entrypoint": entrypoint,
             "notebook": self._is_notebook_run,
             "artifact": f"wandb-artifact://_id/{self._logged_code_artifact['id']}",
+            "build_context": metadata.get("build_context"),
+            "dockerfile": metadata.get("dockerfile"),
         }
         name = self._make_job_name(self._logged_code_artifact["name"])
 
@@ -276,7 +341,7 @@ class JobBuilder:
         return source, name
 
     def _make_job_name(self, input_str: str) -> str:
-        """Use job name from settings if provided, else use programatic name."""
+        """Use job name from settings if provided, else use programmatic name."""
         if self._settings.job_name:
             return self._settings.job_name
 
@@ -289,21 +354,10 @@ class JobBuilder:
     ) -> List[str]:
         # if building a partial job from CLI, overwrite entrypoint and notebook
         # should already be in metadata from create_job
-        if metadata.get("_partial"):
+        if self._partial:
             if metadata.get("entrypoint"):
                 entrypoint: List[str] = metadata["entrypoint"]
                 return entrypoint
-
-            # if entrypoint is not in metadata, then construct from python
-            assert metadata.get("python")
-
-            python = metadata["python"]
-            if python.count(".") > 1:
-                python = ".".join(python.split(".")[:2])
-
-            entrypoint = [f"python{python}", program_relpath]
-            return entrypoint
-
         # job is being built from a run
         entrypoint = [os.path.basename(sys.executable), program_relpath]
 
@@ -315,8 +369,89 @@ class JobBuilder:
     def _is_colab_run(self) -> bool:
         return hasattr(self._settings, "_colab") and bool(self._settings._colab)
 
-    def build(self) -> Optional[Artifact]:
+    def _build_job_source(
+        self,
+        source_type: str,
+        program_relpath: Optional[str],
+        metadata: Dict[str, Any],
+    ) -> Tuple[
+        Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict, None],
+        Optional[str],
+    ]:
+        """Construct a job source dict and name from the current run.
+
+        Arguments:
+            source_type (str): The type of source to build the job from. One of
+                "repo", "artifact", or "image".
+        """
+        source: Union[
+            GitSourceDict,
+            ArtifactSourceDict,
+            ImageSourceDict,
+            None,
+        ] = None
+
+        if source_type == "repo":
+            source, name = self._build_repo_job_source(
+                program_relpath or "",
+                metadata,
+            )
+        elif source_type == "artifact":
+            source, name = self._build_artifact_job_source(
+                program_relpath or "",
+                metadata,
+            )
+        elif source_type == "image" and self._has_image_job_ingredients(metadata):
+            source, name = self._build_image_job_source(metadata)
+        else:
+            source = None
+
+        if source is None:
+            if source_type:
+                self._log_if_verbose(
+                    f"Source type is set to '{source_type}' but some required information is missing "
+                    "from the environment. A job will not be created from this run. See "
+                    "https://docs.wandb.ai/guides/launch/create-job",
+                    "warn",
+                )
+            return None, None
+
+        return source, name
+
+    def build(
+        self,
+        api: Api,
+        build_context: Optional[str] = None,
+        dockerfile: Optional[str] = None,
+    ) -> Optional[Artifact]:
+        """Build a job artifact from the current run.
+
+        Arguments:
+            build_context (Optional[str]): Path within the job source code to
+                the image build context. Saved as part of the job for future
+                builds.
+            dockerfile (Optional[str]): Path within the build context the
+                Dockerfile. Saved as part of the job for future builds.
+
+        Returns:
+            Optional[Artifact]: The job artifact if it was successfully built,
+            otherwise None.
+        """
         _logger.info("Attempting to build job artifact")
+
+        # If a partial job was used, write the input/output types to the metadata
+        # rather than building a new job version.
+        if self._partial_source_id is not None:
+            new_metadata = {
+                "input_types": {"@wandb.config": self.input_types},
+                "output_types": self.output_types,
+            }
+            api.update_artifact_metadata(
+                self._partial_source_id,
+                new_metadata,
+            )
+            return None
+
         if not os.path.exists(
             os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)
         ):
@@ -332,6 +467,8 @@ class JobBuilder:
                 "warn",
             )
             return None
+        metadata["dockerfile"] = dockerfile
+        metadata["build_context"] = build_context
 
         runtime: Optional[str] = metadata.get("python")
         # can't build a job without a python version
@@ -348,84 +485,60 @@ class JobBuilder:
         name: Optional[str] = None
         source_info: Optional[JobSourceDict] = None
 
-        if self._partial_source is not None:
-            # construct source from downloaded partial job metadata
-            name = self._partial_source["job_name"]
-            source_info = self._partial_source["job_source_info"]
-            # add input/output types now that we are actually running a run
-            source_info.update(
-                {"input_types": input_types, "output_types": output_types}
-            )
-            # set source_type to determine whether to add diff file to artifact
-            source_type = source_info.get("source_type")
-        else:
-            # configure job from environment
-            source_type = self._get_source_type(metadata)
-            if not source_type:
-                # if source_type is None, then we don't have enough information to build a job
-                # if the user intended to create a job, warn.
-                if (
-                    self._settings.job_name
-                    or self._settings.job_source
-                    or self._source_type
-                ):
-                    self._log_if_verbose(
-                        "No source type found, not creating job artifact", "warn"
-                    )
-                return None
-
-            program_relpath = self._get_program_relpath(source_type, metadata)
-            if source_type != "image" and not program_relpath:
+        # configure job from environment
+        source_type = self._get_source_type(metadata)
+        if not source_type:
+            # if source_type is None, then we don't have enough information to build a job
+            # if the user intended to create a job, warn.
+            if (
+                self._settings.job_name
+                or self._settings.job_source
+                or self._source_type
+            ):
                 self._log_if_verbose(
-                    "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
-                    "warn",
+                    "No source type found, not creating job artifact", "warn"
                 )
-                return None
-
-            source: Union[
-                Optional[GitSourceDict],
-                Optional[ArtifactSourceDict],
-                Optional[ImageSourceDict],
-            ] = None
-
-            # make source dict
-            if source_type == "repo":
-                assert program_relpath
-                source, name = self._build_repo_job_source(program_relpath, metadata)
-            elif source_type == "artifact":
-                assert program_relpath
-                source, name = self._build_artifact_job_source(
-                    program_relpath, metadata
-                )
-            elif source_type == "image" and self._has_image_job_ingredients(metadata):
-                source, name = self._build_image_job_source(metadata)
-            else:
-                source = None
+            return None
 
-            if source is None:
-                if source_type:
-                    self._log_if_verbose(
-                        f"Source type is set to '{source_type}' but some required information is missing "
-                        "from the environment. A job will not be created from this run. See "
-                        "https://docs.wandb.ai/guides/launch/create-job",
-                        "warn",
-                    )
-                return None
-
-            source_info = {
-                "_version": "v0",
-                "source_type": source_type,
-                "source": source,
-                "input_types": input_types,
-                "output_types": output_types,
-                "runtime": runtime,
-            }
+        program_relpath = self._get_program_relpath(source_type, metadata)
+        if not self._partial and source_type != "image" and not program_relpath:
+            self._log_if_verbose(
+                "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                "warn",
+            )
+            return None
+
+        source, name = self._build_job_source(
+            source_type,
+            program_relpath,
+            metadata,
+        )
+        if source is None:
+            return None
+
+        if build_context:
+            source["build_context"] = build_context  # type: ignore[typeddict-item]
+        if dockerfile:
+            source["dockerfile"] = dockerfile  # type: ignore[typeddict-item]
+
+        # Pop any keys that are initialized to None. The current TypedDict
+        # system for source dicts requires all keys to be present, but we
+        # don't want to include keys that are None in the final dict.
+        for key in list(source.keys()):
+            if source[key] is None:  # type: ignore[literal-required]
+                source.pop(key)  # type: ignore[literal-require,misc]
+
+        source_info = {
+            "_version": str(get_min_supported_for_source_dict(source) or "v0"),
+            "source_type": source_type,
+            "source": source,
+            "input_types": input_types,
+            "output_types": output_types,
+            "runtime": runtime,
+        }
 
         assert source_info is not None
         assert name is not None
-        if metadata.get("_partial"):
-            assert not self._partial_source, "partial job has partial output"
-            source_info.update({"_partial": metadata["_partial"]})
 
         artifact = JobArtifact(name)
 
@@ -511,44 +624,3 @@ class JobBuilder:
 
     def _has_image_job_ingredients(self, metadata: Dict[str, Any]) -> bool:
         return metadata.get("docker") is not None
-
-
-def convert_use_artifact_to_job_source(
-    use_artifact: "UseArtifactRecord",
-) -> PartialJobSourceDict:
-    source_info = use_artifact.partial.source_info
-    source_info_dict: JobSourceDict = {
-        "_version": "v0",
-        "source_type": source_info.source_type,
-        "runtime": source_info.runtime,
-    }
-    if source_info.source_type == "repo":
-        entrypoint = [str(x) for x in source_info.source.git.entrypoint]
-        git_source: GitSourceDict = {
-            "git": {
-                "remote": source_info.source.git.git_info.remote,
-                "commit": source_info.source.git.git_info.commit,
-            },
-            "entrypoint": entrypoint,
-            "notebook": source_info.source.git.notebook,
-        }
-        source_info_dict.update({"source": git_source})
-    elif source_info.source_type == "artifact":
-        entrypoint = [str(x) for x in source_info.source.artifact.entrypoint]
-        artifact_source: ArtifactSourceDict = {
-            "artifact": source_info.source.artifact.artifact,
-            "entrypoint": entrypoint,
-            "notebook": source_info.source.artifact.notebook,
-        }
-        source_info_dict.update({"source": artifact_source})
-    elif source_info.source_type == "image":
-        image_source: ImageSourceDict = {
-            "image": source_info.source.image.image,
-        }
-        source_info_dict.update({"source": image_source})
-
-    partal_job_source_dict: PartialJobSourceDict = {
-        "job_name": use_artifact.partial.job_name.split(":")[0],
-        "job_source_info": source_info_dict,
-    }
-    return partal_job_source_dict

wandb/sdk/internal/progress.py

@@ -81,31 +81,3 @@ class Progress:
         return self.len
 
     next = __next__
-
-
-class AsyncProgress:
-    """Wrapper around Progress, to make it async iterable.
-
-    httpx, for streaming uploads, requires the data source to be an async iterable.
-    If we pass in a sync iterable (like a bare `Progress` instance), httpx will
-    get confused, think we're trying to make a synchronous request, and raise.
-    So we need this wrapper class to be an async iterable but *not* a sync iterable.
-    """
-
-    def __init__(self, progress: Progress) -> None:
-        self._progress = progress
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        try:
-            return next(self._progress)
-        except StopIteration:
-            raise StopAsyncIteration
-
-    def __len__(self):
-        return len(self._progress)
-
-    def rewind(self) -> None:
-        self._progress.rewind()