optimum-rbln 0.8.1a6__tar.gz → 0.8.1rc0__tar.gz

optimum_rbln-0.8.1rc0/.github/scripts/validate_docstrings.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+"""
+Simple and reliable docstring validator for CI/CD.
+Focuses on catching issues that would break mkdocstrings parsing.
+"""
+
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+
+def infer_module_name(file_path: Path) -> str:
+    """
+    Infer module name from file path.
+
+    Example:
+        src/optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py
+        -> optimum.rbln.transformers.models.decoderonly.modeling_decoderonly
+    """
+    if not any(keyword in file_path.name for keyword in ["configuration", "modeling"]):
+        return None
+
+    parts = file_path.parts
+
+    # Find 'src' directory index
+    try:
+        src_index = parts.index("src")
+    except ValueError:
+        # If no 'src' directory found, try common patterns
+        for common_src in ["lib", "package", file_path.parts[0]]:
+            if common_src in parts:
+                src_index = parts.index(common_src)
+                break
+        else:
+            # Fallback: use the first directory
+            src_index = 0
+
+    # Extract module path parts (everything after src, excluding .py)
+    module_parts = parts[src_index + 1 :]
+
+    # Remove .py extension from the last part
+    if module_parts and module_parts[-1].endswith(".py"):
+        module_parts = module_parts[:-1] + (module_parts[-1][:-3],)
+
+    # Join with dots
+    module_name = ".".join(module_parts)
+    return module_name
+
+
+def test_mkdocstrings_parsing(file_path: Path, module_name: str = None) -> bool:
+    """Test if mkdocstrings can parse the module documentation."""
+    try:
+        if module_name:
+            # Test with specific module name
+            test_content = f"""# Test Documentation
+
+:::{module_name}
+    options:
+        show_source: false
+        filters: ["!^_"]
+"""
+        else:
+            # Auto-infer module name from file path
+            inferred_module = infer_module_name(file_path)
+            if inferred_module is None:
+                print(
+                    "⚠️ No module name provided, and the file is not a configuration or modeling file, skipping mkdocstrings test"
+                )
+                return True
+            print(f"ℹ️  Auto-inferred module name: {inferred_module}")
+            test_content = f"""# Test Documentation
+
+:::{inferred_module}
+    options:
+        show_source: false
+        filters: ["!^_"]
+"""
+            module_name = inferred_module
+
+        config = f"""site_name: Test
+plugins:
+  - mkdocstrings:
+      handlers:
+        python:
+          paths: ["{Path(__file__).parent.parent.parent / "src"}"]
+          options:
+            show_source: false
+            filters: ["!^_"]
+"""
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir)
+            (tmp_path / "mkdocs.yml").write_text(config)
+            docs_dir = tmp_path / "docs"
+            docs_dir.mkdir()
+            (docs_dir / "index.md").write_text(test_content)
+
+            result = subprocess.run(
+                ["mkdocs", "build", "--strict"], cwd=tmp_path, capture_output=True, text=True, timeout=10
+            )
+
+            if result.returncode != 0:
+                print(f"❌ mkdocstrings parsing failed for {file_path}:")
+                print(f"   Module: {module_name}")
+                # Show relevant errors and warnings
+                lines = result.stderr.split("\n")
+                relevant = [
+                    l for l in lines if any(keyword in l for keyword in ["griffe", "WARNING", "ERROR", "CRITICAL"])
+                ]
+                for line in relevant[:5]:  # Show first 5 relevant lines
+                    if line.strip():
+                        print(f"    {line.strip()}")
+                return False
+
+            return True
+
+    except subprocess.TimeoutExpired:
+        print(f"⚠️  Timeout testing {file_path} (non-fatal)")
+        return True  # Don't fail on timeout
+    except Exception as e:
+        print(f"⚠️  Could not test mkdocstrings for {file_path}: {e} (non-fatal)")
+        return True  # Don't fail on test errors
+
+
+def main():
+    """Main entry point."""
+    if len(sys.argv) < 2:
+        print("Usage: python validate_docstrings.py <python_file> [module_name]")
+        print()
+        print("Examples:")
+        print("  python validate_docstrings.py src/my_module.py")
+        print("  python validate_docstrings.py src/my_module.py my.module.name")
+        print()
+        print("Note: If module_name is not provided, it will be auto-inferred from file path")
+        sys.exit(1)
+
+    file_path = Path(sys.argv[1])
+    module_name = sys.argv[2] if len(sys.argv) > 2 else None
+
+    if not file_path.exists():
+        print(f"❌ File not found: {file_path}")
+        sys.exit(1)
+
+    print(f"🔍 Testing mkdocstrings parsing: {file_path}")
+    if module_name:
+        print(f"📦 Module: {module_name}")
+
+    # Test mkdocstrings compatibility
+    mkdocs_ok = test_mkdocstrings_parsing(file_path, module_name)
+
+    if mkdocs_ok:
+        print("✅ mkdocstrings parsing test passed!")
+        sys.exit(0)
+    else:
+        print("❌ mkdocstrings parsing test failed!")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/.github/workflows/check_code_quality.yml

@@ -1,19 +1,7 @@
 name: check_code_quality
 
 on:
-  push:
-    branches: [main]
-    paths:
-      - "src/**/*.py"
-      - "tests/**/*.py"
-      - "examples/**/*.py"
-
-  pull_request:
-    branches: [main]
-    paths:
-      - "src/**/*.py"
-      - "tests/**/*.py"
-      - "examples/**/*.py"
+  workflow_call:
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/.github/workflows/rbln_scheduled_test.yaml

@@ -10,6 +10,14 @@ env:
   HF_AUTH_TOKEN: ${{ secrets.HF_AUTH_TOKEN }}
 
 jobs:
+  check-code-quality:
+    uses: ./.github/workflows/check_code_quality.yml
+    
+  test-docstrings:
+    uses: ./.github/workflows/test-docstrings.yml
+    with:
+      test_all_files: true
+
   load-version:
     runs-on: rebel-k8s-runner
     outputs:

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/.github/workflows/rbln_trigger_on_pr.yaml

@@ -24,9 +24,19 @@ jobs:
             echo "No [skip ci] found, continuing with CI"
           fi
 
+  check-code-quality:
+    needs: check-skip-ci
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' }}
+    uses: ./.github/workflows/check_code_quality.yml
+    
+  test-docstrings:
+    needs: check-skip-ci
+    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' }}
+    uses: ./.github/workflows/test-docstrings.yml
+
   load-version:
     runs-on: rebel-k8s-runner
-    needs: check-skip-ci
+    needs: [check-skip-ci, check-code-quality, test-docstrings]
     if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' }}
     outputs:
       compiler_version: ${{ steps.get_version.outputs.compiler_version }}
@@ -41,7 +51,7 @@ jobs:
           echo "compiler_version=$VERSION" >> $GITHUB_OUTPUT
 
   check-compiler:
-    needs: [check-skip-ci, load-version]
+    needs: [check-skip-ci, check-code-quality, test-docstrings, load-version]
     if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' }}
     uses: ./.github/workflows/rbln_check_compiler.yaml
     with:
@@ -50,7 +60,7 @@ jobs:
 
   check-team-member:
     runs-on: rebel-k8s-runner
-    needs: [check-skip-ci, check-compiler]
+    needs: [check-skip-ci, check-code-quality, test-docstrings, check-compiler]
     if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-compiler.outputs.is_deploy_version == 'true' }}
     outputs:
       is_team_member: ${{ steps.check_member.outputs.IS_TEAM_MEMBER }}
@@ -75,7 +85,7 @@ jobs:
           fi
 
   optimum-rbln-pytest:
-    needs: [check-skip-ci, check-compiler, check-team-member]
+    needs: [check-skip-ci, check-code-quality, test-docstrings, check-compiler, check-team-member]
     if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-team-member.outputs.is_team_member == 'true' }}
     uses: ./.github/workflows/rbln_optimum_pytest.yaml
     with:
@@ -85,7 +95,7 @@ jobs:
     secrets: inherit
 
   optimum-rbln-inference-test:
-    needs: [check-skip-ci, check-compiler, check-team-member]
+    needs: [check-skip-ci, check-code-quality, test-docstrings, check-compiler, check-team-member]
     if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-team-member.outputs.is_team_member == 'true' }}
     uses: ./.github/workflows/rbln_optimum_inference_test.yaml
     with:

optimum_rbln-0.8.1rc0/.github/workflows/test-docstrings.yml

@@ -0,0 +1,137 @@
+name: Test Docstring Parsing
+
+on:
+  workflow_call:
+    inputs:
+      test_all_files:
+        description: "Test all files in src/ instead of just changed files"
+        required: false
+        type: boolean
+        default: false
+
+jobs:
+  test-docstrings:
+    runs-on: ubuntu-latest-rbln
+    
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0  # Get full git history for file diff
+    
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.9'
+    
+    - name: Install dependencies
+      run: |
+        pip install mkdocs mkdocs-material mkdocstrings mkdocstrings-python
+    
+    - name: Get Python files to test
+      id: changed-files
+      run: |
+        if [ "${{ inputs.test_all_files }}" = "true" ]; then
+          # For scheduled test, get all Python files in src/
+          ALL_FILES=$(find src/ -name '*.py' -type f || true)
+          echo "Testing all Python files in src/:"
+          echo "$ALL_FILES"
+          
+          if [ -n "$ALL_FILES" ]; then
+            FILES_JSON=$(echo "$ALL_FILES" | jq -R -s -c 'split("\n")[:-1]')
+            echo "files=$FILES_JSON" >> $GITHUB_OUTPUT
+            echo "has_changes=true" >> $GITHUB_OUTPUT
+          else
+            echo "has_changes=false" >> $GITHUB_OUTPUT
+          fi
+        else
+          # For PR, compare with base branch
+          if [ "${{ github.event_name }}" = "pull_request" ]; then
+            CHANGED_FILES=$(git diff --name-only origin/${{ github.base_ref }}...HEAD | grep '\.py$' | grep 'src/' || true)
+          else
+            # For push, compare with previous commit
+            CHANGED_FILES=$(git diff --name-only HEAD~1 HEAD | grep '\.py$' | grep 'src/' || true)
+          fi
+          
+          echo "Changed Python files:"
+          echo "$CHANGED_FILES"
+          
+          # Convert to JSON array for matrix
+          if [ -n "$CHANGED_FILES" ]; then
+            FILES_JSON=$(echo "$CHANGED_FILES" | jq -R -s -c 'split("\n")[:-1]')
+            echo "files=$FILES_JSON" >> $GITHUB_OUTPUT
+            echo "has_changes=true" >> $GITHUB_OUTPUT
+          else
+            echo "has_changes=false" >> $GITHUB_OUTPUT
+          fi
+        fi
+    
+    - name: Test docstring parsing
+      if: steps.changed-files.outputs.has_changes == 'true'
+      run: |
+        FILES='${{ steps.changed-files.outputs.files }}'
+        
+        # Exclude files (hardcoded)
+        EXCLUDE_FILES=(
+          "src/optimum/rbln/transformers/models/llava_next/modeling_llava_next.py"
+          "src/optimum/rbln/transformers/models/t5/modeling_t5.py"
+        )
+        
+        # Filter out excluded files
+        FILTERED_FILES=""
+        echo "$FILES" | jq -r '.[]' | while read -r file; do
+          should_exclude=false
+          for exclude_file in "${EXCLUDE_FILES[@]}"; do
+            if [[ "$file" == "$exclude_file" ]]; then
+              echo "⏭️  Skipping excluded file: $file"
+              should_exclude=true
+              break
+            fi
+          done
+          
+          if [[ "$should_exclude" == "false" ]]; then
+            echo "$file" >> /tmp/filtered_files.txt
+          fi
+        done
+        
+        # Check if we have files to process
+        if [[ ! -f /tmp/filtered_files.txt ]] || [[ ! -s /tmp/filtered_files.txt ]]; then
+          echo "No files to process after excluding files"
+          exit 0
+        fi
+        
+        TOTAL_FILES=$(wc -l < /tmp/filtered_files.txt)
+        echo "Processing $TOTAL_FILES files with 16 parallel processes..."
+        
+        # Function to validate a single file
+        validate_file() {
+          local file=$1
+          echo "Testing: $file"
+          python .github/scripts/validate_docstrings.py "$file"
+          if [ $? -ne 0 ]; then
+            echo "❌ Docstring test failed for $file"
+            return 1
+          fi
+          echo "✅ Docstring test passed for $file"
+          return 0
+        }
+        
+        # Export function for use with xargs
+        export -f validate_file
+        
+        # Process files in parallel with 16 processes
+        cat /tmp/filtered_files.txt | xargs -I {} -P 16 -n 1 bash -c 'validate_file "$@"' _ {}
+        
+        # Check if any process failed
+        if [ $? -ne 0 ]; then
+          echo "❌ Some docstring tests failed"
+          exit 1
+        fi
+        
+        echo "🎉 All docstring tests passed!"
+        
+        # Clean up
+        rm -f /tmp/filtered_files.txt
+    
+    - name: No changes detected
+      if: steps.changed-files.outputs.has_changes == 'false'
+      run: echo "No Python files to test" 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a6
+Version: 0.8.1rc0
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.1a6'
-__version_tuple__ = version_tuple = (0, 8, 1, 'a6')
+__version__ = version = '0.8.1rc0'
+__version_tuple__ = version_tuple = (0, 8, 1, 'rc0')

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNControlNetModelConfig, RBLNUNe
 
 
 class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion ControlNet pipelines.
-    """
-
     submodules = ["text_encoder", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py

@@ -21,8 +21,6 @@ from ..models.configuration_prior_transformer import RBLNPriorTransformerConfig
 
 
 class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
-    """Base configuration class for Kandinsky V2.2 decoder pipelines."""
-
     submodules = ["unet", "movq"]
     _movq_uses_encoder = False
 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
 class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion pipelines.
-    """
-
     submodules = ["text_encoder", "unet", "vae"]
     _vae_uses_encoder = False
 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNSD3Transformer2DModelConfig
 
 
 class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion 3 pipelines.
-    """
-
     submodules = ["transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "vae"]
     _vae_uses_encoder = False
 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py

@@ -20,10 +20,6 @@ from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
 class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
-    """
-    Base configuration for Stable Diffusion XL pipelines.
-    """
-
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]
     _vae_uses_encoder = False
 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/modeling_diffusers.py

@@ -412,24 +412,22 @@ class RBLNDiffusionMixin:
         return compiled_image_size
 
     def handle_additional_kwargs(self, **kwargs):
-        """
-        Function to handle additional compile-time parameters during inference.
-
-        If the additional variable is determined by another module, this method should be overrided.
-
-        Example:
-            ```python
-            if hasattr(self, "movq"):
-                compiled_image_size = self.movq.image_size
-                kwargs["height"] = compiled_image_size[0]
-                kwargs["width"] = compiled_image_size[1]
-
-            compiled_num_frames = self.unet.rbln_config.num_frames
-            if compiled_num_frames is not None:
-                kwargs["num_frames"] = compiled_num_frames
-            return kwargs
-            ```
-        """
+        # Function to handle additional compile-time parameters during inference.
+
+        # If the additional variable is determined by another module, this method should be overrided.
+
+        # Example:
+        #     ```python
+        #     if hasattr(self, "movq"):
+        #         compiled_image_size = self.movq.image_size
+        #         kwargs["height"] = compiled_image_size[0]
+        #         kwargs["width"] = compiled_image_size[1]
+
+        #     compiled_num_frames = self.unet.rbln_config.num_frames
+        #     if compiled_num_frames is not None:
+        #         kwargs["num_frames"] = compiled_num_frames
+        #     return kwargs
+        #     ```
         return kwargs
 
     @remove_compile_time_kwargs

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py

@@ -80,7 +80,12 @@ class RBLNAutoencoderKL(RBLNModel):
 
             wrapped_model.eval()
 
-            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+            compiled_models[model_name] = cls.compile(
+                wrapped_model,
+                rbln_compile_config=rbln_config.compile_cfgs[i],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map[model_name],
+            )
 
         return compiled_models
 

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

@@ -99,11 +99,21 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
             compiled_models = {}
             if rbln_config.uses_encoder:
                 encoder_model, decoder_model = cls.wrap_model_if_needed(model, rbln_config)
-                enc_compiled_model = cls.compile(encoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
+                enc_compiled_model = cls.compile(
+                    encoder_model,
+                    rbln_compile_config=rbln_config.compile_cfgs[0],
+                    create_runtimes=rbln_config.create_runtimes,
+                    device=rbln_config.device_map["encoder"],
+                )
                 compiled_models["encoder"] = enc_compiled_model
             else:
                 decoder_model = cls.wrap_model_if_needed(model, rbln_config)
-            dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[-1])
+            dec_compiled_model = cls.compile(
+                decoder_model,
+                rbln_compile_config=rbln_config.compile_cfgs[-1],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map["decoder"],
+            )
             compiled_models["decoder"] = dec_compiled_model
 
         finally:
@@ -115,7 +125,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
     def update_rbln_config_using_pipe(
         cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
-        rbln_config.vae.num_channels_latents = pipe.transformer.config.in_channels
+        rbln_config.vae.num_channels_latents = pipe.transformer.config.out_channels
         rbln_config.vae.vae_scale_factor_temporal = pipe.vae_scale_factor_temporal
         rbln_config.vae.vae_scale_factor_spatial = pipe.vae_scale_factor_spatial
         return rbln_config

{optimum_rbln-0.8.1a6 → optimum_rbln-0.8.1rc0}/src/optimum/rbln/diffusers/models/autoencoders/vq_model.py

@@ -78,7 +78,12 @@ class RBLNVQModel(RBLNModel):
 
             wrapped_model.eval()
 
-            compiled_models[model_name] = cls.compile(wrapped_model, rbln_compile_config=rbln_config.compile_cfgs[i])
+            compiled_models[model_name] = cls.compile(
+                wrapped_model,
+                rbln_compile_config=rbln_config.compile_cfgs[i],
+                create_runtimes=rbln_config.create_runtimes,
+                device=rbln_config.device_map[model_name],
+            )
 
         return compiled_models