class Arena(object):
"""A wrapper around arena that adds W&B config options"""
def __init__(self,
args,
wandb_project=None,
wandb_api_key=None,
wandb_run_id=None,
timeout_minutes=10):
self.api = InternalApi()
self.args = args
self.wandb_run_id = wandb_run_id or _short_id()
self.wandb_project = wandb_project
self.wandb_api_key = wandb_api_key or self.api.api_key
self.timeout_minutes = timeout_minutes
self.workers = int(self._parse_flag("--workers", 1)[1] or "0")
self.entity = None
def _parse_flag(self, flag, default=-1):
index = next((i for i, arg in enumerate(self.args)
if re.match(r"{}[= ]".format(flag), arg)), default)
if index > -1 and len(self.args) > index:
if "=" in self.args[index]:
val = self.args[index].split("=", 1)[1]
elif " " in self.args[index]:
val = self.args[index].split(" ", 1)[1]
else:
val = True
else:
val = None
return index, val
def submit(self):
try:
from minio import Minio
from google.cloud import storage
except ImportError:
raise ValueError(
"Required libraries for kubeflow aren't installed, run `pip install wandb[kubeflow]`"
)
print('Submitting arena {} job 🚀'.format(self.args[0]))
# TODO: require command?
opt_index, _ = self._parse_flag("--", len(self.args) - 1)
name_index, name = self._parse_flag("--name")
if name_index == -1:
name = "wandb"
name_index = len(self.args) - 1
self.args.insert(name_index, None)
name = '-'.join([name, _short_id(5)])
self.args[name_index] = "--name=" + name
projo = self.wandb_project or self.api.settings("project")
if projo:
if "/" in projo:
ent, projo = projo.split("/")
self.args.insert(opt_index,
"--env=WANDB_ENTITY={}".format(ent))
else:
_, git = self._parse_flag("--syncSource")
_, image = self._parse_flag("--image")
if git:
projo = git.split("/")[-1].replace(".git", "")
elif image:
projo = image.split(":")[0]
if projo:
projo = self.api.format_project(projo)
self.args.insert(opt_index,
"--env=WANDB_PROJECT={}".format(projo))
if self.wandb_api_key:
self.args.insert(
opt_index, "--env=WANDB_API_KEY={}".format(self.wandb_api_key))
else:
# Extract the secret, ideally this would be a secret env in the TFjob YAML
try:
kube_args = {"o": "json"}
index, namespace = self._parse_flag("--namespace")
if namespace:
kube_args["namespace"] = namespace
secret = json.loads(
str(sh.kubectl("get", "secret", "wandb", **kube_args)))
except sh.ErrorReturnCode:
secret = {}
if secret.get("data"):
print("Found wandb k8s secret, adding to environment")
api_key = secret["data"].get("api_key")
if api_key:
self.args.insert(
opt_index, "--env=WANDB_API_KEY=" +
base64.b64decode(api_key).decode("utf8"))
self.wandb_api_key = api_key
if self.wandb_api_key:
try:
# TODO: support someone overriding entity
if self.workers <= 1:
res = self.api.upsert_run(name=self.wandb_run_id,
project=projo)
wandb_run_path = os.path.join(
res["project"]["entity"]["name"],
res["project"]["name"], "runs", res["name"])
print(
'Run configured with W&B\nview live results here: {}'.
format("https://app.wandb.ai/" + wandb_run_path))
self.args.insert(
opt_index, "--env=WANDB_RUN_ID={}".format(res["name"]))
self.args.insert(opt_index, "--env=WANDB_RESUME=allow")
else:
res = self.api.viewer()
self.args.insert(opt_index,
"--env=WANDB_RUN_GROUP=" + name)
wandb_run_path = os.path.join(res["entity"], projo,
"groups", name)
print(
'Distributed run configured with W&B\nview live results here: {}'
.format("https://app.wandb.ai/" + wandb_run_path))
except CommError:
print("Failed to talk to W&B")
else:
print(
"Couldn't authenticate with W&B, run `wandb login` on your local machine"
)
index, gcs_url = self._parse_flag("--logdir")
tensorboard = self._parse_flag("--tensorboard")[0] > -1
if gcs_url and wandb_run_path:
pipeline_metadata(gcs_url, wandb_run_path, tensorboard)
elif wandb_run_path:
print("--logdir isn't set, skipping pipeline asset saving.")
cmd = arena(["submit"] + self.args)
print("Arena job {} submitted, watching state for upto {} minutes".
format(name, self.timeout_minutes))
total_time = 0
poll_rate = 10
while True:
# TODO: parse JSON when it's supported
status = str(arena("get", name)).split("\n")
rows = [
row for row in (re.split(r"\s+", row) for row in status)
if len(row) == 6 and "s" in row[3]
]
if len(rows) <= 1:
print("Final status: ", rows)
break
status = [row[1] for row in rows[1:]]
runtime = [row[3] for row in rows[1:]]
print("Status: {} {}".format(status, runtime))
if not all([s in ("PENDING", "RUNNING") for s in status]):
if not any([s in ("PENDING", "RUNNING") for s in status]):
print("Job finished with statuses: {}".format(status))
if any([s == "FAILED" for s in status]):
arena("logs", name, _fg=True)
break
time.sleep(poll_rate)
total_time += 10
if total_time > 90:
poll_rate = 30
if total_time > self.timeout_minutes * 60:
print("Timeout exceeded")
class Artifact(object):
"""An artifact object you can write files into, and pass to log_artifact."""
def __init__(self, name, type, description=None, metadata=None):
if not re.match("^[a-zA-Z0-9_\-.]+$", name):
raise ValueError(
'Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. Invalid name: "%s"'
% name
)
# TODO: this shouldn't be a property of the artifact. It's a more like an
# argument to log_artifact.
storage_layout = StorageLayout.V2
if env.get_use_v1_artifacts():
storage_layout = StorageLayout.V1
self._storage_policy = WandbStoragePolicy(
config={
"storageLayout": storage_layout,
# TODO: storage region
}
)
self._api = InternalApi()
self._final = False
self._digest = None
self._file_entries = None
self._manifest = ArtifactManifestV1(self, self._storage_policy)
self._cache = get_artifacts_cache()
self._added_new = False
self._added_objs = {}
# You can write into this directory when creating artifact files
self._artifact_dir = compat_tempfile.TemporaryDirectory(
missing_ok_on_cleanup=True
)
self.type = type
self.name = name
self.description = description
self.metadata = metadata
@property
def id(self):
# The artifact hasn't been saved so an ID doesn't exist yet.
return None
@property
def entity(self):
# TODO: querying for default entity a good idea here?
return self._api.settings("entity") or self._api.viewer().get("entity")
@property
def project(self):
return self._api.settings("project")
@property
def manifest(self):
self.finalize()
return self._manifest
@property
def digest(self):
self.finalize()
# Digest will be none if the artifact hasn't been saved yet.
return self._digest
def _ensure_can_add(self):
if self._final:
raise ValueError("Can't add to finalized artifact.")
def new_file(self, name, mode="w"):
self._ensure_can_add()
path = os.path.join(self._artifact_dir.name, name.lstrip("/"))
if os.path.exists(path):
raise ValueError('File with name "%s" already exists' % name)
util.mkdir_exists_ok(os.path.dirname(path))
self._added_new = True
return open(path, mode)
def add_file(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isfile(local_path):
raise ValueError("Path is not a file: %s" % local_path)
name = name or os.path.basename(local_path)
entry = ArtifactManifestEntry(
name,
None,
digest=md5_file_b64(local_path),
size=os.path.getsize(local_path),
local_path=local_path,
)
self._manifest.add_entry(entry)
return entry
def add_dir(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isdir(local_path):
raise ValueError("Path is not a directory: %s" % local_path)
termlog(
"Adding directory to artifact (%s)... "
% os.path.join(".", os.path.normpath(local_path)),
newline=False,
)
start_time = time.time()
paths = []
for dirpath, _, filenames in os.walk(local_path, followlinks=True):
for fname in filenames:
physical_path = os.path.join(dirpath, fname)
logical_path = os.path.relpath(physical_path, start=local_path)
if name is not None:
logical_path = os.path.join(name, logical_path)
paths.append((logical_path, physical_path))
def add_manifest_file(log_phy_path):
logical_path, physical_path = log_phy_path
self._manifest.add_entry(
ArtifactManifestEntry(
logical_path,
None,
digest=md5_file_b64(physical_path),
size=os.path.getsize(physical_path),
local_path=physical_path,
)
)
import multiprocessing.dummy # this uses threads
NUM_THREADS = 8
pool = multiprocessing.dummy.Pool(NUM_THREADS)
pool.map(add_manifest_file, paths)
pool.close()
pool.join()
termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
def add_reference(self, uri, name=None, checksum=True, max_objects=None):
if (
isinstance(uri, object)
and hasattr(uri, "parent_artifact")
and uri.parent_artifact != self
):
ref_url_fn = getattr(uri, "ref_url")
if callable(ref_url_fn):
uri = ref_url_fn()
url = urlparse(uri)
if not url.scheme:
raise ValueError(
"References must be URIs. To reference a local file, use file://"
)
if self._final:
raise ValueError("Can't add to finalized artifact.")
manifest_entries = self._storage_policy.store_reference(
self, uri, name=name, checksum=checksum, max_objects=max_objects
)
for entry in manifest_entries:
self._manifest.add_entry(entry)
return manifest_entries
# TODO: name this add_obj?
def add(self, obj, name):
"""adds `obj` to the artifact, located at `name`. You can use Artifact#get(`name`) after downloading
the artifact to retrieve this object.
Arguments:
- `obj`:wandb.Media - the object to save
- `name`:str - the path to save
"""
if isinstance(obj, Media):
# If the object is coming from another artifact, save it as a reference
if hasattr(obj, "_source") and obj._source is not None:
suffix = "." + obj.get_json_suffix() + ".json"
ref_path = obj._source["artifact"].get_path(
obj._source["name"] + suffix
)
path = name + suffix
return self.add_reference(ref_path, path)[0]
# Otherwise, save the object directly via json
elif type(obj) in JSONABLE_MEDIA_CLASSES:
obj_id = id(obj)
if obj_id in self._added_objs:
return self._added_objs[obj_id]
val = obj.to_json(self)
if "_type" not in val:
val["_type"] = obj.get_json_suffix()
suffix = val["_type"] + ".json"
if not name.endswith(suffix):
name = name + "." + suffix
entry = self._manifest.get_entry_by_path(name)
if entry is not None:
return entry
with self.new_file(name) as f:
import json
# TODO: Do we need to open with utf-8 codec?
f.write(json.dumps(obj.to_json(self), sort_keys=True))
# Note, we add the file from our temp directory.
# It will be added again later on finalize, but succeed since
# the checksum should match
entry = self.add_file(os.path.join(self._artifact_dir.name, name), name)
self._added_objs[obj_id] = entry
return entry
else:
ValueError("Can't add obj of type {} to artifact".format(type(obj)))
else:
raise ValueError("Can't add obj to artifact")
def get_added_local_path_name(self, local_path):
"""If local_path was already added to artifact, return its internal name."""
entry = self._manifest.get_entry_by_local_path(local_path)
if entry is None:
return None
return entry.path
def get_path(self, name):
raise ValueError("Cannot load paths from an artifact before it has been saved")
def download(self):
raise ValueError("Cannot call download on an artifact before it has been saved")
def finalize(self):
if self._final:
return self._file_entries
# Record any created files in the manifest.
if self._added_new:
self.add_dir(self._artifact_dir.name)
# mark final after all files are added
self._final = True
self._digest = self._manifest.digest()
# If there are new files, move them into the artifact cache now. Our temp
# self._artifact_dir may not be available by the time file pusher syncs
# these files.
if self._added_new:
# Update the file entries for new files to point at their new location.
def remap_entry(entry):
if entry.local_path is None or not entry.local_path.startswith(
self._artifact_dir.name
):
return entry
rel_path = os.path.relpath(
entry.local_path, start=self._artifact_dir.name
)
local_path = os.path.join(self._artifact_dir.name, rel_path)
cache_path, hit = self._cache.check_md5_obj_path(
entry.digest, entry.size
)
if not hit:
shutil.copyfile(local_path, cache_path)
entry.local_path = cache_path
for entry in self._manifest.entries.values():
remap_entry(entry)
class Artifact(object):
"""An artifact object you can write files into, and pass to log_artifact."""
def __init__(self, name, type, description=None, metadata=None):
if not re.match("^[a-zA-Z0-9_\-.]+$", name):
raise ValueError(
'Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. Invalid name: "%s"'
% name)
# TODO: this shouldn't be a property of the artifact. It's a more like an
# argument to log_artifact.
storage_layout = StorageLayout.V2
if env.get_use_v1_artifacts():
storage_layout = StorageLayout.V1
self._storage_policy = WandbStoragePolicy(
config={
"storageLayout": storage_layout,
# TODO: storage region
})
self._api = InternalApi()
self._final = False
self._digest = None
self._file_entries = None
self._manifest = ArtifactManifestV1(self, self._storage_policy)
self._cache = get_artifacts_cache()
self._added_new = False
# You can write into this directory when creating artifact files
self._artifact_dir = compat_tempfile.TemporaryDirectory(
missing_ok_on_cleanup=True)
self.type = type
self.name = name
self.description = description
self.metadata = metadata
@property
def id(self):
# The artifact hasn't been saved so an ID doesn't exist yet.
return None
@property
def entity(self):
# TODO: querying for default entity a good idea here?
return self._api.settings("entity") or self._api.viewer().get("entity")
@property
def project(self):
return self._api.settings("project")
@property
def manifest(self):
self.finalize()
return self._manifest
@property
def digest(self):
self.finalize()
# Digest will be none if the artifact hasn't been saved yet.
return self._digest
def _ensure_can_add(self):
if self._final:
raise ValueError("Can't add to finalized artifact.")
def new_file(self, name, mode="w"):
self._ensure_can_add()
path = os.path.join(self._artifact_dir.name, name.lstrip("/"))
if os.path.exists(path):
raise ValueError('File with name "%s" already exists' % name)
util.mkdir_exists_ok(os.path.dirname(path))
self._added_new = True
return open(path, mode)
def add_file(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isfile(local_path):
raise ValueError("Path is not a file: %s" % local_path)
name = name or os.path.basename(local_path)
entry = ArtifactManifestEntry(
name,
None,
digest=md5_file_b64(local_path),
size=os.path.getsize(local_path),
local_path=local_path,
)
self._manifest.add_entry(entry)
def add_dir(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isdir(local_path):
raise ValueError("Path is not a directory: %s" % local_path)
termlog(
"Adding directory to artifact (%s)... " %
os.path.join(".", os.path.normpath(local_path)),
newline=False,
)
start_time = time.time()
paths = []
for dirpath, _, filenames in os.walk(local_path, followlinks=True):
for fname in filenames:
physical_path = os.path.join(dirpath, fname)
logical_path = os.path.relpath(physical_path, start=local_path)
if name is not None:
logical_path = os.path.join(name, logical_path)
paths.append((logical_path, physical_path))
def add_manifest_file(log_phy_path):
logical_path, physical_path = log_phy_path
self._manifest.add_entry(
ArtifactManifestEntry(
logical_path,
None,
digest=md5_file_b64(physical_path),
size=os.path.getsize(physical_path),
local_path=physical_path,
))
import multiprocessing.dummy # this uses threads
NUM_THREADS = 8
pool = multiprocessing.dummy.Pool(NUM_THREADS)
pool.map(add_manifest_file, paths)
pool.close()
pool.join()
termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
def add_reference(self, uri, name=None, checksum=True, max_objects=None):
url = urlparse(uri)
if not url.scheme:
raise ValueError(
"References must be URIs. To reference a local file, use file://"
)
if self._final:
raise ValueError("Can't add to finalized artifact.")
manifest_entries = self._storage_policy.store_reference(
self, uri, name=name, checksum=checksum, max_objects=max_objects)
for entry in manifest_entries:
self._manifest.add_entry(entry)
def get_path(self, name):
raise ValueError(
"Cannot load paths from an artifact before it has been saved")
def download(self):
raise ValueError(
"Cannot call download on an artifact before it has been saved")
def finalize(self):
if self._final:
return self._file_entries
# Record any created files in the manifest.
if self._added_new:
self.add_dir(self._artifact_dir.name)
# mark final after all files are added
self._final = True
self._digest = self._manifest.digest()
# If there are new files, move them into the artifact cache now. Our temp
# self._artifact_dir may not be available by the time file pusher syncs
# these files.
if self._added_new:
# Update the file entries for new files to point at their new location.
def remap_entry(entry):
if entry.local_path is None or not entry.local_path.startswith(
self._artifact_dir.name):
return entry
rel_path = os.path.relpath(entry.local_path,
start=self._artifact_dir.name)
local_path = os.path.join(self._artifact_dir.name, rel_path)
cache_path, hit = self._cache.check_md5_obj_path(
entry.digest, entry.size)
if not hit:
shutil.copyfile(local_path, cache_path)
entry.local_path = cache_path
for entry in self._manifest.entries.values():
remap_entry(entry)
class Artifact(object):
"""An artifact object you can write files into, and pass to log_artifact."""
def __init__(self, name, type, description=None, metadata=None):
if not re.match(r"^[a-zA-Z0-9_\-.]+$", name):
raise ValueError(
'Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. Invalid name: "%s"'
% name)
# TODO: this shouldn't be a property of the artifact. It's a more like an
# argument to log_artifact.
storage_layout = StorageLayout.V2
if env.get_use_v1_artifacts():
storage_layout = StorageLayout.V1
self._storage_policy = WandbStoragePolicy(
config={
"storageLayout": storage_layout,
# TODO: storage region
})
self._api = InternalApi()
self._final = False
self._digest = None
self._file_entries = None
self._manifest = ArtifactManifestV1(self, self._storage_policy)
self._cache = get_artifacts_cache()
self._added_new = False
self._added_objs = {}
# You can write into this directory when creating artifact files
self._artifact_dir = compat_tempfile.TemporaryDirectory(
missing_ok_on_cleanup=True)
self.type = type
self.name = name
self.description = description
self.metadata = metadata
@property
def id(self):
# The artifact hasn't been saved so an ID doesn't exist yet.
return None
@property
def entity(self):
# TODO: querying for default entity a good idea here?
return self._api.settings("entity") or self._api.viewer().get("entity")
@property
def project(self):
return self._api.settings("project")
@property
def manifest(self):
self.finalize()
return self._manifest
@property
def digest(self):
self.finalize()
# Digest will be none if the artifact hasn't been saved yet.
return self._digest
def _ensure_can_add(self):
if self._final:
raise ValueError("Can't add to finalized artifact.")
def new_file(self, name, mode="w"):
self._ensure_can_add()
path = os.path.join(self._artifact_dir.name, name.lstrip("/"))
if os.path.exists(path):
raise ValueError('File with name "%s" already exists at "%s"' %
(name, path))
util.mkdir_exists_ok(os.path.dirname(path))
self._added_new = True
return open(path, mode)
def add_file(self, local_path, name=None, is_tmp=False):
"""Adds a local file to the artifact
Args:
local_path (str): path to the file
name (str, optional): new path and filename to assign inside artifact. Defaults to None.
is_tmp (bool, optional): If true, then the file is renamed deterministically. Defaults to False.
Returns:
ArtifactManifestEntry: the added entry
"""
self._ensure_can_add()
if not os.path.isfile(local_path):
raise ValueError("Path is not a file: %s" % local_path)
name = name or os.path.basename(local_path)
digest = md5_file_b64(local_path)
if is_tmp:
file_path, file_name = os.path.split(name)
file_name_parts = file_name.split(".")
file_name_parts[0] = b64_string_to_hex(digest)[:8]
name = os.path.join(file_path, ".".join(file_name_parts))
entry = ArtifactManifestEntry(
name,
None,
digest=digest,
size=os.path.getsize(local_path),
local_path=local_path,
)
self._manifest.add_entry(entry)
return entry
def add_dir(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isdir(local_path):
raise ValueError("Path is not a directory: %s" % local_path)
termlog(
"Adding directory to artifact (%s)... " %
os.path.join(".", os.path.normpath(local_path)),
newline=False,
)
start_time = time.time()
paths = []
for dirpath, _, filenames in os.walk(local_path, followlinks=True):
for fname in filenames:
physical_path = os.path.join(dirpath, fname)
logical_path = os.path.relpath(physical_path, start=local_path)
if name is not None:
logical_path = os.path.join(name, logical_path)
paths.append((logical_path, physical_path))
def add_manifest_file(log_phy_path):
logical_path, physical_path = log_phy_path
self._manifest.add_entry(
ArtifactManifestEntry(
logical_path,
None,
digest=md5_file_b64(physical_path),
size=os.path.getsize(physical_path),
local_path=physical_path,
))
import multiprocessing.dummy # this uses threads
NUM_THREADS = 8
pool = multiprocessing.dummy.Pool(NUM_THREADS)
pool.map(add_manifest_file, paths)
pool.close()
pool.join()
termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
def add_reference(self, uri, name=None, checksum=True, max_objects=None):
"""adds `uri` to the artifact via a reference, located at `name`.
You can use `Artifact.get_path(name)` to retrieve this object.
Arguments:
uri (str) - the URI path of the reference to add. Can be an object returned from
Artifact.get_path to store a reference to another artifact's entry.
name (str) - the path to save
"""
# This is a bit of a hack, we want to check if the uri is a of the type
# ArtifactEntry which is a private class returned by Artifact.get_path in
# wandb/apis/public.py. If so, then recover the reference URL.
if (isinstance(uri, object) and hasattr(uri, "parent_artifact")
and uri.parent_artifact != self):
ref_url_fn = getattr(uri, "ref_url")
uri = ref_url_fn()
url = urlparse(uri)
if not url.scheme:
raise ValueError(
"References must be URIs. To reference a local file, use file://"
)
if self._final:
raise ValueError("Can't add to finalized artifact.")
manifest_entries = self._storage_policy.store_reference(
self, uri, name=name, checksum=checksum, max_objects=max_objects)
for entry in manifest_entries:
self._manifest.add_entry(entry)
return manifest_entries
def add(self, obj, name):
"""Adds `obj` to the artifact, located at `name`. You can
use `Artifact.get(name)` after downloading the artifact to retrieve this object.
Arguments:
obj (wandb.WBValue): The object to save in an artifact
name (str): The path to save
"""
# Validate that the object is wandb.Media type
if not isinstance(obj, WBValue):
raise ValueError("Can only add `obj` which subclass wandb.WBValue")
obj_id = id(obj)
if obj_id in self._added_objs:
return self._added_objs[obj_id]
# If the object is coming from another artifact, save it as a reference
if obj.artifact_source is not None:
ref_path = obj.artifact_source["artifact"].get_path(
type(obj).with_suffix(obj.artifact_source["name"]))
return self.add_reference(ref_path, type(obj).with_suffix(name))[0]
val = obj.to_json(self)
name = obj.with_suffix(name)
entry = self._manifest.get_entry_by_path(name)
if entry is not None:
return entry
with self.new_file(name) as f:
import json
# TODO: Do we need to open with utf-8 codec?
f.write(json.dumps(val, sort_keys=True))
# Note, we add the file from our temp directory.
# It will be added again later on finalize, but succeed since
# the checksum should match
entry = self.add_file(os.path.join(self._artifact_dir.name, name),
name)
self._added_objs[obj_id] = entry
return entry
def get_added_local_path_name(self, local_path):
"""If local_path was already added to artifact, return its internal name."""
entry = self._manifest.get_entry_by_local_path(local_path)
if entry is None:
return None
return entry.path
def get_path(self, name):
raise ValueError(
"Cannot load paths from an artifact before it has been saved")
def download(self):
raise ValueError(
"Cannot call download on an artifact before it has been saved")
def get(self):
raise ValueError(
"Cannot call get on an artifact before it has been saved")
def finalize(self):
if self._final:
return self._file_entries
# Record any created files in the manifest.
if self._added_new:
self.add_dir(self._artifact_dir.name)
# mark final after all files are added
self._final = True
self._digest = self._manifest.digest()
# If there are new files, move them into the artifact cache now. Our temp
# self._artifact_dir may not be available by the time file pusher syncs
# these files.
if self._added_new:
# Update the file entries for new files to point at their new location.
def remap_entry(entry):
if entry.local_path is None or not entry.local_path.startswith(
self._artifact_dir.name):
return entry
rel_path = os.path.relpath(entry.local_path,
start=self._artifact_dir.name)
local_path = os.path.join(self._artifact_dir.name, rel_path)
cache_path, hit = self._cache.check_md5_obj_path(
entry.digest, entry.size)
if not hit:
shutil.copyfile(local_path, cache_path)
entry.local_path = cache_path
for entry in self._manifest.entries.values():
remap_entry(entry)
class Artifact(ArtifactInterface):
"""
Constructs an empty artifact whose contents can be populated using its
`add` family of functions. Once the artifact has all the desired files,
you can call `wandb.log_artifact()` to log it.
Arguments:
name: (str) A human-readable name for this artifact, which is how you
can identify this artifact in the UI or reference it in `use_artifact`
calls. Names can contain letters, numbers, underscores, hyphens, and
dots. The name must be unique across a project.
type: (str) The type of the artifact, which is used to organize and differentiate
artifacts. Common types include `dataset` or `model`, but you can use any string
containing letters, numbers, underscores, hyphens, and dots.
description: (str, optional) Free text that offers a description of the artifact. The
description is markdown rendered in the UI, so this is a good place to place tables,
links, etc.
metadata: (dict, optional) Structured data associated with the artifact,
for example class distribution of a dataset. This will eventually be queryable
and plottable in the UI. There is a hard limit of 100 total keys.
Examples:
Basic usage
```
wandb.init()
artifact = wandb.Artifact('mnist', type='dataset')
artifact.add_dir('mnist/')
wandb.log_artifact(artifact)
```
Raises:
Exception: if problem.
Returns:
An `Artifact` object.
"""
_added_objs: dict
_added_local_paths: dict
_distributed_id: Optional[str]
_metadata: dict
def __init__(
self,
name: str,
type: str,
description: Optional[str] = None,
metadata: Optional[dict] = None,
) -> None:
if not re.match(r"^[a-zA-Z0-9_\-.]+$", name):
raise ValueError(
"Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. "
'Invalid name: "%s"' % name)
# TODO: this shouldn't be a property of the artifact. It's a more like an
# argument to log_artifact.
storage_layout = StorageLayout.V2
if env.get_use_v1_artifacts():
storage_layout = StorageLayout.V1
self._storage_policy = WandbStoragePolicy(
config={
"storageLayout": storage_layout,
# TODO: storage region
})
self._api = InternalApi()
self._final = False
self._digest = ""
self._file_entries = None
self._manifest = ArtifactManifestV1(self, self._storage_policy)
self._cache = get_artifacts_cache()
self._added_objs = {}
self._added_local_paths = {}
# You can write into this directory when creating artifact files
self._artifact_dir = compat_tempfile.TemporaryDirectory(
missing_ok_on_cleanup=True)
self._type = type
self._name = name
self._description = description
self._metadata = metadata or {}
self._distributed_id = None
@property
def id(self) -> Optional[str]:
# The artifact hasn't been saved so an ID doesn't exist yet.
return None
@property
def entity(self) -> str:
# TODO: querying for default entity a good idea here?
return self._api.settings("entity") or self._api.viewer().get("entity")
@property
def project(self) -> str:
return self._api.settings("project")
@property
def manifest(self) -> ArtifactManifest:
self.finalize()
return self._manifest
@property
def digest(self) -> str:
self.finalize()
# Digest will be none if the artifact hasn't been saved yet.
return self._digest
@property
def description(self) -> Optional[str]:
return self._description
@description.setter
def description(self, desc: Optional[str]) -> None:
self._description = desc
@property
def metadata(self) -> dict:
return self._metadata
@metadata.setter
def metadata(self, metadata: dict) -> None:
self._metadata = metadata
@property
def type(self) -> str:
return self._type
@property
def name(self) -> str:
return self._name
@property
def state(self) -> str:
return "PENDING"
@property
def size(self) -> int:
return sum([entry.size for entry in self._manifest.entries])
@property
def distributed_id(self) -> Optional[str]:
return self._distributed_id
@distributed_id.setter
def distributed_id(self, distributed_id: Optional[str]) -> None:
self._distributed_id = distributed_id
def _ensure_can_add(self):
if self._final:
raise ValueError("Can't add to finalized artifact.")
@contextlib.contextmanager
def new_file(self, name: str, mode: str = "w"):
self._ensure_can_add()
path = os.path.join(self._artifact_dir.name, name.lstrip("/"))
if os.path.exists(path):
raise ValueError('File with name "%s" already exists at "%s"' %
(name, path))
util.mkdir_exists_ok(os.path.dirname(path))
with util.fsync_open(path, mode) as f:
yield f
self.add_file(path, name=name)
def add_file(
self,
local_path: str,
name: Optional[str] = None,
is_tmp: Optional[bool] = False,
):
self._ensure_can_add()
if not os.path.isfile(local_path):
raise ValueError("Path is not a file: %s" % local_path)
name = name or os.path.basename(local_path)
digest = md5_file_b64(local_path)
if is_tmp:
file_path, file_name = os.path.split(name)
file_name_parts = file_name.split(".")
file_name_parts[0] = b64_string_to_hex(digest)[:8]
name = os.path.join(file_path, ".".join(file_name_parts))
return self._add_local_file(name, local_path, digest=digest)
def add_dir(self, local_path: str, name: Optional[str] = None):
self._ensure_can_add()
if not os.path.isdir(local_path):
raise ValueError("Path is not a directory: %s" % local_path)
termlog(
"Adding directory to artifact (%s)... " %
os.path.join(".", os.path.normpath(local_path)),
newline=False,
)
start_time = time.time()
paths = []
for dirpath, _, filenames in os.walk(local_path, followlinks=True):
for fname in filenames:
physical_path = os.path.join(dirpath, fname)
logical_path = os.path.relpath(physical_path, start=local_path)
if name is not None:
logical_path = os.path.join(name, logical_path)
paths.append((logical_path, physical_path))
def add_manifest_file(log_phy_path):
logical_path, physical_path = log_phy_path
self._add_local_file(logical_path, physical_path)
import multiprocessing.dummy # this uses threads
NUM_THREADS = 8
pool = multiprocessing.dummy.Pool(NUM_THREADS)
pool.map(add_manifest_file, paths)
pool.close()
pool.join()
termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
def add_reference(
self,
uri: Union[ArtifactEntry, str],
name: Optional[str] = None,
checksum: bool = True,
max_objects: Optional[int] = None,
):
self._ensure_can_add()
# This is a bit of a hack, we want to check if the uri is a of the type
# ArtifactEntry which is a private class returned by Artifact.get_path in
# wandb/apis/public.py. If so, then recover the reference URL.
if isinstance(uri, ArtifactEntry) and uri.parent_artifact() != self:
ref_url_fn = getattr(uri, "ref_url")
uri = ref_url_fn()
url = urlparse(str(uri))
if not url.scheme:
raise ValueError(
"References must be URIs. To reference a local file, use file://"
)
manifest_entries = self._storage_policy.store_reference(
self, uri, name=name, checksum=checksum, max_objects=max_objects)
for entry in manifest_entries:
self._manifest.add_entry(entry)
return manifest_entries
def add(self, obj: WBValue, name: str):
self._ensure_can_add()
# Validate that the object is wandb.Media type
if not isinstance(obj, WBValue):
raise ValueError("Can only add `obj` which subclass wandb.WBValue")
obj_id = id(obj)
if obj_id in self._added_objs:
return self._added_objs[obj_id]["entry"]
# If the object is coming from another artifact, save it as a reference
if obj.artifact_source and obj.artifact_source.name:
ref_path = obj.artifact_source.artifact.get_path(
type(obj).with_suffix(obj.artifact_source.name))
return self.add_reference(ref_path, type(obj).with_suffix(name))[0]
val = obj.to_json(self)
name = obj.with_suffix(name)
entry = self._manifest.get_entry_by_path(name)
if entry is not None:
return entry
with self.new_file(name) as f:
import json
# TODO: Do we need to open with utf-8 codec?
f.write(json.dumps(val, sort_keys=True))
# Note, we add the file from our temp directory.
# It will be added again later on finalize, but succeed since
# the checksum should match
entry = self.add_file(os.path.join(self._artifact_dir.name, name),
name)
self._added_objs[obj_id] = {"entry": entry, "obj": obj}
return entry
def get_path(self, name: str):
raise ValueError(
"Cannot load paths from an artifact before it has been saved")
def get(self, name: str):
raise ValueError(
"Cannot call get on an artifact before it has been saved")
def download(self, root: str = None, recursive: bool = False):
raise ValueError(
"Cannot call download on an artifact before it has been saved")
def get_added_local_path_name(self, local_path: str):
"""
Get the artifact relative name of a file added by a local filesystem path.
Arguments:
local_path: (str) The local path to resolve into an artifact relative name.
Returns:
str: The artifact relative name.
Examples:
Basic usage
```
artifact = wandb.Artifact('my_dataset', type='dataset')
artifact.add_file('path/to/file.txt', name='artifact/path/file.txt')
# Returns `artifact/path/file.txt`:
name = artifact.get_added_local_path_name('path/to/file.txt')
```
"""
entry = self._added_local_paths.get(local_path, None)
if entry is None:
return None
return entry.path
def finalize(self):
"""
Marks this artifact as final, which disallows further additions to the artifact.
This happens automatically when calling `log_artifact`.
Returns:
None
"""
if self._final:
return self._file_entries
# mark final after all files are added
self._final = True
self._digest = self._manifest.digest()
def _add_local_file(self, name, path, digest=None):
digest = digest or md5_file_b64(path)
size = os.path.getsize(path)
cache_path, hit = self._cache.check_md5_obj_path(digest, size)
if not hit:
shutil.copyfile(path, cache_path)
entry = ArtifactManifestEntry(
name,
None,
digest=digest,
size=size,
local_path=cache_path,
)
self._manifest.add_entry(entry)
self._added_local_paths[path] = entry
return entry
class Artifact(object):
"""An artifact object you can write files into, and pass to log_artifact."""
# A local manifest contains the path to the local file in addition to the path within
# the artifact.
LocalArtifactManifestEntry = collections.namedtuple(
'LocalArtifactManifestEntry', ('path', 'hash', 'local_path'))
def __init__(self, name, type, description=None, metadata=None):
if not re.match('^[a-zA-Z0-9_\-.]+$', name):
raise ValueError(
'Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. Invalid name: "%s"'
% name)
if type is None:
raise ValueError(
"type is required when logging artifacts, specify \"dataset\", \"model\", or a custom type"
)
# TODO: this shouldn't be a property of the artifact. It's a more like an
# argument to log_artifact.
self._storage_policy = WandbStoragePolicy()
self._file_specs = {}
self._api = InternalApi() # TODO: persist project in settings?
self._final = False
self._digest = None
self._file_entries = None
self._manifest = ArtifactManifestV1(self, self._storage_policy)
self._cache = artifacts_cache.get_artifacts_cache()
self._added_new = False
# You can write into this directory when creating artifact files
self._artifact_dir = compat_tempfile.TemporaryDirectory(
missing_ok_on_cleanup=True)
self.server_manifest = None
self.type = type
self.name = name
self.description = description
self.metadata = metadata
@property
def id(self):
# The artifact hasn't been saved so an ID doesn't exist yet.
return None
@property
def entity(self):
# TODO: querying for default entity a good idea here?
return self._api.settings('entity') or self._api.viewer().get("entity")
@property
def project(self):
return self._api.settings('project')
@property
def manifest(self):
self.finalize()
return self._manifest
# TODO: Currently this returns the L0 digest. Is this what we want?
@property
def digest(self):
self.finalize()
return self._digest
def _ensure_can_add(self):
if self._final:
raise ValueError('Can\'t add to finalized artifact.')
def new_file(self, name, mode='w'):
self._ensure_can_add()
path = os.path.join(self._artifact_dir.name, name.lstrip('/'))
if os.path.exists(path):
raise ValueError('File with name "%s" already exists' % name)
util.mkdir_exists_ok(os.path.dirname(path))
self._added_new = True
return open(path, mode)
def add_file(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isfile(local_path):
raise ValueError('Path is not a file: %s' % local_path)
name = name or os.path.basename(local_path)
entry = ArtifactManifestEntry(name,
None,
digest=md5_file_b64(local_path),
size=os.path.getsize(local_path),
local_path=local_path)
self._manifest.add_entry(entry)
def add_dir(self, local_path, name=None):
self._ensure_can_add()
if not os.path.isdir(local_path):
raise ValueError('Path is not a directory: %s' % local_path)
termlog('Adding directory to artifact (%s)... ' %
os.path.join('.', os.path.normpath(local_path)),
newline=False)
start_time = time.time()
paths = []
for dirpath, _, filenames in os.walk(local_path, followlinks=True):
for fname in filenames:
physical_path = os.path.join(dirpath, fname)
logical_path = os.path.relpath(physical_path, start=local_path)
if name is not None:
logical_path = os.path.join(name, logical_path)
paths.append((logical_path, physical_path))
def add_manifest_file(log_phy_path):
logical_path, physical_path = log_phy_path
self._manifest.add_entry(
ArtifactManifestEntry(logical_path,
None,
digest=md5_file_b64(physical_path),
size=os.path.getsize(physical_path),
local_path=physical_path))
import multiprocessing.dummy # this uses threads
NUM_THREADS = 8
pool = multiprocessing.dummy.Pool(NUM_THREADS)
pool.map(add_manifest_file, paths)
pool.close()
pool.join()
termlog('Done. %.1fs' % (time.time() - start_time), prefix=False)
def add_reference(self, uri, name=None, checksum=True, max_objects=None):
url = urlparse(uri)
if not url.scheme:
raise ValueError(
'References must be URIs. To reference a local file, use file://'
)
if self._final:
raise ValueError('Can\'t add to finalized artifact.')
manifest_entries = self._storage_policy.store_reference(
self, uri, name=name, checksum=checksum, max_objects=max_objects)
for entry in manifest_entries:
self._manifest.add_entry(entry)
def get_path(self, name):
raise ValueError(
'Cannot load paths from an artifact before it has been saved')
def download(self):
raise ValueError(
'Cannot call download on an artifact before it has been saved')
def finalize(self):
if self._final:
return self._file_entries
# Record any created files in the manifest.
if self._added_new:
self.add_dir(self._artifact_dir.name)
# mark final after all files are added
self._final = True
# Add the manifest itself as a file.
with tempfile.NamedTemporaryFile('w+', suffix=".json",
delete=False) as fp:
json.dump(self._manifest.to_manifest_json(), fp, indent=4)
self._file_specs['wandb_manifest.json'] = fp.name
manifest_file = fp.name
# Calculate the server manifest
file_entries = []
for name, local_path in self._file_specs.items():
file_entries.append(
self.LocalArtifactManifestEntry(name, md5_file_b64(local_path),
os.path.abspath(local_path)))
self.server_manifest = ServerManifestV1(file_entries)
self._digest = self._manifest.digest()
# If there are new files, move them into the artifact cache now. Our temp
# self._artifact_dir may not be available by the time file pusher syncs
# these files.
if self._added_new:
# Update the file entries for new files to point at their new location.
def remap_entry(entry):
if entry.local_path is None or not entry.local_path.startswith(
self._artifact_dir.name):
return entry
rel_path = os.path.relpath(entry.local_path,
start=self._artifact_dir.name)
local_path = os.path.join(self._artifact_dir.name, rel_path)
cache_path, hit = self._cache.check_md5_obj_path(
entry.digest, entry.size)
if not hit:
shutil.copyfile(local_path, cache_path)
entry.local_path = cache_path
for entry in self._manifest.entries.values():
remap_entry(entry)
def init(ctx, project, entity, reset):
from wandb.old.core import _set_stage_dir, __stage_dir__, wandb_dir
if __stage_dir__ is None:
_set_stage_dir("wandb")
# non interactive init
if reset or project or entity:
api = InternalApi()
if reset:
api.clear_setting("entity", persist=True)
api.clear_setting("project", persist=True)
# TODO(jhr): clear more settings?
if entity:
api.set_setting("entity", entity, persist=True)
if project:
api.set_setting("project", project, persist=True)
return
if os.path.isdir(wandb_dir()) and os.path.exists(
os.path.join(wandb_dir(), "settings")
):
click.confirm(
click.style(
"This directory has been configured previously, should we re-configure it?",
bold=True,
),
abort=True,
)
else:
click.echo(
click.style("Let's setup this directory for W&B!", fg="green", bold=True)
)
api = InternalApi()
if api.api_key is None:
ctx.invoke(login)
viewer = api.viewer()
# Viewer can be `None` in case your API information became invalid, or
# in testing if you switch hosts.
if not viewer:
click.echo(
click.style(
"Your login information seems to be invalid: can you log in again please?",
fg="red",
bold=True,
)
)
ctx.invoke(login)
# This shouldn't happen.
viewer = api.viewer()
if not viewer:
click.echo(
click.style(
"We're sorry, there was a problem logging you in. Please send us a note at [email protected] and tell us how this happened.",
fg="red",
bold=True,
)
)
sys.exit(1)
# At this point we should be logged in successfully.
if len(viewer["teams"]["edges"]) > 1:
team_names = [e["node"]["name"] for e in viewer["teams"]["edges"]]
question = {
"type": "list",
"name": "team_name",
"message": "Which team should we use?",
"choices": team_names
# TODO(jhr): disabling manual entry for cling
# 'choices': team_names + ["Manual Entry"]
}
result = whaaaaat.prompt([question])
# result can be empty on click
if result:
entity = result["team_name"]
else:
entity = "Manual Entry"
if entity == "Manual Entry":
entity = click.prompt("Enter the name of the team you want to use")
else:
entity = viewer.get("entity") or click.prompt(
"What username or team should we use?"
)
# TODO: this error handling sucks and the output isn't pretty
try:
project = prompt_for_project(ctx, entity)
except ClickWandbException:
raise ClickException("Could not find team: %s" % entity)
api.set_setting("entity", entity, persist=True)
api.set_setting("project", project, persist=True)
api.set_setting("base_url", api.settings().get("base_url"), persist=True)
util.mkdir_exists_ok(wandb_dir())
with open(os.path.join(wandb_dir(), ".gitignore"), "w") as file:
file.write("*\n!settings")
click.echo(
click.style("This directory is configured! Next, track a run:\n", fg="green")
+ textwrap.dedent(
"""\
* In your training script:
{code1}
{code2}
* then `{run}`.
"""
).format(
code1=click.style("import wandb", bold=True),
code2=click.style('wandb.init(project="%s")' % project, bold=True),
run=click.style("python <train.py>", bold=True),
)
)