Revert "Support torchbind in OSS proxy executor (pytorch#149747)"

This reverts commit aa70d62.

Reverted pytorch#149747 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally ([comment](pytorch#149747 (comment)))

Author: pytorchmergebot

test/inductor/test_torchbind.py

@@ -128,7 +128,7 @@ def test_torchbind_hop_schema(self):
         schema = CallTorchBind.schema(foo_ir, "add")
         self.assertEqual(
             str(schema),
-            "call_torchbind(__torch__.torch.classes._TorchScriptTesting._Foo _0, str method, int _1) -> int _0",
+            "call_torchbind(__torch__.torch.classes._TorchScriptTesting._Foo obj, str method, int _1) -> int _0",
         )
 
     def test_torchbind_config_not_generated(self):
@@ -146,7 +146,7 @@ def test_torchbind_hop_schema_no_input(self):
         schema = CallTorchBind.schema(q_ir, "pop")
         self.assertEqual(
             str(schema),
-            "call_torchbind(__torch__.torch.classes._TorchScriptTesting._TensorQueue _0, str method) -> Tensor _0",
+            "call_torchbind(__torch__.torch.classes._TorchScriptTesting._TensorQueue obj, str method) -> Tensor _0",
         )
 
     def test_torchbind_hop_schema_no_output(self):
@@ -155,7 +155,7 @@ def test_torchbind_hop_schema_no_output(self):
         schema = CallTorchBind.schema(q_ir, "push")
         self.assertEqual(
             str(schema),
-            "call_torchbind(__torch__.torch.classes._TorchScriptTesting._TensorQueue _0, str method, Tensor _1) -> NoneType _0",
+            "call_torchbind(__torch__.torch.classes._TorchScriptTesting._TensorQueue obj, str method, Tensor _1) -> NoneType _0",
         )
 
     def test_torchbind_aot_compile(self):
@@ -250,7 +250,7 @@ def test_torchbind_aot_compile(self):
                                 "target": "call_torchbind",
                                 "inputs": [
                                     {
-                                        "name": "_0",
+                                        "name": "obj",
                                         "arg": {
                                             "as_custom_obj": {
                                                 "name": "_torchbind_obj0",
@@ -293,20 +293,15 @@ def test_torchbind_aot_compile(self):
                 self.assertTrue((tmp_path_model / "custom_objs_config.json").exists())
                 self.assertTrue((tmp_path_constants / "custom_obj_0").exists())
 
-    def test_torchbind_aoti(self):
-        ep, inputs, orig_res, _ = self.get_exported_model()
-        pt2_path = torch._inductor.aoti_compile_and_package(ep)
-        optimized = torch._inductor.aoti_load_package(pt2_path)
-        result = optimized(*inputs)
-        self.assertEqual(result, orig_res)
+        # TODO: add accuracy test after we support loading and running compiled models with
+        # torchbind objects.
 
     @torch._inductor.config.patch("aot_inductor.use_runtime_constant_folding", True)
     def test_torchbind_aot_compile_constant_folding(self):
-        ep, inputs, orig_res, _ = self.get_exported_model()
-        pt2_path = torch._inductor.aoti_compile_and_package(ep)
-        optimized = torch._inductor.aoti_load_package(pt2_path)
-        result = optimized(*inputs)
-        self.assertEqual(result, orig_res)
+        ep, inputs, _, _ = self.get_exported_model()
+        aot_compile(ep.module(), inputs, options={"aot_inductor.package": True})
+        # TODO: add accuracy test after we support loading and running compiled models with
+        # torchbind objects.
 
     def test_torchbind_list_return_aot_compile(self):
         class M(torch.nn.Module):
@@ -322,48 +317,15 @@ def forward(self, x):
 
         m = M()
         inputs = (torch.ones(2, 3),)
-        orig_res = m(*inputs)
 
         # We can't directly torch.compile because dynamo doesn't trace ScriptObjects yet
         with enable_torchbind_tracing():
             ep = torch.export.export(m, inputs, strict=False)
 
-        pt2_path = torch._inductor.aoti_compile_and_package(ep)
-        optimized = torch._inductor.aoti_load_package(pt2_path)
-        result = optimized(*inputs)
-        self.assertEqual(result, orig_res)
-
-    def test_torchbind_queue(self):
-        class Foo(torch.nn.Module):
-            def __init__(self, tq) -> None:
-                super().__init__()
-                self.tq = tq
-
-            def forward(self, x):
-                self.tq.push(x.cos())
-                self.tq.push(x.sin())
-                # TODO: int return type in fallback kernel not support yet
-                x_cos = self.tq.pop()  # + self.tq.size()
-                x_sin = self.tq.pop()  # - self.tq.size()
-                return x_sin, x_cos
-
-        inputs = (torch.randn(3, 2),)
-
-        q = _empty_tensor_queue()
-        m = Foo(q)
-        orig_res = m(*inputs)
-
-        q2 = _empty_tensor_queue()
-        m2 = Foo(q2)
-
-        # We can't directly torch.compile because dynamo doesn't trace ScriptObjects yet
-        with enable_torchbind_tracing():
-            ep = torch.export.export(m2, inputs, strict=False)
+        aot_compile(ep.module(), inputs, options={"aot_inductor.package": True})
 
-        pt2_path = torch._inductor.aoti_compile_and_package(ep)
-        optimized = torch._inductor.aoti_load_package(pt2_path)
-        result = optimized(*inputs)
-        self.assertEqual(result, orig_res)
+        # TODO: add accuracy test after we support loading and running compiled models with
+        # torchbind objects.
 
     @requires_gpu()
     @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)

torch/_higher_order_ops/torchbind.py

@@ -42,7 +42,9 @@ def schema(obj, method) -> torch.FunctionSchema:
         val = obj.get_real_obj()
         schema = val._get_method(method).schema
         schema_str = str(schema)
-        new_schema_str = f"call_torchbind({str(schema.arguments[0].real_type)} {schema.arguments[0].name},"
+        new_schema_str = (
+            "call_torchbind(" + str(schema.arguments[0].real_type) + " obj,"
+        )
         first_comma_index = schema_str.find(",")
         if first_comma_index == -1:
             # If no comma is found, find the last closing parenthesis

torch/csrc/inductor/aoti_package/model_package_loader.cpp

@@ -1,7 +1,6 @@
 #if !defined(C10_MOBILE) && !defined(ANDROID)
 
 #include <c10/util/error.h>
-#include <c10/util/string_view.h>
 #include <torch/csrc/inductor/aoti_package/model_package_loader.h>
 #include <torch/csrc/inductor/aoti_runner/model_container_runner.h>
 #include <torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h>
@@ -63,6 +62,7 @@ const std::string k_separator = "\\";
 #else
 const std::string k_separator = "/";
 #endif
+
 } // namespace
 
 namespace torch::inductor {
@@ -187,7 +187,7 @@ bool recursive_mkdir(const std::string& dir) {
   }
 
   // Find folder separator and check if we are at the top
-  auto pos = dir.find_last_of(k_separator);
+  auto pos = dir.find_last_of("/\\");
   if (pos == std::string::npos) {
     return false;
   }
@@ -372,7 +372,6 @@ AOTIModelPackageLoader::AOTIModelPackageLoader(
   std::string found_filenames; // Saving for bookkeeping
   std::string model_directory =
       "data" + k_separator + "aotinductor" + k_separator + model_name;
-  std::string const_directory = "data" + k_separator + "constants";
 
   for (uint32_t i = 0; i < zip_archive.m_total_files; i++) {
     uint32_t filename_len =
@@ -390,30 +389,14 @@ AOTIModelPackageLoader::AOTIModelPackageLoader(
     found_filenames += " ";
 
     // Only compile files in the specified model directory
-    if (c10::starts_with(filename_str, model_directory) ||
-        c10::starts_with(filename_str, const_directory)) {
+    if (filename_str.length() >= model_directory.length() &&
+        filename_str.substr(0, model_directory.length()) == model_directory) {
       std::string output_path_str = temp_dir_;
-
-      if (c10::starts_with(filename_str, model_directory)) {
-        output_path_str += k_separator;
-        output_path_str += filename_str;
-      } else { // startsWith(filename_str, const_directory)
-        // Extract constants to the same directory as the rest of the files
-        // to be consistent with internal implementation
-        size_t lastSlash = filename_str.find_last_of(k_separator);
-        std::string filename = filename_str;
-        if (lastSlash != std::string::npos) {
-          filename = filename_str.substr(lastSlash + 1);
-        }
-        output_path_str +=
-            k_separator + model_directory + k_separator + filename;
-      }
-
-      LOG(INFO) << "Extract file: " << filename_str << " to "
-                << output_path_str;
+      output_path_str += k_separator;
+      output_path_str += filename_str;
 
       // Create the parent directory if it doesn't exist
-      size_t parent_path_idx = output_path_str.find_last_of(k_separator);
+      size_t parent_path_idx = output_path_str.find_last_of("/\\");
       if (parent_path_idx == std::string::npos) {
         throw std::runtime_error(
             "Failed to find parent path in " + output_path_str);