Managing Training Artifacts

# Download artifacts from a job
wave storage download j-xyz789 --output ./results
 
# List artifacts in a job
wave storage list j-xyz789
 
# Clean up old artifacts
wave storage cleanup --older-than 30d

/workspace/
├── artifacts/              # Main artifacts directory
│   ├── models/            # Trained models
│   ├── checkpoints/       # Training checkpoints
│   ├── logs/             # Training logs
│   └── results/          # Evaluation results
├── data/                  # Input data
└── src/                  # Your source code

# PyTorch example
import torch
 
# Save model
torch.save(model.state_dict(), '/workspace/artifacts/models/model.pt')
 
# Save checkpoint
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, '/workspace/artifacts/checkpoints/checkpoint.pt')

# TensorFlow example
import tensorflow as tf
 
# Save model
model.save('/workspace/artifacts/models/model')
 
# Save checkpoint
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
checkpoint.save('/workspace/artifacts/checkpoints/ckpt')

# List artifacts
wave storage list j-xyz789
 
# Download specific artifacts
wave storage download j-xyz789 \
  --include "*.pt" \
  --output ./models
 
# Download all artifacts
wave storage download j-xyz789 \
  --output ./results
 
# Stream logs in real-time
wave storage logs j-xyz789 -f
 
# Clean up old artifacts
wave storage cleanup \
  --older-than 30d \
  --exclude "*.pt"

import torch
from pathlib import Path
 
class ModelCheckpoint:
    def __init__(self, model, optimizer, save_dir):
        self.model = model
        self.optimizer = optimizer
        self.save_dir = Path('/workspace/artifacts/checkpoints') / save_dir
        self.save_dir.mkdir(parents=True, exist_ok=True)
 
    def save(self, epoch, loss):
        checkpoint_path = self.save_dir / f'checkpoint_epoch_{epoch}.pt'
        torch.save({
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'loss': loss,
        }, checkpoint_path)
 
    def load(self, checkpoint_path):
        checkpoint = torch.load(checkpoint_path)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        return checkpoint['epoch'], checkpoint['loss']

import tensorflow as tf
import os
 
class TrainingCallback(tf.keras.callbacks.Callback):
    def __init__(self, checkpoint_dir):
        super().__init__()
        self.checkpoint_dir = os.path.join('/workspace/artifacts/checkpoints', checkpoint_dir)
        os.makedirs(self.checkpoint_dir, exist_ok=True)
 
    def on_epoch_end(self, epoch, logs=None):
        # Save checkpoint
        checkpoint_path = os.path.join(self.checkpoint_dir, f'checkpoint_epoch_{epoch}')
        self.model.save_weights(checkpoint_path)
 
        # Save metrics
        with open(os.path.join(self.checkpoint_dir, 'metrics.txt'), 'a') as f:
            f.write(f'Epoch {epoch}: {logs}\n')

import boto3
from pathlib import Path
 
def upload_artifacts(job_id, local_dir='/workspace/artifacts'):
    s3 = boto3.client('s3')
    bucket = 'your-artifact-bucket'
 
    for path in Path(local_dir).rglob('*'):
        if path.is_file():
            s3_key = f'jobs/{job_id}/{path.relative_to(local_dir)}'
            s3.upload_file(str(path), bucket, s3_key)

from google.cloud import storage
from pathlib import Path
 
def upload_artifacts(job_id, local_dir='/workspace/artifacts'):
    client = storage.Client()
    bucket = client.get_bucket('your-artifact-bucket')
 
    for path in Path(local_dir).rglob('*'):
        if path.is_file():
            blob_name = f'jobs/{job_id}/{path.relative_to(local_dir)}'
            blob = bucket.blob(blob_name)
            blob.upload_from_filename(str(path))

import torch
from pathlib import Path
 
def save_training_artifacts(model, metrics, config, run_dir):
    # Create artifact directories
    artifact_dir = Path('/workspace/artifacts') / run_dir
    model_dir = artifact_dir / 'models'
    metrics_dir = artifact_dir / 'metrics'
    config_dir = artifact_dir / 'config'
 
    # Create directories
    for d in [model_dir, metrics_dir, config_dir]:
        d.mkdir(parents=True, exist_ok=True)
 
    # Save model
    torch.save(model.state_dict(), model_dir / 'model.pt')
 
    # Save metrics
    with open(metrics_dir / 'metrics.json', 'w') as f:
        json.dump(metrics, f)
 
    # Save config
    with open(config_dir / 'config.json', 'w') as f:
        json.dump(config, f)

import wandb
from pathlib import Path
 
def track_experiment(model, config, artifacts_dir):
    run = wandb.init(project="my-project", config=config)
 
    # Log model artifacts
    model_path = Path(artifacts_dir) / 'models' / 'model.pt'
    wandb.save(str(model_path))
 
    # Log metrics
    wandb.log({
        "train_loss": train_loss,
        "val_loss": val_loss,
        "learning_rate": config["lr"]
    })

import torch.distributed as dist
from pathlib import Path
 
def save_distributed_checkpoint(model, optimizer, epoch, rank):
    if rank == 0:  # Save only on main process
        checkpoint_dir = Path('/workspace/artifacts/checkpoints')
        checkpoint_dir.mkdir(parents=True, exist_ok=True)
 
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, checkpoint_dir / f'checkpoint_epoch_{epoch}.pt')
 
    dist.barrier()  # Wait for save to complete