Swift for TensorFlow Resurrection: Swift running on Colab again

This weekend, I've been working on how to sideload Swift on Google Colab (repo: philipturner/swift-colab). Eventually, this will turn into loading Swift for TensorFlow as a Swift package, pre-compiled as a binary target instead of a toolchain. I got the point where I can pass an arbitrary string of Swift code as a Python string, then compile and run it.

The next step is to install PythonKit into the Colab virtual machine and support SwiftPM. I was exploring the Swift compiler and Swift Package Manager, and came across this tutorial.

After getting it to work on my Mac, I decided to try it out again on Linux. I realized that having the ability to carry out the commands in Swift instead of Shell had some advantages (I could write a string literal to a file without learning the Shell command for that :smiley:), and Google Colab is easier to setup than Docker. Also, it was exciting to see Swift code execute something generally useful in Google Colab!

Below is the link to the Colab notebook. Anyone can make a copy and run the program. The first code block is modified to pull from the save-1 branch of the GitHub repository, which will stay stable unlike the main branch.

import swift

swift.run('''
import Foundation
let fm = FileManager.default

... // do some stuff

func doCommand(args: [String]) throws {
    let command = Process()
    command.executableURL = .init(fileURLWithPath: "/usr/bin/env")
    command.arguments = args
    try command.run()
    command.waitUntilExit()
}

try doCommand(args: ["swiftc", "-D", "DEBUG", "point.swift", "main.swift", "-o", "point-app"])
try doCommand(args: ["./point-app"])
''')

Output:

debug mode
Hello world! 4 20

I can't guarantee that side-loading will break out of passing strings to Python and support the full Jupyter notebook experience. Regardless, sideloading makes it possible to utilize cloud GPUs and TPUs in the new Swift for TensorFlow. I can validate that the upcoming Metal backend doesn't break CUDA support.

For more context about the Swift for TensorFlow resurrection:

24 Likes

I should be able to call Swift functions from Python. I can take the memory address of an Python object via id(object) and spawn a process using the address as an argument. Then, a pre-compiled Swift executable can transform the memory address into a PythonObject using a modified fork of PythonKit.

This form of bridging is needed to allow subclassing the Jupyter Kernel class, while implementing the kernel's logic in Swift instead of Python. Currently, the Swift Jupyter kernel from @marcrasi is written almost entirely with Python (in the swift_kernel.py file).

4 Likes

I got to the point where I can call a Swift function from Python! I made a C-compatible Swift function, compiled that into a dynamic library (.so file), then loaded that file using the Python ctypes library. I passed the id of a Python string into the C-like Swift function, and verified that it was a memory address!

The next step is to compile my fork of PythonKit and convert that reference back into a Python string! I'll post the source code for this once I have the two-way interface between Swift and Python fully working.

3 Likes

I got PythonKit to compile on Google Colab. All I have left is to make Swift executables accessible to Python, then I can subclass the Jupyter kernel and (hopefully) resurrect full Swift support in Colab.

5 Likes

@philipturner I have a question for you/sort of a challenge/I don't know if it is possible... but I would think about how the swift project can have a test to make sure that whatever you are fixing still works.

Do you have any thoughts on how we could do that? Otherwise there isn't a guarantee that the breakage will not reappear.

I am eager to help you out. Colab support is still unstable, and I'm nowhere near finished. The first paragraph of your reply wasn't very clear due to how it was worded. Could you reiterate that?

Sorry, my first sentence means that I am unfamiliar with what you are exactly doing and what that would mean in terms of testing it.

For instance, I am not sure if it is possible to test on swift's CI whatever you are fixing. The reason I am saying it is a challenge is that I am challenging you to think if it is possible = p.

3 Likes

It is possible to set up a semi-automated test. Since Colab requires a Google account and authorization to run, I don't know whether a GitHub YAML bot could use it. However, you could automate most of the workflow and manually trigger the test on a personal Colab notebook, much like you manually trigger a @swift-ci workflow.

Modify my build script to pull from the nightly build, and add some Swift scripts that you can download into a Colab notebook. You need to be cautious about build times, as you will be kicked off of Colab if one task takes too long. However, what you are thinking of is entirely doable.

Once I have Colab fully functional again (or the best it can possibly get), I'm going to switch to MetalXLA and be fully occupied with finishing it before PyTorch releases their Metal backend. When that happens, you could work on some testing scripts and try them out in Colab.

1 Like

I got libSwiftPythonBridge.so and libPythonKit.so to successfully link to a Swift script, then compiled and executed that script. Soon, I should be able to replace the Shell command for compiling a Swift string with a call into SwiftPythonBridge's C-compatible interface from Python. This might make the output from print(...) in Swift synchronize with the Jupyter kernel.

3 Likes

I got around a major bug with the Global Interpreter Lock, which had me really scared for a moment (I need to use PyDLL instead of CDLL in ctypes). I can now execute a Swift script by calling directly into the dynamic library's C interface. However, every call to print(...) goes to a hidden system output, instead of the Jupyter notebook's output. Previously, this is what happened, but I manually extracted that output and logged it to Jupyter's output.

Edit: Using Wurlitzer, I can restore output synchronization to how it was before, although you still have to wait until all of the code executes before reading any output.

1 Like

I can now subclass Python objects and coordinate the logic of their methods to Swift (repository save point #3). The next step will be subclassing the Jupyter kernel and finding what restrictions Google added to it in March.

Python code temporarily included in the swift Python package:

class SwiftInteropTestSuperclass:
    pass
    
class SwiftInteropTest(SwiftInteropTestSuperclass): 
    def __init__(self):
        self.swift_delegate = SwiftDelegate()
        
    def example_func(self, string_param):
        return self.swift_delegate.call("example_func", [self, string_param])
    
    def example_func_2(self, string_param):
        return self.swift_delegate.call("example_func_2", { string_param: self })

Swift counterpart:

import PythonKit
import SwiftPythonBridge // this module is internal to Swift-Colab
let swiftModule = Python.import("swift")

let interopTest = swiftModule.SwiftInteropTest()

interopTest.registerFunction(name: "example_func") { param -> Void in
    print("example_func called from Python with param \(param)")
}
            
interopTest.registerFunction(name: "example_func_2") { param -> PythonConvertible in
    print("example_func_2 called from Python with param \(param)")
    return String("return value")
}

print(interopTest.example_func("Input string for example_func"))
print(interopTest.example_func_2("Input string for example_func_2"))

Output:

example_func called from Python with param [<swift.SwiftInteropTest object at 0x7f6c20490a90>, 'Input string for example_func']
None
example_func_2 called from Python with param {'Input string for example_func_2': <swift.SwiftInteropTest object at 0x7f6c20490a90>}
return value
2 Likes

I translated the register.py file in google/swift-jupyter to Swift and it runs without crashing. I still need to translate swift_kernel.py, which is larger and likely what's affected by Google's restrictions.

1 Like

I got to the point where I can alter the behavior of Google Colab, making it output whatever the code cell puts as input. I had to manually overwrite some Python code, then restart the Jupyter runtime. Now that I got to this point, I'm very confident I can follow through all the way and bring back Swift support to its state before the death of S4TF. Code completion, syntax coloring, everything.

5 Likes

Syntax coloring is working. All I had to do was clone an S4TF tutorial, inspect its metadata, and copy that over to a blank notebook.

Open the notebook template in Colab, and the text is syntax-colored like Swift instead of Python. For example, an import statement is blue and green (Swift) instead of purple and white (Python).

3 Likes

Swift on Google Colab has entered the beta stage! Executing code is still in the works, but you can follow the steps of side-loading and prepare for when it's feature-complete. Check out the README or the Colab notebook:

8 Likes

Translated the StdoutHandler from the original Jupyter kernel's swift_kernel.py to Swift (won't get around to testing it for quite a while):

Just finished translating the entire Swift kernel from Python to Swift. Now, it's time for heavy testing and experiencing many painful bugs :frowning:.

7 Likes

The S4TF team was right. There were major restrictions on LLDB, but I just bypassed them!

2 Likes

Due to some major additions to the Swift Package manager since S4TF died, I need to rework the Jupyter kernel's Swift package loader. I'm aiming to remove the restriction that you can't execute % commands outside of the first cell.

By the way, I got a very basic line of Swift code to execute in Colab.

Int.bitWidth
64
4 Likes

Swift-Colab is complete! Several tutorials from Swift for TensorFlow have been tested on it, and the Python unit testing suite has been transformed into a series of Colab notebooks. These future-proof it by allowing you to test specific Swift versions, ensuring Colab support is never dropped again. Furthermore, I can test Swift 5.3 (the last version S4TF worked on) and the toolchain only takes 30 seconds to download because of how fast Google's internal servers are :smiley:.

Unlike before March 2021, this version of the Swift Jupyter kernel does not come with any libraries built in. You must explicitly import PythonKit and Differentiation. Soon, I'll patch up TensorFlow and allow that too. It will require a special installation command because it takes very long to compile and SwiftPM doesn't support pre-compiled Swift binaries yet.

Thanks @Michael_Gottesman for the suggestion to future-proof it. These tests must be manually run, but they don't take too long and I don't expect them to break often.

If anyone is interested in catching up on the effort to resurrect S4TF, here's a good repository for reference: GitHub - philipturner/resurrection-of-s4tf: Chronicling the resurrection of Swift for TensorFlow

14 Likes
Terms of Service

Privacy Policy

Cookie Policy