@ancient-application-36103 Having one issue actually that is somewhat related to the PR but I don’t think should affect the validity of this PR. Basically I am running some inference code where we have these basic steps: 1. Preprocess 2. Compile model for inferentia instance type 3. Do our inference 4. Other steps In step 2, we use something that was suggested by you or ville before:

class InferentiaCompiledModel:
    @classmethod
    def save_model(cls, model):
        with tempfile.NamedTemporaryFile() as f:
            torch.jit.save(model, f.name)
            return f.read()

    @classmethod
    def load_model(cls, blob):
        with tempfile.NamedTemporaryFile() as f:
            f.write(blob)
            f.flush()
            return torch.jit.load(f.name)

So we can do (at the end of step 2 above)

self.model_obj = InferentiaCompiledModel(compiled_model)

. This can be read in step 3 similarly (which is a parallel AWS Batch step) :

model = InferentiaCompiledModel(self.model_obj)

and this works very nicely when orchestrating our code locally. However, when using step functions, this fails weirdly on the load step (step 3) with some failure to initialise neuron error that comes from this loading line. The docker images seem fine, the inferentia capable machine is running the step and the docker image itself has neuron enabled, as I say everything is the same as the one where we orchestrate this locally. Is there anything you can think that would be messing this up. Note, there is a small ‘launching’ intermediate step which is sandwiched between the compile step (which runs on AWS batch but only one machine) and the inference step (which runs on AWS batch in a parallel way). This is the only difference between these two methods, where in the local orchestrator, this step receives the self. model_obj from AWS and passes it to the parallel step whereas on Step Functions, this is all handled by AWS. The flow graph is attached as reference if any of this is unclear. I will say this is, in my opinion shouldn’t affect the PR as the container with inf enabled fires up correctly according to the decorator.

stepfunctions_graph.png

@dry-beach-38304 Would be good if you could re-add the okay-to-test label on this/also help to approve this

given a lot seem to be due to formatting, you might have an easier time cherry-picking yoru chnages back onto master and pushing that instead.

so it should be generally similar

iirc, we use everything default.

If you have pre-commit installed, it should run it for you.