#21: ds.model() and ds.add_model(), tests thereof, and docs.

Author: alaric-dotmesh

README.md

@@ -329,6 +329,50 @@ ds.add_parameters(prefilter=True, smooth=True, smoothing_factor=12)
 ds.publish('Did some awesome data science!')
 ```
 
+### Models
+
+If your run has generated a Tensorflow model, you can declare it as such. This will load the model into the Model Library on the Dotscience Hub, and will enable automated deployment and model tracking features.
+
+This can be done with `model()` or `add_model()`:
+
+```python
+import dotscience as ds
+import tensorflow as tf
+
+ds.script()
+ds.start()
+
+...
+
+tf.saved_model.simple_save(
+    tf.keras.backend.get_session(),
+    ds.model(tf, "potatoes", "./model"),        # <---
+    inputs={'input_image_bytes': model.input},
+    outputs={t.name:t for t in model.outputs})
+
+...or...
+
+tf.saved_model.simple_save(
+    tf.keras.backend.get_session(),
+    "./model",
+    inputs={'input_image_bytes': model.input},
+    outputs={t.name:t for t in model.outputs})
+
+ds.add_model(tf, "potatoes", "./model")         # <---
+
+ds.publish('Trained the potato classifier')
+```
+
+The first argument to `model` or `add_model` should be the Tensorflow module itself, as imported by `import tensorflow as tf` in our example. This is used to identify it as a Tensorflow model (other types of model will be used in future), and to record the Tensorflow version used to generate it.
+
+The second argument is the model name for the Model Library. In this case, we called it `potatoes`, as our model is a potato classifier.
+
+The third argument is the path to the directory we're saving the Tensorflow model in, in this case `./model`. If called as `model()` rather than `add_model()`, this path is returned, so that it can be used to wrap the output path argument to `simple_save` in our example.
+
+For classifier models, an optional keyword argument is supported in both `model()` and `add_model()`: `classes` can be provided as a path to a CSV file listing your classes, to enable automatic model metric tracking in deployment.
+
+Note that we don't need to call `output()` for the paths passed to `model()` and `add_model()`; they automatically declare the files as outputs from this run.
+
 ## Multiple runs
 
 There's nothing to stop you from doing more than one "run" in one go; just call `start()` at the beginning and `publish()` at the end of each.

dotscience/__init__.py

@@ -124,6 +124,43 @@ def label(self, label, value):
         self.add_label(label, value)
         return value
 
+    def model(self, kind, name, *args, **kwargs):
+        artefact_type = None
+        try:
+            if kind.__name__ == "tensorflow":
+                artefact_type = "tensorflow-model"
+        except:
+            pass
+
+        if artefact_type == None:
+            raise RuntimeError('Unknown model type %r' % (kind,))
+
+        aj = {"type": artefact_type}
+        files = {}
+        return_value = None
+        if artefact_type == "tensorflow-model":
+            aj["version"] = kind.__version__
+            if len(args) != 1:
+                raise RuntimeError('Tensorflow models require a path to the model as the third argument')
+            files["model"] = args[0]
+            return_value = args[0]
+            if "classes" in kwargs:
+                files["classes"] = kwargs["classes"]
+
+        relative_files = {}
+        for key in files:
+            self.add_output(files[key])
+            relative_files[key] = os.path.relpath(str(files[key]),start=self._root)
+        aj["files"] = relative_files
+
+        self.add_label("artefact:" + name, json.dumps(aj, sort_keys=True, separators=(',', ':')))
+
+        return return_value
+
+    def add_model(self, kind, name, *args, **kwargs):
+        self.model(kind, name, *args, **kwards)
+        return None
+
     def add_summary(self, label, value):
         self._summary[str(label)] = str(value)
 

dotscience/test_dotscience.py

@@ -210,6 +210,45 @@ def test_run_output_recursive():
     finally:
         shutil.rmtree("test_run_output_recursive.tmp")
 
+class MockTensorflow:
+    def __init__(self):
+        self.__name__ = "tensorflow"
+        self.__version__ = "1.2.3.4"
+
+@given(text(),sampled_from(output_files))
+def test_run_tensorflow_model(name,x):
+    r = dotscience.Run("/workspace-root")
+    xpath = tidy_path("/workspace-root/" + x)
+    relxpath = os.path.relpath(xpath,start="/workspace-root")
+    assert r.model(MockTensorflow(), name, xpath) == xpath
+    assert str(r) == """[[DOTSCIENCE-RUN:%s]]%s[[/DOTSCIENCE-RUN:%s]]""" % \
+    (r._id, json.dumps({"version": "1",
+                        "summary": {},
+                        "parameters": {},
+                        "input": [],
+                        "output": [relxpath],
+                        "labels": {"artefact:"+name: "{\"files\":{\"model\":\"" + relxpath + "\"},\"type\":\"tensorflow-model\",\"version\":\"1.2.3.4\"}"},
+    }, sort_keys=True, indent=4), r._id)
+
+@given(text(),sampled_from(output_files),sampled_from(output_files))
+def test_run_tensorflow_model_with_classes(name,x,c):
+    assume(x != c)
+    r = dotscience.Run("/workspace-root")
+    xpath = tidy_path("/workspace-root/" + x)
+    relxpath = os.path.relpath(xpath,start="/workspace-root")
+    cpath = tidy_path("/workspace-root/" + c)
+    relcpath = os.path.relpath(cpath,start="/workspace-root")
+
+    assert r.model(MockTensorflow(), name, xpath, classes=cpath) == xpath
+    assert str(r) == """[[DOTSCIENCE-RUN:%s]]%s[[/DOTSCIENCE-RUN:%s]]""" % \
+    (r._id, json.dumps({"version": "1",
+                        "summary": {},
+                        "parameters": {},
+                        "input": [],
+                        "output": sorted([relxpath,relcpath]),
+                        "labels": {"artefact:"+name: "{\"files\":{\"classes\":\"" + relcpath + "\",\"model\":\"" + relxpath + "\"},\"type\":\"tensorflow-model\",\"version\":\"1.2.3.4\"}"},
+    }, sort_keys=True, indent=4), r._id)
+
 @given(text())
 def test_run_labels_1(x):
     r = dotscience.Run("/workspace-root")

test.sh

@@ -14,4 +14,4 @@ BASE=dotscience-python-test:$CI_DOCKER_TAG
 
 docker build -t $BASE -f Dockerfile.test .
 
-docker run -v dotscience-python-test-examples:/dsbuild/.hypothesis/examples $BASE /bin/bash -c "cd dsbuild ; pytest dotscience"
+docker run -v dotscience-python-test-examples:/dsbuild/.hypothesis/examples $BASE /bin/bash -c "cd dsbuild ; pytest dotscience $@"