search

philipturner / swift-colab

Swift kernel for Google Colaboratory
MIT License
105 stars 9 forks source link

Provides a JupyterDisplay protocol package? #21

Open liuliu opened 2 years ago

liuliu commented 2 years ago

Seems we are the only two people use Swift with notebooks. Anyway, I think we need to collaborate on a new header-like library called JupyterDisplay. I have a prototype: https://github.com/liuliu/swift-jupyter/blob/main/display/JupyterDisplay.swift

It is really simple. The idea is that any other libraries can depend on this library and provide notebook rich display support. If you looked at Python packages, many of them "battery-included" for notebook displays either inside the main package or in the support package.

So far, the EnableJupyterDisplay.swift implementation took hands into itself and provides backends for matplotlib or Python or SwiftPlot. That is really limited.

On the technical side, this package simply exposes a JupyterDisplay instance that EnableJupyterDisplay.swift can switch to a different implementation, thus, enabling the said behavior.

Let me know what do you think. I mainly want to propose moving this into a "neutral" org and package it lightly (like package to SwiftPM and Bazel) and everyone can use.

philipturner commented 2 years ago

This sounds like a good idea. I can look into it when I get around to the next major update to Swift-Colab. But probably for the next few months, I’ll be occupied with the Metal/GPU side of my project. We will probably need to wait until changes can be made to Swift-Colab, because Jupyter display stuff is deeply integrated into the kernel. I removed the EnableJupyterDisplay file a long time ago and only EnableIPythonDisplay is maintained. That demonstrates the crux of the issue here, lack of local Jupyter notebook support.

philipturner commented 2 years ago

I saw that your fork supports local Jupyter notebooks, a feature I would like to integrate into Swift-Colab via JupyterLab (Swift-CoJupyterLab). For example, I would like to use my upcoming Metal backend for S4TF in a familiar Colab-like environment. I made some highly optimized SwiftPM caching stuff for Colab, which would fit perfectly into JupyterLab like a Swift alternative to Python pip.

Could you provide a quick explanation of what’s been going on with the swift-jupyter fork lately? What bug fixes, changes to the README, etc. have you added? I’d like to help your contributions reach a larger audience, and I can leverage the large traffic Swift-Colab gets (48 vs 9 stars). We can start off with a link in my README directed to your fork, but eventually I might merge your work in entirely into Swift-Colab 3.0.

As a time frame of what I’ve said previously when conversing with you:

philipturner commented 2 years ago

Instead of a Swift package dependency, this library could be a system module imported like Apple’s Foundation library. You don’t compile that with SwiftPM; it just comes with the toolchain. Same thing with this. In v2.3, I’m planning a way to link modules into each other like they’re system modules, enabling the compilation of https://github.com/s4tf/models. This idea can be repurposed for the package you propose, and I think it will be easier to maintain that way. Plus, more flexibility because I can swap out the backend implementation at any time. Similar to my idea for the Metal and OpenCL backends for S4TF, which use an API-stable C-like interface and virtual function calls.

liuliu commented 2 years ago

This sounds like a good idea. I can look into it when I get around to the next major update to Swift-Colab. But probably for the next few months, I’ll be occupied with the Metal/GPU side of my project. We will probably need to wait until changes can be made to Swift-Colab, because Jupyter display stuff is deeply integrated into the kernel. I removed the EnableJupyterDisplay file a long time ago and only EnableIPythonDisplay is maintained. That demonstrates the crux of the issue here, lack of local Jupyter notebook support.

Interesting. Yeah, I went to a different route by not using EnableIPythonDisplay (as that internally will initiate (as in google/swift-jupyter, don't know if you did any improvements) ipython kernel, and that for a particular Swift version at the time (5.5.x), it will silently kill the swift kernel for unknown reason (the LLDB process exit normally if you attach that to a debugger) after a few cell executions. Without ipython kernel, everything works fine) and just add necessary setup for pandas and matplotlib to work (the latter is more involved).

Could you provide a quick explanation of what’s been going on with the swift-jupyter fork lately? What bug fixes, changes to the README, etc. have you added? I’d like to help your contributions reach a larger audience, and I can leverage the large traffic Swift-Colab gets (48 vs 9 stars). We can start off with a link in my README directed to your fork, but eventually I might merge your work in entirely into Swift-Colab 3.0.

I recently become interested in using notebooks to do more GUI heavy stuff besides some data crunching. Most work in swift-jupyter now revolves around that and making development with Bazel easier (and correct). A short list:

One thing I liked yours is the independency LLDB integration without Python. The CI is flaky recently and packages were published without the Python LLDB plugin (I guess that will be more of a problem for VSCode Swift extension than notebooks? Why no one report that earlier is super strange).

Are you using Bazel for internal development too? Curious why it is slated to v2.3. There are definitely quite a bit more weirdness with Bazel than SwiftPM (which natively supports compile framework) than I anticipated when I did it. I am not certain that I have solved it all though. Seems work for my repo now (which is nontrivial, as it has interactions with several C heavy / Swift heavy libraries), but in the past every time I use it extensively, there are bugs pop up somehow (mostly around compilation flags, linker flags or missing module.modulemap for C libs).

philipturner commented 2 years ago

Are you using Bazel for internal development too? Curious why it is slated to v2.3.

I have only ever used Bazel to compile tensorflow/tensorflow, but not to compile Swift code. I am not very familiar with it and do not use it internally. I wanted to put it v2.3 because that's less difficult than adding local notebook support. I was also going to change the inline SwiftPM package spec parser, supporting multiple types of syntax. I prefer the terse .bazel instead of .bazelBuild, which is less mentally demanding to read.

%install '.package(url: "https://...", .branch("main")' // 4.2 style, currently only supported
%install '.package(url: "https://...", branch: "main")' // 5.0 style, coming in v2.3
%install '.bazel(url: "https://...", anotherArgument: "@PythonKit//:PythonKit")' // coming in v2.3
%install '.cmake(...)' // distant future; Swift-Colab is easily extensible to more build systems

This differs form your fork's new %bazel magic command. I avoid adding magic commands unless absolutely necessary, because it breaks backward compatibility if you change your mind and want to rename it. But that also goes for how you process arguments to magic commands, hence I must retain support for Swift 4.2-style manifests.

One thing I liked yours is the independency LLDB integration without Python.

That was one of the big reasons why I rewrote Swift-Colab to create v2.0. Development toolchains kept crashing, and I didn't know why. I assumed that only release toolchains could be compatible with Swift-Colab. Apparently, dev toolchains were compatible but I had to rewrite the LLDB dependency from scratch. I had a long-term vision of compiling S4TF on Colab months down the road, which only supported dev toolchains at the time. That was a good choice, because I think Swift 5.6 dropped the Python bindings.

There are definitely quite a bit more weirdness with Bazel than SwiftPM

From my experience, SwiftPM is quite buggy. I reported two bugs on the swift-package-manager repo, along with an issue in swift-colab.

Working with buggy build systems is nothing new to me, but I am concerned about immediate download time. Google Colab has an incredibly fast network bandwidth, so downloading Swift toolchains is fine. I found that downloading via apt-get can take a few seconds, as can compiling a Swift package like swift-crypto. Any optional delay is unacceptable in Colab, so I will lazily install Bazel via apt the moment it's requested.

Default to optimized version of the final .so library for importing, this helps quite a bit for Swift dependencies that is computationally-heavy;

I want to adapt the mindset of "zero tolerance for latency" to local notebook support, which should create the best possible user experience. For example, users should not need to compile LLDB Python bindings themselves; that takes minutes. The SwiftPM engine is really powerful in this respect, because it compiles in debug mode by default. Bazel should not compile in fully-optimized mode by default, because that increases compile time. However, we should provide a well-documented mechanism to enable release-mode compilation with Bazel. The README can yell out "enable this option because it's fastest!"

philipturner commented 2 years ago

Is there any chance we can translate SwiftPM flags into Bazel build options? For example, -c release could translate into fully-optimized .so/.dylib/.dll binaries. I would like to reuse %install-swiftpm-flags for both SwiftPM and Bazel. The alternative is adding a %install-bazel-flags command, but:

When I designed Swift-Colab, I thought from the end-user's experience more than anything else. They are often impatient and don't know as much about SwiftPM/command-line stuff as we do. However, I can't change the semantic purpose of certain features just for the sake of ergonomics. The idea above would change the semantic purpose of %install-swiftpm-flags.

Balancing theoretical correctness with ergonomics led to some other wierd decisions previously. For example, I included a Python code snippet for restarting the runtime, but only on the README. It is not part of template notebooks. You can expect a lot of deliberation and caution over little changes like this to Swift-Colab. This quality assurance stems from me trying to create the "authentic"/"official" continuation of Swift for TensorFlow. What would Google do if they were making this repository?

liuliu commented 2 years ago

This differs form your fork's new %bazel magic command. I avoid adding magic commands unless absolutely necessary, because it breaks backward compatibility if you change your mind and want to rename it. But that also goes for how you process arguments to magic commands, hence I must retain support for Swift 4.2-style manifests.

I can go more detail into how I use Bazel, and in general how it works if you want. Basically, it is unlike SwiftPM which has a vibrant GitHub presence that you can simply include a URL and pulling in dependencies. However, it is great if you have a projects that has mixed-language dependencies, such as mine, which mixed C libraries, Python libraries and tied together into a Swift binary (or notebook).

That was a good choice, because I think Swift 5.6 dropped the Python bindings.

TBH, this is probably an oversight. Last time this happened for 5.4, I pinged in #swift channel and the CI were fixed, so we have it in 5.5 :)

From my experience, SwiftPM is quite buggy. I reported two bugs on the swift-package-manager repo, along with an issue in swift-colab.

Ha, nothing is easy then. I suspect REPL mode is used by very few especially the library support. But hopefully Swift Playground on various platforms can help this.

I want to adapt the mindset of "zero tolerance for latency" to local notebook support, which should create the best possible user experience.

Yeah, at that point, probably have a common cache for people to mount :) TBH, the speed improvements when linking to optimized version v.s. debug version is night-and-day.

Is there any chance we can translate SwiftPM flags into Bazel build options? For example, -c release could translate into fully-optimized .so/.dylib/.dll binaries. I would like to reuse %install-swiftpm-flags for both SwiftPM and Bazel. The alternative is adding a %install-bazel-flags command, but:

Yeah, Bazel have 3 build modes for all its supported languages by default (that includes Rust, Go, C, C++ and Swift): debug, fastbuild and opt, you can switch between them with --compilation_mode=opt for example. The only difference between debug and fastbuild is that latter strips debug symbols so it is leaner.

Balancing theoretical correctness with ergonomics led to some other wierd decisions previously.

Agreed. Hence I asked why Bazel. It is not exactly a popular choice. However, it does give you some advantages, such as setup a consistent environment.

Now, if you want a quick starter on Bazel as a build system for Swift code:

Your repo needs to have a WORKSPACE file at the root, it contains references to other repos, including repos that provides Bazel rules (more on this later).

You specify the code to build with BUILD file in directories, and use Python-like language (it is imperative, but looks declarative) to compile, for example:

swift_library(
  name = "SomePackageName",
  srcs = ["some.swift", "package.swift"],
  deps = ["OtherPackages"]
)

You can build, or run (if it is a executable target) with bazel command-line:

bazel run //Directory1/Directory2:TargetName

Here is the thing that probably doesn't map well mentally between Bazel / SwiftPM: SwiftPM has only one Package.swift file, Bazel can have as many BUILD files as you can in one repo. That also means it needs to be efficient to reference dependencies. Thus, all dependencies from other repos (or places), needs to be first defined in WORKSPACE and reference through the syntax like @RepoName//Directory:TargetName. The RepoName is what you defined in WORKSPACE. That explains the @PythonKit//:PythonKit I referenced so much about in the notebook.

Minor but often important, you can put some bazel build parameters in .bazelrc file to avoid typing it every time when bazel run or bazel build.

Here is what my setup looks like when using Bazel + swift-jupyter:

Once setup, this is great for me. However, as you can see, it differs from mental model when operating with SwiftPM. Most notably: with SwiftPM, the notebook can operate almost independently, it doesn't need a repo with WORKSPACE file.

Although you can, in theory, just create a temporary directory with WORKSPACE file and host the notebook there. It is actually not ideal for me, because I need to reference to libraries from within the repo to use in the notebook. If I made modifications to the library, I just need to restart the kernel. If I use SwiftPM setup (or go with separate WORKSPACE route), I have to setup these libraries in a separate repo and push a new commit to use.

All these, not to mention that not many Swift projects have no Bazel BUILD files, and the support to directly use SwiftPM Package.swift through https://github.com/cgrindel/rules_spm is not well-tested for non-Apple platforms.

philipturner commented 2 years ago

The only difference between debug and fastbuild is that latter strips debug symbols so it is leaner.

That's basically -c release -Xswiftc -Onone, which is debug mode without debug symbols. That shortens S4TF's compile time from 3 to 2 minutes, and I tried really hard to get that build option supported in a Colab notebook.

It definitely seems like Bazel is different enough from SwiftPM to warrant a %bazel magic command. I wasn't thinking with that in mind when saying I would support Bazel in Swift-Colab. It seems like a lot of work to support, and the benefits aren't that significant - you stated that it's not popular.

I think I'll leave the work of Bazel support to your fork for the time being. I'm not ruling it out; I just have finite time and other priorities right now. Swift-Colab v2.3 can just focus on the newest build system modifications to support compiling S4TF on release toolchains.

philipturner commented 2 years ago

I linked your repository in the README - would you mind reading my additions (+ drop-down) and telling me if there's anything that can be written better? The new text appears between the Table of Contents and Getting Started.

liuliu commented 2 years ago

Thank you! It reads great!

philipturner commented 2 years ago

I'd prefer if we interface with libraries using virtual function calls. For example, a top-level function could take an any JupyterDisplayProtocol as a parameter. The library would implement the header that defines the protocol, but the notebook would instantiate the conforming instance. This is similar to how I'm implementing Metal and OpenCL backends for S4TF, using any PluggableDevice as the main unit of computation.

For example, we might extend the JupyterDisplay protocol to something that can submit streamed images for video processing.

// Encoding happens behind the scenes; the calling library submits uncompressed frames.
// Always let the user decide the codec.
enum JupyterVideoCodec {
  case motionJPEG
  case automatic // backend chooses the fastest codec
  // case mpeg4 - might be added in the future
}

struct JupyterVideoConfig {
  var width: Int
  var height: Int
  var codec: JupyterVideoCodec
  // Best practice: rendering library intercepts this variable and sets the timestamp.
  var timestamp: Double?
}

protocol JupyterDisplayProtocol {
  ...
  /*static*/ func submitVideoFrame(bytes: Data, config: JupyterVideoConfig)

  // Perhaps also a feature for accepting keyboard/mouse input?
  /*static*/ func addEventListener(_ body: @escaping (???) -> Void)
  /*static*/ func setEventListener(for event: ???, _ body: @escaping (???) -> Void)
}

A library function for rendering might do this:

import JupyterDisplayModule // pre-installed into the system

func renderGLVideo(display: any JupyterDisplayProtocol, config: JupyterVideoConfig) {
  var receivedEvent = false
  display.addEventListender { receivedEvent = true }

  // Loop ends when the user manually interrupts cell execution, so the program should save game simulation progress accordingly.
  while true {
    usleep(16_000) // Interrupt might stop the program here.
    if receivedEvent { /*Do something*/ }

    // Base its size on the video config's `width` and `height`.
    let rawFrameBuffer = Data(...)

    var currentFrameConfig = config
    currentFrameConfig.timestamp = ... // Optional; just enhances encoder's ability to synchronize frames.
    display.submitVideoFrame(bytes: rawFrameBuffer, config: currentFrameConfig)
  }
}

Usage in a Jupyter notebook:


import SomeLibrary
setupSimulation()
// Allow interop between IPythonDisplay <-> JupyterDisplay either by:
// (a) Making them typealiases of each other.
// (b) Using conditional compilation for Colab vs. local notebook.
// (c) Retaining `IPythonDisplay` for backward-compatibility, but introducing a new
//     `JupyterDisplay` type to Swift-Colab.

let config = JupyterDisplayConfig(
  width: 720,
  height: 1280,
  codec: .motionJPEG)
// Several possibilities for how to input the local display variable:
renderGLVideo(display: IPythonDisplay.current, config: config)
renderGLVideo(display: JupyterDisplay.current, config: config)
renderGLVideo(display: jupyterDisplay, config: config)
renderGLVideo(display: JupyterDisplay.self, config: config)

// Alternatively, the user can render frames with custom code:
while true {
  usleep(...)
  JupyterDisplay.current.submitVideoFrame(...)
}
liuliu commented 2 years ago

Yes. The inspiration is mostly from Python side, where things like matplotlib or tdqm can identify Jupyter environment automagically and act accordingly.

Having explicit pass in from library point view is great. But I maintain that also provides a singleton access would enable that "automagic" feeling for library users (they don't need to pass in anything to get the benefit).

Thinking about the tdqm use cases, if we have a Swift equivalent:

for i in tdqm(someCollection) {
}

that can render progress bar in Jupyter correctly, it is automagic. If we require a pass in:

for i in tdqm(someCollection, display: IPythonDisplay.current) {
}

it is less so magical, but OK. However, if we wrap it one layer deeper:

final class SomeTrainer {
  public func train() {
    for i in tdqm(data) {
    }
  }
}

let trainer = SomeTrainer()
trainer.fit()

comparing with:

final class SomeTrainer {
  public func train(display: any JupyterDisplayProtocol) {
    for i in tdqm(data, display: display) {
    }
  }
}

let trainer = SomeTrainer()
trainer.fit(display: IPythonDisplay.current)

Much more needle to threading through now.

I think for libraries that require such rigor, having dependency passed in is great. But for libraries that don't, have an escape hatch (a global singleton access) would be good too.

philipturner commented 2 years ago

In that case, I could just create a system library with a concrete type called JupyterDisplayManager or JupyterContext. Either a package or the notebook itself can import the dependency, and everyone accesses the same singleton variable. I should probably pipe the framebuffer through a file in case the package and notebook open two different .so instances, or to offload serialization/networking to a background process. I'll have to find some good way to set settings like resolution, if that's even needed.

We may not need a protocol to encapsulate the functionality because there's one concrete type. Would you be fine if your current rendering notebooks are not backward-compatible, and have to be written after the final iteration of this design?

Also, I should just typealias JupyterDisplay == IPythonDisplay in Colab and IPythonDisplay == JupyterDisplay in local notebooks so that notebooks run everywhere, despite the fragmented history. That necessitates creating something new for the functionality we'll add, hence alternative names like JupyterContext.

liuliu commented 2 years ago

I think you focused too much on streaming aspect of it :) The proposal for protocolize it is to encapsulate notebook differences, if there are any, when displaying images or text or htmls. For example, the concrete implementation in mine EnableJupyterDisplay.swift involves format the message and pipe it through stdout, it will then be retrieved by the Swift kernel and send to the notebook. It may not be the same for your implementation or IPythonDisplay.

philipturner commented 2 years ago

Then we can just add static methods to the JupyterDisplay type object from the old EnableJupyterDisplay.swift, then do typealias IPythonDisplay == JupyterDisplay. That solves API differences between Colab and local notebooks, and it's even source-compatible with your existing Jupyter notebooks for simulation.

We already have:

https://github.com/philipturner/swift-colab/blob/405e27197757f545457fb0bd8e049666ffd6b40c/Sources/include/EnableIPythonDisplay.swift#L81-L93

How about we change your JupyterDisplay.swift to the following, then seeing whether your test notebooks still work? You may have to fill in empty methods with the actual implementation logic.

Click to expand ```swift /// A display manager for Jupyter notebook. /// /// This provides necessary methods for other Swift libraries to introspect whether you can /// display from within a Jupyter notebook, and methods to send HTML, images or texts over /// to the notebook. public enum JupyterDisplay { /// flush display messages to the notebook immediately. By default, these messages are /// processed at the end of cell execution. public func flush() {} /// Display a base64 encoded image in the notebook. public func display(base64EncodedPNG: String, metadata: String?) {} /// Display a snippet of HTML in the notebook. public func display(html: String, metadata: String?) {} /// Display plain text in the notebook. public func display(text: String, metadata: String?) {} /// Check whether the notebook is active and accepting display messages. public private(set) var isEnabled: Bool = false /// Get the cell number. public var executionCount: Int { 0 /*TODO: Return actual cell count*/ } public func enable() { // TODO: Some other logic JupyterDisplay.isEnabled = true } } extension JupyterDisplay { /// Display a base64 encoded image in the notebook. public func display(base64EncodedPNG: String) { display(base64EncodedPNG: base64EncodedPNG, metadata: nil) } /// Display a snippet of HTML in the notebook. public func display(html: String) { display(html: html, metadata: nil) } /// Display plain text in the notebook. public func display(text: String) { display(text: text, metadata: nil) } } public typealias IPythonDisplay = JupyterDisplay ```

Then %include JupyterDisplay.swift becomes:

import JupyterDisplay

extension JupyterDisplay {
  static func foo() {} /* methods for communicating with the kernel */
}

JupyterDisplay.enable()

Also, unify IPythonDisplay.swift and JupyterDisplay.swift by parsing the filename before executing the %include command. If it matches IPythonDisplay.swift, redirect it to JupyterDisplay.swift. I can easily change Swift-Colab to use a new file for IPythonDisplay.

philipturner commented 2 years ago

Also, would you mind reimplementing your logic for HTTP streaming using Python libraries, with PythonKit installed as a system library (or regular package dependency for now) for all conformant packages? Then source code would work across both Colab and local notebooks, swapping the two different implementations at runtime. You could try different build configurations if this doesn't work out verbatim.

#if canImport(PythonKit)
let jpegDependency1 = try? Python.attemptImport("jpegDependency1")
#endif
#if canImport(NIOCore)
// alternative implementation
#endif

That puts the burden of streaming logic on the package designer, not the Jupyter kernel. Maybe the kernel can include streaming logic in the future, but this approach provides more flexibility at the moment.

liuliu commented 2 years ago

Also, unify IPythonDisplay.swift and JupyterDisplay.swift by parsing the filename before executing the %include command. If it matches IPythonDisplay.swift, redirect it to JupyterDisplay.swift. I can easily change Swift-Colab to use a new file for IPythonDisplay.

We may talk about slightly different things. I am not interested in how the file we use for %include "EnableJupyterDisplay.swift" was written. I am interested in having a "protocol" library that other libraries can depend on to provide display support. For example, now, the MuJoCoViewer.swift inside my Gym package https://github.com/liuliu/s4nnc/blob/main/gym/renders/MuJoCoViewer.swift#L32 can detect whether Jupyter notebook is on or not and choose to display differently.

Thus, when you use it in the notebook, it looks like this:

import Gym

%include "EnableJupyterDisplay.swift"

let viewer = MuJoCoViewer()

viewer.render()

If it is not in notebook, this will create a window and render a image. If it is in notebook, this will pipe an image through the notebook.

To implement this kind of support, JupyterDisplay token has to be a global instance, rather than a enum. Like you said, it is possible still to not provide the global instance, and just pass the instance in:

viewer.render(with: IPythonDisplay.current)

In this way, we still provides a protocol, but library doesn't need to talk to the global instance, and feels "better engineered". As I elaborated earlier, this may not feel as magical as Python, so it is a balanced act.

philipturner commented 2 years ago

So what you're saying is:

That package has a mutable variable called JupyterDisplay. By default, it's an instance that doesn't do anything. At runtime, the file EnableJupyterDisplay.swift will mutate the module's variable, changing it to something that is enabled. For example:

// EnableJupyterDisplay.swift
import JupyterDisplay

// Renamed the existing `JupyterDisplay` type to `_JupyterDisplay`
enum _JupyterDisplay {
  ... // already existing code
}
extension _JupyterDisplay: JupyterDisplayProtocol {}
JupyterDisplay.JupyterDisplay = // something that delegates to `_JupyterDisplay`
// EnableIPythonDisplay.swift
import JupyterDisplay

enum IPythonDisplay {
  ... // already existing code
}
extension IPythonDisplay: JupyterDisplayProtocol {}
JupyterDisplay.JupyterDisplay = // something that delegates to `IPythonDisplay`

That would be very flexible, although I'm not sure we should expose a global variable like that. Perhaps we at least make it immutable at face-value, but provide an API to mutate it. For example:

extension JupyterDisplayProtocol {
  public func _setJupyterDisplayInstance(_ instance: any JupyterDisplayProtocol) {
    // Optionally do some cleanup on the existing instance.
    // Optionally do some validation on the input instance.
    JupyterDisplay = instance
  }
}

I think I'm getting it now, but I'd like to consider alternative approaches as well. For example, we could substitute a protocol with an API that just swaps out function pointers, which a static JupyterDisplay type object calls into. Capitalizing something and then making it a variable goes against traditional naming conventions (it should start with lowercase). If we can accomplish the same library interface while sticking more to the semantic meaning, I would like that.

For example, we could do the following:

protocol JupyterDisplayProtocol {
  // non-static function
  func display(base64EncodedPNG: ...)
}

enum JupyterDisplay {
  private var delegate: JupyterDisplayProtocol = EmptyJupyterDisplay()
  public func _setCurrentInstance(_ instance: JupyterDisplayProtocol) {
    Self.delegate = instance
  }

  // static function
  static func display(base64EncodedPNG: ...) {
    delegate.display(base64EncodedPNG: base64EncodedPNG)
  }

  static var isEnabled: Bool {
    delegate.isEnabled
  }
}
philipturner commented 2 years ago

You stated above that we could make a neutral organization for hosting the Swift package. I'm not yet sure that's the best path forward. If it's specific to the Jupyter kernel and maintained alongside the Jupyter kernel, perhaps it should be embedded inside each Jupyter kernel's repository. Creating an external repo is an irreversible decision; if we change our mind later, people now have Swift package dependencies to a stale repository (or API breakage). Would you mind trying out alternative options first?

In previous comments, I had the ideas of embedding JupyterDisplay as a system library. That means the JuypterDisplay module doesn't exist in a non-Jupyter environment. You need to doubly check that the JupyterDisplay module exists, and the current JupyterDisplay instance is enabled.

#if canImport(JupyterDisplay)
import JupyterDisplay
#endif

extension MuJoCoViewer: Renderable {
  public func render() {
    #if canImport(JupyterDisplay)
    if JupyterDisplay.isEnabled {
      // Check to see if we launched the render server yet.
      if httpChannel == nil {
        httpChannel = try? renderServer.bind(host: "0.0.0.0", port: .random(in: 10_000..<20_000))
          .wait()
      }
      if JupyterDisplay.executionCount != renderCell {
        JupyterDisplay.display(html: renderServer.html)
        JupyterDisplay.flush()
        renderCell = JupyterDisplay.executionCount
      }
    }
    #endif
    ...
  }
}

Another concern is how to detect a module's dependency to the JupyterDisplay repository. My SwiftPM engine doesn't yet have introspection capabilities. Also, if two separate libraries depend on the same JupyterDisplay repo (possibly with different package versions), which one do I import into LLDB? Or what if someone installs a Swift package after including "EnableJupyterDisplay.swift"? The file can only run once in Swift-Colab, and it ran when the JupyterDisplay module couldn't be imported. It lost its chance to mutate the JupyterDisplay.JupyterDisplay variable. This problem would be solved by making JupyterDisplay an internal rather than external module.

philipturner commented 2 years ago

Also, why can't we use IPythonDisplay in local Jupyter notebooks? Your README said it crashes after a few dozen cell executions, but I couldn't find any more information on the crash. EnableIPythonDisplay.swift is simpler to implement than EnableJupyterDisplay.swift, which requires cryptography and compiling large external dependencies. Perhaps we could replicate the crash in Colab, and I could determine what's causing it.

liuliu commented 2 years ago

I think I'm getting it now, but I'd like to consider alternative approaches as well. For example, we could substitute a protocol with an API that just swaps out function pointers, which a static JupyterDisplay type object calls into. Capitalizing something and then making it a variable goes against traditional naming conventions (it should start with lowercase). If we can accomplish the same library interface while sticking more to the semantic meaning, I would like that.

Yes, there are several ways to implement this. It is a kind of dependency injection pattern. The important thing is to expose a protocol-only interface module to other libraries and the heavy-weight dependency can be injected from notebook end. Your enum example would be an alternative way to implement the same thing.

Another concern is how to detect a module's dependency to the JupyterDisplay repository. My SwiftPM engine doesn't yet have introspection capabilities. Also, if two separate libraries depend on the same JupyterDisplay repo (possibly with different package versions), which one do I import into LLDB? Or what if someone installs a Swift package after including "EnableJupyterDisplay.swift"? The file can only run once in Swift-Colab, and it ran when the JupyterDisplay module couldn't be imported. It lost its chance to mutate the JupyterDisplay.JupyterDisplay variable. This problem would be solved by making JupyterDisplay an internal rather than external module.

That's correct, SwiftPM doesn't support system module detection, due to it is a very cross-compilation friendly nature (iOS always cross-compiled on a macOS system). That's why have a protocol-only library as the root dependency is ideal actually. The order of importing and include "EnableJupyterDisplay.swift" is not important because the injected dependency only become relevant upon use. Luckily, in Swift, there is no eager initialization for library level code, unlike Python or C++ (where you can run code when import __init__.py, or at startup for your static initializers (C++)). In Swift, static initializers are all lazy. Now I see what you mean, yes, importing the protocol-only library needs to happen before "EnableJupyterDisplay.swift". That's why I made the EnableJupyterDisplay.swift call automatically after the library import.

Also, why can't we use IPythonDisplay in local Jupyter notebooks? Your README said it crashes after a few dozen cell executions, but I couldn't find any more information on the crash. EnableIPythonDisplay.swift is simpler to implement than EnableJupyterDisplay.swift, which requires cryptography and compiling large external dependencies. Perhaps we could replicate the crash in Colab, and I could determine what's causing it.

TBH, I am a bit burned by all weirdness of Swift REPL. The last shenanigan I found is that by linking to liblinear, for loop will crash with Swift 5.6.x REPL. The EnableIPythonDisplay.swift problem with Swift 5.5.x (I haven't tried with 5.6.x) is not exactly crash. After a few cell executions, the Swift kernel will die if EnableIPythonDisplay.swift is called prior. The kernel is dead with the lldb exited with code 0. Thus, if you attach a debugger, it doesn't really show why. At that time, I guess it may due to the fact that it doesn't like a IPython kernel (launched by the EnableIPythonDisplay.swift code) connected along with a Swift kernel, but because it doesn't really "crash", I have limited way to debug.

philipturner commented 2 years ago

I'm deciding to take a break from my Metal backend for S4TF, and complete Swift-Colab v2.3 now. For starters, I resolved #17. I'm also going through the documentation and slowly completing it, which I've been too lazy to do for a while. I plan to add two new magic commands:

The current documentation of %install-swiftpm-import is in MagicCommands.md. Behind the scenes, I'm going to prepend a call to this command for JupyterDisplay (with a special execution path), before compiling every Swift package. The user will not have to (and cannot) declare the import in the Jupyter notebook. This makes logical sense because you never exported JupyterDisplay with an %include command. Are you fine with this being how we inject a dependency to JupyterDisplay, once liuliu/swift-jupyter and philipturner/swift-colab converge on the matter?

liuliu commented 2 years ago

The current documentation of %install-swiftpm-import is in MagicCommands.md. Behind the scenes, I'm going to prepend a call to this command for JupyterDisplay (with a special execution path), before compiling every Swift package. The user will not have to (and cannot) declare the import in the Jupyter notebook. This makes logical sense because you never exported JupyterDisplay with an %include command. Are you fine with this being how we inject a dependency to JupyterDisplay, once liuliu/swift-jupyter and philipturner/swift-colab converge on the matter?

I need to learn more about differences between what you called "environment" and "import". import JupyterDisplay automagically sounds good (from my understanding, this will make it simpler later when we %include EnableXXXDisplay.swift?). This change doesn't change the dependency tree for other SwiftPM packages right? (If they don't use JupyterDisplay, it is not one of their dependencies, and if they use, it doesn't matter if we import JupyterDisplay for they automagically or not?).

I see what you mean now. This allow Colab to inject a concrete implementation directly (or not, can still be a protocol-only library and inject later, but it seems to be the purpose). I need to think a bit to weigh the pros and cons for this approach. It definitely have tensions with how traditionally dependency-injection was done, or with how Bazel / SwiftPM manages its dependencies (both prefers concrete dependencies, rather than phantom ones (i.e. the dependencies declared in Package.swift, rather than implicit ones doesn't show up anywhere). For one, it makes unit tests your code that specializes with JupyterDisplay harder to do (previously, you can inject your FakeJupyterDisplay and run swift test just fine, now it is not obvious how you can do it with SwiftPM alone).

philipturner commented 2 years ago

It definitely have tensions with how traditionally dependency-injection was done

If possible, I encourage the user to create a concrete dependency in Package.swift. This is a solution to a problem that cannot be solved any other way. When compiling S4TF, it takes 3 minutes. When compiling both s4tf/s4tf and s4tf/models, you would compile S4TF two times, because each repository makes their own copy of the TensorFlow module.

This feature augments the "clever idea" to bypass S4TF's long compile time, where you rewrite a notebook's contents without recompiling S4TF. A user would have just looked through several tutorials with their one S4TF instance compiled. Now they encounter the tutorial that runs s4tf/models, and have to wait 3 extra minutes because that GitHub repo has its own copy of s4tf/s4tf. With %install-swiftpm-import, someone could reuse the existing TensorFlow module inside s4tf/models and skip the 3 minutes of wait time.

Other than this specific use case, and JupyterDisplay, I can't see a reason someone would use the feature.

For one, it makes unit tests your code that specializes with JupyterDisplay harder to do

My unit tests are a bunch of Colab notebooks (see Testing). You test that the code works by firing up Swift-Colab and running the notebook cells. This setup does not permit running swift test on any dependencies downloaded by a notebook. Instead, you copy the function bodies of each XCTestCase member into its own Jupyter cell, encapsulated in a do {} statement for safe measure. I'm oversimplifying the process, but this is the general idea.

However, I would like to add a dedicated SwiftPM testing feature. I often set certain environment variables before running SwiftPM tests, which %install-swiftpm-environment would help with.

liuliu commented 2 years ago

This is a solution to a problem that cannot be solved any other way. When compiling S4TF, it takes 3 minutes. When compiling both s4tf/s4tf and s4tf/models, you would compile S4TF two times, because each repository makes their own copy of the TensorFlow module.

You probably already thought through this many times, but have you tried to just support a generic cache directive? This is often supported in CI environments to make CI jobs faster. SwiftPM have everything checked out and compiled under .build directory. Since Colab has a pretty consistent environment, you can provide a GDrive directory with compiled .o / .a files for S4TF that people can do %checkout-cache SomeGDriveLink .build to populate .build directory.

philipturner commented 2 years ago

I recently though of using Google Drive to cache build products, but it doesn't solve the problem that necessitates %install-swiftpm-import. My SwiftPM engine allows for loading each package dynamically because they all have separate .build directories. It would cache the S4TF builds, but the user still has to wait 6 minutes on their very first time ever using Swift-Colab, instead of 3. Also, not everyone's going to make an alt Google account just to create a scratch Google Drive (connecting to your actual Google account is a very bad security risk).

This is unlike the swift-jupyter or Swift-Colab pre-v2.0 mechanism, which I assume your fork is still using. With that mechanism, all packages compile into one mega-package, so no dependencies duplicate. There is one .build directory for everything combined.

Google Drive caching seems like a good analogue to local filesystem caching, because it stores data over long periods of time. I am revamping the experience of mounting Google Drives, and just bypassed a major restriction around calling google.colab.drive.mount (still only a POC). Perhaps that feature can become part of Swift-Colab v2.3, and help me plan for local notebook support.

philipturner commented 2 years ago

The restriction on Google Drive was that I couldn't mount the drive from Swift code. I had to switch over to the Python Jupyter kernel, mount the drive, then switch back. I am bypassing this mechanism by forwarding stdin requests from the LLDB process to the SwiftKernel process and back - through a custom piping mechanism. This removed the need to cache and asynchronously execute send_multipart's in KernelCommunicator, thus I may delete the kernel communicator soon. Would your fork be able to remove the kernel communicator and directly copy image data into a Jupyter display?

For reference, here is the serialization mechanism. Encode 8 bytes, which are an Int64 containing how large the payload is. Then, encode the payload. The recipient deserializes the data and moves the file's cursor upward. At the end of each Jupyter cell, delete and recreate the files that facilitate inter-process communication. Otherwise, they may grow infinitely large.

liuliu commented 2 years ago

I recently though of using Google Drive to cache build products, but it doesn't solve the problem that necessitates %install-swiftpm-import. My SwiftPM engine allows for loading each package dynamically because they all have separate .build directories. It would cache the S4TF builds, but the user still has to wait 6 minutes on their very first time ever using Swift-Colab, instead of 3. Also, not everyone's going to make an alt Google account just to create a scratch Google Drive (connecting to your actual Google account is a very bad security risk).

I was actually talking about provide a public accessible cache, so whether it is compiled separately or compiled into a megapackage will be the same. And for first-time user, it will be as fast because you pull that from public cache everyone uses. Whether it is from GDrive or a S3 bucket is an implementation detail.

I am bypassing this mechanism by forwarding stdin requests from the LLDB process to the SwiftKernel process and back - through a custom piping mechanism. This removed the need to cache and asynchronously execute send_multipart's in KernelCommunicator, thus I may delete the kernel communicator soon. Would your fork be able to remove the kernel communicator and directly copy image data into a Jupyter display?

Similar, I haven't removed the KernelCommunicator method yet, as it seems more efficient (avoid formatting and copying). But I do provided an alternative flush mechanism to enable send display messages to notebook in the middle of execution. The implementation is similar, I wrapped the JSON payload with ---SOMEBOUNDARYSTRINGS and the stdout reader (from Swift kernel, reads LLDB process) will parse this particular stdout and treat it as a notebook message instead.

philipturner commented 2 years ago

I was actually talking about provide a public accessible cache,

I don't think that's a very good idea. You would have to cache exponential variations of different exact Swift toolchain + different build flags + different version of the external Swift package. I'm not even sure you can distinguish each unique environment. If one factor is different, the .build folder could have different contents. That's why I'm so hesitant to make a public cache as described in this thread.

However, making a cache just for the environment combos found in tutorial notebooks could artificially boost performance. It would make tutorials go faster for beginners, but it doesn't augment power users' workflows. In other words, it's optimizing for benchmark performance.

Similar, I haven't removed the KernelCommunicator method yet, as it seems more efficient (avoid formatting and copying).

KernelCommunicator pipes the data through the file system, which has the same overhead as the new mechanism I propose (which also pipes). I think it also serializes the data before copying, which is what my new mechanism does. KernelCommunicator got painful when I had to marshal the data through LLDB's C++ API, rather than its Python API. I had to make a recursive serialization protocol just for one Swift <-> C++ bridging function, which used exponentially expanding buffers while encoding - written in straight C.

liuliu commented 2 years ago

I don't think that's a very good idea. You would have to cache exponential variations of different exact Swift toolchain + different build flags + different version of the external Swift package. I'm not even sure you can distinguish each unique environment. If one factor is different, the .build folder could have different contents. That's why I'm so hesitant to make a public cache as described in this thread.

Depends on how SwiftPM's .build directory structured, I didn't read too much into their code to know how they keep these directories up-to-date. Sharing cache for Bazel is safe and efficient (due to content hashing). If the cache is on object / library level, it will benefit both beginners and power users because the later can still reuse big chunk of compiled code. Anyway, SwiftPM is out of my expertise, just want to clarify some of these claims might be build system specific, not an inherited limitation of doing public read-only cache itself.

KernelCommunicator pipes the data through the file system, which has the same overhead as the new mechanism I propose (which also pipes). I think it also serializes the data before copying, which is what my new mechanism does. KernelCommunicator got painful when I had to marshal the data through LLDB's C++ API, rather than its Python API. I had to make a recursive serialization protocol just for one Swift <-> C++ bridging function, which used exponentially expanding buffers while encoding - written in straight C.

Hmm, that doesn't match what I saw in KernelCommunicator.swift inherited from google/swift-jupyter (https://github.com/google/swift-jupyter/blob/main/KernelCommunicator.swift#L60)? For that file, it serializes into JSON and pass the pointer of that serialized JSON back to LLDB. In LLDB, it can use read memory method to read it back. By saving one memcpy, I meant that. But I guess it may not worth it though.

philipturner commented 2 years ago

In LLDB, it can use read memory method to read it back

I was thinking that the LLDB process is a separate system-level process, which you spawn and communicate with over pipes. I wasn't thinking correctly, because LLDB does run everything in the current process; it's just a bunch of library calls. That means you could access stuff stored in raw memory. However, the lifetime overhead of each data segment is probably dominated by disk writing/piping in kernel.iopub_socket.send_multipart. The worst I could do is double the overhead (roughly), but that's outweighed by having much greater flexibility.

liuliu commented 2 years ago

I was thinking that the LLDB process is a separate system-level process, which you spawn and communicate with over pipes. I wasn't thinking correctly, because LLDB does run everything in the current process; it's just a bunch of library calls. That means you could access stuff stored in raw memory. However, the lifetime overhead of each data segment is probably dominated by disk writing/piping in kernel.iopub_socket.send_multipart. The worst I could do is double the overhead (roughly), but that's outweighed by having much greater flexibility.

Yeah, once it is to iopub_socket.send_multipart, I am oblivious what kind of communication mechanism they use (I think they use ZeroMQ encoded named pipes, but whatever). I am more interested in how we can get from notebook code (in Swift) to the kernel (in my case, Python, in your case, Swift with PythonKit). Currently I am happy with my ad-hoc stdout mechanism, but definitely can learn if there are better ways.

philipturner commented 2 years ago

I plan to add a %test directive in this round of changes to Swift-Colab. The documentation link is here. It allows running SwiftPM tests, but you can't specify any [PRODUCT ...] to export. The test build products will be cached between executions in a containerized folder. Let me know what you think!

I should have explained this, but I like to hop around between components of the S4TF project. It's so massive in scope, that I can get bored of one part and move on to another. When I work on Swift-Colab, I usually spend a week or two making several disparate changes, culminating in a minor release. The changes I'm discussing in this thread are achievable within that time frame, but local notebook support is not. Thus, local notebook support will come in the round after this, circa early 2023.

I'm also simultaneously building SwiftOpenCL* in staggered intervals (April 2022, June/July 2022, probably September 2022), patching up s4tf/s4tf gradually, and creating the Metal backend for S4TF. More info here, which is linked on my GitHub profile. A few days ago, I switched gears from the Metal backend to Swift-Colab.

*This project reused the dynamic loading mechanism from PythonKit quite masterfully.