torchx/notebook: added workspacefile magic for use with notebooks

Author: d4l3k

dev-requirements.txt

@@ -19,3 +19,4 @@ torchtext>=0.11.0
 torchvision>=0.11.1
 ts>=0.5.1
 usort==0.6.4
+ipython

torchx/notebook.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+This contains TorchX utilities for creating and running components and apps from
+an Jupyter/IPython Notebook.
+"""
+
+import posixpath
+
+import fsspec
+from IPython.core.magic import register_cell_magic
+
+
+def get_workspace() -> str:
+    """
+    get_workspace returns the TorchX notebook workspace fsspec path.
+    """
+    return "memory://torchx-workspace/"
+
+
+@register_cell_magic
+def workspacefile(line: str, cell: str) -> None:
+    workspace = get_workspace()
+    fs, path = fsspec.core.url_to_fs(workspace)
+    path = posixpath.join(path, line)
+
+    base = posixpath.dirname(path)
+    if not fs.exists(base):
+        fs.mkdirs(base, exist_ok=True)
+
+    with fs.open(path, "wt") as f:
+        f.write(cell)
+
+    print(f"Added {line} to workspace {workspace}")

torchx/test/notebook_test.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import io
+import posixpath
+import unittest
+from unittest.mock import MagicMock, patch
+
+import fsspec
+from IPython.testing.globalipapp import get_ipython
+
+
+class VersionTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.ip = get_ipython()
+
+    def test_get_workspace(self) -> None:
+        from torchx.notebook import get_workspace
+
+        self.assertEqual(get_workspace(), "memory://torchx-workspace/")
+
+    @patch("sys.stdout", new_callable=io.StringIO)
+    def test_workspacefile(self, stdout: MagicMock) -> None:
+        from torchx.notebook import get_workspace
+
+        cell = "print('Arrakis')"
+        file_path = "foo/bar.py"
+        out = self.ip.run_cell_magic("workspacefile", file_path, cell)
+        self.assertEqual(out, None)
+        self.assertEqual(
+            stdout.getvalue(),
+            "Added foo/bar.py to workspace memory://torchx-workspace/\n",
+        )
+        fs, workspace_path = fsspec.core.url_to_fs(get_workspace())
+        path = posixpath.join(workspace_path, file_path)
+        with fs.open(path, "rt") as f:
+            self.assertEqual(f.read(), cell)