Commit 3b39cb99 authored by Pierre Augier's avatar Pierre Augier
Browse files

Solve path_jit issues

parent 2ddadc7efcf0
......@@ -41,8 +41,9 @@ from transonic.util import (
has_to_build,
modification_date,
is_method,
path_jit_classes,
write_if_has_to_write,
find_module_name_from_path,
path_root,
)
if mpi.nb_proc == 1:
......@@ -53,9 +54,6 @@ is_transpiling = False
modules_backends = {backend_name: {} for backend_name in backends.keys()}
modules = modules_backends[backend_default]
# TODO: warning backend_default
path_jit_classes = mpi.Path(path_jit_classes)
def _get_transonic_calling_module(backend_name: str = None, index_frame: int = 2):
"""Get the Transonic instance corresponding to the calling module
......@@ -548,11 +546,15 @@ def jit_class(cls, jit_methods, backend):
cls_name = cls.__name__
mod_name = cls.__module__
module = sys.modules[mod_name]
if mod_name == "__main__":
mod_name = find_module_name_from_path(module.__file__)
path_jit_class = mpi.Path(backend.jit.path_class)
# 1. create a Python file with @jit functions and methods
python_path_dir = path_jit_classes / mod_name.replace(".", os.path.sep)
python_path_dir = path_jit_class / mod_name.replace(".", os.path.sep)
python_path = python_path_dir / (cls_name + ".py")
if mpi.has_to_build(python_path, module.__file__):
......@@ -570,7 +572,7 @@ def jit_class(cls, jit_methods, backend):
mpi.barrier()
# 2. import the file
python_mod_name = path_jit_classes.name + "." + mod_name + "." + cls_name
python_mod_name = path_jit_class.name + "." + mod_name + "." + cls_name
module = import_from_path(python_path, python_mod_name)
# 3. replace the methods
......
......@@ -11,6 +11,7 @@ Internal API
"""
import re
from pathlib import Path
try:
import numpy as np
......@@ -22,7 +23,7 @@ from transonic.signatures import make_signatures_from_typehinted_func
from transonic.log import logger
from transonic import mpi
from transonic.typing import format_type_as_backend_type, typeof
from transonic.util import get_source_without_decorator
from transonic.util import get_source_without_decorator, path_root
from .for_classes import produce_code_class
......@@ -35,8 +36,13 @@ class SubBackendJIT:
self.name_capitalized = name.capitalize()
self.type_formatter = type_formatter
def get_path_jit(self):
pass
self.path_base = Path(path_root) / self.name / "__jit__"
self.path_class = self.path_base.parent / "__jit_class__"
if mpi.rank == 0:
self.path_base.mkdir(parents=True, exist_ok=True)
self.path_class.mkdir(parents=True, exist_ok=True)
mpi.barrier()
def make_backend_source(self, info_analysis, func, path_backend):
func_name = func.__name__
......
......@@ -77,13 +77,6 @@ from transonic.util import (
modules_backends = {backend_name: {} for backend_name in backends.keys()}
modules = modules_backends[backend_default]
# TODO: warning backend_default
path_jit = mpi.Path(path_root) / backend_default / "__jit__"
if mpi.rank == 0:
path_jit.mkdir(parents=True, exist_ok=True)
mpi.barrier()
_COMPILE_JIT = strtobool(os.environ.get("TRANSONIC_COMPILE_JIT", "True"))
......@@ -128,6 +121,10 @@ class ModuleJIT:
"special": special,
}
backend = backends[backend_name]
path_jit = mpi.Path(backend.jit.path_base)
path_jit_class = mpi.Path(backend.jit.path_class)
# TODO: check if these files have to be written here...
# Write exterior code for functions
for file_name, code in code_ext["function"].items():
......@@ -137,14 +134,7 @@ class ModuleJIT:
# Write exterior code for classes
for file_name, code in code_ext["class"].items():
print("DEBUG JIT", path_jit)
print("self.module_name", self.module_name)
path_ext = (
path_jit.parent
/ "__jit_classes__"
/ self.module_name.replace(".", os.path.sep)
)
print("DEBUG", path_ext)
path_ext = path_jit_class / self.module_name.replace(".", os.path.sep)
path_ext_file = path_ext / (file_name + ".py")
write_if_has_to_write(path_ext_file, format_str(code), logger.info)
......@@ -254,6 +244,7 @@ class JIT:
mod.jit_functions[func_name] = self
module_name = mod.module_name
path_jit = mpi.Path(backend.jit.path_base)
path_backend = path_jit / module_name.replace(".", os.path.sep)
if mpi.rank == 0:
......
......@@ -6,19 +6,22 @@ from shutil import rmtree
import numpy as np
import pytest
from transonic.backends import backends
from transonic.compiler import scheduler, wait_for_all_extensions
from transonic.justintime import path_jit, modules
from transonic.justintime import modules
from transonic import mpi
from transonic.util import path_jit_classes, can_import_accelerator
from transonic.util import can_import_accelerator
from transonic.config import backend_default
backend = backends[backend_default]
scheduler.nb_cpus = 2
module_name = "transonic.for_test_justintime"
str_relative_path = module_name.replace(".", os.path.sep)
path_jit = mpi.PathSeq(path_jit)
path_jit = backend.jit.path_base
path_jit_classes = backend.jit.path_class
path_jit_dir = path_jit / str_relative_path
path_classes_dir = path_jit_classes / str_relative_path
......
......@@ -127,28 +127,23 @@ def can_import_accelerator(backend: str = backend_default):
return True
path_jit_classes = path_root / backend_default / "__jit_classes__"
def find_module_name_from_path(path_py: Path):
"""Find the module name from the path of a Python file
It is done by looking to ``sys.path`` to see how the module can be imported.
"""
path_py = Path(path_py)
cwd = Path.cwd()
path = path_py.absolute().parent
module_name = path_py.stem
# special case for jit_classes
try:
path_rel = path.relative_to(path_jit_classes)
except ValueError:
pass
else:
tmp = [path_jit_classes.name]
name_pack = str(path_rel).replace(os.path.sep, ".")
special_dir = "__jit_class__"
if special_dir in path.parts:
tmp = [special_dir]
name_pack = ".".join(path.parts[path.parts.index(special_dir) + 1 :])
if name_pack:
tmp.append(name_pack)
tmp.append(module_name)
......@@ -353,10 +348,11 @@ def clear_cached_extensions(module_name: str, force: bool = False):
"""
from transonic.justintime import path_jit
from transonic.backends import backends
from transonic import mpi
backend = backends[backend_default]
path_jit = mpi.Path(backend.jit.path_base)
if module_name.endswith(".py"):
module_name = module_name[:-3]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment