Merge pull request #49 from simvue-io/pickled_objects

Support Python objects as artifacts

Author: alahiff

.github/workflows/python-app.yml

@@ -28,7 +28,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install flake8 pytest
         pip install -e .
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+        if [ -f test-requirements.txt ]; then pip install -r test-requirements.txt; fi
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Change log
 
+## v0.8.0
+
+* Support NumPy arrays, PyTorch tensors, Matplotlib and Plotly plots and picklable Python objects as artifacts.
+* (Bug fix) Events in offline mode didn't work.
+
 ## v0.7.2
 
 * Pydantic model is used for input validation.

examples/PyTorch/main.py

@@ -0,0 +1,158 @@
+# Taken from https://github.com/pytorch/examples/blob/main/mnist/main.py
+from __future__ import print_function
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+from simvue import Run
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        output = F.log_softmax(x, dim=1)
+        return output
+
+
+def train(args, model, device, train_loader, optimizer, epoch, run):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+            run.log_metrics({"train.loss.%d" % epoch: float(loss.item())}, step=batch_idx)
+            if args.dry_run:
+                break
+
+
+def test(model, device, test_loader, epoch, run):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
+            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+    test_accuracy = 100. * correct / len(test_loader.dataset)
+
+    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset),
+        test_accuracy))
+    run.log_metrics({'test.loss': test_loss,
+                     'test.accuracy': test_accuracy}, step=epoch)
+
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
+                        help='input batch size for testing (default: 1000)')
+    parser.add_argument('--epochs', type=int, default=14, metavar='N',
+                        help='number of epochs to train (default: 14)')
+    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
+                        help='learning rate (default: 1.0)')
+    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
+                        help='Learning rate step gamma (default: 0.7)')
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--no-mps', action='store_true', default=False,
+                        help='disables macOS GPU training')
+    parser.add_argument('--dry-run', action='store_true', default=False,
+                        help='quickly check a single pass')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+    parser.add_argument('--save-model', action='store_true', default=False,
+                        help='For Saving the current Model')
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+    use_mps = not args.no_mps and torch.backends.mps.is_available()
+
+    torch.manual_seed(args.seed)
+
+    if use_cuda:
+        device = torch.device("cuda")
+    elif use_mps:
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    train_kwargs = {'batch_size': args.batch_size}
+    test_kwargs = {'batch_size': args.test_batch_size}
+    if use_cuda:
+        cuda_kwargs = {'num_workers': 1,
+                       'pin_memory': True,
+                       'shuffle': True}
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    transform=transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,))
+        ])
+    dataset1 = datasets.MNIST('../data', train=True, download=True,
+                       transform=transform)
+    dataset2 = datasets.MNIST('../data', train=False,
+                       transform=transform)
+    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
+    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
+
+    model = Net().to(device)
+    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
+
+    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
+
+    run = Run()
+    run.init(tags=['PyTorch'])
+
+    for epoch in range(1, args.epochs + 1):
+        train(args, model, device, train_loader, optimizer, epoch, run)
+        test(model, device, test_loader, epoch, run)
+        scheduler.step()
+
+    if args.save_model:
+        run.save(model.state_dict(), "output", name="mnist_cnn.pt")
+
+    run.close()
+
+
+if __name__ == '__main__':
+    main()
+

examples/PyTorch/requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+simvue

setup.py

@@ -16,7 +16,7 @@
     long_description_content_type="text/markdown",
     url="https://simvue.io",
     platforms=["any"],
-    install_requires=["requests", "msgpack", "tenacity", "pyjwt", "psutil", "pydantic"],
+    install_requires=["dill", "requests", "msgpack", "tenacity", "pyjwt", "psutil", "pydantic", "plotly"],
     package_dir={'': '.'},
     packages=["simvue"],
     package_data={"": ["README.md"]},

simvue/__init__.py

@@ -2,4 +2,4 @@
 from simvue.client import Client
 from simvue.handler import Handler
 from simvue.models import RunInput
-__version__ = '0.7.2'
+__version__ = '0.8.0'

simvue/client.py

@@ -3,6 +3,7 @@
 import pickle
 import requests
 
+from .serialization import Deserializer
 from .utilities import get_auth
 
 CONCURRENT_DOWNLOADS = 10
@@ -51,7 +52,7 @@ def list_artifacts(self, run, category=None):
 
         return None
 
-    def get_artifact(self, run, name):
+    def get_artifact(self, run, name, allow_pickle=False):
         """
         Return the contents of the specified artifact
         """
@@ -62,23 +63,23 @@ def get_artifact(self, run, name):
         except requests.exceptions.RequestException:
             return None
 
-        if response.status_code == 200 and response.json():
-            url = response.json()[0]['url']
-
-            try:
-                response = requests.get(url, timeout=DOWNLOAD_TIMEOUT)
-            except requests.exceptions.RequestException:
-                return None
-        else:
+        if response.status_code != 200:
             return None
 
+        url = response.json()[0]['url']
+        mimetype = response.json()[0]['type']
+
         try:
-            content = pickle.loads(response.content)
-        except:
-            return response.content
-        else:
+            response = requests.get(url, timeout=DOWNLOAD_TIMEOUT)
+        except requests.exceptions.RequestException:
+            return None
+
+        content = Deserializer().deserialize(response.content, mimetype, allow_pickle)
+        if content is not None:
             return content
 
+        return response.content
+
     def get_artifact_as_file(self, run, name, path='./'):
         """
         Download an artifact

simvue/offline.py

@@ -1,9 +1,11 @@
+import codecs
 import json
 import logging
 import os
 import time
+import uuid
 
-from .utilities import get_offline_directory, create_file
+from .utilities import get_offline_directory, create_file, prepare_for_api
 
 logger = logging.getLogger(__name__)
 
@@ -90,9 +92,14 @@ def save_file(self, data):
         """
         Save file
         """
+        if 'pickled' in data:
+            temp_file = f"{self._directory}/temp-{str(uuid.uuid4())}.pickle"
+            with open(temp_file, 'wb') as fh:
+                fh.write(data['pickled'])
+            data['pickledFile'] = temp_file
         unique_id = time.time()
         filename = f"{self._directory}/file-{unique_id}.json"
-        self._write_json(filename, data)
+        self._write_json(filename, prepare_for_api(data, False))
         return True
 
     def add_alert(self, data):

simvue/remote.py

@@ -1,9 +1,8 @@
 import logging
 import time
-import requests
 
 from .api import post, put
-from .utilities import get_auth, get_expiry
+from .utilities import get_auth, get_expiry, prepare_for_api
 
 logger = logging.getLogger(__name__)
 
@@ -53,10 +52,13 @@ def create_run(self, data):
 
         return self._name
 
-    def update(self, data):
+    def update(self, data, run=None):
         """
         Update metadata, tags or status
         """
+        if run is not None:
+            data['name'] = run
+
         try:
             response = put(f"{self._url}/api/runs", self._headers, data)
         except Exception as err:
@@ -69,10 +71,13 @@ def update(self, data):
         self._error(f"Got status code {response.status_code} when updating run")
         return False
 
-    def set_folder_details(self, data):
+    def set_folder_details(self, data, run=None):
         """
         Set folder details
         """
+        if run is not None:
+            data['run'] = run
+
         try:
             response = put(f"{self._url}/api/folders", self._headers, data)
         except Exception as err:
@@ -85,13 +90,16 @@ def set_folder_details(self, data):
         self._error(f"Got status code {response.status_code} when updating folder details")
         return False
 
-    def save_file(self, data):
+    def save_file(self, data, run=None):
         """
         Save file
         """
+        if run is not None:
+            data['run'] = run
+
         # Get presigned URL
         try:
-            response = post(f"{self._url}/api/data", self._headers, data)
+            response = post(f"{self._url}/api/data", self._headers, prepare_for_api(data))
         except Exception as err:
             self._error(f"Got exception when preparing to upload file {data['name']} to object storage: {str(err)}")
             return False
@@ -105,22 +113,40 @@ def save_file(self, data):
 
         if 'url' in response.json():
             url = response.json()['url']
-            try:
-                with open(data['originalPath'], 'rb') as fh:
-                    response = put(url, {}, fh, is_json=False, timeout=UPLOAD_TIMEOUT)
+            if 'pickled' in data and 'pickledFile' not in data:
+                try:
+                    response = put(url, {}, data['pickled'], is_json=False, timeout=UPLOAD_TIMEOUT)
                     if response.status_code != 200:
-                        self._error(f"Got status code {response.status_code} when uploading file {data['name']} to object storage")
+                        self._error(f"Got status code {response.status_code} when uploading object {data['name']} to object storage")
                         return None
-            except Exception as err:
-                self._error(f"Got exception when uploading file {data['name']} to object storage: {str(err)}")
-                return None
+                except Exception as err:
+                    self._error(f"Got exception when uploading object {data['name']} to object storage: {str(err)}")
+                    return None
+            else:
+                if 'pickledFile' in data:
+                    use_filename = data['pickledFile']
+                else:
+                    use_filename = data['originalPath']
+
+                try:
+                    with open(use_filename, 'rb') as fh:
+                        response = put(url, {}, fh, is_json=False, timeout=UPLOAD_TIMEOUT)
+                        if response.status_code != 200:
+                            self._error(f"Got status code {response.status_code} when uploading file {data['name']} to object storage")
+                            return None
+                except Exception as err:
+                    self._error(f"Got exception when uploading file {data['name']} to object storage: {str(err)}")
+                    return None
 
         return True
 
-    def add_alert(self, data):
+    def add_alert(self, data, run=None):
         """
         Add an alert
         """
+        if run is not None:
+            data['run'] = run
+
         try:
             response = post(f"{self._url}/api/alerts", self._headers, data)
         except Exception as err: